Skip to main content

memra_engine/
hybrid_forward.rs

1//! Hybrid forward pass (Stage-1, f32, prefill, single sequence). Per layer dispatches to a
2//! linear-attention (Gated DeltaNet) or full-attention mixer, then SwiGLU FFN. Matches
3//! llama.cpp src/models/qwen35.cpp node-for-node.
4
5use crate::Engine;
6use crate::cache::Cache;
7use cudarc::driver::CudaSlice;
8use memra_gguf::config::ModelConfig;
9
10/// Resident trunk transients for the eager prime (piecewise-graph foundation; see
11/// HybridModel::prime_slabs). Every buffer is fully overwritten before use per prime.
12pub struct PrimeSlabs {
13    pub t_cap: usize,
14    pub h: CudaSlice<f32>,
15    pub x1: CudaSlice<f32>,
16    pub z: CudaSlice<f32>,
17    pub act: CudaSlice<f32>,
18    pub xa: CudaSlice<f32>,
19    pub xb: CudaSlice<f32>,
20    pub h16: CudaSlice<u8>,
21    pub z16: CudaSlice<u8>,
22    /// piecewise boundary slabs (increment 2): GEMM outputs land here so the
23    /// downstream captured segments see fixed addresses.
24    pub gate: CudaSlice<f32>, // t * n_ff_max
25    pub up: CudaSlice<f32>,      // t * n_ff_max
26    pub ffn_out: CudaSlice<f32>, // t * n_embd
27    /// piecewise increment 3: per-layer S-glue segment graphs (down-add + next
28    /// attn-norm, ALL-slab IO, zero in-graph allocations -> keeperless capture is
29    /// clean). Baked at this t_cap; replay only when t == t_cap. seg_glue[il] fires
30    /// between layer il and il+1 (ping-pong parity is deterministic per il).
31    pub seg_glue: Vec<Option<cudarc::driver::CudaGraph>>,
32    /// increment 5 (core-split edition): the mixer out-GEMM writes _into_ `mixed`
33    /// directly (no staging copy — the increment-4 copy route was refuted), making
34    /// S-mid [add + post-norm] all-slab and capturable.
35    pub mixed: CudaSlice<f32>,
36    pub seg_mid: Vec<Option<cudarc::driver::CudaGraph>>,
37    pub seg_t: usize,
38}
39
40// Split prime ranges cannot enter the full-range segment-graph arm, and every slab access
41// is serialized by its device mutex after binding that device's CUDA context on the thread.
42unsafe impl Send for PrimeSlabs {}
43
44fn empty_cache_layers<T>(n: usize) -> Vec<Option<T>> {
45    std::iter::repeat_with(|| None).take(n).collect()
46}
47
48/// Temporarily move a PP-2 cache's layer state into two independently-owned cache shells.
49/// The stage walkers then receive disjoint `&mut Cache` values and can run on separate host
50/// threads without aliasing. GPU buffers are moved, not copied; Drop restores every layer
51/// and publishes the last position completed by both stages.
52struct PrimeCacheStages<'a> {
53    parent: &'a mut Cache,
54    cut: usize,
55    stage0: Cache,
56    stage1: Cache,
57}
58
59impl<'a> PrimeCacheStages<'a> {
60    fn new(parent: &'a mut Cache, cut: usize) -> Self {
61        let n = parent.kv.len();
62        assert_eq!(parent.recur.len(), n, "cache layer vectors disagree");
63        assert!(cut <= n, "PP-2 cache cut {cut} exceeds {n} layers");
64        let mut kv0 = empty_cache_layers(n);
65        let mut kv1 = empty_cache_layers(n);
66        let mut recur0 = empty_cache_layers(n);
67        let mut recur1 = empty_cache_layers(n);
68        for i in 0..cut {
69            kv0[i] = parent.kv[i].take();
70            recur0[i] = parent.recur[i].take();
71        }
72        for i in cut..n {
73            kv1[i] = parent.kv[i].take();
74            recur1[i] = parent.recur[i].take();
75        }
76        let pos = parent.pos;
77        let max_ctx = parent.max_ctx;
78        Self {
79            parent,
80            cut,
81            stage0: Cache {
82                kv: kv0,
83                recur: recur0,
84                pos,
85                max_ctx,
86                last_logits_dev: None,
87                dflash_taps: None,
88            },
89            stage1: Cache {
90                kv: kv1,
91                recur: recur1,
92                pos,
93                max_ctx,
94                last_logits_dev: None,
95                dflash_taps: None,
96            },
97        }
98    }
99
100    fn parts(&mut self) -> (&mut Cache, &mut Cache) {
101        (&mut self.stage0, &mut self.stage1)
102    }
103}
104
105impl Drop for PrimeCacheStages<'_> {
106    fn drop(&mut self) {
107        let n = self.parent.kv.len();
108        for i in 0..n {
109            let source = if i < self.cut {
110                &mut self.stage0
111            } else {
112                &mut self.stage1
113            };
114            debug_assert!(self.parent.kv[i].is_none());
115            debug_assert!(self.parent.recur[i].is_none());
116            self.parent.kv[i] = source.kv[i].take();
117            self.parent.recur[i] = source.recur[i].take();
118        }
119        self.parent.pos = self.stage0.pos.min(self.stage1.pos);
120    }
121}
122
123/// task #18 (attn side): one sequence's pre-attention outputs (post-rope q/k, v, out-gate).
124pub(crate) struct AttnPre {
125    pub q: cudarc::driver::CudaSlice<f32>,
126    pub k: cudarc::driver::CudaSlice<f32>,
127    pub v: cudarc::driver::CudaSlice<f32>,
128    pub gate: Option<cudarc::driver::CudaSlice<f32>>,
129}
130
131/// task #18: one sequence's GDN prep outputs (the scan inputs).
132pub(crate) struct GdnPrep {
133    pub hk: usize,
134    pub q_l2: cudarc::driver::CudaSlice<f32>,
135    pub k_l2: cudarc::driver::CudaSlice<f32>,
136    pub v_g: cudarc::driver::CudaSlice<f32>,
137    pub beta: cudarc::driver::CudaSlice<f32>,
138    pub g_log: cudarc::driver::CudaSlice<f32>,
139    pub kb16: Option<cudarc::driver::CudaSlice<u8>>,
140    pub qb16: Option<cudarc::driver::CudaSlice<u8>>,
141}
142
143/// Device scratch for the burst verify stream (see `verify_stream_scratch`).
144pub(crate) struct VerifyStreamScratch {
145    pub pos_d: CudaSlice<i32>,
146    pub row_ctrs: Vec<CudaSlice<i32>>,
147}
148use crate::hybrid::{FullAttnLayer, HybridModel, LinearAttnLayer, Mixer, MoeWeights};
149
150struct MoeInputTraceWriter {
151    dir: std::path::PathBuf,
152    index: std::fs::File,
153    payloads: std::collections::HashMap<u16, (std::fs::File, u64)>,
154}
155
156static MOE_INPUT_TRACE_WRITER: std::sync::OnceLock<std::sync::Mutex<Option<MoeInputTraceWriter>>> =
157    std::sync::OnceLock::new();
158
159/// STAGE-2 GROUPED DECODE gate (MEMRA_MOE_GDEC, default ON; `=0` restores the sequential
160/// per-expert launch chain). See `moe_gdec_token`.
161fn gdec_enabled() -> bool {
162    static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
163    *E.get_or_init(|| {
164        std::env::var("MEMRA_MOE_GDEC")
165            .map(|v| v != "0")
166            .unwrap_or(true)
167    })
168}
169
170/// SLAB-LOCAL RESIDENT ARM gate (lane/pp-leverb 2026-08-08, MEMRA_MOE_SLAB, default ON;
171/// `=0` restores the SLRU dispatch even when resident slabs exist). Read PER CALL, never
172/// memoized — probes A/B the two provenances in one process (the MEMRA_PRIME_PP pattern).
173/// See `moe_ffn_sequential_zq8`'s slab_local arm: the sigmoid-router archs (step35/M3/Hy3)
174/// are denied every `dev_exps` consumer (pairs/dev route softmax), so before this arm the
175/// fits-VRAM resident slabs were UPLOADED for them but never READ — the SLRU kept staging
176/// the same bytes beside a dead copy (37 GB H2D per pp4096 prime on the Step SKU, anatomy
177/// receipt). The arm reads the SAME bytes through the SAME kernels; only the pointer
178/// PROVENANCE changes (slab base + ex*stride vs SLRU slot address) — the bit-identity class
179/// `moe_ffn_dev`'s resident arm already documents against its SLRU arm.
180fn moe_slab_enabled() -> bool {
181    std::env::var("MEMRA_MOE_SLAB").as_deref() != Ok("0")
182}
183
184/// Expert-grouped dispatch remains opt-in after the local 5090 transfer gate rejected the
185/// default flip. `=0` selects the established path, while any other explicit value enables the
186/// grouped research arm for the current call.
187fn moe_grouped_enabled(_cfg: &ModelConfig, _prefill: bool) -> bool {
188    std::env::var("MEMRA_MOE_GROUPED")
189        .map(|value| value != "0")
190        .unwrap_or(false)
191}
192
193/// Deterministic in-token expert prefetch. `MEMRA_MOE_PREFETCH=1` overlaps memory-source H2D on the
194/// copy stream; selecting the opt-in worker spill backend enables the same known-next hook for disk.
195fn moe_prefetch_enabled() -> bool {
196    static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
197    *E.get_or_init(|| {
198        std::env::var("MEMRA_MOE_PREFETCH").as_deref() == Ok("1")
199            || crate::spill_pread::worker_enabled()
200    })
201}
202
203/// Best-effort OS page-cache prefetch distance for mmap-backed expert ranges. Independent of the
204/// H2D copy-stream experiment so storage->RAM and RAM->HBM overlap can be measured separately.
205/// The opt-in default stays one expert to preserve the original experiment; spill rigs can widen
206/// it with `MEMRA_MOE_PAGE_PREFETCH_WINDOW` to cover NVMe latency.
207fn moe_page_prefetch_window() -> usize {
208    static W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
209    *W.get_or_init(|| {
210        page_prefetch_window_from_values(
211            std::env::var("MEMRA_MOE_PAGE_PREFETCH").as_deref() == Ok("1"),
212            std::env::var("MEMRA_MOE_PAGE_PREFETCH_WINDOW")
213                .ok()
214                .as_deref(),
215        )
216    })
217}
218
219fn page_prefetch_window_from_values(enabled: bool, raw_window: Option<&str>) -> usize {
220    if !enabled {
221        return 0;
222    }
223    raw_window.and_then(|value| value.parse().ok()).unwrap_or(1)
224}
225
226/// Return only the newly exposed positions in a rolling lookahead window. Position zero seeds the
227/// full window; each later position adds one expert at the far edge. Thus widening the window does
228/// not repeatedly issue `MADV_WILLNEED` for the same range.
229fn page_prefetch_positions(position: usize, len: usize, window: usize) -> std::ops::Range<usize> {
230    if window == 0 || position >= len {
231        return len..len;
232    }
233    let (start, count) = if position == 0 {
234        (1, window)
235    } else {
236        (position.saturating_add(window), 1)
237    };
238    let start = start.min(len);
239    start..start.saturating_add(count).min(len)
240}
241
242/// Grouped worker-I/O schedule: prime the first active expert before the loop, then queue exactly
243/// one known-next expert at each iteration. Returning positions keeps expert ordering authoritative.
244fn grouped_worker_prefetch_position(order_len: usize, current: Option<usize>) -> Option<usize> {
245    let position = current.map_or(0, |position| position.saturating_add(1));
246    (position < order_len).then_some(position)
247}
248
249/// Fill the worker ring with complete experts, retaining one pinned buffer for an unexpected
250/// demand miss. Each expert has gate/up/down extents, so depth 16 admits a rolling five-expert
251/// window. Position zero primes the current expert too: its three independent reads can run in
252/// parallel instead of demand-serializing gate, up, and down before any useful GPU work exists.
253fn worker_prefetch_window() -> usize {
254    static WINDOW: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
255    *WINDOW.get_or_init(|| {
256        let automatic = crate::spill_pread::configured_depth().saturating_sub(1) / 3;
257        std::env::var("MEMRA_SPILL_WORKER_EXPERT_WINDOW")
258            .ok()
259            .and_then(|value| value.parse::<usize>().ok())
260            .unwrap_or(automatic.max(1))
261    })
262}
263
264/// Return only positions newly exposed by a rolling worker-I/O window. Unlike mmap page advice,
265/// this includes the current expert when the window is seeded so all three current projections
266/// enter the CPU pool together.
267fn worker_prefetch_positions(position: usize, len: usize, window: usize) -> std::ops::Range<usize> {
268    if window == 0 || position >= len {
269        return len..len;
270    }
271    let (start, count) = if position == 0 {
272        (0, window)
273    } else {
274        (position.saturating_add(window).saturating_sub(1), 1)
275    };
276    let start = start.min(len);
277    start..start.saturating_add(count).min(len)
278}
279
280/// LAUNCH-STRUCTURE STAGE 3 gate (MEMRA_MOE_DEV, default ON; `=0` restores host routing). The
281/// zero-DtoH device-dispatch path for fully-resident layers: router top-k output stays on device,
282/// expert weight pointers come from the per-layer device table. Requires the fused router (the
283/// dev path consumes the device sel/w directly), so MEMRA_FUSED_ROUTER=0 also disables it.
284fn moe_dev_enabled() -> bool {
285    static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
286    *E.get_or_init(|| {
287        std::env::var("MEMRA_MOE_DEV")
288            .map(|v| v != "0")
289            .unwrap_or(true)
290            && !matches!(std::env::var("MEMRA_FUSED_ROUTER").as_deref(), Ok("0"))
291    })
292}
293
294/// Device sigmoid top-k is the default for Step-3.7 / M3 / Hy3 / GLM-DSA. `MEMRA_SIG_ROUTER=0` restores
295/// the full-logit DtoH plus `moe_route_sigmoid_host` oracle without changing expert dispatch.
296fn sigmoid_router_enabled() -> bool {
297    static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
298    *E.get_or_init(|| {
299        std::env::var("MEMRA_SIG_ROUTER")
300            .map(|v| v != "0")
301            .unwrap_or(true)
302    })
303}
304
305/// MoE EXPERT dp4a gate (MEMRA_MOE_Q8, default ON; `=0` restores the Stage-A f32-dequant expert
306/// kernels). Applies when gate/up/down expert qtypes are all in the dp4a body set (IQ3_S/IQ4_XS).
307/// FP-order differs from Stage-A (int dp4a + warp tree) — argmax/run-gen/stream-identity gates
308/// arbitrate; the sequential and fused q8 paths ship as a matched pair (MEMRA_MOE_GATE contract).
309fn moe_q8_enabled() -> bool {
310    static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
311    *E.get_or_init(|| {
312        std::env::var("MEMRA_MOE_Q8")
313            .map(|v| v != "0")
314            .unwrap_or(true)
315    })
316}
317
318/// gemma4 fast-arm gate: qtypes with an `expert_dot_g` dp4a body (superset used by the gelu
319/// dev arm; the qwen q8 arms keep their own battery-gated q8_expert_supported policy).
320fn expert_dp4a_supported(qt: i32) -> bool {
321    qt == crate::QT_Q4_0
322        || qt == crate::QT_IQ3_S
323        || qt == crate::QT_IQ4_XS
324        || qt == crate::QT_Q3_K
325        || qt == crate::QT_Q4_K
326        || qt == crate::QT_Q6_K
327}
328
329fn q8_expert_supported(qt: i32) -> bool {
330    // k-quant arms added 2026-07-06 (Q3_K/Q4_K/Q6_K bodies for the UD tail layers). Briefly
331    // default-excluded the same day when they appeared to break 35B real-prompt spec — the
332    // ACTUAL culprit was the MoE router's cuBLASLt n-dependence (d994271); with the router
333    // decode-exact at verify t, the k-quant arms pass the full spec battery (p1/p2/p3 + raw
334    // K=1..8) and are DEFAULT ON again (+9 tok/s: 148.9 -> 157.9). MEMRA_MOE_Q8_KQ=0 excludes.
335    static KQ: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
336    let kq = *KQ.get_or_init(|| {
337        std::env::var("MEMRA_MOE_Q8_KQ")
338            .map(|v| v != "0")
339            .unwrap_or(true)
340    });
341    // NVFP4 experts: DEFAULT ON (2026-07-17). The M3-era "decode-vs-verify MISMATCH 3.4e1"
342    // that had this excluded was the missing per-expert macro-scale fold, fixed in the
343    // dev-kernel epilogues + moe_w_scale_by_expert; the 35B ct-NVFP4 artifact now runs the
344    // q8 arm at parity with the IQ4_XS daily (174-178 tok/s, spec K=1..8 exact). M3/Hy3
345    // never reach the q8 arms regardless (sigmoid-router cfg gates on pairs/dev/gdec).
346    // MEMRA_MOE_Q8_NVFP4=0 restores the f32 arm.
347    let nvfp4_q8 = std::env::var("MEMRA_MOE_Q8_NVFP4")
348        .map(|v| v != "0")
349        .unwrap_or(true);
350    qt == crate::QT_IQ3_S
351        || qt == crate::QT_IQ4_XS
352        || (nvfp4_q8 && qt == crate::QT_NVFP4)
353        || (kq && (qt == crate::QT_Q3_K || qt == crate::QT_Q4_K || qt == crate::QT_Q6_K))
354}
355
356/// The decode-once (_dec) and IQ-MMA expert kernels dequant via IQ-specific extractors —
357/// k-quant tensors must fall to the _em dot path instead.
358fn q8_expert_dec_supported(qt: i32) -> bool {
359    qt == crate::QT_IQ3_S || qt == crate::QT_IQ4_XS || qt == crate::QT_Q4_0
360}
361
362/// Grouped-f16 door (MEMRA_MOE_F16G) per-projection admission: the qtype has a dequant-to-f16
363/// kernel in cu/moe_f16_grouped.cu AND the projection's k dimension tiles its block size.
364/// Round 49 widened coverage to q35's UD mix (gate/up IQ3_S x39 + Q3_K x1 + IQ4_XS x1; down
365/// IQ4_XS x37 + Q6_K x3 + Q4_K x1) — the round-47 IQ4_XS/Q4_0-only table admitted ~1 of 41
366/// q35 layers, which is why that cell measured FLAT.
367fn f16g_proj_ok(qt: i32, in_f: usize) -> bool {
368    match qt {
369        crate::QT_Q4_0 => in_f % 32 == 0,
370        crate::QT_IQ4_XS | crate::QT_IQ3_S | crate::QT_Q3_K | crate::QT_Q4_K | crate::QT_Q6_K => {
371            in_f % 256 == 0
372        }
373        _ => false,
374    }
375}
376
377/// STAGE 3 prewarm gate (MEMRA_MOE_PREWARM, default ON; `=0` leaves residency organic). One-shot
378/// per layer: force-admit every block while FREE slots cover the whole layer (never evicts).
379fn moe_prewarm_enabled() -> bool {
380    static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
381    *E.get_or_init(|| {
382        std::env::var("MEMRA_MOE_PREWARM")
383            .map(|v| v != "0")
384            .unwrap_or(true)
385    })
386}
387
388/// During a discarded fixed-residency profile, admit CPU-routed misses after their current-token
389/// CPU result is complete. The current result and numeric path are unchanged; later warmup tokens
390/// can then vote for and exercise those experts on GPU before the cache is frozen.
391fn cpu_expert_profile_admit_enabled() -> bool {
392    static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
393    *E.get_or_init(|| std::env::var("MEMRA_CPU_EXPERT_FREEZE_PROFILE_ADMIT").as_deref() == Ok("1"))
394}
395
396/// Minimum prompt length for the BATCHED cache prime (`prime_cache`). Below this the tokenwise
397/// decode loop wins anyway (the batched path's GEMM dispatch needs m>=16, and the stateful conv
398/// kernel needs T >= d_conv-1). Callers: generate / generate_spec.
399pub const PRIME_MIN_T: usize = 16;
400const PRIME_PIPE_MICROBATCHES: usize = 8;
401const PRIME_PIPE_MIN_CHUNK: usize = 128;
402const PRIME_PIPE_EDGE_MIN_CHUNK: usize = 64;
403const PRIME_PIPE_LINEAR_WORK: usize = 8;
404
405fn prime_pp2_auto_geometry(n_layers: usize) -> bool {
406    crate::pp::prime_pp_on()
407        && !crate::pp::pp2_streams_off()
408        && crate::pp::pp_cuts(n_layers).is_some_and(|cuts| cuts.len() == 3)
409}
410
411/// Effective internal prime chunk. An explicit MEMRA_PRIME_CHUNK is authoritative.
412/// Naked PP-2 primes use the measured pipeline geometry: up to eight microchunks, never
413/// below 128 tokens, while the legacy 4096-token cap remains the long-context bound.
414pub fn prime_chunk_tokens(t: usize, n_layers: usize) -> usize {
415    if let Ok(value) = std::env::var("MEMRA_PRIME_CHUNK") {
416        let parsed = value
417            .parse::<usize>()
418            .unwrap_or(crate::cache::PRIME_CHUNK_MAX_TOKENS);
419        return if crate::cache::swa_ring_on() {
420            if parsed == 0 {
421                crate::cache::PRIME_CHUNK_MAX_TOKENS
422            } else {
423                parsed.min(crate::cache::PRIME_CHUNK_MAX_TOKENS)
424            }
425        } else {
426            parsed
427        };
428    }
429    let chunk = crate::cache::PRIME_CHUNK_MAX_TOKENS;
430    if prime_pp2_auto_geometry(n_layers) && t >= 2 * PRIME_PIPE_MIN_CHUNK {
431        chunk.min(
432            t.div_ceil(PRIME_PIPE_MICROBATCHES)
433                .max(PRIME_PIPE_MIN_CHUNK),
434        )
435    } else {
436        chunk
437    }
438}
439
440fn fixed_prime_chunk_ranges(t: usize, chunk: usize) -> Vec<(usize, usize)> {
441    fixed_prime_chunk_ranges_for_ring(t, chunk, crate::cache::swa_ring_on())
442}
443
444fn fixed_prime_chunk_ranges_for_ring(t: usize, chunk: usize, ring_on: bool) -> Vec<(usize, usize)> {
445    if chunk == 0 || t <= chunk {
446        return vec![(0, t)];
447    }
448    let mut ranges = Vec::with_capacity(t.div_ceil(chunk));
449    let mut start = 0usize;
450    while start < t {
451        let mut end = (start + chunk).min(t);
452        if t - end > 0 && t - end < PRIME_MIN_T {
453            if ring_on {
454                let shifted = t - PRIME_MIN_T;
455                end = if shifted > start { shifted } else { t };
456            } else {
457                end = t;
458            }
459        }
460        ranges.push((start, end));
461        start = end;
462    }
463    ranges
464}
465
466fn prime_chunk_work(prefix: usize, total: usize) -> u128 {
467    let prefix = prefix as u128;
468    prefix * (prefix + (PRIME_PIPE_LINEAR_WORK as u128) * (total as u128))
469}
470
471fn dynamic_prime_chunk_ranges(
472    t: usize,
473    fixed_chunk: usize,
474    fixed: &[(usize, usize)],
475) -> Vec<(usize, usize)> {
476    let n = fixed.len();
477    if n < 3 {
478        return fixed.to_vec();
479    }
480
481    let max_first = t - (n - 1) * PRIME_MIN_T;
482    let first = fixed_chunk
483        .div_ceil(2)
484        .max(PRIME_PIPE_EDGE_MIN_CHUNK)
485        .min(max_first);
486    let mut ranges = Vec::with_capacity(n);
487    ranges.push((0, first));
488
489    let first_work = prime_chunk_work(first, t);
490    let work_span = prime_chunk_work(t, t) - first_work;
491    let denominator = (n - 1) as u128;
492    let mut previous = first;
493    for boundary in 1..n - 1 {
494        let target = first_work * denominator + work_span * (boundary as u128);
495        let remaining = n - 1 - boundary;
496        let mut low = previous + PRIME_MIN_T;
497        let mut high = t - remaining * PRIME_MIN_T;
498        while low < high {
499            let mid = low + (high - low) / 2;
500            if prime_chunk_work(mid, t) * denominator >= target {
501                high = mid;
502            } else {
503                low = mid + 1;
504            }
505        }
506        ranges.push((previous, low));
507        previous = low;
508    }
509    ranges.push((previous, t));
510    ranges
511}
512
513/// Internal prime ranges. The naked PP-2 pipeline defaults to a short-fill,
514/// equal-modeled-time schedule; MEMRA_PRIME_CHUNK_SCHED=fixed restores the measured
515/// equal-token ranges. An explicit MEMRA_PRIME_CHUNK always retains fixed semantics.
516pub fn prime_chunk_ranges(t: usize, n_layers: usize) -> Vec<(usize, usize)> {
517    let explicit_chunk = std::env::var_os("MEMRA_PRIME_CHUNK").is_some();
518    let chunk = prime_chunk_tokens(t, n_layers);
519    let fixed = fixed_prime_chunk_ranges(t, chunk);
520    let dynamic = match std::env::var("MEMRA_PRIME_CHUNK_SCHED") {
521        Ok(value) => value == "dynamic",
522        Err(_) => true,
523    };
524    if explicit_chunk || !dynamic || !prime_pp2_auto_geometry(n_layers) {
525        fixed
526    } else {
527        dynamic_prime_chunk_ranges(t, chunk, &fixed)
528    }
529}
530
531impl HybridModel {
532    /// CHUNK-INVARIANCE BISECT SEAM (lane/chunk-invariance): `MEMRA_PRIME_TRACE=<path>`
533    /// appends one JSONL row per (chunk, layer) with a hash of that layer's last-row
534    /// post-residual hidden. Diagnostic only — never on in a measured or gated run
535    /// (it forces a dtoh + host hash per layer).
536    fn prime_trace_path() -> Option<&'static str> {
537        static P: std::sync::OnceLock<Option<String>> = std::sync::OnceLock::new();
538        P.get_or_init(|| std::env::var("MEMRA_PRIME_TRACE").ok())
539            .as_deref()
540    }
541
542    /// Prefill forward over `tokens`; returns logits [T, n_vocab] (host f32).
543    pub fn forward(
544        &self,
545        e: &Engine,
546        tokens: &[u32],
547    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
548        if self.is_gemma4_e4b() {
549            return self.gemma4_e4b_forward(e, tokens, false);
550        }
551        if self.cfg.gemma4.is_some() {
552            return self.gemma4_forward(e, tokens, false);
553        }
554        let cfg = &self.cfg;
555        let n_embd = cfg.n_embd as usize;
556        let t = tokens.len();
557        let eps = cfg.rms_eps;
558        let pos: Vec<i32> = (0..t as i32).collect();
559        let pos_d = e.htod_i32(&pos)?;
560
561        let mut x = self.embed(e, tokens)?; // [T, n_embd]
562
563        for (il, layer) in self.layers.iter().enumerate() {
564            // attn_norm
565            let mut h = e.uninit(t * n_embd)?;
566            e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
567
568            let mixed = match &layer.mixer {
569                Mixer::Full(fa) => self.full_attn(e, fa, &h, &pos_d, t, il)?,
570                Mixer::Linear(la) => self.linear_attn(e, la, &h, t)?,
571                Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
572            };
573
574            // residual 1
575            let mut x1 = e.uninit(t * n_embd)?;
576            e.add(&x, &mixed, &mut x1, t * n_embd)?;
577
578            // pre-FFN norm (post_attention_norm), FFN (Dense or MoE), residual 2
579            let mut z = e.uninit(t * n_embd)?;
580            e.rms_norm(
581                &x1,
582                layer.post_attn_norm.float_data(),
583                &mut z,
584                n_embd,
585                t,
586                eps,
587            )?;
588            let ffn_out = match &layer.ffn {
589                crate::hybrid::Ffn::Dense {
590                    ffn_gate,
591                    ffn_up,
592                    ffn_down,
593                } => {
594                    let n_ff = ffn_gate.out_features();
595                    let mut g2 = e.matmul_group(&[ffn_gate, ffn_up], &z, t)?;
596                    let up = g2.pop().unwrap();
597                    let gate = g2.pop().unwrap();
598                    let mut act = e.uninit(t * n_ff)?;
599                    // A DENSE FFN reads the SHEXP clamp array: upstream's one `build_ffn` serves
600                    // both the dense MLP and the shared expert, and its limit is
601                    // `swiglu_clamp_shexp[il]` (llama-graph.cpp:1751). step35's leading dense
602                    // blocks 0-2 therefore key off clamp_shexp, not clamp_exp.
603                    Self::ffn_act_lim(
604                        e,
605                        &self.cfg,
606                        &gate,
607                        &up,
608                        1.0,
609                        1.0,
610                        self.cfg.clamp_shexp_at(il as u32),
611                        &mut act,
612                        t * n_ff,
613                    )?;
614                    e.matmul(ffn_down, &act, t)?
615                }
616                crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_prefill(e, m, &z, t, il as u16)?,
617            };
618            let mut x2 = e.uninit(t * n_embd)?;
619            e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
620            x = x2;
621        }
622
623        let mut hn = e.uninit(t * n_embd)?;
624        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
625        let logits = e.matmul(&self.output, &hn, t)?;
626        Ok(e.dtoh(&logits)?)
627    }
628
629    /// Prefill that returns ONLY the last token's logits — the common case (greedy/sample needs
630    /// just the final position to start decode). Runs the trunk over all T, then the lm_head
631    /// (output.weight, the largest matrix — 248320 rows) on the LAST hidden row ONLY, not all T.
632    /// On a 512-token prompt this turns a [512,248320] GEMM into [1,248320] — the dominant prefill
633    /// cost (nsys: ~99ms when done for all T). Bit-identical last-row logits to forward()[last].
634    pub fn forward_last(
635        &self,
636        e: &Engine,
637        tokens: &[u32],
638    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
639        if self.cfg.gemma4.is_some() {
640            return self.gemma4_forward(e, tokens, true);
641        }
642        let cfg = &self.cfg;
643        let n_embd = cfg.n_embd as usize;
644        let t = tokens.len();
645        let eps = cfg.rms_eps;
646        let pos: Vec<i32> = (0..t as i32).collect();
647        let pos_d = e.htod_i32(&pos)?;
648
649        let mut x = self.embed(e, tokens)?; // [T, n_embd]
650        // MEMRA_LAYER_PROBE=1: synchronize + print after every stage — bisects an in-graph
651        // ILLEGAL_ADDRESS to (layer, stage) at ~1 line of output per layer (M3 bring-up tool).
652        let probe = std::env::var("MEMRA_LAYER_PROBE").is_ok();
653        for (il, layer) in self.layers.iter().enumerate() {
654            let mut h = e.uninit(t * n_embd)?;
655            e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
656            if probe {
657                e.stream().synchronize()?;
658                eprintln!("[probe] L{il} norm ok");
659            }
660            let mixed = match &layer.mixer {
661                Mixer::Full(fa) => self.full_attn(e, fa, &h, &pos_d, t, il)?,
662                Mixer::Linear(la) => self.linear_attn(e, la, &h, t)?,
663                Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
664            };
665            if probe {
666                e.stream().synchronize()?;
667                eprintln!("[probe] L{il} mixer ok");
668            }
669            let mut x1 = e.uninit(t * n_embd)?;
670            e.add(&x, &mixed, &mut x1, t * n_embd)?;
671            let mut z = e.uninit(t * n_embd)?;
672            e.rms_norm(
673                &x1,
674                layer.post_attn_norm.float_data(),
675                &mut z,
676                n_embd,
677                t,
678                eps,
679            )?;
680            let ffn_out = match &layer.ffn {
681                crate::hybrid::Ffn::Dense {
682                    ffn_gate,
683                    ffn_up,
684                    ffn_down,
685                } => {
686                    let n_ff = ffn_gate.out_features();
687                    let mut g2 = e.matmul_group(&[ffn_gate, ffn_up], &z, t)?;
688                    let up = g2.pop().unwrap();
689                    let gate = g2.pop().unwrap();
690                    let mut act = e.uninit(t * n_ff)?;
691                    // dense FFN keys off the SHEXP clamp array — see forward()'s note.
692                    Self::ffn_act_lim(
693                        e,
694                        &self.cfg,
695                        &gate,
696                        &up,
697                        1.0,
698                        1.0,
699                        self.cfg.clamp_shexp_at(il as u32),
700                        &mut act,
701                        t * n_ff,
702                    )?;
703                    e.matmul(ffn_down, &act, t)?
704                }
705                crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_prefill(e, m, &z, t, il as u16)?,
706            };
707            if probe {
708                e.stream().synchronize()?;
709                eprintln!("[probe] L{il} ffn ok");
710            }
711            let mut x2 = e.uninit(t * n_embd)?;
712            e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
713            x = x2;
714        }
715        // norm over all T, then slice the LAST row and run lm_head on that single row.
716        let mut hn = e.uninit(t * n_embd)?;
717        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
718        let last = e.view(&hn, t * n_embd); // [T, n_embd]
719        let last_row = last.slice((t - 1) * n_embd..t * n_embd); // [1, n_embd]
720        let mut hlast = e.uninit(n_embd)?;
721        e.copy_view_into(&mut hlast, 0, &last_row, n_embd)?;
722        let logits = e.matmul(&self.output, &hlast, 1)?; // [1, n_vocab] — lm_head on ONE row
723        Ok(e.dtoh(&logits)?)
724    }
725
726    /// BATCHED PROMPT PRIME (the measured #1 e2e gap, e2e-image-1): `forward_last`'s batched
727    /// prefill body EXTENDED to leave a DECODE-READY cache behind — vs the tokenwise prime's
728    /// ~102/38 tok/s (9B/27B) decode_step loop, this runs the whole prompt at prefill throughput.
729    ///   (a) full-attn layers append their T post-RoPE K/V rows into `cache.kv[il]` via the SAME
730    ///       per-row quantize kernel as the decode append (bit-identical cache bytes per row);
731    ///   (b) linear layers run STATEFULLY from the cache's current recurrent state (zero at a
732    ///       fresh prime): carried-ring conv (ssm_conv1d_tm_state) + ONE gdn_scan(state_in,
733    ///       state_out) whose internal sequential t-loop equals T chained T=1 steps — but with
734    ///       the NORMAL prefill matmul dispatch (GEMM at m>=16), NOT the decode-exact MMVQ the
735    ///       spec verify uses (prime is a prefill-regime pass; the run-gen prefill==decode
736    ///       argmax gate is the accuracy authority, exactly as for forward_last);
737    ///   (c) `cache.pos`/KV len/len_d advance by T.
738    /// Returns (last-row logits host, h_seed = last-row PRE-output_norm hidden [n_embd],
739    /// hiddens = the full pre-output_norm hidden stack [T, n_embd] — generate_spec's prompt_h).
740    /// FRESH-PROMPT ONLY (cache.pos == 0): the fa_prefill tiles attend within `tokens` alone.
741    /// forward_last itself stays untouched (kernel-check / run-gen gate on it).
742    ///
743    /// `queued_after` (lane/tick-seg, 2026-08-07): the number of prompt tokens of the SAME
744    /// REQUEST that the caller will prime in LATER calls — 0 when this call is the whole
745    /// request (every single-shot caller). Serve splits a long prompt across SEVERAL
746    /// prime_cache calls (one per scheduler tick, plus the prefix-cache LCP split), and the
747    /// request's absolute end position `seq_end = cache.pos + t + queued_after` steers step35's
748    /// SWA prefill arm — computing it per CALL made the arm a function of the tick budget
749    /// (budgets 512/256/64 DIFFER 1.813e0 vs monolithic, greedy diverging at step 6; dark lanes
750    /// default to 256 AND cap by live SLO headroom, so identical judge requests primed
751    /// differently under load — research/tick-seg-20260807, receipt in
752    /// research/step35-chunkfix-20260807 §9). The parameter is what prime_cache structurally
753    /// lacked: it cannot know from `tokens` and `cache` alone whether more of the request is
754    /// coming. A SESSION CONTINUATION (a NEW user turn primed onto a live cache) is a NEW
755    /// request — its arithmetic is keyed to its own extent, so those callers pass 0; only a
756    /// caller that SPLITS one request across calls passes the remainder.
757    pub fn prime_cache(
758        &self,
759        e: &Engine,
760        tokens: &[u32],
761        cache: &mut Cache,
762        queued_after: usize,
763    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
764        self.prime_cache_overlaid(e, tokens, cache, queued_after, None)
765    }
766
767    /// `prime_cache` with a vision embedding overlay (lane/vision): image merger outputs
768    /// replace the `<|image_pad|>` token embeddings at prompt-relative positions before the
769    /// trunk walk — the mixed-embedding prime. Text-only callers use `prime_cache` (overlay
770    /// None, byte-identical path). v1 scope: the serial chunk walk only — PP prime arms and
771    /// gemma4 refuse loudly (the vision serving box is single-GPU).
772    pub fn prime_cache_overlaid(
773        &self,
774        e: &Engine,
775        tokens: &[u32],
776        cache: &mut Cache,
777        queued_after: usize,
778        overlay: Option<&crate::vision::EmbedOverlay>,
779    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
780        let n_embd = self.cfg.n_embd as usize;
781        let t = tokens.len();
782        // SESSION CONTINUATION (2026-07-05): cache.pos > 0 = priming a NEW SUFFIX onto a live
783        // session cache — every chunk (including the first) takes the continuation arm
784        // (fa_prefill_view over the quantized past + this chunk). Fresh prime (pos==0) unchanged.
785        assert!(
786            t >= PRIME_MIN_T,
787            "prime_cache needs T >= {PRIME_MIN_T} (caller gates)"
788        );
789        assert!(
790            cache.pos + t <= cache.max_ctx,
791            "prime_cache: prompt exceeds cache max_ctx"
792        );
793
794        // CHUNKED PRIME (2026-07-05, the long-ctx OOM fix): the monolithic prime allocates
795        // per-layer transients proportional to T (gate/up/act = T*n_ff*4B EACH — 1.5GB apiece at
796        // 16k on the 27B), which OOMs a 24GB card around 16k prompt tokens. Chunk the prompt:
797        // each chunk runs the full layer stack with transients sized to the chunk, appending its
798        // K/V to the resident quantized cache and carrying the GDN conv-ring + recurrent state
799        // through `cache.recur` (linear_attn_prime is already stateful — a chunk boundary is
800        // exactly the state carry it was built for). Full-attn chunks after the first attend to
801        // the QUANTIZED past KV via fa_prefill_view (the spec-verify pattern) — same numeric
802        // class as decode reading the cache. Prompts <= one chunk take the ORIGINAL monolithic
803        // body byte-for-byte (chunk 0 short-circuits to the f32 fa_prefill path).
804        // MEMRA_PRIME_CHUNK sets the chunk size (tokens); 0 disables chunking (monolithic).
805        if self.is_gemma4_e4b() || self.cfg.gemma4.is_some() {
806            if self.is_gemma4_e4b() {
807                if overlay.is_some() {
808                    return Err(
809                        "vision embedding overlay is unsupported on gemma4 E4B (PLE prime)".into(),
810                    );
811                }
812                return self.gemma4_e4b_prime(e, tokens, cache);
813            }
814            // gemma4 v0: monolithic fresh-prompt prime (chunked/continuation arms later).
815            // An overlay takes the masked-prefill arm: image rows splice in unscaled
816            // (gemma4.cpp:182 — embd batches skip the sqrt(n_embd) scale) and the image
817            // spans become bidirectional attention islands (lane/gemma-vision).
818            return self.gemma4_prime(e, tokens, cache, overlay);
819        }
820        let ranges = prime_chunk_ranges(t, self.layers.len());
821        // CHUNK-ORDER INVARIANCE (lane/chunk-invariance, 2026-08-05; vLLM #38561 shape).
822        // MEMRA_PRIME_CHUNK is documented as a memory-transient knob, but it also decides
823        // the prefill's ARITHMETIC, so two rigs with different values produced different
824        // greedy text for the same prompt (research/session-affinity-20260805: 97- and
825        // 149-token prompts). ROOT CAUSE, measured in research/chunk-invariance-20260805
826        // (VERDICT.md) — and it is NOT what docs originally said:
827        //   * NOT trunk GEMM m / reduction order. REFUTED: the prefill GEMM is m-INVARIANT
828        //     (rows [0,32) bit-identical at m=32 vs m=33..80, both quantized wq and the
829        //     output head), so growing a chunk cannot move an existing row's value.
830        //   * NOT the GDN scan segmentation. REFUTED: MEMRA_GDN_CHUNKED=0 (sequential scan,
831        //     no WY segmentation at all) still diverges, so vLLM's mamba-boundary fix does
832        //     not describe our leak.
833        //   * IT IS the attention numeric CLASS edge in full_attn_prime_fa_dispatch, which
834        //     selects on `base_len == 0`: chunk 0 attends over this batch's f32 K/V
835        //     (fa_prefill) while every later chunk attends over the q8_0/q5_1 quantized KV
836        //     cache (fa_prefill_view_ws). The chunk size therefore decides WHERE in the
837        //     prompt that precision edge falls. Signature: per-row maxdiff is exactly 0.0
838        //     before the first boundary and O(1) right after, first_div_pos == chunk size.
839        // The grain-free fix (full_attn_prime_fa_dispatch below) removed that class edge at
840        // the source — every row is in one numeric class, so the chunk size no longer steers
841        // arithmetic and MEMRA_PRIME_CHUNK is a pure memory/transient knob again. The interim
842        // MEMRA_PRIME_INVARIANT/MEMRA_PRIME_GRAIN pin-the-boundary door was superseded by that
843        // fix and KILLED at v0.71 per the flags doctrine (the jsonl + VERDICT.md are the
844        // record); the chunkinv gate asserts byte-identity across chunk sizes naked.
845        // The REQUEST's absolute end position, computed ONCE before the loop: every chunk sees the
846        // same value, whatever the chunk size. step35's SWA arm selects on it so that kernel
847        // selection — and therefore the logits — cannot depend on MEMRA_PRIME_CHUNK
848        // (research/step35-chunkfix-20260807; see step35_attn_pre_wo's doc note).
849        // `+ queued_after` closes the SECOND axis (lane/tick-seg): when serve splits the request
850        // across calls, the request still ends at the same absolute position, whatever the tick
851        // budget or LCP split point. MEMRA_PRIME_CALLLOCAL=1 is the ROLLBACK SEAM to the FULL
852        // pre-fix arithmetic: the per-call value here AND the unaligned FA view offset in
853        // step35_attn_pre_wo. Both halves are required for tickinv35's canary teeth under the FA
854        // default. Read per call, not cached (the probe flips it in-process between arms). Never
855        // on in a measured default run.
856        let legacy_calllocal = std::env::var("MEMRA_PRIME_CALLLOCAL").as_deref() == Ok("1");
857        let seq_end = if legacy_calllocal {
858            cache.pos + t
859        } else {
860            cache.pos + t + queued_after
861        };
862        if ranges.len() == 1 {
863            return self.prime_chunk(e, tokens, cache, seq_end, 0, overlay);
864        }
865        // PIPELINED PP-2 PRIME (lane/cx-pipeline-prime, 2026-08-08): overlap stage 0 of
866        // chunk N+1 with stage 1 of chunk N. The serial split stays reachable through
867        // MEMRA_PRIME_PIPE=0 and is the exactness oracle. N>2 keeps the serial walker;
868        // this lane owns the balanced two-stage schedule only.
869        if crate::pp::prime_pipe_on() && crate::pp::prime_pp_on() && !crate::pp::pp2_streams_off() {
870            if let Some(fence) = crate::pp::pp_cuts(self.layers.len()).filter(|f| f.len() == 3) {
871                if overlay.is_some() {
872                    return Err(
873                        "vision embedding overlay + pipelined PP prime unsupported (v1); \
874                         run the serial prime (single device or MEMRA_PRIME_PIPE=0)"
875                            .into(),
876                    );
877                }
878                if crate::pp::pp_multi_stream_same_device() {
879                    return Err(
880                        "prime chunk pipeline refused with 2 stage streams on one device — \
881                         that concurrent-stream placement remains quarantined by the deferred \
882                         pp flake record. Use one device per stage or MEMRA_PRIME_PIPE=0 for \
883                         the serial split."
884                            .into(),
885                    );
886                }
887                return self.prime_cache_pp2_pipelined(e, tokens, cache, seq_end, &ranges, &fence);
888            }
889        }
890        let mut hiddens = e.uninit(t * n_embd)?;
891        let mut last: Option<(Vec<f32>, CudaSlice<f32>)> = None;
892        for &(start, end) in &ranges {
893            let (l, hs, x) =
894                self.prime_chunk(e, &tokens[start..end], cache, seq_end, start, overlay)?;
895            e.copy_into(&mut hiddens, start * n_embd, &x, (end - start) * n_embd)?;
896            last = Some((l, hs));
897        }
898        let (logits, h_seed) = last.unwrap();
899        Ok((logits, h_seed, hiddens))
900    }
901
902    /// PP-2 chunk scheduler: stage 1 of chunk N and stage 0 of chunk N+1 are both queued
903    /// before N's epilogue D2H drains the last-stage stream. Arithmetic is unchanged:
904    /// every chunk still runs the same two `prime_layers` ranges, boundary copy, output
905    /// norm, lm head, and caller hidden-stack copy as the serial split.
906    fn prime_cache_pp2_pipelined(
907        &self,
908        e: &Engine,
909        tokens: &[u32],
910        cache: &mut Cache,
911        seq_end: usize,
912        ranges: &[(usize, usize)],
913        fence: &[usize],
914    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
915        debug_assert_eq!(fence.len(), 3);
916        debug_assert!(ranges.len() >= 2);
917        let rt = crate::pp::PpNRt::get(e)?;
918        assert_eq!(
919            rt.n_stages(),
920            2,
921            "prime pipeline requires exactly two PP stages"
922        );
923        let n_embd = self.cfg.n_embd as usize;
924        let t = tokens.len();
925        let initial_base = cache.pos;
926        let caller_stream = e.stream();
927
928        // #87 reverse publication before any new stage allocation, then prewarm both
929        // boundary slots while the stage streams are otherwise empty. Lazy-growing slot B
930        // after stage 1(N) is queued would synchronize that stream and erase the first
931        // overlap on a two-chunk prompt.
932        rt.fence_stages_behind(&caller_stream)?;
933        let max_payload = ranges.iter().map(|(s, e)| (e - s) * n_embd).max().unwrap();
934        rt.prepare_overlap_slots(0, max_payload)?;
935
936        let mut hiddens = e.uninit(t * n_embd)?;
937        let mut last: Option<(Vec<f32>, CudaSlice<f32>)> = None;
938        let mut stage_caches = PrimeCacheStages::new(cache, fence[1]);
939        let (cache0, cache1) = stage_caches.parts();
940        let (first_start, first_end) = ranges[0];
941        let mut slot = self.prime_pp2_stage0_enqueue(
942            e,
943            rt,
944            &tokens[first_start..first_end],
945            cache0,
946            seq_end,
947            fence,
948            initial_base + first_start,
949            true,
950        )?;
951        cache0.pos = initial_base + first_end;
952
953        for (i, &(start, end)) in ranges.iter().enumerate() {
954            let base = initial_base + start;
955            debug_assert_eq!(
956                cache1.pos, base,
957                "stage 1 must drain chunks in original position order"
958            );
959            let (out, next_slot) = if let Some(&(next_start, next_end)) = ranges.get(i + 1) {
960                let next_base = initial_base + next_start;
961                debug_assert_eq!(
962                    cache0.pos, next_base,
963                    "stage 0 must issue chunks in original position order"
964                );
965                let cache0_stage = &mut *cache0;
966                // Step's MoE router readback synchronizes once per layer. Two CUDA streams
967                // on one host thread therefore serialize even if the calls are ordered as
968                // a pipeline. Drive the disjoint stage caches from two scoped host threads:
969                // stage 1 consumes slot N while stage 0 produces slot N+1.
970                std::thread::scope(|scope| -> Result<_, Box<dyn std::error::Error>> {
971                    let stage0 = scope.spawn(move || -> Result<usize, String> {
972                        let next = self
973                            .prime_pp2_stage0_enqueue(
974                                e,
975                                rt,
976                                &tokens[next_start..next_end],
977                                cache0_stage,
978                                seq_end,
979                                fence,
980                                next_base,
981                                true,
982                            )
983                            .map_err(|err| err.to_string())?;
984                        cache0_stage.pos = initial_base + next_end;
985                        Ok(next)
986                    });
987                    let x = self.prime_pp2_stage1_enqueue(
988                        e,
989                        rt,
990                        slot,
991                        end - start,
992                        cache1,
993                        seq_end,
994                        fence,
995                        base,
996                        true,
997                    )?;
998                    let out = {
999                        rt.bind_stage(1)?;
1000                        let _st1 = rt.enter(1);
1001                        let e1 = rt.engine(1, e);
1002                        self.prime_chunk_epilogue(e1, x, end - start, cache1)?
1003                    };
1004                    let next = stage0
1005                        .join()
1006                        .map_err(|_| "pipeprime stage-0 host walker panicked")?
1007                        .map_err(|err| -> Box<dyn std::error::Error> { err.into() })?;
1008                    Ok((out, Some(next)))
1009                })?
1010            } else {
1011                let x = self.prime_pp2_stage1_enqueue(
1012                    e,
1013                    rt,
1014                    slot,
1015                    end - start,
1016                    cache1,
1017                    seq_end,
1018                    fence,
1019                    base,
1020                    true,
1021                )?;
1022                let out = {
1023                    rt.bind_stage(1)?;
1024                    let _st1 = rt.enter(1);
1025                    let e1 = rt.engine(1, e);
1026                    self.prime_chunk_epilogue(e1, x, end - start, cache1)?
1027                };
1028                (out, None)
1029            };
1030
1031            rt.publish_to(1, &caller_stream)?;
1032            e.copy_into(&mut hiddens, start * n_embd, &out.2, (end - start) * n_embd)?;
1033            last = Some((out.0, out.1));
1034            crate::pp::PRIME_SPLIT_CHUNKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1035
1036            if let Some(next) = next_slot {
1037                // The caller copy above reads a stage-1 allocation. Before stage 1 of the
1038                // next chunk can allocate/reuse blocks, mirror #87's body-entry fence.
1039                // Stage 0(N+1) is already queued before this wait is appended, so its
1040                // overlap with stage 1(N) is preserved.
1041                rt.fence_stages_behind(&caller_stream)?;
1042                slot = next;
1043            }
1044        }
1045
1046        debug_assert_eq!(cache0.pos, initial_base + t);
1047        debug_assert_eq!(cache1.pos, initial_base + t);
1048        let (logits, h_seed) = last.unwrap();
1049        Ok((logits, h_seed, hiddens))
1050    }
1051
1052    /// One prime chunk: the full layer stack over `tokens`, continuing from the cache's current
1053    /// state (`cache.pos` = tokens already primed; 0 = fresh). Positions/RoPE are absolute
1054    /// task #21 de-broadcast: the q/k head count PREP must emit so the scan's consumers
1055    /// agree. Compact (num_k) ONLY when the scan will take the chunked+mma route AND
1056    /// num_k == num_v/2 (the engine-side hint mirrors this exact formula); everything
1057    /// else (s128 verify tier, chunked-off, mma-off) keeps the broadcast layout.
1058    fn gdn_hk(e: &Engine, t: usize, num_v: usize, num_k: usize) -> usize {
1059        if Engine::gdn_db_on()
1060            && Engine::gdn_chunked_enabled()
1061            && t >= 16
1062            && e.gdn_mma_enabled(Engine::gdn_chunk_size())
1063            && num_k * 2 == num_v
1064        {
1065            num_k
1066        } else {
1067            num_v
1068        }
1069    }
1070
1071    /// task #17: gate for the fused fp16-operand epilogues (silu_mul/gated_rmsnorm/sig_mul
1072    /// `_f16out` twins). Bit-identical class (the twins emit the cvt kernel's exact halves),
1073    /// but seam-gated (MEMRA_F16OUT=0) so a bit-check can arbitrate, and OFF under verify-exact
1074    /// (matmul_group skips the f16 lane there; the twins must not resurrect it).
1075    fn f16out_on(e: &Engine, t: usize) -> bool {
1076        crate::f16_ffi::pp_f16_enabled()
1077            && t >= 16
1078            && !e.verify_exact_on()
1079            && std::env::var("MEMRA_F16OUT").as_deref() != Ok("0")
1080    }
1081
1082    /// (cache.pos + i). Returns (last-row logits, h_seed, this chunk's hidden stack [T, n_embd]).
1083    /// See HybridModel::prime_slabs — the eager prime's resident trunk transients.
1084    /// PER-DEVICE since lane/pp-leverb (2026-08-08): the map is keyed by the allocating
1085    /// engine's CUDA ordinal — under the prime stage split each stage's range walks through
1086    /// its OWN slabs on its own device (a dev0 slab dereferenced by a dev1 kernel would be
1087    /// a peer read per GEMM operand, the exact class Lever B removes). Single-device rigs
1088    /// see one entry, byte-identical behavior.
1089    pub fn prime_slabs_get(
1090        &self,
1091        e: &Engine,
1092        t: usize,
1093        n_embd: usize,
1094        n_ff_max: usize,
1095    ) -> Result<std::sync::Arc<std::sync::Mutex<PrimeSlabs>>, Box<dyn std::error::Error>> {
1096        let mut slabs = self.prime_slabs.lock().unwrap();
1097        let dev = e.ctx().ordinal();
1098        let need_new = match slabs.get(&dev) {
1099            None => true,
1100            Some(sl) => sl.lock().unwrap().t_cap < t,
1101        };
1102        if need_new {
1103            slabs.insert(
1104                dev,
1105                std::sync::Arc::new(std::sync::Mutex::new(PrimeSlabs {
1106                    t_cap: t,
1107                    h: e.uninit(t * n_embd)?,
1108                    x1: e.uninit(t * n_embd)?,
1109                    z: e.uninit(t * n_embd)?,
1110                    act: e.uninit(t * n_ff_max)?,
1111                    xa: e.uninit(t * n_embd)?,
1112                    xb: e.uninit(t * n_embd)?,
1113                    h16: e.alloc_u8_uninit(t * n_embd * 2)?,
1114                    z16: e.alloc_u8_uninit(t * n_embd * 2)?,
1115                    gate: e.uninit(t * n_ff_max)?,
1116                    up: e.uninit(t * n_ff_max)?,
1117                    ffn_out: e.uninit(t * n_embd)?,
1118                    seg_glue: Vec::new(),
1119                    mixed: e.uninit(t * n_embd)?,
1120                    seg_mid: Vec::new(),
1121                    seg_t: 0,
1122                })),
1123            );
1124        }
1125        Ok(slabs.get(&dev).expect("prime slab inserted").clone())
1126    }
1127
1128    /// `seq_end` = the whole REQUEST's absolute end position (`cache.pos + prompt_len` at
1129    /// `prime_cache` entry), NOT this chunk's end. Chunk-size-invariant by construction; step35's
1130    /// SWA arm selects on it (see `step35_attn_pre_wo`). Every other arch ignores it.
1131    fn prime_chunk(
1132        &self,
1133        e: &Engine,
1134        tokens: &[u32],
1135        cache: &mut Cache,
1136        seq_end: usize,
1137        chunk_off: usize,
1138        overlay: Option<&crate::vision::EmbedOverlay>,
1139    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
1140        if crate::pp::pp_host_bounce_active()
1141            && (self.cfg.gemma4.is_some() || !crate::pp::prime_pp_on())
1142        {
1143            return Err(
1144                "prime_chunk: refused with MEMRA_PP_HOST_BOUNCE=1 because this configuration \
1145                 has no active prime stage split and would peer-read remote weights; keep \
1146                 MEMRA_PRIME_PP enabled and use a PP-prime-supported model"
1147                    .into(),
1148            );
1149        }
1150        // LEVER B (lane/pp-leverb, 2026-08-08): the ppN door for the chunked prime. With the
1151        // door open + per-stage streams, each chunk walks its layer ranges on the OWNING
1152        // stage's engine/device (the anatomy receipt this kills: dev1 ran ZERO prefill
1153        // kernels; stage-1 trunk weights were peer-read = 22% of the pp4096 wall).
1154        // MEMRA_PRIME_PP=0 = the unsplit rollback (also the prime-split-gate's reference
1155        // arm — prime deliberately keeps NO refuse_unsplit_if_remote, see pp.rs). The
1156        // MEMRA_PP_STREAMS=0 seam keeps the unsplit walk too: in that regime the sharded
1157        // loader is off and there is nothing remote to split for.
1158        if self.cfg.gemma4.is_none() && !crate::pp::pp2_streams_off() && crate::pp::prime_pp_on() {
1159            if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
1160                if overlay.is_some() {
1161                    return Err("vision embedding overlay + PP prime unsupported (v1); \
1162                         run single-device or MEMRA_PRIME_PP=0"
1163                        .into());
1164                }
1165                return self.prime_chunk_ppn(e, tokens, cache, seq_end, &fence);
1166            }
1167        }
1168        if crate::pp::pp_host_bounce_active() {
1169            return Err(
1170                "prime_chunk: MEMRA_PP_HOST_BOUNCE=1 found no valid prime stage split; \
1171                 refusing an unsplit remote-weight walk"
1172                    .into(),
1173            );
1174        }
1175        let t = tokens.len();
1176        let base = cache.pos;
1177        debug_assert!(
1178            seq_end >= base + t,
1179            "prime_chunk: seq_end must cover this chunk"
1180        );
1181        let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
1182        let pos_d = e.htod_i32(&pos)?;
1183
1184        let mut x_embed = self.embed(e, tokens)?; // [T, n_embd]
1185        if let Some(ov) = overlay {
1186            // Mixed-embedding splice: image rows overwrite the pad-token embeddings that
1187            // fall inside this chunk's prompt-relative window [chunk_off, chunk_off+t).
1188            // Images larger than one prime chunk straddle boundaries, hence the clipping.
1189            let n_embd = self.cfg.n_embd as usize;
1190            for &(pos, row_off, n_rows) in &ov.spans {
1191                let lo = pos.max(chunk_off);
1192                let hi = (pos + n_rows).min(chunk_off + t);
1193                if lo < hi {
1194                    let src_row = row_off + (lo - pos);
1195                    let view = ov
1196                        .rows
1197                        .slice(src_row * n_embd..(src_row + (hi - lo)) * n_embd);
1198                    e.copy_view_into(
1199                        &mut x_embed,
1200                        (lo - chunk_off) * n_embd,
1201                        &view,
1202                        (hi - lo) * n_embd,
1203                    )?;
1204                }
1205            }
1206        }
1207        let x = self.prime_layers(
1208            e,
1209            x_embed,
1210            0,
1211            self.layers.len(),
1212            &pos_d,
1213            t,
1214            base,
1215            cache,
1216            seq_end,
1217        )?;
1218        self.prime_chunk_epilogue(e, x, t, cache)
1219    }
1220
1221    /// PRIME RANGE SUBGRAPH (lane/pp-leverb, 2026-08-08): layers `[lo, hi)` of the chunked
1222    /// prime walk — `prime_chunk`'s trunk loop extracted verbatim to the
1223    /// `decode_layers_eager(lo, hi)` / `verify_layers(lo, hi)` contract: enters with a
1224    /// MATERIALIZED `[T, n_embd]` residual, exits with the range's final residual
1225    /// materialized (cloned out of the slab). At `lo=0, hi=n_layers` — the unsplit call —
1226    /// the launch sequence is byte-identical to the pre-extraction body. Range semantics:
1227    ///   - the cross-layer [down-add + NEXT attn-norm] fusion is range-LOCAL (`il + 1 < hi`):
1228    ///     layer `hi`'s attn_norm belongs to the NEXT stage's device, so the range ends with
1229    ///     the plain add (materialize) and the next stage hoists its own first norm — the
1230    ///     kernel-check-pinned `add_rms_norm == add then rms_norm` identity, the same law
1231    ///     the decode split rests on (`prime-split-gate` arbitrates end-to-end);
1232    ///   - prime slabs are PER-DEVICE (`prime_slabs_get` keys on the engine ordinal), so
1233    ///     each stage walks through its own resident transients;
1234    ///   - the S-glue/S-mid capture path requires the FULL range (its lookahead fuses
1235    ///     `self.layers[il+1]` unconditionally) — `use_seg` gains `lo == 0 && hi == n_layers`.
1236    #[allow(clippy::too_many_arguments)]
1237    fn prime_layers(
1238        &self,
1239        e: &Engine,
1240        x_in: CudaSlice<f32>,
1241        lo: usize,
1242        hi: usize,
1243        pos_d: &CudaSlice<i32>,
1244        t: usize,
1245        base: usize,
1246        cache: &mut Cache,
1247        seq_end: usize,
1248    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1249        let cfg = &self.cfg;
1250        let n_embd = cfg.n_embd as usize;
1251        let eps = cfg.rms_eps;
1252        // task #14: fuse the fp16 GEMM-operand emission into the trunk norms (kills the
1253        // standalone convert launches). Only when the f16 lane serves and T reaches the
1254        // GEMM tier; bit-identical either way.
1255        let f16fuse = crate::f16_ffi::pp_f16_enabled() && t >= 16;
1256        // PRIME SLABS (piecewise foundation): trunk transients in resident buffers —
1257        // ~224 fewer alloc/free calls per prime and FROZEN Lt operand addresses
1258        // (nvjet's alignment-variant selection becomes run-to-run stable). Every slab is
1259        // fully overwritten before use; x ping-pongs xa<->xb; the hidden-stack return
1260        // clones the final x (the slab cannot leave). Non-slab fallback: MEMRA_PRIME_SLABS=0.
1261        let n_ff_max = self
1262            .layers
1263            .iter()
1264            .map(|l| match &l.ffn {
1265                crate::hybrid::Ffn::Dense { ffn_gate, .. } => ffn_gate.out_features(),
1266                _ => n_embd,
1267            })
1268            .max()
1269            .unwrap_or(n_embd)
1270            .max(n_embd);
1271        let use_slabs = std::env::var("MEMRA_PRIME_SLABS").as_deref() != Ok("0");
1272        let slab = if use_slabs {
1273            Some(self.prime_slabs_get(e, t, n_embd, n_ff_max)?)
1274        } else {
1275            None
1276        };
1277        let mut slab_guard = slab.as_ref().map(|sl| sl.lock().unwrap());
1278        let mut x_own; // fallback storage when slabs are off
1279        type SlabRefs<'a> = (
1280            &'a mut CudaSlice<f32>,
1281            &'a mut CudaSlice<f32>,
1282            &'a mut CudaSlice<f32>,
1283            &'a mut CudaSlice<f32>,
1284            &'a mut CudaSlice<u8>,
1285            &'a mut CudaSlice<u8>,
1286            &'a mut CudaSlice<f32>,
1287            &'a mut CudaSlice<f32>,
1288            &'a mut CudaSlice<f32>,
1289        );
1290        let (mut x_cur, mut x_nxt, sl): (
1291            &mut CudaSlice<f32>,
1292            &mut CudaSlice<f32>,
1293            Option<SlabRefs>,
1294        );
1295        let mut seg: Option<(
1296            &mut Vec<Option<cudarc::driver::CudaGraph>>,
1297            &mut Vec<Option<cudarc::driver::CudaGraph>>,
1298            &mut CudaSlice<f32>,
1299            &mut usize,
1300        )> = None;
1301        let mut x_own2;
1302        match slab_guard.as_mut() {
1303            Some(g) => {
1304                let slabs = &mut **g;
1305                e.copy_into(&mut slabs.xa, 0, &x_in, t * n_embd)?;
1306                let PrimeSlabs {
1307                    xa,
1308                    xb,
1309                    h,
1310                    x1,
1311                    z,
1312                    act,
1313                    h16,
1314                    z16,
1315                    gate,
1316                    up,
1317                    ffn_out,
1318                    seg_glue,
1319                    mixed,
1320                    seg_mid,
1321                    seg_t,
1322                    ..
1323                } = slabs;
1324                x_cur = xa;
1325                x_nxt = xb;
1326                seg = Some((seg_glue, seg_mid, mixed, seg_t));
1327                sl = Some((h, x1, z, act, h16, z16, gate, up, ffn_out));
1328            }
1329            None => {
1330                x_own = x_in;
1331                x_own2 = e.uninit(t * n_embd)?;
1332                x_cur = &mut x_own;
1333                x_nxt = &mut x_own2;
1334                sl = None;
1335            }
1336        }
1337        let mut alloc_h;
1338        let mut alloc_x1;
1339        let mut alloc_z;
1340        let mut alloc_act;
1341        let mut alloc_h16;
1342        let mut alloc_z16;
1343        let mut alloc_gate;
1344        let mut alloc_up;
1345        let mut alloc_fo;
1346        let (h, x1, z, act): (
1347            &mut CudaSlice<f32>,
1348            &mut CudaSlice<f32>,
1349            &mut CudaSlice<f32>,
1350            &mut CudaSlice<f32>,
1351        );
1352        let (h16, z16): (&mut CudaSlice<u8>, &mut CudaSlice<u8>);
1353        let (sl_gate, sl_up, sl_fo): (
1354            &mut CudaSlice<f32>,
1355            &mut CudaSlice<f32>,
1356            &mut CudaSlice<f32>,
1357        );
1358        match sl {
1359            Some((a, b, c, d, e16, f16b, g, u, fo)) => {
1360                h = a;
1361                x1 = b;
1362                z = c;
1363                act = d;
1364                h16 = e16;
1365                z16 = f16b;
1366                sl_gate = g;
1367                sl_up = u;
1368                sl_fo = fo;
1369            }
1370            None => {
1371                alloc_h = e.uninit(t * n_embd)?;
1372                alloc_x1 = e.uninit(t * n_embd)?;
1373                alloc_z = e.uninit(t * n_embd)?;
1374                alloc_act = e.uninit(t * n_ff_max)?;
1375                alloc_h16 = e.alloc_u8_uninit(t * n_embd * 2)?;
1376                alloc_z16 = e.alloc_u8_uninit(t * n_embd * 2)?;
1377                alloc_gate = e.uninit(t * n_ff_max)?;
1378                alloc_up = e.uninit(t * n_ff_max)?;
1379                alloc_fo = e.uninit(t * n_embd)?;
1380                h = &mut alloc_h;
1381                x1 = &mut alloc_x1;
1382                z = &mut alloc_z;
1383                act = &mut alloc_act;
1384                h16 = &mut alloc_h16;
1385                z16 = &mut alloc_z16;
1386                sl_gate = &mut alloc_gate;
1387                sl_up = &mut alloc_up;
1388                sl_fo = &mut alloc_fo;
1389            }
1390        }
1391        // piecewise increment 3: layer-0 norm hoisted; the per-layer tail fuses
1392        // [down-add + NEXT layer's attn-norm] into one captured S-glue segment
1393        // (all-slab IO, zero in-graph allocations). Capture happens lazily on the
1394        // first prime at this t (capture does not execute -> launch right after).
1395        let n_layers = self.layers.len();
1396        // OPT-IN (2026-07-26 interleaved verdict): 2-kernel segments measured NET
1397        // -0.7%-to-neutral under the interleaved A/B protocol — one cuGraphLaunch costs
1398        // about what two kernel submissions do. The earlier "+0.9%" was cross-run clock
1399        // drift (the repo's interleaved-A/B law exists for exactly this). Larger segments
1400        // (S-prep/S-attn, 7-9 kernels) remain the open hypothesis; the core-split
1401        // machinery stays (byte-identical) as their foundation.
1402        // step35 is excluded: the core-split path calls `full_attn_prime_core_inner`, which is
1403        // the GENERIC attn core (uniform n_head, rope_dim_count, no window, no head-wise gate).
1404        // step35 rides its own mixer through the normal per-layer arm below.
1405        let use_seg = f16fuse
1406            && seg.is_some()
1407            && self.cfg.step35.is_none()
1408            && lo == 0
1409            && hi == n_layers
1410            && std::env::var("MEMRA_PRIME_SEG").as_deref() == Ok("1");
1411        if let Some((sg, sm, _, st)) = seg.as_mut() {
1412            if **st != t {
1413                sg.clear();
1414                sg.extend((0..n_layers).map(|_| None));
1415                sm.clear();
1416                sm.extend((0..n_layers).map(|_| None));
1417                **st = t;
1418            }
1419        }
1420        {
1421            let layer_lo = &self.layers[lo];
1422            if f16fuse {
1423                e.rms_norm_f16out(
1424                    x_cur,
1425                    layer_lo.attn_norm.float_data(),
1426                    h,
1427                    h16,
1428                    n_embd,
1429                    t,
1430                    eps,
1431                )?;
1432            } else {
1433                e.rms_norm(x_cur, layer_lo.attn_norm.float_data(), h, n_embd, t, eps)?;
1434            }
1435        }
1436        for il in lo..hi {
1437            let layer = &self.layers[il];
1438            let hx16 = if f16fuse { Some(&*h16) } else { None };
1439            if use_seg {
1440                // core-split path: projections -> _inner core -> out-GEMM INTO the mixed
1441                // slab (no copies) -> S-mid segment [add + post-norm] as one graph launch.
1442                let (pre, pre16, w_out) = match &layer.mixer {
1443                    Mixer::Full(fa) => {
1444                        let g3 = match hx16 {
1445                            Some(xh) => e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], h, xh, t)?,
1446                            None => e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?,
1447                        };
1448                        let (pre, pre16) =
1449                            self.full_attn_prime_core_inner(e, fa, g3, &pos_d, t, cache, il)?;
1450                        (pre, pre16, &fa.wo)
1451                    }
1452                    Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
1453                    Mixer::Linear(la) => {
1454                        let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
1455                        let g4 = match hx16 {
1456                            Some(xh) => e.matmul_group_xh(&ws, h, xh, t)?,
1457                            None => e.matmul_group(&ws, h, t)?,
1458                        };
1459                        let (pre, pre16) =
1460                            self.linear_attn_prime_core_pad_inner(e, la, g4, t, cache, il, None)?;
1461                        (pre, pre16, &la.ssm_out)
1462                    }
1463                };
1464                {
1465                    let (_, sm, mslab, _) = seg.as_mut().unwrap();
1466                    let pre_n = pre.len() / t;
1467                    let xh_pre = match pre16 {
1468                        Some(x) => x,
1469                        None => e.f16_act(&pre, t * pre_n, pre_n)?,
1470                    };
1471                    if !e.try_f16_gemm_pre_into(w_out, &xh_pre, t, mslab)? {
1472                        let y = e.matmul(w_out, &pre, t)?;
1473                        e.copy_into(mslab, 0, &y, t * n_embd)?;
1474                    }
1475                    if sm[il].is_none() {
1476                        use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
1477                        let w_post = layer.post_attn_norm.float_data();
1478                        e.stream().synchronize()?;
1479                        e.stream()
1480                            .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
1481                        let r = (|| -> Result<(), Box<dyn std::error::Error>> {
1482                            e.add(x_cur, mslab, x1, t * n_embd)?;
1483                            e.rms_norm_f16out(x1, w_post, z, z16, n_embd, t, eps)?;
1484                            Ok(())
1485                        })();
1486                        let g = e.stream().end_capture(
1487                            CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH);
1488                        r?;
1489                        sm[il] = Some(g?.ok_or("S-mid capture produced no graph")?);
1490                    }
1491                    sm[il].as_ref().unwrap().launch()?;
1492                }
1493            } else {
1494                let mixed = match &layer.mixer {
1495                    Mixer::Full(fa) => {
1496                        self.full_attn_prime(e, fa, h, hx16, &pos_d, t, cache, il, seq_end)?
1497                    }
1498                    Mixer::Linear(la) => self.linear_attn_prime(e, la, h, hx16, t, cache, il)?,
1499                    Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
1500                };
1501                if f16fuse {
1502                    // round 28: residual+norm in ONE kernel (add_rms_norm precedent,
1503                    // bit-identical) — the standalone add pass disappears.
1504                    e.add_rms_norm_f16out(
1505                        x_cur,
1506                        &mixed,
1507                        layer.post_attn_norm.float_data(),
1508                        x1,
1509                        z,
1510                        z16,
1511                        n_embd,
1512                        t,
1513                        eps,
1514                    )?;
1515                } else {
1516                    e.add(x_cur, &mixed, x1, t * n_embd)?;
1517                    e.rms_norm(x1, layer.post_attn_norm.float_data(), z, n_embd, t, eps)?;
1518                }
1519            }
1520            let zx16 = if f16fuse { Some(&*z16) } else { None };
1521            match &layer.ffn {
1522                crate::hybrid::Ffn::Dense {
1523                    ffn_gate,
1524                    ffn_up,
1525                    ffn_down,
1526                } => {
1527                    let n_ff = ffn_gate.out_features();
1528                    // gate/up INTO boundary slabs (piecewise increment 2); fall back to
1529                    // the allocating group + copy when a mirror is missing.
1530                    let mut into_ok = false;
1531                    if let Some(xh) = zx16 {
1532                        into_ok = e.try_f16_gemm_pre_into(ffn_gate, xh, t, sl_gate)?
1533                            && e.try_f16_gemm_pre_into(ffn_up, xh, t, sl_up)?;
1534                    }
1535                    if !into_ok {
1536                        let mut g2 = match zx16 {
1537                            Some(xh) => e.matmul_group_xh(&[ffn_gate, ffn_up], z, xh, t)?,
1538                            None => e.matmul_group(&[ffn_gate, ffn_up], z, t)?,
1539                        };
1540                        let up_y = g2.pop().unwrap();
1541                        let gate_y = g2.pop().unwrap();
1542                        e.copy_into(sl_gate, 0, &gate_y, t * n_ff)?;
1543                        e.copy_into(sl_up, 0, &up_y, t * n_ff)?;
1544                    }
1545                    // task #17: the silu arm's f16out twin emits the down GEMM's fp16
1546                    // operand in-epilogue; non-silu activations keep the standalone convert.
1547                    // silu_mul_f16out is PLAIN silu(gate)*up — a clamped layer (step35 dense
1548                    // blocks under a live swiglu_clamp_shexp) must take the ffn_act_lim arm.
1549                    let d_lim = self.cfg.clamp_shexp_at(il as u32);
1550                    let act16 = if Self::f16out_on(e, t) && self.cfg.m3.is_none() && d_lim.is_none()
1551                    {
1552                        let mut a16 = e.alloc_u8_uninit(t * n_ff * 2)?;
1553                        e.silu_mul_f16out(sl_gate, sl_up, act, &mut a16, t * n_ff)?;
1554                        Some(a16)
1555                    } else {
1556                        Self::ffn_act_lim(
1557                            e,
1558                            &self.cfg,
1559                            sl_gate,
1560                            sl_up,
1561                            1.0,
1562                            1.0,
1563                            d_lim,
1564                            act,
1565                            t * n_ff,
1566                        )?;
1567                        None
1568                    };
1569                    // down GEMM into the ffn_out slab (f16 arm; fallback copies)
1570                    let xh_act = match act16 {
1571                        Some(x) => x,
1572                        None => e.f16_act(act, t * n_ff, n_ff)?,
1573                    };
1574                    if !e.try_f16_gemm_pre_into(ffn_down, &xh_act, t, sl_fo)? {
1575                        let y = e.matmul(ffn_down, &*act, t)?;
1576                        e.copy_into(sl_fo, 0, &y, t * n_embd)?;
1577                    }
1578                }
1579                crate::hybrid::Ffn::Moe(m) => {
1580                    let y = self.moe_ffn_il_prefill(e, m, z, t, il as u16)?;
1581                    e.copy_into(sl_fo, 0, &y, t * n_embd)?;
1582                }
1583            }
1584            if use_seg && il + 1 < hi {
1585                // S-glue segment: [add + next attn-norm(+f16out)] — one cuGraphLaunch
1586                let w_next = self.layers[il + 1].attn_norm.float_data();
1587                let (sg, _, _, _) = seg.as_mut().unwrap();
1588                if sg[il].is_none() {
1589                    use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
1590                    e.stream().synchronize()?;
1591                    e.stream()
1592                        .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
1593                    let r = (|| -> Result<(), Box<dyn std::error::Error>> {
1594                        e.add(x1, sl_fo, x_nxt, t * n_embd)?;
1595                        e.rms_norm_f16out(x_nxt, w_next, h, h16, n_embd, t, eps)?;
1596                        Ok(())
1597                    })();
1598                    let g = e.stream().end_capture(
1599                        CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
1600                    );
1601                    r?;
1602                    sg[il] = Some(g?.ok_or("S-glue capture produced no graph")?);
1603                }
1604                sg[il].as_ref().unwrap().launch()?;
1605            } else {
1606                if il + 1 < hi {
1607                    let w_next = self.layers[il + 1].attn_norm.float_data();
1608                    if f16fuse {
1609                        e.add_rms_norm_f16out(x1, sl_fo, w_next, x_nxt, h, h16, n_embd, t, eps)?;
1610                    } else {
1611                        e.add(x1, sl_fo, x_nxt, t * n_embd)?;
1612                        e.rms_norm(x_nxt, w_next, h, n_embd, t, eps)?;
1613                    }
1614                } else {
1615                    e.add(x1, sl_fo, x_nxt, t * n_embd)?;
1616                }
1617            }
1618            // CHUNK-INVARIANCE BISECT SEAM (lane/chunk-invariance, 2026-08-05): dump the
1619            // per-layer post-residual hidden for the LAST row of this chunk, keyed by
1620            // absolute position, so two runs at different MEMRA_PRIME_CHUNK can be diffed
1621            // layer-by-layer to find the FIRST diverging layer. Diagnostic only —
1622            // unset (the default) costs one OnceLock read per layer.
1623            if let Some(path) = Self::prime_trace_path() {
1624                let row = (base + t - 1) as usize;
1625                let host = e.dtoh(x_nxt)?;
1626                let last = &host[(t - 1) * n_embd..t * n_embd];
1627                use std::io::Write as _;
1628                let mut f = std::fs::OpenOptions::new()
1629                    .create(true)
1630                    .append(true)
1631                    .open(path)?;
1632                let mut h64: u64 = 0xcbf29ce484222325;
1633                for v in last {
1634                    h64 ^= v.to_bits() as u64;
1635                    h64 = h64.wrapping_mul(0x100000001b3);
1636                }
1637                writeln!(
1638                    f,
1639                    "{{\"pos\":{row},\"layer\":{il},\"t\":{t},\"base\":{base},\
1640                             \"hash\":\"{h64:016x}\",\"v0\":{:.9e},\"v1\":{:.9e},\"v2\":{:.9e}}}",
1641                    last[0], last[1], last[2]
1642                )?;
1643            }
1644            std::mem::swap(&mut x_cur, &mut x_nxt);
1645        }
1646        // hidden-stack return: clone the final x out of the slab
1647        let mut x = e.uninit(t * n_embd)?;
1648        e.copy_into(&mut x, 0, x_cur, t * n_embd)?;
1649        drop(slab_guard);
1650        Ok(x)
1651    }
1652
1653    /// The prime chunk's tail — h_seed + output_norm + one-row lm head + cache.pos advance —
1654    /// shared verbatim by the unsplit walk and the last stage of the ppN walk (`e` = the
1655    /// engine that produced `x`, i.e. the last stage's under the split; output_norm/output
1656    /// were loaded through that engine by the sharded loader, hybrid.rs `e_head`).
1657    fn prime_chunk_epilogue(
1658        &self,
1659        e: &Engine,
1660        x: CudaSlice<f32>,
1661        t: usize,
1662        cache: &mut Cache,
1663    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
1664        let n_embd = self.cfg.n_embd as usize;
1665        let eps = self.cfg.rms_eps;
1666        // h_seed = LAST row of x BEFORE output_norm (MTP-PLAN §A default) or AFTER it
1667        // (MEMRA_SPEC_HPOST — reference convention; hn is computed just below either way, so
1668        // the post-norm copy happens after hn exists).
1669        let mut h_seed = e.uninit(n_embd)?;
1670        if !crate::spec::spec_hpost() {
1671            e.copy_view_into(
1672                &mut h_seed,
1673                0,
1674                &x.slice((t - 1) * n_embd..t * n_embd),
1675                n_embd,
1676            )?;
1677        }
1678        // last-row logits, exactly like forward_last (norm all T — per-row op — then lm_head on 1 row).
1679        let mut hn = e.uninit(t * n_embd)?;
1680        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
1681        if crate::spec::spec_hpost() {
1682            e.copy_view_into(
1683                &mut h_seed,
1684                0,
1685                &hn.slice((t - 1) * n_embd..t * n_embd),
1686                n_embd,
1687            )?;
1688        }
1689        let last = e.view(&hn, t * n_embd);
1690        let last_row = last.slice((t - 1) * n_embd..t * n_embd);
1691        let mut hlast = e.uninit(n_embd)?;
1692        e.copy_view_into(&mut hlast, 0, &last_row, n_embd)?;
1693        let logits = e.matmul(&self.output, &hlast, 1)?;
1694        cache.pos += t;
1695        // Hidden stack handed to generate_spec as prompt_h: pre-norm x (default) or the full
1696        // post-norm stack hn (MEMRA_SPEC_HPOST).
1697        Ok((
1698            e.dtoh(&logits)?,
1699            h_seed,
1700            if crate::spec::spec_hpost() { hn } else { x },
1701        ))
1702    }
1703
1704    /// THE PRIME STAGE SPLIT (lane/pp-leverb, 2026-08-08): one prime chunk as N stage
1705    /// subgraphs, each range on ITS OWN engine/stream/device, with the `[T, n_embd]`
1706    /// boundary handoff at every fence cut — the prime twin of `decode_step_h_ppn` /
1707    /// `decode_step_t_core_ppn`, and the kill for the anatomy's two receipts: stage-1 trunk
1708    /// weights stop being peer-read (22% of the pp4096 wall) and dev1 stops running zero
1709    /// prefill kernels. Structure mirrors the verify split exactly:
1710    ///   stage 0        `rt.enter(0)` → per-stage pos_d + embed (the table lives with
1711    ///                  stage 0) → `prime_layers(fence[0], fence[1])` → `rt.tx`
1712    ///   middle stages  `rt.rx` → per-stage pos_d → range → `rt.tx`
1713    ///   last stage     `rt.rx` → range → the shared epilogue (output_norm + head live
1714    ///                  there via the sharded loader) → `publish_to`
1715    /// Laws inherited (not relearned): `fence_stages_behind` at entry (#87 — a previous
1716    /// round's stage-freed buffers must not be reused under the caller's queued reads);
1717    /// per-stage `pos_d` (allocated/consumed/freed on one stream); per-stage Engines from
1718    /// `PpNRt` (shared-scratch race); EXIT PUBLICATION for the device-resident returns
1719    /// (h_seed + hidden stack live on the last stage — the caller's stream must wait).
1720    /// KV/MoE locality falls out: each range's KV appends run on the owning stage
1721    /// (`pp::new_cache` placed the buffers there), each stage engine owns its own SLRU
1722    /// pool (per-Engine `moe_cache`, sized on ITS device — the SGLang #33666 law), and the
1723    /// slab-local MoE arm's `DevExps.dev` gate now matches on stage-1 layers too.
1724    /// EXACTNESS: the split adds zero deviation by construction (same kernels, same bytes,
1725    /// boundary = straight f32 copy); `prime-split-gate` (ppsplit) arbitrates bit-for-bit
1726    /// and its liveness counter is bumped here — the gate goes green with this function.
1727    fn prime_chunk_ppn(
1728        &self,
1729        e: &Engine,
1730        tokens: &[u32],
1731        cache: &mut Cache,
1732        seq_end: usize,
1733        fence: &[usize],
1734    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
1735        let rt = crate::pp::PpNRt::get(e)?;
1736        let n_st = fence.len() - 1;
1737        assert_eq!(
1738            rt.n_stages(),
1739            n_st,
1740            "PpNRt stage count {} != fence stages {n_st}",
1741            rt.n_stages()
1742        );
1743        let n_embd = self.cfg.n_embd as usize;
1744        let t = tokens.len();
1745        let base = cache.pos;
1746        debug_assert!(
1747            seq_end >= base + t,
1748            "prime_chunk_ppn: seq_end must cover this chunk"
1749        );
1750        let payload = t * n_embd;
1751        // The caller's ambient stream, captured BEFORE any enter() pushes a stage stream
1752        // (inside a stage scope e.stream() IS the stage stream and the exit wait would
1753        // self-order into a no-op) — the decode_step_t_core_ppn pattern.
1754        let caller_stream = e.stream();
1755        rt.fence_stages_behind(&caller_stream)?;
1756
1757        if n_st == 2 {
1758            let slot =
1759                self.prime_pp2_stage0_enqueue(e, rt, tokens, cache, seq_end, fence, base, false)?;
1760            let x =
1761                self.prime_pp2_stage1_enqueue(e, rt, slot, t, cache, seq_end, fence, base, false)?;
1762            let out = {
1763                rt.bind_stage(1)?;
1764                let _st1 = rt.enter(1);
1765                let e1 = rt.engine(1, e);
1766                self.prime_chunk_epilogue(e1, x, t, cache)?
1767            };
1768            rt.publish_to(1, &caller_stream)?;
1769            crate::pp::PRIME_SPLIT_CHUNKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1770            return Ok(out);
1771        }
1772
1773        let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
1774
1775        // ---- STAGE 0: embed + layers [0, fence[1]) + boundary-0 TX ----
1776        let mut slot = {
1777            let _st0 = rt.enter(0);
1778            let e0 = rt.engine(0, e);
1779            let pos_d = e0.htod_i32(&pos)?;
1780            let x = self.embed(e0, tokens)?;
1781            let x =
1782                self.prime_layers(e0, x, fence[0], fence[1], &pos_d, t, base, cache, seq_end)?;
1783            rt.tx(0, &x, payload)?
1784            // x + pos_d drop here: freed stream-ordered on stage-0's stream after use.
1785        };
1786
1787        // ---- MIDDLE STAGES: RX boundary s-1 -> range -> TX boundary s ----
1788        for s in 1..n_st - 1 {
1789            let _st = rt.enter(s);
1790            let es = rt.engine(s, e);
1791            let pos_d = es.htod_i32(&pos)?;
1792            let x = rt.rx(s - 1, slot, payload)?;
1793            let x = self.prime_layers(
1794                es,
1795                x,
1796                fence[s],
1797                fence[s + 1],
1798                &pos_d,
1799                t,
1800                base,
1801                cache,
1802                seq_end,
1803            )?;
1804            slot = rt.tx(s, &x, payload)?;
1805        }
1806
1807        // ---- LAST STAGE: RX + final range + the shared epilogue ----
1808        let _stl = rt.enter(n_st - 1);
1809        let el = rt.engine(n_st - 1, e);
1810        let pos_d = el.htod_i32(&pos)?;
1811        let x = rt.rx(n_st - 2, slot, payload)?;
1812        let x = self.prime_layers(
1813            el,
1814            x,
1815            fence[n_st - 1],
1816            fence[n_st],
1817            &pos_d,
1818            t,
1819            base,
1820            cache,
1821            seq_end,
1822        )?;
1823        let out = self.prime_chunk_epilogue(el, x, t, cache)?;
1824        // EXIT PUBLICATION: h_seed + the hidden stack are device-resident on the last
1825        // stage's stream; the caller resumes on its own stream (chunk-loop copy_into /
1826        // generate_spec's prompt_h consumer). The logits dtoh above already drained the
1827        // stage stream host-side, but the law is stated in events, not in a dtoh side
1828        // effect a later deferred form would remove.
1829        rt.publish_to(n_st - 1, &caller_stream)?;
1830        crate::pp::PRIME_SPLIT_CHUNKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1831        Ok(out)
1832    }
1833
1834    fn prime_pp2_stage0_enqueue(
1835        &self,
1836        e: &Engine,
1837        rt: &crate::pp::PpNRt,
1838        tokens: &[u32],
1839        cache: &mut Cache,
1840        seq_end: usize,
1841        fence: &[usize],
1842        base: usize,
1843        pipelined: bool,
1844    ) -> Result<usize, Box<dyn std::error::Error>> {
1845        let t = tokens.len();
1846        let n_embd = self.cfg.n_embd as usize;
1847        let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
1848        rt.bind_stage(0)?;
1849        let _st0 = rt.enter(0);
1850        let e0 = rt.engine(0, e);
1851        let pos_d = e0.htod_i32(&pos)?;
1852        let x = self.embed(e0, tokens)?;
1853        let _overlap = pipelined.then(crate::pp::enter_prime_pipe_stage);
1854        let x = self.prime_layers(e0, x, fence[0], fence[1], &pos_d, t, base, cache, seq_end)?;
1855        if pipelined {
1856            rt.tx_pipelined(0, &x, t * n_embd)
1857        } else {
1858            rt.tx(0, &x, t * n_embd)
1859        }
1860    }
1861
1862    fn prime_pp2_stage1_enqueue(
1863        &self,
1864        e: &Engine,
1865        rt: &crate::pp::PpNRt,
1866        slot: usize,
1867        t: usize,
1868        cache: &mut Cache,
1869        seq_end: usize,
1870        fence: &[usize],
1871        base: usize,
1872        pipelined: bool,
1873    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1874        let n_embd = self.cfg.n_embd as usize;
1875        let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
1876        rt.bind_stage(1)?;
1877        let _st1 = rt.enter(1);
1878        let e1 = rt.engine(1, e);
1879        let pos_d = e1.htod_i32(&pos)?;
1880        let x = rt.rx(0, slot, t * n_embd)?;
1881        let _overlap = pipelined.then(crate::pp::enter_prime_pipe_stage);
1882        self.prime_layers(e1, x, fence[1], fence[2], &pos_d, t, base, cache, seq_end)
1883    }
1884
1885    /// CAPTURE-SAFE prime trunk (task #14 increment 1, ARCHITECTURE-H100.md design v2):
1886    /// prime_chunk's layer stack with every capture hazard hoisted — `x` is the
1887    /// PRE-EMBEDDED input in a STABLE graph-input buffer (embed + its 8MB htod stay
1888    /// eager, one launch), `pos_d` is a baked device param (fresh prime = 0..T, constant
1889    /// per bucket), logits/h_seed/hidden stay DEVICE-resident (no dtoh), and cache.pos is
1890    /// NOT advanced (host state — the replay wrapper owns it). Body mirrors prime_chunk
1891    /// (the prime-graph-gate pins them together). KNOWN smoke-scope gap: append host-len
1892    /// bookkeeping still runs on the host per call — the real replay path moves the write
1893    /// slot to the len_d device counter (increment 3).
1894    /// GRAPH-OUTPUT CONTRACT: results are COPIED into caller-provided stable buffers
1895    /// (`logits_out` [n_vocab], `h_seed_out` [n_embd]) — every internal allocation drops
1896    /// INSIDE the capture region (alloc+free node pairs). Retaining an in-capture
1897    /// allocation across end_capture makes instantiate throw INVALID_VALUE (smoke finding
1898    /// 2026-07-26), and under AUTO_FREE_ON_LAUNCH its address wouldn't survive a launch
1899    /// anyway — the decode GraphSession's pre-allocated-output pattern is the law here.
1900    pub fn prime_chunk_captured(
1901        &self,
1902        e: &Engine,
1903        x_in: &CudaSlice<f32>,
1904        pos_d: &CudaSlice<i32>,
1905        t: usize,
1906        cache: &mut Cache,
1907        len_d: &CudaSlice<i32>,
1908        logits_out: &mut CudaSlice<f32>,
1909        h_seed_out: &mut CudaSlice<f32>,
1910    ) -> Result<(), Box<dyn std::error::Error>> {
1911        let cfg = &self.cfg;
1912        let n_embd = cfg.n_embd as usize;
1913        let eps = cfg.rms_eps;
1914        let f16fuse = crate::f16_ffi::pp_f16_enabled() && t >= 16;
1915        let mut x = e.uninit(t * n_embd)?;
1916        e.copy_into(&mut x, 0, x_in, t * n_embd)?;
1917        for (il, layer) in self.layers.iter().enumerate() {
1918            let mut h = e.uninit(t * n_embd)?;
1919            let mut hx16: Option<CudaSlice<u8>> = None;
1920            if f16fuse {
1921                let mut b16 = e.alloc_u8_uninit(t * n_embd * 2)?;
1922                e.rms_norm_f16out(
1923                    &x,
1924                    layer.attn_norm.float_data(),
1925                    &mut h,
1926                    &mut b16,
1927                    n_embd,
1928                    t,
1929                    eps,
1930                )?;
1931                hx16 = Some(b16);
1932            } else {
1933                e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
1934            }
1935            let mixed = match &layer.mixer {
1936                // The captured prime is ONE unchunked bucket over a FRESH cache (pos == 0), so the
1937                // request ends at t. If bucketed capture ever composes with chunking, seq_end must
1938                // come from the caller (see step35_attn_pre_wo's doc note).
1939                Mixer::Full(fa) => {
1940                    self.full_attn_prime(e, fa, &h, hx16.as_ref(), pos_d, t, cache, il, t)?
1941                }
1942                Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
1943                Mixer::Linear(la) => {
1944                    let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
1945                    let g4 = match hx16.as_ref() {
1946                        Some(xh) => e.matmul_group_xh(&ws, &h, xh, t)?,
1947                        None => e.matmul_group(&ws, &h, t)?,
1948                    };
1949                    self.linear_attn_prime_core_pad(e, la, g4, t, cache, il, Some(len_d))?
1950                }
1951            };
1952            let mut x1 = e.uninit(t * n_embd)?;
1953            e.add(&x, &mixed, &mut x1, t * n_embd)?;
1954            let mut z = e.uninit(t * n_embd)?;
1955            let mut zx16: Option<CudaSlice<u8>> = None;
1956            if f16fuse {
1957                let mut b16 = e.alloc_u8_uninit(t * n_embd * 2)?;
1958                e.rms_norm_f16out(
1959                    &x1,
1960                    layer.post_attn_norm.float_data(),
1961                    &mut z,
1962                    &mut b16,
1963                    n_embd,
1964                    t,
1965                    eps,
1966                )?;
1967                zx16 = Some(b16);
1968            } else {
1969                e.rms_norm(
1970                    &x1,
1971                    layer.post_attn_norm.float_data(),
1972                    &mut z,
1973                    n_embd,
1974                    t,
1975                    eps,
1976                )?;
1977            }
1978            let ffn_out = match &layer.ffn {
1979                crate::hybrid::Ffn::Dense {
1980                    ffn_gate,
1981                    ffn_up,
1982                    ffn_down,
1983                } => {
1984                    let n_ff = ffn_gate.out_features();
1985                    let mut g2 = match &zx16 {
1986                        Some(xh) => e.matmul_group_xh(&[ffn_gate, ffn_up], &z, xh, t)?,
1987                        None => e.matmul_group(&[ffn_gate, ffn_up], &z, t)?,
1988                    };
1989                    let up = g2.pop().unwrap();
1990                    let gate = g2.pop().unwrap();
1991                    let mut act = e.uninit(t * n_ff)?;
1992                    // dense FFN keys off the SHEXP clamp array — see forward()'s note.
1993                    Self::ffn_act_lim(
1994                        e,
1995                        &self.cfg,
1996                        &gate,
1997                        &up,
1998                        1.0,
1999                        1.0,
2000                        self.cfg.clamp_shexp_at(il as u32),
2001                        &mut act,
2002                        t * n_ff,
2003                    )?;
2004                    e.matmul(ffn_down, &act, t)?
2005                }
2006                crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_prefill(e, m, &z, t, il as u16)?,
2007            };
2008            let mut x2 = e.uninit(t * n_embd)?;
2009            e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
2010            x = x2;
2011        }
2012        // device-indexed TRUE last row (pads sit past it in a bucketed graph)
2013        if !crate::spec::spec_hpost() {
2014            e.row_gather_dev(&x, h_seed_out, len_d, n_embd)?;
2015        }
2016        let mut hn = e.uninit(t * n_embd)?;
2017        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
2018        if crate::spec::spec_hpost() {
2019            e.row_gather_dev(&hn, h_seed_out, len_d, n_embd)?;
2020        }
2021        let mut hlast = e.uninit(n_embd)?;
2022        e.row_gather_dev(&hn, &mut hlast, len_d, n_embd)?;
2023        let logits = e.matmul(&self.output, &hlast, 1)?;
2024        let nv = logits.len();
2025        e.copy_into(logits_out, 0, &logits, nv)?;
2026        Ok(())
2027    }
2028
2029    fn step35_prime_batch_on() -> bool {
2030        std::env::var("MEMRA_STEP35_PRIME_BATCH").as_deref() != Ok("0")
2031    }
2032
2033    /// Step35 cross-request prime range: weight-streaming work runs once at `m=sum(T)`;
2034    /// sequence-scoped attention/KV work stays on each request's own cache and positions.
2035    #[allow(clippy::too_many_arguments)]
2036    fn step35_prime_batch_layers(
2037        &self,
2038        e: &Engine,
2039        mut x: CudaSlice<f32>,
2040        lo: usize,
2041        hi: usize,
2042        ts: &[usize],
2043        offs: &[usize],
2044        pos_ds: &[CudaSlice<i32>],
2045        caches: &mut [&mut Cache],
2046    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2047        let cfg = &self.cfg;
2048        let n_embd = cfg.n_embd as usize;
2049        let eps = cfg.rms_eps;
2050        let b = ts.len();
2051        let total: usize = ts.iter().sum();
2052        let f16fuse = crate::f16_ffi::pp_f16_enabled() && total >= 16;
2053
2054        let split = |e: &Engine,
2055                     y: &CudaSlice<f32>,
2056                     dim: usize|
2057         -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
2058            let mut out = Vec::with_capacity(b);
2059            for s in 0..b {
2060                let mut ys = e.uninit(ts[s] * dim)?;
2061                e.copy_view_into(
2062                    &mut ys,
2063                    0,
2064                    &y.slice(offs[s] * dim..(offs[s] + ts[s]) * dim),
2065                    ts[s] * dim,
2066                )?;
2067                out.push(ys);
2068            }
2069            Ok(out)
2070        };
2071
2072        for il in lo..hi {
2073            let layer = &self.layers[il];
2074            let Mixer::Full(fa) = &layer.mixer else {
2075                return Err(format!("step35 layer {il} is not full-attn — corrupt config").into());
2076            };
2077
2078            let mut h = e.uninit(total * n_embd)?;
2079            let mut hx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
2080            if f16fuse {
2081                e.rms_norm_f16out(
2082                    &x,
2083                    layer.attn_norm.float_data(),
2084                    &mut h,
2085                    &mut hx16,
2086                    n_embd,
2087                    total,
2088                    eps,
2089                )?;
2090            } else {
2091                e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, total, eps)?;
2092            }
2093
2094            // Q/K/V + gate projections and wo are batched weight streams. The step35 core
2095            // remains per sequence, so its partial RoPE, SWA view, KV append, and gate
2096            // application stay verbatim.
2097            let gate_w = fa
2098                .attn_gate
2099                .as_ref()
2100                .ok_or("step35 layer is missing attn_gate.weight (head-wise attention gate)")?;
2101            let mut g4 = if f16fuse {
2102                e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv, gate_w], &h, &hx16, total)?
2103            } else {
2104                e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv, gate_w], &h, total)?
2105            };
2106            let gate = g4.pop().unwrap();
2107            let mut parts: Vec<Vec<CudaSlice<f32>>> =
2108                (0..b).map(|_| Vec::with_capacity(3)).collect();
2109            for (w, y) in [&fa.wq, &fa.wk, &fa.wv].iter().zip(g4) {
2110                for (s, ys) in split(e, &y, w.out_features())?.into_iter().enumerate() {
2111                    parts[s].push(ys);
2112                }
2113            }
2114            let gates = split(e, &gate, gate_w.out_features())?;
2115            let geometry = self.step35_geom(il);
2116            let hd = geometry.head_dim_k as usize;
2117            let nh = geometry.n_head as usize;
2118            let mut ag_cat = e.uninit(total * nh * hd)?;
2119            for (s, (g3s, gate)) in parts.into_iter().zip(gates).enumerate() {
2120                let ag = self.step35_attn_pre_wo(
2121                    e,
2122                    fa,
2123                    g3s,
2124                    None,
2125                    Some(&gate),
2126                    &pos_ds[s],
2127                    ts[s],
2128                    Some(&mut *caches[s]),
2129                    il,
2130                    ts[s],
2131                )?;
2132                e.copy_into(&mut ag_cat, offs[s] * nh * hd, &ag, ts[s] * nh * hd)?;
2133            }
2134            let mixed = e.matmul(&fa.wo, &ag_cat, total)?;
2135
2136            let mut x1 = e.uninit(total * n_embd)?;
2137            let mut z = e.uninit(total * n_embd)?;
2138            let mut zx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
2139            if f16fuse {
2140                e.add_rms_norm_f16out(
2141                    &x,
2142                    &mixed,
2143                    layer.post_attn_norm.float_data(),
2144                    &mut x1,
2145                    &mut z,
2146                    &mut zx16,
2147                    n_embd,
2148                    total,
2149                    eps,
2150                )?;
2151            } else {
2152                e.add(&x, &mixed, &mut x1, total * n_embd)?;
2153                e.rms_norm(
2154                    &x1,
2155                    layer.post_attn_norm.float_data(),
2156                    &mut z,
2157                    n_embd,
2158                    total,
2159                    eps,
2160                )?;
2161            }
2162
2163            let ffn_out = match &layer.ffn {
2164                crate::hybrid::Ffn::Dense {
2165                    ffn_gate,
2166                    ffn_up,
2167                    ffn_down,
2168                } => {
2169                    let n_ff = ffn_gate.out_features();
2170                    let mut g2 = if f16fuse {
2171                        e.matmul_group_xh(&[ffn_gate, ffn_up], &z, &zx16, total)?
2172                    } else {
2173                        e.matmul_group(&[ffn_gate, ffn_up], &z, total)?
2174                    };
2175                    let up = g2.pop().unwrap();
2176                    let gate = g2.pop().unwrap();
2177                    let mut act = e.uninit(total * n_ff)?;
2178                    let d_lim = cfg.clamp_shexp_at(il as u32);
2179                    if Self::f16out_on(e, total) && cfg.m3.is_none() && d_lim.is_none() {
2180                        let mut a16 = e.alloc_u8_uninit(total * n_ff * 2)?;
2181                        e.silu_mul_f16out(&gate, &up, &mut act, &mut a16, total * n_ff)?;
2182                        match e.try_f16_gemm_pre(ffn_down, &a16, total)? {
2183                            Some(y) => y,
2184                            None => e.matmul(ffn_down, &act, total)?,
2185                        }
2186                    } else {
2187                        Self::ffn_act_lim(
2188                            e,
2189                            cfg,
2190                            &gate,
2191                            &up,
2192                            1.0,
2193                            1.0,
2194                            d_lim,
2195                            &mut act,
2196                            total * n_ff,
2197                        )?;
2198                        e.matmul(ffn_down, &act, total)?
2199                    }
2200                }
2201                crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, total, il as u16)?,
2202            };
2203            let mut x2 = e.uninit(total * n_embd)?;
2204            e.add(&x1, &ffn_out, &mut x2, total * n_embd)?;
2205            x = x2;
2206        }
2207        Ok(x)
2208    }
2209
2210    fn step35_prime_batch_epilogue(
2211        &self,
2212        e: &Engine,
2213        x: CudaSlice<f32>,
2214        ts: &[usize],
2215        offs: &[usize],
2216        caches: &mut [&mut Cache],
2217    ) -> Result<Vec<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
2218        let n_embd = self.cfg.n_embd as usize;
2219        let total: usize = ts.iter().sum();
2220        let mut hn = e.uninit(total * n_embd)?;
2221        e.rms_norm(
2222            &x,
2223            self.output_norm.float_data(),
2224            &mut hn,
2225            n_embd,
2226            total,
2227            self.cfg.rms_eps,
2228        )?;
2229
2230        let hidden_src = if crate::spec::spec_hpost() { &hn } else { &x };
2231        let mut out = Vec::with_capacity(ts.len());
2232        for s in 0..ts.len() {
2233            let mut hidden = e.uninit(ts[s] * n_embd)?;
2234            e.copy_view_into(
2235                &mut hidden,
2236                0,
2237                &hidden_src.slice(offs[s] * n_embd..(offs[s] + ts[s]) * n_embd),
2238                ts[s] * n_embd,
2239            )?;
2240            let last0 = (offs[s] + ts[s] - 1) * n_embd;
2241            let mut h_seed = e.uninit(n_embd)?;
2242            e.copy_view_into(
2243                &mut h_seed,
2244                0,
2245                &hidden_src.slice(last0..last0 + n_embd),
2246                n_embd,
2247            )?;
2248            // Exactness-first: the serial reference runs the output head at m=1.
2249            let mut hlast = e.uninit(n_embd)?;
2250            e.copy_view_into(&mut hlast, 0, &hn.slice(last0..last0 + n_embd), n_embd)?;
2251            let logits = e.dtoh(&e.matmul(&self.output, &hlast, 1)?)?;
2252            caches[s].pos += ts[s];
2253            out.push((logits, h_seed, hidden));
2254        }
2255        Ok(out)
2256    }
2257
2258    fn step35_prime_cache_batch(
2259        &self,
2260        e: &Engine,
2261        prompts: &[&[u32]],
2262        caches: &mut [&mut Cache],
2263    ) -> Result<Vec<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
2264        if crate::pp::pp_host_bounce_active()
2265            && (!crate::pp::prime_pp_on() || crate::pp::pp_cuts(self.layers.len()).is_none())
2266        {
2267            return Err(
2268                "step35_prime_cache_batch: MEMRA_PP_HOST_BOUNCE=1 requires a valid prime \
2269                 stage split; refusing an unsplit remote-weight walk"
2270                    .into(),
2271            );
2272        }
2273        if !Self::step35_prime_batch_on() {
2274            return Err("step35 batched prime is disabled (MEMRA_STEP35_PRIME_BATCH=0)".into());
2275        }
2276        if caches.iter().any(|c| c.pos != 0) {
2277            return Err(
2278                "step35 batched prime currently supports complete fresh prompts only; \
2279                 continuation/tick chunks require per-request queued_after"
2280                    .into(),
2281            );
2282        }
2283
2284        let ts: Vec<usize> = prompts.iter().map(|p| p.len()).collect();
2285        for &t in &ts {
2286            assert!(
2287                t >= PRIME_MIN_T,
2288                "step35 batched prime needs T >= {PRIME_MIN_T}"
2289            );
2290        }
2291        for (s, c) in caches.iter().enumerate() {
2292            assert!(
2293                ts[s] <= c.max_ctx,
2294                "step35 batched prime exceeds cache max_ctx"
2295            );
2296        }
2297        let offs: Vec<usize> = ts
2298            .iter()
2299            .scan(0usize, |a, &t| {
2300                let o = *a;
2301                *a += t;
2302                Some(o)
2303            })
2304            .collect();
2305        let total: usize = ts.iter().sum();
2306        let payload = total * self.cfg.n_embd as usize;
2307        let cat_tokens: Vec<u32> = prompts.iter().flat_map(|p| p.iter().copied()).collect();
2308        let positions: Vec<Vec<i32>> = ts.iter().map(|&t| (0..t as i32).collect()).collect();
2309        let upload_positions =
2310            |e: &Engine| -> Result<Vec<CudaSlice<i32>>, Box<dyn std::error::Error>> {
2311                positions
2312                    .iter()
2313                    .map(|p| e.htod_i32(p))
2314                    .collect::<Result<_, _>>()
2315            };
2316
2317        static ONCE: std::sync::Once = std::sync::Once::new();
2318        ONCE.call_once(|| {
2319            eprintln!(
2320                "[step35-prime-batch] first concat prime: B={} tokens={total}",
2321                prompts.len()
2322            );
2323        });
2324
2325        let out = if !crate::pp::pp2_streams_off() && crate::pp::prime_pp_on() {
2326            if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
2327                let rt = crate::pp::PpNRt::get(e)?;
2328                let n_st = fence.len() - 1;
2329                assert_eq!(
2330                    rt.n_stages(),
2331                    n_st,
2332                    "step35 prime batch stage count mismatch"
2333                );
2334                let caller_stream = e.stream();
2335                rt.fence_stages_behind(&caller_stream)?;
2336
2337                let mut slot = {
2338                    let _st0 = rt.enter(0);
2339                    let e0 = rt.engine(0, e);
2340                    let pos_ds = upload_positions(e0)?;
2341                    let x = self.embed(e0, &cat_tokens)?;
2342                    let x = self.step35_prime_batch_layers(
2343                        e0, x, fence[0], fence[1], &ts, &offs, &pos_ds, caches,
2344                    )?;
2345                    rt.tx(0, &x, payload)?
2346                };
2347                for s in 1..n_st - 1 {
2348                    let _st = rt.enter(s);
2349                    let es = rt.engine(s, e);
2350                    let pos_ds = upload_positions(es)?;
2351                    let x = rt.rx(s - 1, slot, payload)?;
2352                    let x = self.step35_prime_batch_layers(
2353                        es,
2354                        x,
2355                        fence[s],
2356                        fence[s + 1],
2357                        &ts,
2358                        &offs,
2359                        &pos_ds,
2360                        caches,
2361                    )?;
2362                    slot = rt.tx(s, &x, payload)?;
2363                }
2364
2365                let _stl = rt.enter(n_st - 1);
2366                let el = rt.engine(n_st - 1, e);
2367                let pos_ds = upload_positions(el)?;
2368                let x = rt.rx(n_st - 2, slot, payload)?;
2369                let x = self.step35_prime_batch_layers(
2370                    el,
2371                    x,
2372                    fence[n_st - 1],
2373                    fence[n_st],
2374                    &ts,
2375                    &offs,
2376                    &pos_ds,
2377                    caches,
2378                )?;
2379                let out = self.step35_prime_batch_epilogue(el, x, &ts, &offs, caches)?;
2380                rt.publish_to(n_st - 1, &caller_stream)?;
2381                crate::pp::STEP35_PRIME_BATCH_SPLITS
2382                    .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2383                out
2384            } else {
2385                let pos_ds = upload_positions(e)?;
2386                let x = self.embed(e, &cat_tokens)?;
2387                let x = self.step35_prime_batch_layers(
2388                    e,
2389                    x,
2390                    0,
2391                    self.layers.len(),
2392                    &ts,
2393                    &offs,
2394                    &pos_ds,
2395                    caches,
2396                )?;
2397                self.step35_prime_batch_epilogue(e, x, &ts, &offs, caches)?
2398            }
2399        } else {
2400            let pos_ds = upload_positions(e)?;
2401            let x = self.embed(e, &cat_tokens)?;
2402            let x = self.step35_prime_batch_layers(
2403                e,
2404                x,
2405                0,
2406                self.layers.len(),
2407                &ts,
2408                &offs,
2409                &pos_ds,
2410                caches,
2411            )?;
2412            self.step35_prime_batch_epilogue(e, x, &ts, &offs, caches)?
2413        };
2414        crate::pp::STEP35_PRIME_BATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2415        Ok(out)
2416    }
2417
2418    /// Cross-request BATCHED fresh prime (task #13, design in ARCHITECTURE-H100.md): the
2419    /// trunk's token-parallel ops (embed, norms, adds, ffn, projection GROUPS) run once on
2420    /// the CONCATENATION of B sequences — GEMMs at m = sum_T, the continuous-batching win
2421    /// the serving-lane bench attributed the remaining vLLM gap to. The stateful mixer
2422    /// CORES (QK-norm/RoPE/FA/append, conv/GDN scans) run per sequence on split projection
2423    /// buffers (D2D row copies; the mixers' own out-projections stay per-seq this
2424    /// increment). CONTINUATION primes (increment (b), 2026-07-30): cache.pos > 0 seqs
2425    /// batch the projections/FFN/lm_head exactly like fresh; the mixer cores take the
2426    /// per-seq CONTINUATION arms (Full: core_inner with carried pos_d + fa_prefill_view
2427    /// over the quantized past; Linear: the stateful pad_view twin — the same state
2428    /// carry the chunked single-seq prime rides). The fresh-only favl/gdn-vl fast paths
2429    /// stay byte-identical (gated on !carried). gemma4 models have no continuation
2430    /// prime (v0 monolithic fresh) — carried gemma4 batches return Err (caller falls
2431    /// back to single-chunk serving).
2432    /// NUMERIC CONFIG: a concat GEMM tiles K differently than per-seq GEMMs — same class
2433    /// as every prefill GEMM change; prime_batch_gate arbitrates (argmax + stream battery).
2434    pub fn prime_cache_batch(
2435        &self,
2436        e: &Engine,
2437        prompts: &[&[u32]],
2438        caches: &mut [&mut Cache],
2439    ) -> Result<Vec<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
2440        let cfg = &self.cfg;
2441        let n_embd = cfg.n_embd as usize;
2442        let eps = cfg.rms_eps;
2443        let b = prompts.len();
2444        assert!(b >= 1 && b == caches.len());
2445        let pos0s: Vec<usize> = caches.iter().map(|c| c.pos).collect();
2446        let carried = pos0s.iter().any(|&p| p > 0);
2447        // gemma4: refuse UNCONDITIONALLY (2026-08-07, lane/gemma4-serve-gaps). The old guard
2448        // covered only `carried` — two concurrent FRESH gemma4 prompts batched into the
2449        // generic concat attn core below (uniform geometry, no per-layer swa window, no
2450        // softcapped head): compiles, runs, wrong logits. Same silent-wrong class as the
2451        // step35 refusal beneath. The per-sequence `gemma4_prime` is the supported prefill.
2452        if cfg.gemma4.is_some() {
2453            return Err(
2454                "prime_cache_batch: gemma4 has no batched prime core (per-layer \
2455                        swa/global geometry, softcapped head) — use gemma4_prime per sequence"
2456                    .into(),
2457            );
2458        }
2459        // Step35 has a dedicated concat walk: the generic core below cannot express its
2460        // per-layer geometry/SWA/head gate, and under PP the dedicated path is stage-scoped.
2461        if cfg.step35.is_some() {
2462            return self.step35_prime_cache_batch(e, prompts, caches);
2463        }
2464        let ts: Vec<usize> = prompts.iter().map(|p| p.len()).collect();
2465        for &t in &ts {
2466            assert!(
2467                t >= PRIME_MIN_T,
2468                "prime_cache_batch needs T >= {PRIME_MIN_T}"
2469            );
2470        }
2471        for (s, c) in caches.iter().enumerate() {
2472            assert!(
2473                c.pos + ts[s] <= c.max_ctx,
2474                "prime_cache_batch: prompt exceeds cache max_ctx"
2475            );
2476        }
2477        let total: usize = ts.iter().sum();
2478        let offs: Vec<usize> = ts
2479            .iter()
2480            .scan(0usize, |a, &t| {
2481                let o = *a;
2482                *a += t;
2483                Some(o)
2484            })
2485            .collect();
2486        // per-seq positions (fresh: 0..T_s; continuation: pos0..pos0+T_s)
2487        let pos_ds: Vec<CudaSlice<i32>> = ts
2488            .iter()
2489            .zip(&pos0s)
2490            .map(|(&t, &p0)| e.htod_i32(&(p0 as i32..(p0 + t) as i32).collect::<Vec<_>>()))
2491            .collect::<Result<_, _>>()?;
2492        // split a concat [total, dim] buffer into per-seq copies
2493        let split = |e: &Engine,
2494                     y: &CudaSlice<f32>,
2495                     dim: usize|
2496         -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
2497            let mut out = Vec::with_capacity(b);
2498            for s in 0..b {
2499                let mut ys = e.uninit(ts[s] * dim)?;
2500                e.copy_view_into(
2501                    &mut ys,
2502                    0,
2503                    &y.slice(offs[s] * dim..(offs[s] + ts[s]) * dim),
2504                    ts[s] * dim,
2505                )?;
2506                out.push(ys);
2507            }
2508            Ok(out)
2509        };
2510
2511        let cat_tokens: Vec<u32> = prompts.iter().flat_map(|p| p.iter().copied()).collect();
2512        let mut x = self.embed(e, &cat_tokens)?; // [total, n_embd]
2513        for (il, layer) in self.layers.iter().enumerate() {
2514            let mut h = e.uninit(total * n_embd)?;
2515            let mut hx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
2516            e.rms_norm_f16out(
2517                &x,
2518                layer.attn_norm.float_data(),
2519                &mut h,
2520                &mut hx16,
2521                n_embd,
2522                total,
2523                eps,
2524            )?;
2525            // mixer: projection GROUP on the concat (m = total), stateful core per seq
2526            let mut mixed = e.uninit(total * n_embd)?;
2527            match &layer.mixer {
2528                Mixer::Full(fa) => {
2529                    let g3 = e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], &h, &hx16, total)?;
2530                    // task #18 (attn side): the WHOLE attn core is varlen for fresh gated
2531                    // batches — split/QK-norm/RoPE/append (attn_pre_vl8, view inputs: the
2532                    // q/k/v split copies vanish) + ONE varlen FA. Per-block math identical
2533                    // everywhere (bit-gateable); MEMRA_FA_VL=0 or a non-bf16kv config falls
2534                    // back to the per-seq dispatch.
2535                    let geometry = self.cfg.full_attention_geometry_at(il as u32);
2536                    let (n_head, n_head_kv, head_dim) = (
2537                        geometry.n_head as usize,
2538                        geometry.n_head_kv as usize,
2539                        geometry.head_dim_k as usize,
2540                    );
2541                    let fa_scale = geometry.attention_scale();
2542                    let use_favl = !carried
2543                        && (2..=8).contains(&b)
2544                        && (head_dim == 256 || head_dim == 128)
2545                        && geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ
2546                        && std::env::var("MEMRA_NOFA").is_err()
2547                        && std::env::var("MEMRA_FA_FLOOR").is_err()
2548                        && std::env::var("MEMRA_FA_PP_W2").as_deref() != Ok("1")
2549                        && std::env::var("MEMRA_FA_BF16KV").as_deref() != Ok("0")
2550                        && std::env::var("MEMRA_FA_VL").as_deref() != Ok("0");
2551                    if use_favl {
2552                        let (qf_w, kf_w, vf_w) = (
2553                            fa.wq.out_features(),
2554                            fa.wk.out_features(),
2555                            fa.wv.out_features(),
2556                        );
2557                        struct APre {
2558                            q: CudaSlice<f32>,
2559                            gate: Option<CudaSlice<f32>>,
2560                            qn: CudaSlice<f32>,
2561                            kn: CudaSlice<f32>,
2562                        }
2563                        let mut aps = Vec::with_capacity(b);
2564                        for &t in ts.iter().take(b) {
2565                            aps.push(APre {
2566                                q: e.uninit(t * n_head * head_dim)?,
2567                                gate: Some(e.uninit(t * n_head * head_dim)?),
2568                                qn: e.uninit(t * n_head * head_dim)?,
2569                                kn: e.uninit(t * n_head_kv * head_dim)?,
2570                            });
2571                        }
2572                        let (kv_dim_k, kv_dim_v, ktb, vtb) = {
2573                            let kvl = caches[0].kv[il].as_ref().unwrap();
2574                            (kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes)
2575                        };
2576                        let pargs: Vec<crate::AttnPreVl> = (0..b)
2577                            .map(|s| {
2578                                let (o, t) = (offs[s], ts[s]);
2579                                let kvl = caches[s].kv[il].as_ref().unwrap();
2580                                assert!(
2581                                    kvl.len == 0 && kvl.len + t <= caches[s].max_ctx,
2582                                    "prime_cache_batch attn vl: fresh + capacity"
2583                                );
2584                                crate::AttnPreVl {
2585                                    qf: e.addr_f32v(&g3[0].slice(o * qf_w..(o + t) * qf_w)),
2586                                    kf: e.addr_f32v(&g3[1].slice(o * kf_w..(o + t) * kf_w)),
2587                                    vf: e.addr_f32v(&g3[2].slice(o * vf_w..(o + t) * vf_w)),
2588                                    q: e.addr_f32(&aps[s].q),
2589                                    gate: e.addr_f32(aps[s].gate.as_ref().unwrap()),
2590                                    qn: e.addr_f32(&aps[s].qn),
2591                                    kn: e.addr_f32(&aps[s].kn),
2592                                    kc: e.addr_u8(&kvl.k),
2593                                    vc: e.addr_u8(&kvl.v),
2594                                    t: t as i32,
2595                                    pad: 0,
2596                                }
2597                            })
2598                            .collect();
2599                        e.attn_pre_vl8(
2600                            &pargs,
2601                            fa.q_norm.float_data(),
2602                            fa.k_norm.float_data(),
2603                            head_dim,
2604                            geometry.n_rot as usize,
2605                            n_head,
2606                            n_head_kv,
2607                            self.cfg.rms_eps,
2608                            geometry.rope_base,
2609                            1.0,
2610                            kv_dim_k,
2611                            kv_dim_v,
2612                            ktb,
2613                            vtb,
2614                        )?;
2615                        for s in 0..b {
2616                            let kvl = caches[s].kv[il].as_mut().unwrap();
2617                            kvl.len += ts[s];
2618                            let new_len = kvl.len as i32;
2619                            e.set_i32_one(&mut kvl.len_d, new_len)?;
2620                        }
2621                        let mut attns = Vec::with_capacity(b);
2622                        let mut mirrors = Vec::with_capacity(b);
2623                        for &t in ts.iter().take(b) {
2624                            attns.push(e.uninit(t * n_head * head_dim)?);
2625                            let n = t * n_head_kv * head_dim;
2626                            mirrors.push((e.alloc_u8_uninit(n * 2)?, e.alloc_u8_uninit(n * 2)?));
2627                        }
2628                        // FA3 batched twin (round 31): TMA-swizzled wgmma vl when the
2629                        // promoted single-seq config is on; else the mma favl.
2630                        let fa3_on = match std::env::var("MEMRA_FA3").as_deref() {
2631                            Ok("0") => false,
2632                            Ok("1") => true,
2633                            _ => cfg!(memra_hopper_mma),
2634                        };
2635                        if fa3_on {
2636                            let mut q16s = Vec::with_capacity(b);
2637                            let mut v16s = Vec::with_capacity(b);
2638                            for s in 0..b {
2639                                let t = ts[s];
2640                                let mut q16 = e.alloc_u8_uninit(t * n_head * head_dim * 2)?;
2641                                e.f32_to_bf16_into(&aps[s].qn, &mut q16, t * n_head * head_dim)?;
2642                                let mut k16 = e.alloc_u8_uninit(t * n_head_kv * head_dim * 2)?;
2643                                e.f32_to_bf16_into(&aps[s].kn, &mut k16, t * n_head_kv * head_dim)?;
2644                                let mut v16 = e.alloc_u8_uninit(t * n_head_kv * head_dim * 2)?;
2645                                e.f32_to_bf16_v(
2646                                    &g3[2].slice(offs[s] * vf_w..(offs[s] + t) * vf_w),
2647                                    &mut v16,
2648                                    t * n_head_kv * head_dim,
2649                                )?;
2650                                q16s.push(q16);
2651                                v16s.push((k16, v16));
2652                            }
2653                            let mut qp = [core::ptr::null::<core::ffi::c_void>(); 8];
2654                            let mut kp = qp;
2655                            let mut vp = qp;
2656                            let mut op = [core::ptr::null_mut::<f32>(); 8];
2657                            let mut tsv = [0i32; 8];
2658                            for s in 0..b {
2659                                qp[s] = e.addr_u8(&q16s[s]) as *const core::ffi::c_void;
2660                                kp[s] = e.addr_u8(&v16s[s].0) as *const core::ffi::c_void;
2661                                vp[s] = e.addr_u8(&v16s[s].1) as *const core::ffi::c_void;
2662                                op[s] = e.addr_f32(&attns[s]) as *mut f32;
2663                                tsv[s] = ts[s] as i32;
2664                            }
2665                            let rc = unsafe {
2666                                crate::fa3_vl_raw(
2667                                    qp.as_ptr(),
2668                                    kp.as_ptr(),
2669                                    vp.as_ptr(),
2670                                    op.as_ptr(),
2671                                    tsv.as_ptr(),
2672                                    b as i32,
2673                                    n_head as i32,
2674                                    n_head_kv as i32,
2675                                    head_dim as i32,
2676                                    fa_scale,
2677                                    e.stream().cu_stream() as *mut core::ffi::c_void,
2678                                )
2679                            };
2680                            if rc != 0 {
2681                                return Err(format!("memra_fa3_vl rc={rc}").into());
2682                            }
2683                        } else {
2684                            let fargs: Vec<crate::FaSeqVl> = (0..b)
2685                                .map(|s| crate::FaSeqVl {
2686                                    q: e.addr_f32(&aps[s].qn),
2687                                    k16: e.addr_u8(&mirrors[s].0),
2688                                    v16: e.addr_u8(&mirrors[s].1),
2689                                    o: e.addr_f32(&attns[s]),
2690                                    kf: e.addr_f32(&aps[s].kn),
2691                                    vf: e.addr_f32v(
2692                                        &g3[2].slice(offs[s] * vf_w..(offs[s] + ts[s]) * vf_w),
2693                                    ),
2694                                    t: ts[s] as i32,
2695                                    pad: 0,
2696                                })
2697                                .collect();
2698                            e.fa_prefill_vl8(&fargs, head_dim, n_head, n_head_kv, fa_scale)?;
2699                        }
2700                        for (s, attn) in attns.into_iter().enumerate() {
2701                            let (attn_g, ag16) = self.full_attn_prime_post_fa(
2702                                e,
2703                                attn,
2704                                &aps[s].gate,
2705                                ts[s],
2706                                n_head,
2707                                head_dim,
2708                            )?;
2709                            let mut done = false;
2710                            if let Some(xh) = &ag16 {
2711                                done = e.try_f16_gemm_pre_into_off(
2712                                    &fa.wo,
2713                                    xh,
2714                                    ts[s],
2715                                    &mut mixed,
2716                                    offs[s] * n_embd,
2717                                )?;
2718                            }
2719                            if !done {
2720                                let m = e.matmul(&fa.wo, &attn_g, ts[s])?;
2721                                e.copy_into(&mut mixed, offs[s] * n_embd, &m, ts[s] * n_embd)?;
2722                            }
2723                        }
2724                    } else {
2725                        let mut parts: Vec<Vec<CudaSlice<f32>>> =
2726                            (0..b).map(|_| Vec::new()).collect();
2727                        for (w, y) in [&fa.wq, &fa.wk, &fa.wv].iter().zip(g3) {
2728                            for (s, ys) in split(e, &y, w.out_features())?.into_iter().enumerate() {
2729                                parts[s].push(ys);
2730                            }
2731                        }
2732                        for (s, g3s) in parts.into_iter().enumerate() {
2733                            // task #16 gather removal: wo writes into `mixed` at offs[s] directly.
2734                            let (attn_g, ag16) = self.full_attn_prime_core_inner(
2735                                e, fa, g3s, &pos_ds[s], ts[s], caches[s], il,
2736                            )?;
2737                            let mut done = false;
2738                            if let Some(xh) = &ag16 {
2739                                done = e.try_f16_gemm_pre_into_off(
2740                                    &fa.wo,
2741                                    xh,
2742                                    ts[s],
2743                                    &mut mixed,
2744                                    offs[s] * n_embd,
2745                                )?;
2746                            }
2747                            if !done {
2748                                let m = e.matmul(&fa.wo, &attn_g, ts[s])?;
2749                                e.copy_into(&mut mixed, offs[s] * n_embd, &m, ts[s] * n_embd)?;
2750                            }
2751                        }
2752                    }
2753                }
2754                Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
2755                Mixer::Linear(la) => {
2756                    // task #16: NO split copies (cores read row-offset views of the concat
2757                    // outputs; out-GEMMs write into `mixed` at offs[s]). task #18: the core
2758                    // itself is BATCHED — per-seq prep/K1-K3, then ONE varlen K4 + ONE
2759                    // varlen K5 launch for all sequences.
2760                    let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
2761                    let g4 = e.matmul_group_xh(&ws, &h, &hx16, total)?;
2762                    let outs =
2763                        self.linear_attn_prime_core_batch(e, la, &g4, &offs, &ts, caches, il)?;
2764                    for (s, (gn, gn16)) in outs.into_iter().enumerate() {
2765                        let (o, t) = (offs[s], ts[s]);
2766                        let mut done = false;
2767                        if let Some(xh) = &gn16 {
2768                            done = e.try_f16_gemm_pre_into_off(
2769                                &la.ssm_out,
2770                                xh,
2771                                t,
2772                                &mut mixed,
2773                                o * n_embd,
2774                            )?;
2775                        }
2776                        if !done {
2777                            let m = e.matmul(&la.ssm_out, &gn, t)?;
2778                            e.copy_into(&mut mixed, o * n_embd, &m, t * n_embd)?;
2779                        }
2780                    }
2781                }
2782            }
2783            let mut x1 = e.uninit(total * n_embd)?;
2784            let mut z = e.uninit(total * n_embd)?;
2785            let mut zx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
2786            e.add_rms_norm_f16out(
2787                &x,
2788                &mixed,
2789                layer.post_attn_norm.float_data(),
2790                &mut x1,
2791                &mut z,
2792                &mut zx16,
2793                n_embd,
2794                total,
2795                eps,
2796            )?;
2797            let ffn_out = match &layer.ffn {
2798                crate::hybrid::Ffn::Dense {
2799                    ffn_gate,
2800                    ffn_up,
2801                    ffn_down,
2802                } => {
2803                    let n_ff = ffn_gate.out_features();
2804                    let mut g2 = e.matmul_group_xh(&[ffn_gate, ffn_up], &z, &zx16, total)?;
2805                    let up = g2.pop().unwrap();
2806                    let gate = g2.pop().unwrap();
2807                    let mut act = e.uninit(total * n_ff)?;
2808                    // task #17 (batch trunk): silu twin emits the down GEMM's fp16 operand
2809                    // in-epilogue (nsys round-26: this trunk still paid 32 cvt passes).
2810                    // A clamped layer must skip the plain-SiLU twin (see prime_chunk's note).
2811                    let d_lim = self.cfg.clamp_shexp_at(il as u32);
2812                    if Self::f16out_on(e, total) && self.cfg.m3.is_none() && d_lim.is_none() {
2813                        let mut a16 = e.alloc_u8_uninit(total * n_ff * 2)?;
2814                        e.silu_mul_f16out(&gate, &up, &mut act, &mut a16, total * n_ff)?;
2815                        match e.try_f16_gemm_pre(ffn_down, &a16, total)? {
2816                            Some(y) => y,
2817                            None => e.matmul(ffn_down, &act, total)?,
2818                        }
2819                    } else {
2820                        Self::ffn_act_lim(
2821                            e,
2822                            &self.cfg,
2823                            &gate,
2824                            &up,
2825                            1.0,
2826                            1.0,
2827                            d_lim,
2828                            &mut act,
2829                            total * n_ff,
2830                        )?;
2831                        e.matmul(ffn_down, &act, total)?
2832                    }
2833                }
2834                crate::hybrid::Ffn::Moe(m) => {
2835                    self.moe_ffn_il_prefill(e, m, &z, total, il as u16)?
2836                }
2837            };
2838            let mut x2 = e.uninit(total * n_embd)?;
2839            e.add(&x1, &ffn_out, &mut x2, total * n_embd)?;
2840            x = x2;
2841        }
2842        // epilogue per seq (identical math to prime_chunk: norm all rows, lm_head on last row)
2843        let mut hn = e.uninit(total * n_embd)?;
2844        e.rms_norm(
2845            &x,
2846            self.output_norm.float_data(),
2847            &mut hn,
2848            n_embd,
2849            total,
2850            eps,
2851        )?;
2852        // batched lm_head (nsys round-26: B sequential m=1 matvecs re-read the 600MB
2853        // lm_head weight B times — 2.2ms at B=6): gather the B last rows and run ONE
2854        // m=B GEMM (f16 lane; falls back to the per-seq matvec when no mirror). The
2855        // f16-vs-mmvq logits delta is the prefill GEMM numeric class — prime_batch_gate's
2856        // argmax battery arbitrates, same as every other prefill GEMM change.
2857        let mut hcat = e.uninit(b * n_embd)?;
2858        for s in 0..b {
2859            let last0 = (offs[s] + ts[s] - 1) * n_embd;
2860            e.copy_view_into(
2861                &mut hcat,
2862                s * n_embd,
2863                &hn.slice(last0..last0 + n_embd),
2864                n_embd,
2865            )?;
2866        }
2867        let logits_cat = if b >= 2 {
2868            e.try_f16_gemm(&self.output, &hcat, b)?
2869        } else {
2870            None
2871        };
2872        let logits_host: Option<Vec<f32>> = match &logits_cat {
2873            Some(lc) => Some(e.dtoh(lc)?),
2874            None => None,
2875        };
2876        let n_vocab = self.output.out_features();
2877        let mut hidden_all = if crate::spec::spec_hpost() {
2878            split(e, &hn, n_embd)?
2879        } else {
2880            split(e, &x, n_embd)?
2881        };
2882        let mut out = Vec::with_capacity(b);
2883        for s in 0..b {
2884            let last0 = (offs[s] + ts[s] - 1) * n_embd;
2885            let mut h_seed = e.uninit(n_embd)?;
2886            if !crate::spec::spec_hpost() {
2887                e.copy_view_into(&mut h_seed, 0, &x.slice(last0..last0 + n_embd), n_embd)?;
2888            } else {
2889                e.copy_view_into(&mut h_seed, 0, &hn.slice(last0..last0 + n_embd), n_embd)?;
2890            }
2891            let logits = match &logits_host {
2892                Some(lh) => lh[s * n_vocab..(s + 1) * n_vocab].to_vec(),
2893                None => {
2894                    let mut hlast = e.uninit(n_embd)?;
2895                    e.copy_view_into(&mut hlast, 0, &hn.slice(last0..last0 + n_embd), n_embd)?;
2896                    e.dtoh(&e.matmul(&self.output, &hlast, 1)?)?
2897                }
2898            };
2899            caches[s].pos += ts[s];
2900            out.push((logits, h_seed, hidden_all.remove(0)));
2901        }
2902        Ok(out)
2903    }
2904
2905    /// `full_attn` (batched prefill mixer) + the cache side-effect: append the T post-RoPE K/V
2906    /// rows into the resident quantized KV cache (q8_0 K / q5_1 V) and advance len/len_d. Row
2907    /// bytes are BIT-IDENTICAL to the decode append (same per-warp quant kernel per row; the
2908    /// batched `append_kv_quantized_rows` runs that exact warp math on a (block, token) grid).
2909    /// The attention itself is unchanged prefill math (fa_prefill over the f32 K/V).
2910    ///
2911    /// `seq_end` = the ABSOLUTE end position of the WHOLE prime request (`cache.pos + prompt_len`
2912    /// measured BEFORE the chunk loop starts), i.e. a chunk-size-invariant property of the request.
2913    /// Only step35's SWA arm reads it (`step35_attn_pre_wo`'s doc note explains why keying on the
2914    /// chunk's own `t_kv` made the output depend on MEMRA_PRIME_CHUNK); every other arch ignores it.
2915    #[allow(clippy::too_many_arguments)]
2916    fn full_attn_prime(
2917        &self,
2918        e: &Engine,
2919        fa: &FullAttnLayer,
2920        h: &CudaSlice<f32>,
2921        hx: Option<&CudaSlice<u8>>,
2922        pos_d: &CudaSlice<i32>,
2923        t: usize,
2924        cache: &mut Cache,
2925        il: usize,
2926        seq_end: usize,
2927    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2928        if self.cfg.step35.is_some() {
2929            return self.step35_attn_prime(e, fa, h, hx, pos_d, t, cache, il, seq_end);
2930        }
2931        // PROJ/CORE SPLIT (task #13, 2026-07-26): the q/k/v group projection is hoisted so
2932        // the cross-request batch driver can run it at m = sum_T over concatenated tokens;
2933        // this single-seq path composes proj+core identically (byte-for-byte the old body).
2934        // task #14: `hx` = the norm-fused fp16 twin of `h` (skips the convert launch).
2935        let g3 = match hx {
2936            Some(xh) => e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], h, xh, t)?,
2937            None => e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?,
2938        };
2939        self.full_attn_prime_core(e, fa, g3, pos_d, t, cache, il)
2940    }
2941
2942    /// Everything after the q/k/v projections (split/QK-norm/RoPE/FA/gate/append + wo).
2943    /// `g3` = matmul_group([wq, wk, wv]) output rows for THIS sequence.
2944    /// Composes inner + wo (== the pre-split core; every existing caller unchanged).
2945    fn full_attn_prime_core(
2946        &self,
2947        e: &Engine,
2948        fa: &FullAttnLayer,
2949        g3: Vec<CudaSlice<f32>>,
2950        pos_d: &CudaSlice<i32>,
2951        t: usize,
2952        cache: &mut Cache,
2953        il: usize,
2954    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2955        let (attn_g, ag16) = self.full_attn_prime_core_inner(e, fa, g3, pos_d, t, cache, il)?;
2956        if let Some(xh) = &ag16 {
2957            if let Some(y) = e.try_f16_gemm_pre(&fa.wo, xh, t)? {
2958                return Ok(y);
2959            }
2960        }
2961        Ok(e.matmul(&fa.wo, &attn_g, t)?)
2962    }
2963
2964    fn full_attn_prime_core_inner(
2965        &self,
2966        e: &Engine,
2967        fa: &FullAttnLayer,
2968        g3: Vec<CudaSlice<f32>>,
2969        pos_d: &CudaSlice<i32>,
2970        t: usize,
2971        cache: &mut Cache,
2972        il: usize,
2973    ) -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
2974        let cfg = &self.cfg;
2975        let geometry = cfg.full_attention_geometry_at(il as u32);
2976        let n_head = geometry.n_head as usize;
2977        let n_head_kv = geometry.n_head_kv as usize;
2978        let head_dim = geometry.head_dim_k as usize;
2979        let scale = geometry.attention_scale();
2980        let (pre, base_len) = self.full_attn_prime_pre_fa(e, fa, g3, pos_d, t, cache, il)?;
2981        let AttnPre { q, k, v, gate } = pre;
2982        let mut attn = e.uninit(t * n_head * head_dim)?;
2983        self.full_attn_prime_fa_dispatch(
2984            e, &q, &k, &v, &mut attn, base_len, t, cache, il, head_dim, n_head, n_head_kv, scale,
2985        )?;
2986        self.full_attn_prime_post_fa(e, attn, &gate, t, n_head, head_dim)
2987    }
2988
2989    /// task #18 (attn side): projections tail through KV append — everything before the
2990    /// attention kernel. Returns the post-rope q/k, v, optional out-gate, and the KV rows
2991    /// present BEFORE this chunk's append (base_len; 0 == fresh).
2992    #[allow(clippy::type_complexity)]
2993    fn full_attn_prime_pre_fa(
2994        &self,
2995        e: &Engine,
2996        fa: &FullAttnLayer,
2997        mut g3: Vec<CudaSlice<f32>>,
2998        pos_d: &CudaSlice<i32>,
2999        t: usize,
3000        cache: &mut Cache,
3001        il: usize,
3002    ) -> Result<(AttnPre, usize), Box<dyn std::error::Error>> {
3003        let cfg = &self.cfg;
3004        let geometry = cfg.full_attention_geometry_at(il as u32);
3005        let n_head = geometry.n_head as usize;
3006        let n_head_kv = geometry.n_head_kv as usize;
3007        let head_dim = geometry.head_dim_k as usize;
3008        let eps = cfg.rms_eps;
3009
3010        // qwen35 fuses [q|gate] per head in wq (2*head_dim stride); M3/Hy3 have NO output gate
3011        // (attention_output_gate=false) — wq out = n_head*head_dim exactly, and q_gate_split
3012        // would read 2x out of bounds. `gated` keys both the split and the sigmoid epilogue.
3013        let gated = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
3014        let v = g3.pop().unwrap();
3015        let mut k = g3.pop().unwrap();
3016        let qf = g3.pop().unwrap();
3017        let (mut q, gate) = if gated {
3018            let mut q = e.uninit(t * n_head * head_dim)?;
3019            let mut gate = e.uninit(t * n_head * head_dim)?;
3020            e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, t)?;
3021            (q, Some(gate))
3022        } else {
3023            (qf, None)
3024        };
3025
3026        let mut qn = e.uninit(t * n_head * head_dim)?;
3027        e.rms_norm(
3028            &q,
3029            fa.q_norm.float_data(),
3030            &mut qn,
3031            head_dim,
3032            n_head * t,
3033            eps,
3034        )?;
3035        q = qn;
3036        let mut kn = e.uninit(t * n_head_kv * head_dim)?;
3037        e.rms_norm(
3038            &k,
3039            fa.k_norm.float_data(),
3040            &mut kn,
3041            head_dim,
3042            n_head_kv * t,
3043            eps,
3044        )?;
3045        k = kn;
3046        let rope_dims = geometry.n_rot as usize;
3047        e.rope_neox(
3048            &mut q,
3049            pos_d,
3050            head_dim,
3051            rope_dims,
3052            n_head,
3053            t,
3054            geometry.rope_base,
3055            1.0,
3056        )?;
3057        e.rope_neox(
3058            &mut k,
3059            pos_d,
3060            head_dim,
3061            rope_dims,
3062            n_head_kv,
3063            t,
3064            geometry.rope_base,
3065            1.0,
3066        )?;
3067
3068        // CACHE SIDE-EFFECT: append the T post-rope K/V token rows (token-major [T, kv_dim] ==
3069        // the cache row layout) quantized into cache.kv[il], then advance len + device len_d.
3070        {
3071            let kvl = cache.kv[il].as_mut().unwrap();
3072            assert!(kvl.len + t <= cache.max_ctx, "prime_cache: KV overflow");
3073            e.append_kv_quantized_rows(
3074                &k,
3075                &v,
3076                &mut kvl.k,
3077                &mut kvl.v,
3078                kvl.len,
3079                t,
3080                kvl.kv_dim_k,
3081                kvl.kv_dim_v,
3082                kvl.k_tok_bytes,
3083                kvl.v_tok_bytes,
3084                crate::Engine::kv_fp8_on(),
3085            )?;
3086            kvl.len += t;
3087            let new_len = kvl.len as i32;
3088            e.set_i32_one(&mut kvl.len_d, new_len)?;
3089        }
3090
3091        let base_len = {
3092            let kvl = cache.kv[il].as_ref().unwrap();
3093            kvl.len - t // KV rows present BEFORE this chunk's append above
3094        };
3095        Ok((AttnPre { q, k, v, gate }, base_len))
3096    }
3097
3098    /// batched prefill attention. FRESH prime (no past KV): unchanged forward_last math over
3099    /// the f32 K/V of this batch. CONTINUATION chunk (past KV present): the chunk's queries
3100    /// must attend to [0 .. base+t) — run fa_prefill_view over the resident QUANTIZED cache
3101    /// (the spec-verify pattern; kernel's causal mask offsets by T_kv-T). Numerically this
3102    /// reads q8_0/q5_1-dequantized K/V for the past AND the current chunk — the same class as
3103    /// decode reading the cache; the run-gen/first-16 battery is the accuracy authority.
3104    #[allow(clippy::too_many_arguments)]
3105    fn full_attn_prime_fa_dispatch(
3106        &self,
3107        e: &Engine,
3108        q: &CudaSlice<f32>,
3109        k: &CudaSlice<f32>,
3110        v: &CudaSlice<f32>,
3111        attn: &mut CudaSlice<f32>,
3112        base_len: usize,
3113        t: usize,
3114        cache: &mut Cache,
3115        il: usize,
3116        head_dim: usize,
3117        n_head: usize,
3118        n_head_kv: usize,
3119        scale: f32,
3120    ) -> Result<(), Box<dyn std::error::Error>> {
3121        // GRAIN-FREE CHUNK-INVARIANCE FIX (lane/chunkinv-flip, 2026-08-05): the base_len == 0
3122        // f32 special case (fa_prefill over this batch's f32 K/V) is DROPPED. pre_fa appends
3123        // the chunk's quantized rows into cache.kv[il] BEFORE this dispatch, so chunk 0 can
3124        // attend through the quantized cache exactly like every later chunk (quantize-then-
3125        // attend). One numeric class for every row => the chunk size cannot decide where a
3126        // precision edge falls, so chunked prefill is reduction-order-stable with NO grain
3127        // knob (the stronger fix VERDICT.md filed; supersedes the MEMRA_PRIME_INVARIANT door's
3128        // pin-the-boundary approach).
3129        // MEMRA_PRIME_F32CHUNK0=1 is the ROLLBACK SEAM to the old arithmetic (chunk 0 attends
3130        // f32 K/V) — flags-doctrine rollback door AND the chunkinv gate's canary injection:
3131        // with the fix unconditional, only re-introducing the class edge can prove the gate
3132        // still detects the mechanism. Never on in a measured default run.
3133        if base_len == 0 && std::env::var("MEMRA_PRIME_F32CHUNK0").as_deref() == Ok("1") {
3134            if std::env::var("MEMRA_NOFA").is_ok() || !(head_dim == 256 || head_dim == 128) {
3135                e.sdpa_naive(
3136                    q, k, v, attn, head_dim, n_head, n_head_kv, t, t, scale, true,
3137                )?;
3138            } else {
3139                e.fa_prefill(
3140                    q, k, v, attn, head_dim, n_head, n_head_kv, t, t, scale, true,
3141                )?;
3142            }
3143            return Ok(());
3144        }
3145        let kvl = cache.kv[il].as_ref().unwrap();
3146        let t_kv = base_len + t;
3147        let k_view = e.view_u8(&kvl.k, t_kv * kvl.k_tok_bytes);
3148        let v_view = e.view_u8(&kvl.v, t_kv * kvl.v_tok_bytes);
3149        // fa_prefill_q/_qw twins are stamped for head_dim 256 (qwen35) and 128 (M3) only;
3150        // other dims (and MEMRA_NOFA) take the naive quantized-view SDPA — SAME cache bytes,
3151        // same numeric class, so the uniform contract holds on the fallback too.
3152        if std::env::var("MEMRA_NOFA").is_ok() || !(head_dim == 256 || head_dim == 128) {
3153            e.sdpa_naive_quantized_view(
3154                q,
3155                &k_view,
3156                &v_view,
3157                attn,
3158                head_dim,
3159                n_head,
3160                n_head_kv,
3161                t,
3162                t_kv,
3163                scale,
3164                true,
3165                kvl.k_tok_bytes,
3166                kvl.v_tok_bytes,
3167            )?;
3168            return Ok(());
3169        }
3170        // ARC B (2026-07-05): dequant-once workspace, DEFAULT ON. fa_prefill_q's inline
3171        // dequant re-reads+re-dequants the whole quantized KV stream from every one of the
3172        // T/64 x n_head CTAs (64x+ redundant at chunk=4096; 30.5% of the 32k prime wall).
3173        // fa_prefill_view_ws dequants K/V ONCE into a resident bf16 workspace then runs the
3174        // bit-identical bf16 twin (fa_prefill_qw) — same staged values, same FP order, token-
3175        // identical output (gate: MEMRA_PRIME_CHUNK=4096 ws-on vs ws-off vs monolithic).
3176        // MEMRA_PRIME_DEQW=0 reverts to the inline-dequant kernel.
3177        let deqw = std::env::var("MEMRA_PRIME_DEQW")
3178            .map(|v| v != "0")
3179            .unwrap_or(true);
3180        if deqw {
3181            e.fa_prefill_view_ws(
3182                q,
3183                &k_view,
3184                &v_view,
3185                attn,
3186                head_dim,
3187                n_head,
3188                n_head_kv,
3189                t,
3190                t_kv,
3191                scale,
3192                true,
3193                kvl.k_tok_bytes,
3194                kvl.v_tok_bytes,
3195                crate::Engine::kv_fp8_on(),
3196            )?;
3197        } else {
3198            e.fa_prefill_view(
3199                q,
3200                &k_view,
3201                &v_view,
3202                attn,
3203                head_dim,
3204                n_head,
3205                n_head_kv,
3206                t,
3207                t_kv,
3208                scale,
3209                true,
3210                kvl.k_tok_bytes,
3211                kvl.v_tok_bytes,
3212                crate::Engine::kv_fp8_on(),
3213            )?;
3214        }
3215        Ok(())
3216    }
3217
3218    /// task #17: sig_mul_f16out fuses [sigmoid + mul + f16 convert] into one launch
3219    /// (bit-identical composition) and hands wo its fp16 operand directly.
3220    fn full_attn_prime_post_fa(
3221        &self,
3222        e: &Engine,
3223        attn: CudaSlice<f32>,
3224        gate: &Option<CudaSlice<f32>>,
3225        t: usize,
3226        n_head: usize,
3227        head_dim: usize,
3228    ) -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
3229        let (attn_g, ag16) = match gate {
3230            Some(gate) => {
3231                let n = t * n_head * head_dim;
3232                let mut ag = e.uninit(n)?;
3233                if Self::f16out_on(e, t) {
3234                    let mut a16 = e.alloc_u8_uninit(n * 2)?;
3235                    e.sig_mul_f16out(&attn, gate, &mut ag, &mut a16, n)?;
3236                    (ag, Some(a16))
3237                } else {
3238                    let mut gsig = e.uninit(n)?;
3239                    e.sigmoid(gate, &mut gsig, n)?;
3240                    e.mul(&attn, &gsig, &mut ag, n)?;
3241                    (ag, None)
3242                }
3243            }
3244            None => (attn, None),
3245        };
3246        Ok((attn_g, ag16))
3247    }
3248
3249    /// STATEFUL batched linear-attention prime: `linear_attn`'s prefill-dispatch pass (normal
3250    /// `e.matmul` — GEMM at m>=16 — plus the prefill repack/L2/glog kernels) but with the state
3251    /// carried THROUGH the cache like the spec verify does: carried-ring conv
3252    /// (ssm_conv1d_tm_state writes the final ring back) + ONE gdn_scan from cache.recur[il]'s
3253    /// current state (zero at a fresh prime) whose final state ping-pongs back into the cache.
3254    /// Wiring mirrors `linear_attn_verify_t` (spec.rs); dispatch mirrors `linear_attn` (prefill).
3255    fn linear_attn_prime(
3256        &self,
3257        e: &Engine,
3258        la: &LinearAttnLayer,
3259        h: &CudaSlice<f32>,
3260        hx: Option<&CudaSlice<u8>>,
3261        t: usize,
3262        cache: &mut Cache,
3263        il: usize,
3264    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3265        // PROJ/CORE SPLIT (task #13): see full_attn_prime — same hoist for the GDN 4-tuple.
3266        let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
3267        let g4 = match hx {
3268            Some(xh) => e.matmul_group_xh(&ws, h, xh, t)?,
3269            None => e.matmul_group(&ws, h, t)?,
3270        };
3271        self.linear_attn_prime_core(e, la, g4, t, cache, il)
3272    }
3273
3274    /// Everything after the GDN 4-tuple projections (conv/chunk stack/gated norm + ssm_out).
3275    fn linear_attn_prime_core(
3276        &self,
3277        e: &Engine,
3278        la: &LinearAttnLayer,
3279        mut g4: Vec<CudaSlice<f32>>,
3280        t: usize,
3281        cache: &mut Cache,
3282        il: usize,
3283    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3284        self.linear_attn_prime_core_pad(e, la, g4.drain(..).collect(), t, cache, il, None)
3285    }
3286
3287    /// task #14: `pad_len` = device true length for PADDED prime graphs. Pads become
3288    /// identity GDN steps (gdn_pad_mask zeroes beta/g_log past the true length) and the
3289    /// conv ring writes back from the true tail. None = classic path, byte-identical.
3290    #[allow(clippy::too_many_arguments)]
3291    fn linear_attn_prime_core_pad_inner(
3292        &self,
3293        e: &Engine,
3294        la: &LinearAttnLayer,
3295        mut g4: Vec<CudaSlice<f32>>,
3296        t: usize,
3297        cache: &mut Cache,
3298        il: usize,
3299        pad_len: Option<&CudaSlice<i32>>,
3300    ) -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
3301        // shim over the view twin (task #16): full-range views of the owned buffers.
3302        let ssm = self.cfg.ssm.as_ref().unwrap();
3303        let d_state = ssm.state_size as usize;
3304        let num_k = ssm.group_count as usize;
3305        let num_v = ssm.time_step_rank as usize;
3306        let key_dim = d_state * num_k;
3307        let value_dim = d_state * num_v;
3308        let conv_dim = key_dim * 2 + value_dim;
3309        let alpha = g4.pop().unwrap(); // [T, num_v]
3310        let beta_raw = g4.pop().unwrap(); // [T, num_v]
3311        let z = g4.pop().unwrap(); // [T, value_dim]
3312        let qkv_mixed = g4.pop().unwrap(); // [T, conv_dim] token-major
3313        self.linear_attn_prime_core_pad_view(
3314            e,
3315            la,
3316            &qkv_mixed.slice(0..t * conv_dim),
3317            &z.slice(0..t * value_dim),
3318            &beta_raw.slice(0..t * num_v),
3319            &alpha.slice(0..t * num_v),
3320            t,
3321            cache,
3322            il,
3323            pad_len,
3324        )
3325    }
3326
3327    /// task #18: the GDN prep stage (conv + repack + l2 x2 + sigmoid + glog + pad-mask) —
3328    /// shared verbatim by the per-seq scan path and the varlen batched path.
3329    #[allow(clippy::too_many_arguments)]
3330    fn linear_attn_gdn_prep(
3331        &self,
3332        e: &Engine,
3333        la: &LinearAttnLayer,
3334        qkv_mixed: &cudarc::driver::CudaView<f32>,
3335        beta_raw: &cudarc::driver::CudaView<f32>,
3336        alpha: &cudarc::driver::CudaView<f32>,
3337        t: usize,
3338        cache: &mut Cache,
3339        il: usize,
3340        pad_len: Option<&CudaSlice<i32>>,
3341    ) -> Result<GdnPrep, Box<dyn std::error::Error>> {
3342        let cfg = &self.cfg;
3343        let ssm = cfg.ssm.as_ref().unwrap();
3344        let d_state = ssm.state_size as usize; // 128
3345        let num_k = ssm.group_count as usize; // 16
3346        let num_v = ssm.time_step_rank as usize; // 32
3347        let d_conv = ssm.conv_kernel as usize; // 4
3348        let key_dim = d_state * num_k; // 2048
3349        let value_dim = d_state * num_v; // 4096
3350        let conv_dim = key_dim * 2 + value_dim; // 8192
3351        let eps = cfg.rms_eps;
3352        debug_assert!(
3353            t >= d_conv - 1,
3354            "stateful conv needs T >= pad (PRIME_MIN_T gates)"
3355        );
3356
3357        // conv with CARRIED ring state + ring roll (state read + final-window write-back).
3358        // task #18 conv-fuse (default ON): conv + SiLU + repack in one pass — the 67MB
3359        // conv_out intermediate and its transposed re-read disappear (values bit-identical).
3360        // MEMRA_CONV_FUSE=0 reverts to the two-kernel chain.
3361        let rl = cache.recur[il].as_mut().unwrap();
3362        let hk = Self::gdn_hk(e, t, num_v, num_k);
3363        let conv_fuse = std::env::var("MEMRA_CONV_FUSE").as_deref() != Ok("0");
3364        let hk = if conv_fuse { hk } else { num_v }; // de-broadcast rides the fused conv
3365        let mut q_g = e.uninit(d_state * hk * t)?;
3366        let mut k_g = e.uninit(d_state * hk * t)?;
3367        let mut v_g = e.uninit(d_state * num_v * t)?;
3368        if conv_fuse {
3369            e.ssm_conv1d_gdn_state_pad(
3370                qkv_mixed,
3371                &mut rl.conv_state,
3372                la.ssm_conv1d.float_data(),
3373                &mut q_g,
3374                &mut k_g,
3375                &mut v_g,
3376                conv_dim,
3377                t,
3378                d_conv,
3379                d_state,
3380                num_v,
3381                num_k,
3382                key_dim,
3383                hk,
3384                pad_len,
3385            )?;
3386        } else {
3387            let mut conv_out = e.uninit(conv_dim * t)?; // [conv_dim, T] channel-major, SiLU
3388            e.ssm_conv1d_tm_state_pad_v(
3389                qkv_mixed,
3390                &mut rl.conv_state,
3391                la.ssm_conv1d.float_data(),
3392                &mut conv_out,
3393                conv_dim,
3394                t,
3395                d_conv,
3396                pad_len,
3397            )?;
3398            e.qkv_to_gdn_repack(
3399                &conv_out, &mut q_g, &mut k_g, &mut v_g, d_state, num_v, num_k, key_dim, t,
3400            )?;
3401        }
3402        let mut q_l2 = e.uninit(d_state * hk * t)?;
3403        // mirror-fold (round 35): q's bf16 twin (wgmma K45/K2 A-operand) in-epilogue too.
3404        // Emitted only where a consumer exists (the wgmma config) — on other arches the
3405        // alloc + epilogue stores would be pure waste.
3406        let qb16 = if Engine::l2_v2_on(d_state) && e.gdn_wgmma_on(32) {
3407            let mut qb = e.alloc_u8_uninit(d_state * hk * t * 2)?;
3408            e.l2_norm_pp(&q_g, &mut q_l2, Some(&mut qb), d_state, hk * t, eps)?;
3409            Some(qb)
3410        } else {
3411            e.l2_norm_pp(&q_g, &mut q_l2, None, d_state, hk * t, eps)?;
3412            None
3413        };
3414        let mut k_l2 = e.uninit(d_state * hk * t)?;
3415        // mirror-fold: emit k's bf16 twin (the chunked-scan kb16 mirror) in-epilogue
3416        let kb16 = if Engine::l2_v2_on(d_state) {
3417            let mut kb = e.alloc_u8_uninit(d_state * hk * t * 2)?;
3418            e.l2_norm_pp(&k_g, &mut k_l2, Some(&mut kb), d_state, hk * t, eps)?;
3419            Some(kb)
3420        } else {
3421            e.l2_norm_pp(&k_g, &mut k_l2, None, d_state, hk * t, eps)?;
3422            None
3423        };
3424        let mut beta = e.uninit(t * num_v)?;
3425        e.sigmoid_v(beta_raw, &mut beta, t * num_v)?;
3426        let mut g_log = e.uninit(t * num_v)?;
3427        e.gdn_glog_v(
3428            alpha,
3429            la.ssm_dt.float_data(),
3430            la.ssm_a.float_data(),
3431            &mut g_log,
3432            num_v,
3433            t,
3434        )?;
3435        if let Some(len_d) = pad_len {
3436            e.gdn_pad_mask(&mut beta, &mut g_log, len_d, num_v, t)?;
3437        }
3438        Ok(GdnPrep {
3439            hk,
3440            q_l2,
3441            k_l2,
3442            v_g,
3443            beta,
3444            g_log,
3445            kb16,
3446            qb16,
3447        })
3448    }
3449
3450    /// task #18: batched GDN mixer core — per-seq prep + K1-K3, then ONE varlen K4 and
3451    /// ONE varlen K5 launch for all B sequences (full-machine occupancy vs B underfilled
3452    /// per-seq trains). Per-block math identical to the per-seq path (bit-gateable);
3453    /// MEMRA_GDN_VL=0 or a non-mma/non-C32 config falls back to the per-seq loop.
3454    #[allow(clippy::too_many_arguments)]
3455    fn linear_attn_prime_core_batch(
3456        &self,
3457        e: &Engine,
3458        la: &LinearAttnLayer,
3459        g4: &[CudaSlice<f32>],
3460        offs: &[usize],
3461        ts: &[usize],
3462        caches: &mut [&mut Cache],
3463        il: usize,
3464    ) -> Result<Vec<(CudaSlice<f32>, Option<CudaSlice<u8>>)>, Box<dyn std::error::Error>> {
3465        let ssm = self.cfg.ssm.as_ref().unwrap();
3466        let d_state = ssm.state_size as usize;
3467        let num_k = ssm.group_count as usize;
3468        let num_v = ssm.time_step_rank as usize;
3469        let key_dim = d_state * num_k;
3470        let value_dim = d_state * num_v;
3471        let conv_dim = key_dim * 2 + value_dim;
3472        let eps = self.cfg.rms_eps;
3473        let scale = 1.0 / (d_state as f32).sqrt();
3474        let b = ts.len();
3475        let c = Engine::gdn_chunk_size();
3476        // CONTINUATION batches (increment (b)) take the per-seq stateful pad_view twin
3477        // below — the varlen K4/K5 chain is fresh-only (zero initial state assumed).
3478        let carried = caches.iter().any(|c| c.pos > 0);
3479        let use_vl = !carried
3480            && (2..=8).contains(&b)
3481            && Engine::gdn_chunked_enabled()
3482            && ts.iter().all(|&t| t >= 16)
3483            && e.gdn_mma_enabled(c)
3484            && std::env::var("MEMRA_GDN_VL").as_deref() != Ok("0");
3485        if !use_vl {
3486            return (0..b)
3487                .map(|s| {
3488                    let (o, t) = (offs[s], ts[s]);
3489                    self.linear_attn_prime_core_pad_view(
3490                        e,
3491                        la,
3492                        &g4[0].slice(o * conv_dim..(o + t) * conv_dim),
3493                        &g4[1].slice(o * value_dim..(o + t) * value_dim),
3494                        &g4[2].slice(o * num_v..(o + t) * num_v),
3495                        &g4[3].slice(o * num_v..(o + t) * num_v),
3496                        t,
3497                        caches[s],
3498                        il,
3499                        None,
3500                    )
3501                })
3502                .collect();
3503        }
3504        // increment 3: the ENTIRE core is varlen — allocs only per seq, then 13 launches
3505        // total for the whole batch: [conv, ring, repack, l2(q+k), gate-prep] +
3506        // [k-mirror, K1, K2, K3, w-mirror, K4, K5] + [gated-norm tail].
3507        struct SeqBufs {
3508            conv_out: CudaSlice<f32>,
3509            q_g: CudaSlice<f32>,
3510            k_g: CudaSlice<f32>,
3511            v_g: CudaSlice<f32>,
3512            q_l2: CudaSlice<f32>,
3513            k_l2: CudaSlice<f32>,
3514            beta: CudaSlice<f32>,
3515            g_log: CudaSlice<f32>,
3516            gn: CudaSlice<f32>,
3517            gn16: CudaSlice<u8>,
3518        }
3519        let d_conv = ssm.conv_kernel as usize;
3520        let f16o = Self::f16out_on(e, 16);
3521        let hk = Self::gdn_hk(e, 16, num_v, num_k); // vl path is always chunked+mma
3522        let mut sb = Vec::with_capacity(b);
3523        let mut pres = Vec::with_capacity(b);
3524        for &t in ts.iter().take(b) {
3525            sb.push(SeqBufs {
3526                conv_out: e.uninit(conv_dim * t)?,
3527                q_g: e.uninit(d_state * hk * t)?,
3528                k_g: e.uninit(d_state * hk * t)?,
3529                v_g: e.uninit(d_state * num_v * t)?,
3530                q_l2: e.uninit(d_state * hk * t)?,
3531                k_l2: e.uninit(d_state * hk * t)?,
3532                beta: e.uninit(t * num_v)?,
3533                g_log: e.uninit(t * num_v)?,
3534                gn: e.uninit(d_state * num_v * t)?,
3535                gn16: e.alloc_u8_uninit(d_state * num_v * t * 2)?,
3536            });
3537            pres.push(e.gdn_chunk_alloc(num_v, t, c, hk)?);
3538        }
3539        let prep_args: Vec<crate::GdnPrepVl> = (0..b)
3540            .map(|s| {
3541                let (o, t) = (offs[s], ts[s]);
3542                let rl = caches[s].recur[il].as_ref().unwrap();
3543                crate::GdnPrepVl {
3544                    qkv: e.addr_f32v(&g4[0].slice(o * conv_dim..(o + t) * conv_dim)),
3545                    conv_state: e.addr_f32(&rl.conv_state),
3546                    conv_out: e.addr_f32(&sb[s].conv_out),
3547                    q_g: e.addr_f32(&sb[s].q_g),
3548                    k_g: e.addr_f32(&sb[s].k_g),
3549                    v_g: e.addr_f32(&sb[s].v_g),
3550                    q_l2: e.addr_f32(&sb[s].q_l2),
3551                    k_l2: e.addr_f32(&sb[s].k_l2),
3552                    beta_raw: e.addr_f32v(&g4[2].slice(o * num_v..(o + t) * num_v)),
3553                    alpha: e.addr_f32v(&g4[3].slice(o * num_v..(o + t) * num_v)),
3554                    beta: e.addr_f32(&sb[s].beta),
3555                    g_log: e.addr_f32(&sb[s].g_log),
3556                    o: e.addr_f32(&pres[s].o),
3557                    z: e.addr_f32v(&g4[1].slice(o * value_dim..(o + t) * value_dim)),
3558                    gn: e.addr_f32(&sb[s].gn),
3559                    gn16: e.addr_u8(&sb[s].gn16),
3560                    kb16: if Engine::l2_v2_on(d_state) {
3561                        e.addr_u8(&pres[s].kb16)
3562                    } else {
3563                        0
3564                    },
3565                    qb16: if Engine::l2_v2_on(d_state) && e.gdn_wgmma_on(c) {
3566                        e.addr_u8(&pres[s].qb16)
3567                    } else {
3568                        0
3569                    },
3570                    t: t as i32,
3571                    pad: 0,
3572                }
3573            })
3574            .collect();
3575        let args: Vec<crate::GdnSeqVl> = (0..b)
3576            .map(|s| {
3577                let rl = caches[s].recur[il].as_ref().unwrap();
3578                crate::GdnSeqVl {
3579                    kb16: e.addr_u8(&pres[s].kb16),
3580                    gcum: e.addr_f32(&pres[s].gcum),
3581                    beta: e.addr_f32(&sb[s].beta),
3582                    u: e.addr_f32(&pres[s].u),
3583                    wb16: e.addr_u8(&pres[s].wb16),
3584                    y: e.addr_u8(&pres[s].y16),
3585                    ssnap: e.addr_u8(&pres[s].ssnap16),
3586                    state_in: e.addr_f32(&rl.ssm_state),
3587                    state_out: e.addr_f32(&rl.ssm_state_alt),
3588                    q: e.addr_f32(&sb[s].q_l2),
3589                    p: e.addr_f32(&pres[s].p),
3590                    o: e.addr_f32(&pres[s].o),
3591                    k: e.addr_f32(&sb[s].k_l2),
3592                    v: e.addr_f32(&sb[s].v_g),
3593                    g: e.addr_f32(&sb[s].g_log),
3594                    a: e.addr_f32(&pres[s].a),
3595                    w: e.addr_f32(&pres[s].w),
3596                    t: ts[s] as i32,
3597                    nc: pres[s].nc as i32,
3598                }
3599            })
3600            .collect();
3601        e.gdn_prep_vl8(
3602            &prep_args,
3603            la.ssm_conv1d.float_data(),
3604            la.ssm_dt.float_data(),
3605            la.ssm_a.float_data(),
3606            conv_dim,
3607            d_conv,
3608            d_state,
3609            num_v,
3610            num_k,
3611            key_dim,
3612            hk,
3613            eps,
3614        )?;
3615        // mirror-fold (round 27): l2 v2 emits kb16 in-epilogue; K3's store emits wb16 —
3616        // both standalone mirror launches vanish on the default config.
3617        if !Engine::l2_v2_on(d_state) {
3618            e.gdn_mirror_vl8(&args, num_v, 0, hk)?;
3619        }
3620        // task #22: wgmma-fused vl twins — qb16 mirrors + GdnWVl8 side-struct.
3621        let wq8: Option<crate::GdnWVl8> = if e.gdn_wgmma_on(c) {
3622            // qb16 emitted by the vl l2 mirror-fold when l2-v2 serves; bulk cvt otherwise
3623            if !Engine::l2_v2_on(d_state) {
3624                for s in 0..b {
3625                    e.f32_to_bf16_into(&sb[s].q_l2, &mut pres[s].qb16, d_state * hk * ts[s])?;
3626                }
3627            }
3628            let mut wa = [crate::GdnWVl::default(); 8];
3629            for s in 0..b {
3630                wa[s] = crate::GdnWVl {
3631                    qb16: e.addr_u8(&pres[s].qb16),
3632                    pb16: e.addr_u8(&pres[s].pb16),
3633                };
3634            }
3635            Some(crate::GdnWVl8(wa))
3636        } else {
3637            None
3638        };
3639        e.gdn_chunk_k123_vl8(&args, num_v, hk, wq8.as_ref())?;
3640        e.gdn_chunk_vl8(&args, num_v, scale, hk, wq8.as_ref())?;
3641        if f16o {
3642            e.gdn_tail_vl8(&prep_args, la.ssm_norm.float_data(), d_state, num_v, eps)?;
3643        }
3644        // per-seq state swap (+ non-f16out tail fallback)
3645        let mut out = Vec::with_capacity(b);
3646        for (s, bufs) in sb.into_iter().enumerate() {
3647            let rl = caches[s].recur[il].as_mut().unwrap();
3648            std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
3649            let (o, t) = (offs[s], ts[s]);
3650            let SeqBufs { mut gn, gn16, .. } = bufs;
3651            if f16o {
3652                out.push((gn, Some(gn16)));
3653            } else {
3654                let z_v = g4[1].slice(o * value_dim..(o + t) * value_dim);
3655                e.gated_rmsnorm_zv(
3656                    &pres[s].o,
3657                    la.ssm_norm.float_data(),
3658                    &z_v,
3659                    &mut gn,
3660                    d_state,
3661                    num_v * t,
3662                    eps,
3663                )?;
3664                out.push((gn, None));
3665            }
3666        }
3667        Ok(out)
3668    }
3669
3670    /// task #16: view-consuming GDN prime core — the batched prime hands row-offset
3671    /// views of the CONCAT projection outputs directly (no per-seq split copies).
3672    /// Same kernels, same values, byte-identical to the Vec shim above.
3673    #[allow(clippy::too_many_arguments)]
3674    fn linear_attn_prime_core_pad_view(
3675        &self,
3676        e: &Engine,
3677        la: &LinearAttnLayer,
3678        qkv_mixed: &cudarc::driver::CudaView<f32>,
3679        z: &cudarc::driver::CudaView<f32>,
3680        beta_raw: &cudarc::driver::CudaView<f32>,
3681        alpha: &cudarc::driver::CudaView<f32>,
3682        t: usize,
3683        cache: &mut Cache,
3684        il: usize,
3685        pad_len: Option<&CudaSlice<i32>>,
3686    ) -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
3687        let cfg = &self.cfg;
3688        let ssm = cfg.ssm.as_ref().unwrap();
3689        let d_state = ssm.state_size as usize; // 128
3690        let num_v = ssm.time_step_rank as usize; // 32
3691        let eps = cfg.rms_eps;
3692        let scale = 1.0 / (d_state as f32).sqrt();
3693
3694        let prep =
3695            self.linear_attn_gdn_prep(e, la, qkv_mixed, beta_raw, alpha, t, cache, il, pad_len)?;
3696
3697        // ONE gdn_scan over T from the cache's CURRENT state (zero at fresh prime); the final
3698        // state lands in the spare buffer and ping-pongs back (stable resident pointers, the
3699        // decode-determinism discipline from linear_attn_decode_inner). A4: `gdn_scan_prefill`
3700        // dispatches the chunked WY form under MEMRA_GDN_CHUNKED (prefill-only seam; decode +
3701        // verify keep the sequential kernel).
3702        let mut o = e.uninit(d_state * num_v * t)?;
3703        let rl = cache.recur[il].as_mut().unwrap();
3704        {
3705            let crate::cache::RecurLayer {
3706                ssm_state,
3707                ssm_state_alt,
3708                ..
3709            } = rl;
3710            e.gdn_scan_prefill(
3711                &prep.q_l2,
3712                &prep.k_l2,
3713                &prep.v_g,
3714                &prep.g_log,
3715                &prep.beta,
3716                prep.kb16.as_ref(),
3717                prep.qb16.as_ref(),
3718                ssm_state,
3719                ssm_state_alt,
3720                &mut o,
3721                num_v,
3722                t,
3723                scale,
3724                prep.hk,
3725            )?;
3726        }
3727        std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
3728
3729        // gated RMSNorm + out projection (prefill dispatch). task #17: the f16out twin also
3730        // emits the ssm_out GEMM's fp16 operand in-epilogue (kills the standalone convert).
3731        let mut gn = e.uninit(d_state * num_v * t)?;
3732        let gn16 = if Self::f16out_on(e, t) {
3733            let mut g16 = e.alloc_u8_uninit(d_state * num_v * t * 2)?;
3734            e.gated_rmsnorm_f16out_zv(
3735                &o,
3736                la.ssm_norm.float_data(),
3737                z,
3738                &mut gn,
3739                &mut g16,
3740                d_state,
3741                num_v * t,
3742                eps,
3743            )?;
3744            Some(g16)
3745        } else {
3746            e.gated_rmsnorm_zv(
3747                &o,
3748                la.ssm_norm.float_data(),
3749                z,
3750                &mut gn,
3751                d_state,
3752                num_v * t,
3753                eps,
3754            )?;
3755            None
3756        };
3757        Ok((gn, gn16))
3758    }
3759
3760    /// Task #15 core-split wrapper: composes inner + ssm_out (== the old core).
3761    #[allow(clippy::too_many_arguments)]
3762    fn linear_attn_prime_core_pad(
3763        &self,
3764        e: &Engine,
3765        la: &LinearAttnLayer,
3766        g4: Vec<CudaSlice<f32>>,
3767        t: usize,
3768        cache: &mut Cache,
3769        il: usize,
3770        pad_len: Option<&CudaSlice<i32>>,
3771    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3772        let (gn, gn16) = self.linear_attn_prime_core_pad_inner(e, la, g4, t, cache, il, pad_len)?;
3773        if let Some(xh) = &gn16 {
3774            if let Some(y) = e.try_f16_gemm_pre(&la.ssm_out, xh, t)? {
3775                return Ok(y);
3776            }
3777        }
3778        Ok(e.matmul(&la.ssm_out, &gn, t)?)
3779    }
3780
3781    /// Full-attention mixer with QK-norm, partial RoPE, sigmoid output gate (qwen35 :257-336).
3782    ///
3783    /// `il` = layer index: step35 needs it (per-layer n_head / rope width / window / gate) and
3784    /// routes to its own mixer. Every other arch ignores it (uniform geometry).
3785    pub fn full_attn(
3786        &self,
3787        e: &Engine,
3788        fa: &FullAttnLayer,
3789        h: &CudaSlice<f32>,
3790        pos_d: &CudaSlice<i32>,
3791        t: usize,
3792        il: usize,
3793    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3794        if self.cfg.step35.is_some() {
3795            return self.step35_attn(e, fa, h, pos_d, t, il);
3796        }
3797        let cfg = &self.cfg;
3798        let _n_embd = cfg.n_embd as usize;
3799        let geometry = cfg.full_attention_geometry_at(il as u32);
3800        let n_head = geometry.n_head as usize;
3801        let n_head_kv = geometry.n_head_kv as usize;
3802        let head_dim = geometry.head_dim_k as usize;
3803        let eps = cfg.rms_eps;
3804        let scale = geometry.attention_scale();
3805
3806        // qwen35: wq output = head_dim*2*n_head (fused [q|gate] per head). M3/Hy3: NO output
3807        // gate — wq out = n_head*head_dim, no split (see prime-path note).
3808        let gated = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
3809        // grouped: one f16 activation convert feeds q/k/v (matmul_group)
3810        let mut g3 = e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?;
3811        let v = g3.pop().unwrap();
3812        let mut k = g3.pop().unwrap();
3813        let qf = g3.pop().unwrap();
3814        let (mut q, gate) = if gated {
3815            let mut q = e.uninit(t * n_head * head_dim)?;
3816            let mut gate = e.uninit(t * n_head * head_dim)?;
3817            e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, t)?;
3818            (q, Some(gate))
3819        } else {
3820            (qf, None)
3821        };
3822
3823        // QK-norm (per head_dim row), then partial RoPE.
3824        let mut qn = e.uninit(t * n_head * head_dim)?;
3825        e.rms_norm(
3826            &q,
3827            fa.q_norm.float_data(),
3828            &mut qn,
3829            head_dim,
3830            n_head * t,
3831            eps,
3832        )?;
3833        q = qn;
3834        let mut kn = e.uninit(t * n_head_kv * head_dim)?;
3835        e.rms_norm(
3836            &k,
3837            fa.k_norm.float_data(),
3838            &mut kn,
3839            head_dim,
3840            n_head_kv * t,
3841            eps,
3842        )?;
3843        k = kn;
3844        let rope_dims = geometry.n_rot as usize;
3845        e.rope_neox(
3846            &mut q,
3847            pos_d,
3848            head_dim,
3849            rope_dims,
3850            n_head,
3851            t,
3852            geometry.rope_base,
3853            1.0,
3854        )?;
3855        e.rope_neox(
3856            &mut k,
3857            pos_d,
3858            head_dim,
3859            rope_dims,
3860            n_head_kv,
3861            t,
3862            geometry.rope_base,
3863            1.0,
3864        )?;
3865
3866        // SDPA
3867        let mut attn = e.uninit(t * n_head * head_dim)?;
3868        // hand-written FlashAttention prefill (head_dim 256/128 stamped twins). MEMRA_NOFA
3869        // falls back to naive sdpa.
3870        if std::env::var("MEMRA_NOFA").is_ok() || !(head_dim == 256 || head_dim == 128) {
3871            // head_dim gate: see prime-path note (fa_prefill is stamped at 256 and 128 only).
3872            e.sdpa_naive(
3873                &q, &k, &v, &mut attn, head_dim, n_head, n_head_kv, t, t, scale, true,
3874            )?;
3875        } else {
3876            e.fa_prefill(
3877                &q, &k, &v, &mut attn, head_dim, n_head, n_head_kv, t, t, scale, true,
3878            )?;
3879        }
3880
3881        // output gate: attn * sigmoid(gate) — qwen35 only (M3 has no gate).
3882        let attn_g = match &gate {
3883            Some(gate) => {
3884                let mut gsig = e.uninit(t * n_head * head_dim)?;
3885                e.sigmoid(gate, &mut gsig, t * n_head * head_dim)?;
3886                let mut ag = e.uninit(t * n_head * head_dim)?;
3887                e.mul(&attn, &gsig, &mut ag, t * n_head * head_dim)?;
3888                ag
3889            }
3890            None => attn,
3891        };
3892
3893        // o projection
3894        let o = e.matmul(&fa.wo, &attn_g, t)?;
3895        Ok(o)
3896    }
3897
3898    /// Linear-attention (Gated DeltaNet) mixer (qwen35 :338-470).
3899    pub fn linear_attn(
3900        &self,
3901        e: &Engine,
3902        la: &LinearAttnLayer,
3903        h: &CudaSlice<f32>,
3904        t: usize,
3905    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3906        let cfg = &self.cfg;
3907        let _n_embd = cfg.n_embd as usize;
3908        let ssm = cfg.ssm.as_ref().unwrap();
3909        let d_state = ssm.state_size as usize; // 128
3910        let num_k = ssm.group_count as usize; // 16
3911        let num_v = ssm.time_step_rank as usize; // 32
3912        let d_conv = ssm.conv_kernel as usize; // 4
3913        let head_k = d_state;
3914        let head_v = d_state;
3915        let key_dim = head_k * num_k; // 2048
3916        let value_dim = head_v * num_v; // 4096
3917        let conv_dim = key_dim * 2 + value_dim; // 8192
3918        let eps = cfg.rms_eps;
3919        let scale = 1.0 / (d_state as f32).sqrt();
3920
3921        // projections
3922        // grouped: one f16 activation convert feeds all four projections (matmul_group)
3923        let mut g4 = e.matmul_group(
3924            &[&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha],
3925            h,
3926            t,
3927        )?;
3928        let alpha = g4.pop().unwrap(); // [T, num_v]
3929        let beta_raw = g4.pop().unwrap(); // [T, num_v]
3930        let z = g4.pop().unwrap(); // [T, value_dim]
3931        let qkv_mixed = g4.pop().unwrap(); // [T, conv_dim] token-major
3932
3933        // conv + GDN repack, FUSED (2026-07-03): ssm_conv1d_gdn reads qkv_mixed [T, conv_dim]
3934        // token-major DIRECTLY (causal window rows t-pad..t, rows<0 = zero prefill state), applies
3935        // the 8-tap conv + SiLU, and scatters straight into the GDN [d_state, num_v, T] q/k/v
3936        // layout with the modulo head-repeat. Replaces transpose + zeros + conv_left_pad +
3937        // ssm_conv1d + qkv_to_gdn_repack (5 launches, conv_in/conv_out scratch + a 16MB@T=512
3938        // round-trip). BIT-IDENTICAL accumulation and scatter mapping.
3939        let _ = (head_k, head_v);
3940        let mut q_g = e.uninit(d_state * num_v * t)?;
3941        let mut k_g = e.uninit(d_state * num_v * t)?;
3942        let mut v_g = e.uninit(d_state * num_v * t)?;
3943        e.ssm_conv1d_gdn(
3944            &qkv_mixed,
3945            la.ssm_conv1d.float_data(),
3946            &mut q_g,
3947            &mut k_g,
3948            &mut v_g,
3949            conv_dim,
3950            t,
3951            d_conv,
3952            d_state,
3953            num_v,
3954            num_k,
3955            key_dim,
3956        )?;
3957        // L2-norm q,k per (head_dim) row — rows are contiguous d_state in q_g.
3958        let mut q_l2 = e.uninit(d_state * num_v * t)?;
3959        e.l2_norm(&q_g, &mut q_l2, d_state, num_v * t, eps)?;
3960        let mut k_l2 = e.uninit(d_state * num_v * t)?;
3961        e.l2_norm(&k_g, &mut k_l2, d_state, num_v * t, eps)?;
3962        let v_gd = v_g;
3963
3964        // beta = sigmoid(beta_raw) ; g_log = a * softplus(alpha + dt). Both need [num_v, T] layout
3965        // (g[t*num_v + h]). beta_raw/alpha are [T, num_v] token-major == that layout already.
3966        let mut beta = e.uninit(t * num_v)?;
3967        e.sigmoid(&beta_raw, &mut beta, t * num_v)?;
3968        // gdn_glog expects alpha [H,T] with alpha[t*H+h] and dt_bias/a [H] — matches token-major [T,num_v].
3969        let mut g_log = e.uninit(t * num_v)?;
3970        e.gdn_glog(
3971            &alpha,
3972            la.ssm_dt.float_data(),
3973            la.ssm_a.float_data(),
3974            &mut g_log,
3975            num_v,
3976            t,
3977        )?;
3978
3979        // GDN scan (A4: gdn_scan_prefill dispatches chunked WY under MEMRA_GDN_CHUNKED)
3980        let state_in = e.zeros(d_state * d_state * num_v)?; // zero state (prefill)
3981        let mut state_out = e.zeros(d_state * d_state * num_v)?;
3982        let mut o = e.uninit(d_state * num_v * t)?;
3983        e.gdn_scan_prefill(
3984            &q_l2,
3985            &k_l2,
3986            &v_gd,
3987            &g_log,
3988            &beta,
3989            None,
3990            None,
3991            &state_in,
3992            &mut state_out,
3993            &mut o,
3994            num_v,
3995            t,
3996            scale,
3997            num_v,
3998        )?;
3999
4000        // gated RMSNorm: dst = RMSNorm(o, ssm_norm[head_v]) * silu(z). o is [d_state, num_v, T];
4001        // rows of head_v=d_state, nrows = num_v*T. z must match row layout: z is [T, value_dim] token-major
4002        // = [T, num_v*head_v]; per (t, vh) the head_v slice is contiguous -> rows align as (t*num_v+vh).
4003        // o rows are (t*num_v+vh) too. Good.
4004        let mut gn = e.uninit(d_state * num_v * t)?;
4005        e.gated_rmsnorm(
4006            &o,
4007            la.ssm_norm.float_data(),
4008            &z,
4009            &mut gn,
4010            d_state,
4011            num_v * t,
4012            eps,
4013        )?;
4014
4015        // ssm_out projection: gn is [d_state, num_v, T] = [value_dim, T] viewed token-major as [T, value_dim]?
4016        // gn layout: (t*num_v+vh)*d_state + i  == token t, then (vh,i) = channel vh*d_state+i. That's
4017        // token-major [T, value_dim]. linear wants [T, in=value_dim]. Good.
4018        let out = e.matmul(&la.ssm_out, &gn, t)?;
4019        Ok(out)
4020    }
4021}
4022
4023impl HybridModel {
4024    /// MoE FFN (EDGE-1). z: [T, n_embd] (already post-attention-normed). Returns moe_out [T, n_embd].
4025    /// Node-for-node vs llama.cpp build_moe_ffn + qwen35moe::build_layer_ffn.
4026    ///
4027    /// `il` is the trunk layer index — the residency-cache key prefix (a gate-expert of layer 3 is a
4028    /// different 860160-byte block than the same expert of layer 7).
4029    ///
4030    /// Routing: host softmax+sort (default) OR the fused router kernel (MEMRA_FUSED_ROUTER).
4031    /// Dispatch: stage-every-token into 3 scratch slots (default) OR the SLRU residency cache
4032    /// (MEMRA_MOE_CACHE). The cache-HIT weight path is bit-identical to stage-every-token (§B.3).
4033    /// Convenience wrapper used by the hybrid trunk/MTP loops: pulls dims + max-block from `self`.
4034    pub fn moe_ffn_il(
4035        &self,
4036        e: &Engine,
4037        m: &MoeWeights,
4038        z: &CudaSlice<f32>,
4039        t: usize,
4040        il: u16,
4041    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4042        Self::moe_ffn_inner(e, m, z, None, t, &self.cfg, il, self.max_moe_block(), false)
4043    }
4044
4045    /// Prefill twin: Step35 promotes expert-grouped dispatch by default while decode/spec callers
4046    /// keep `moe_ffn_il` and therefore retain their existing dispatch class.
4047    pub fn moe_ffn_il_prefill(
4048        &self,
4049        e: &Engine,
4050        m: &MoeWeights,
4051        z: &CudaSlice<f32>,
4052        t: usize,
4053        il: u16,
4054    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4055        Self::moe_ffn_inner(e, m, z, None, t, &self.cfg, il, self.max_moe_block(), true)
4056    }
4057
4058    /// Decode-path twin with a PRE-QUANTIZED z (from add_rms_norm_zq8): threads (zq, zd) into the
4059    /// t=1 dev arm so the per-layer standalone quantize_q8_1 launch folds away. Identical bytes
4060    /// (the fused kernel reproduces quantize_q8_1 exactly); every other path ignores the pair.
4061    pub fn moe_ffn_il_zq8(
4062        &self,
4063        e: &Engine,
4064        m: &MoeWeights,
4065        z: &CudaSlice<f32>,
4066        zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
4067        t: usize,
4068        il: u16,
4069    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4070        Self::moe_ffn_inner(e, m, z, zq8, t, &self.cfg, il, self.max_moe_block(), false)
4071    }
4072
4073    /// MoE FFN (EDGE-1), source-/model-agnostic. z: [T, n_embd] (already post-attention-normed).
4074    /// Returns moe_out [T, n_embd]. Node-for-node vs llama.cpp build_moe_ffn. Shared by the hybrid
4075    /// (qwen35moe, shared expert present) and the dense-attention MoE (OLMoE, no shared expert) paths;
4076    /// `cfg.moe` supplies the dims and the optional shexp fields decide whether step 3 runs.
4077    ///
4078    /// `il` is the layer index — the residency-cache key prefix. `max_block` is the global max expert
4079    /// stride (fixed cache-slot size); pass `self.max_moe_block()`.
4080    pub(crate) fn moe_ffn(
4081        e: &Engine,
4082        m: &MoeWeights,
4083        z: &CudaSlice<f32>,
4084        t: usize,
4085        cfg: &ModelConfig,
4086        il: u16,
4087        max_block: usize,
4088    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4089        Self::moe_ffn_inner(e, m, z, None, t, cfg, il, max_block, false)
4090    }
4091
4092    #[allow(clippy::too_many_arguments)]
4093    pub(crate) fn moe_ffn_inner(
4094        e: &Engine,
4095        m: &MoeWeights,
4096        z: &CudaSlice<f32>,
4097        zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
4098        t: usize,
4099        cfg: &ModelConfig,
4100        il: u16,
4101        max_block: usize,
4102        prefill: bool,
4103    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4104        let worker_io = crate::spill_pread::worker_enabled();
4105        let epoch_lfu = std::env::var_os("MEMRA_MOE_LFU_DECAY").is_some();
4106        if Engine::moe_cache_enabled() && (worker_io || epoch_lfu) {
4107            e.with_moe_cache(max_block, |cache, _| {
4108                cache.begin_forward_epoch(il, t);
4109                if worker_io {
4110                    cache.begin_worker_scope();
4111                }
4112                Ok(())
4113            })?;
4114        }
4115        if Self::sigmoid_resident_dev_eligible(e, m, cfg) {
4116            let moe = cfg.moe.as_ref().unwrap();
4117            let n_expert = moe.expert_count as usize;
4118            let n_used = moe.expert_used_count as usize;
4119            let sigmoid = cfg.sigmoid_router().unwrap();
4120            let logits = Self::moe_router_logits(e, m, z, t, cfg)?;
4121            Self::trace_sigmoid_router_logits(e, il, t, n_expert, n_used, &logits, m, sigmoid)?;
4122            return Self::moe_ffn_sigmoid_dev(e, m, z, zq8, &logits, t, cfg, il, sigmoid);
4123        }
4124        // Expert-grouped dispatch for prefill (T>1). MEMRA_MOE_GROUPED explicitly opts the
4125        // current caller into this research arm; the naked default stays on the established path.
4126        if t > 1 && moe_grouped_enabled(cfg, prefill) {
4127            let grouped_out = Self::moe_ffn_grouped(e, m, z, t, cfg, il, max_block)?;
4128            // MEMRA_MOE_GATE: byte-identity comparison vs sequential path.
4129            // The grouped q8 path uses the same row-wise quantize/dot/activation/FMA programs as
4130            // sequential dispatch; mixed or q8-disabled layouts fall back to the matching f32
4131            // path. A mismatch is therefore a correctness failure, not an accepted numeric class.
4132            if std::env::var("MEMRA_MOE_GATE").is_ok() {
4133                let seq_out = Self::moe_ffn_sequential(e, m, z, t, cfg, il, max_block)?;
4134                let g_host = e.dtoh(&grouped_out)?;
4135                let s_host = e.dtoh(&seq_out)?;
4136                let g_bytes: &[u8] = unsafe {
4137                    std::slice::from_raw_parts(g_host.as_ptr() as *const u8, g_host.len() * 4)
4138                };
4139                let s_bytes: &[u8] = unsafe {
4140                    std::slice::from_raw_parts(s_host.as_ptr() as *const u8, s_host.len() * 4)
4141                };
4142                if g_bytes == s_bytes {
4143                    println!("moe-gate il={il} t={t} BYTE-IDENTICAL");
4144                } else {
4145                    let diffs = g_host
4146                        .iter()
4147                        .zip(s_host.iter())
4148                        .enumerate()
4149                        .filter(|(_, (a, b))| a != b)
4150                        .count();
4151                    let maxdiff = g_host
4152                        .iter()
4153                        .zip(s_host.iter())
4154                        .map(|(a, b)| (a - b).abs())
4155                        .fold(0.0f32, f32::max);
4156                    panic!(
4157                        "moe-gate il={il} t={t} MISMATCH: {diffs}/{} elems differ, maxdiff={maxdiff:.6e}",
4158                        g_host.len()
4159                    );
4160                }
4161            }
4162            return Ok(grouped_out);
4163        }
4164        Self::moe_ffn_sequential_zq8(e, m, z, zq8, t, cfg, il, max_block)
4165    }
4166
4167    fn sigmoid_resident_dev_eligible(e: &Engine, m: &MoeWeights, cfg: &ModelConfig) -> bool {
4168        let Some(moe) = cfg.moe.as_ref() else {
4169            return false;
4170        };
4171        // Cached once per process: this predicate runs per MoE layer per decode step, and five
4172        // uncached env reads were ~300+ environ scans/step on the exact path this arm de-hosts.
4173        static OBSERVATION_MODE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4174        let observation_mode = *OBSERVATION_MODE.get_or_init(|| {
4175            std::env::var("MEMRA_MOE_STATS").is_ok()
4176                || std::env::var("MEMRA_MOE_TRACE").is_ok()
4177                || std::env::var("MEMRA_MOE_WEIGHT_TRACE").is_ok()
4178                || std::env::var("MEMRA_MOE_INPUT_TRACE_DIR").is_ok()
4179                || std::env::var("MEMRA_MOE_GATE").is_ok()
4180        });
4181        cfg.step35.is_some()
4182            && sigmoid_router_enabled()
4183            && moe_dev_enabled()
4184            && moe_slab_enabled()
4185            && !observation_mode
4186            && moe.expert_used_count <= 8
4187            && m.has_uniform_expert_layout()
4188            && m.gate_exps.macros.is_none()
4189            && m.up_exps.macros.is_none()
4190            && m.down_exps.macros.is_none()
4191            && !m.has_macros
4192            && moe_q8_enabled()
4193            && q8_expert_supported(m.gate_exps.qtype)
4194            && q8_expert_supported(m.up_exps.qtype)
4195            && q8_expert_supported(m.down_exps.qtype)
4196            && m.dev_exps
4197                .as_ref()
4198                .is_some_and(|dev| dev.dev == e.ctx().ordinal())
4199    }
4200
4201    /// Sequential (per-token) MoE FFN -- the original path. Factored out for the gate comparison.
4202    pub(crate) fn moe_ffn_sequential(
4203        e: &Engine,
4204        m: &MoeWeights,
4205        z: &CudaSlice<f32>,
4206        t: usize,
4207        cfg: &ModelConfig,
4208        il: u16,
4209        max_block: usize,
4210    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4211        Self::moe_ffn_sequential_zq8(e, m, z, None, t, cfg, il, max_block)
4212    }
4213
4214    /// One router-logit selector shared by sequential and grouped dispatch. Real prefill must use
4215    /// the row-wise GEMV when exact routing is enabled: cuBLASLt's reduction changes with `m`,
4216    /// which makes expert selection depend on the caller's chunk or concat-batch shape.
4217    fn moe_router_logits(
4218        e: &Engine,
4219        m: &MoeWeights,
4220        z: &CudaSlice<f32>,
4221        t: usize,
4222        cfg: &ModelConfig,
4223    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4224        if t < PRIME_MIN_T {
4225            // Decode and speculative verify use one fixed per-row reduction program.
4226            if crate::router_kernel_on() {
4227                e.router_gemv(
4228                    m.gate_inp.float_data(),
4229                    z,
4230                    cfg.n_embd as usize,
4231                    m.gate_exps.n_expert,
4232                    t,
4233                )
4234            } else {
4235                e.matmul_decode_exact(&m.gate_inp, z, t)
4236            }
4237        } else if crate::router_prefill_exact_on() && crate::router_kernel_on() {
4238            e.router_gemv(
4239                m.gate_inp.float_data(),
4240                z,
4241                cfg.n_embd as usize,
4242                m.gate_exps.n_expert,
4243                t,
4244            )
4245        } else {
4246            e.matmul(&m.gate_inp, z, t)
4247        }
4248    }
4249
4250    /// Append the host-visible router selection for one layer/forward when calibration tracing is
4251    /// enabled. Both sequential and expert-grouped prefill must call this after routing so the
4252    /// trace is independent of the dispatch optimization selected for the forward.
4253    fn trace_moe_routes(
4254        il: u16,
4255        t: usize,
4256        sel_all: &[u32],
4257        weights: &[f32],
4258    ) -> Result<(), Box<dyn std::error::Error>> {
4259        use std::io::Write as _;
4260        if let Ok(path) = std::env::var("MEMRA_MOE_TRACE") {
4261            let mut f = std::fs::OpenOptions::new()
4262                .create(true)
4263                .append(true)
4264                .open(path)?;
4265            let ids: Vec<String> = sel_all.iter().map(|s| s.to_string()).collect();
4266            writeln!(f, "{} {} {}", il, t, ids.join(","))?;
4267        }
4268        if let Ok(path) = std::env::var("MEMRA_MOE_WEIGHT_TRACE") {
4269            let mut f = std::fs::OpenOptions::new()
4270                .create(true)
4271                .append(true)
4272                .open(path)?;
4273            let pairs: Vec<String> = sel_all
4274                .iter()
4275                .zip(weights)
4276                .map(|(expert, weight)| format!("{expert}:{weight:.9}"))
4277                .collect();
4278            writeln!(f, "{} {} {}", il, t, pairs.join(","))?;
4279        }
4280        Ok(())
4281    }
4282
4283    #[allow(clippy::too_many_arguments)]
4284    fn trace_sigmoid_router_logits(
4285        e: &Engine,
4286        il: u16,
4287        t: usize,
4288        n_expert: usize,
4289        n_used: usize,
4290        logits: &CudaSlice<f32>,
4291        m: &MoeWeights,
4292        (scaling_factor, route_norm): (f32, bool),
4293    ) -> Result<(), Box<dyn std::error::Error>> {
4294        if !crate::sigrouter_contract::served_logit_trace_enabled() || t != 1 {
4295            return Ok(());
4296        }
4297        let logits = e.dtoh(logits)?;
4298        let active: Vec<u8> = m
4299            .active_experts
4300            .as_ref()
4301            .map(|mask| mask.iter().map(|&enabled| u8::from(enabled)).collect())
4302            .unwrap_or_else(|| vec![1; n_expert]);
4303        let bias = m.exp_probs_b.clone().unwrap_or_else(|| vec![0.0; n_expert]);
4304        crate::sigrouter_contract::capture_served_logits(
4305            il as u32,
4306            t,
4307            n_expert,
4308            n_used,
4309            scaling_factor,
4310            route_norm,
4311            &active,
4312            &bias,
4313            &logits,
4314        )?;
4315        Ok(())
4316    }
4317
4318    /// Append the f32 input to one MoE layer for offline, layerwise calibration. This diagnostic
4319    /// intentionally performs a DtoH copy and is therefore disabled unless an explicit fresh trace
4320    /// directory is supplied. Each layer owns one payload file; index.jsonl records byte offsets so
4321    /// a validator can prove request/layer coverage before the trace is used for pruning or healing.
4322    fn trace_moe_input(
4323        e: &Engine,
4324        il: u16,
4325        t: usize,
4326        n_embd: usize,
4327        z: &CudaSlice<f32>,
4328    ) -> Result<(), Box<dyn std::error::Error>> {
4329        use std::io::Write as _;
4330        let Ok(dir) = std::env::var("MEMRA_MOE_INPUT_TRACE_DIR") else {
4331            return Ok(());
4332        };
4333        let host = e.dtoh(z)?;
4334        if host.len() != t * n_embd {
4335            return Err(format!(
4336                "MoE input trace shape mismatch at layer {il}: got {} values, expected {}x{}",
4337                host.len(),
4338                t,
4339                n_embd
4340            )
4341            .into());
4342        }
4343        let bytes = unsafe {
4344            std::slice::from_raw_parts(
4345                host.as_ptr().cast::<u8>(),
4346                host.len() * std::mem::size_of::<f32>(),
4347            )
4348        };
4349        let state = MOE_INPUT_TRACE_WRITER.get_or_init(|| std::sync::Mutex::new(None));
4350        let mut state = state
4351            .lock()
4352            .map_err(|_| "MoE input trace writer lock is poisoned")?;
4353        if state.is_none() {
4354            let dir = std::path::PathBuf::from(&dir);
4355            std::fs::create_dir_all(&dir)?;
4356            let index = std::fs::OpenOptions::new()
4357                .create(true)
4358                .append(true)
4359                .open(dir.join("index.jsonl"))?;
4360            *state = Some(MoeInputTraceWriter {
4361                dir,
4362                index,
4363                payloads: std::collections::HashMap::new(),
4364            });
4365        }
4366        let writer = state.as_mut().unwrap();
4367        if writer.dir != std::path::Path::new(&dir) {
4368            return Err("MEMRA_MOE_INPUT_TRACE_DIR changed after capture started".into());
4369        }
4370        let file_name = format!("layer-{il:03}.f32");
4371        if !writer.payloads.contains_key(&il) {
4372            let payload = std::fs::OpenOptions::new()
4373                .create(true)
4374                .append(true)
4375                .open(writer.dir.join(&file_name))?;
4376            let offset = payload.metadata()?.len();
4377            writer.payloads.insert(il, (payload, offset));
4378        }
4379        let (payload, offset) = writer.payloads.get_mut(&il).unwrap();
4380        let row_offset = *offset;
4381        payload.write_all(bytes)?;
4382        *offset += bytes.len() as u64;
4383        writeln!(
4384            writer.index,
4385            "{{\"format\":\"memra-moe-input-trace-v1\",\"layer\":{il},\"tokens\":{t},\
4386             \"hidden_size\":{n_embd},\"file\":\"{file_name}\",\"offset\":{row_offset},\
4387             \"payload_bytes\":{}}}",
4388            bytes.len()
4389        )?;
4390        Ok(())
4391    }
4392
4393    #[allow(clippy::too_many_arguments)]
4394    pub(crate) fn moe_ffn_sequential_zq8(
4395        e: &Engine,
4396        m: &MoeWeights,
4397        z: &CudaSlice<f32>,
4398        zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
4399        t: usize,
4400        cfg: &ModelConfig,
4401        il: u16,
4402        max_block: usize,
4403    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4404        use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
4405        let moe = cfg.moe.as_ref().unwrap();
4406        let n_embd = cfg.n_embd as usize; // 2048 (gate/up in_f, down out_f)
4407        let n_expert = moe.expert_count as usize; // 256
4408        let n_used = moe.expert_used_count as usize; // 8
4409        let n_ff_exp = moe.expert_ff_length as usize; // 512 (gate/up out_f, down in_f)
4410
4411        // verify the HostExps dims match cfg (catches a wrong-file / transpose mixup)
4412        debug_assert_eq!(m.gate_exps.in_f, n_embd);
4413        debug_assert_eq!(m.gate_exps.out_f, n_ff_exp);
4414        debug_assert_eq!(m.down_exps.in_f, n_ff_exp); // down is TRANSPOSED: in=512
4415        debug_assert_eq!(m.down_exps.out_f, n_embd); //                     out=2048
4416        debug_assert_eq!(m.gate_exps.n_expert, n_expert);
4417
4418        // step35 PER-LAYER SwiGLU clamp (None on every other arch and on every unclamped layer).
4419        // Routed experts and the shared expert read SEPARATE arrays — never share one value.
4420        let lim_exp = cfg.clamp_exp_at(il as u32);
4421        let lim_shexp = cfg.clamp_shexp_at(il as u32);
4422        let use_cache = Engine::moe_cache_enabled();
4423        let uniform_experts = m.has_uniform_expert_layout();
4424        let moe_q8 = uniform_experts
4425            && moe_q8_enabled()
4426            && q8_expert_supported(m.gate_exps.qtype)
4427            && q8_expert_supported(m.up_exps.qtype)
4428            && q8_expert_supported(m.down_exps.qtype);
4429        // Experimental secondary backend: complete experts already resident in the SLRU stay on
4430        // CUDA; any expert missing one or more projections runs from the original host GGUF bytes
4431        // through llama.cpp CPU quant dots. Small-t only keeps decode and speculative verification
4432        // in the same numeric/dispatch class; real prefill remains on the established GPU path and
4433        // seeds the residency cache. An explicit library path is the build/runtime gate, so naked
4434        // commands and CI have no llama.cpp or OpenMP dependency.
4435        let cpu_expert_requested = crate::cpu_experts::configured();
4436        if cpu_expert_requested && (cfg.hy3.is_none() || cfg.m3.is_some()) {
4437            return Err(std::io::Error::other(
4438                "MEMRA_CPU_EXPERT_LIB is experimental and currently gated to Hy3",
4439            )
4440            .into());
4441        }
4442        let cpu_hybrid = cpu_expert_requested && t < PRIME_MIN_T && m.dev_exps.is_none();
4443        // A dynamic SLRU changes which complete experts run on CUDA versus llama.cpp CPU dots.
4444        // Those backends are each deterministic but are different numeric configurations, so a
4445        // later prefill eviction can change greedy output. Freeze after the first real prefill;
4446        // decode/spec reads the fixed resident set, while later prefill misses use transient GPU
4447        // staging below and cannot change backend assignment.
4448        let freeze_cpu_residency = cpu_expert_requested
4449            && std::env::var("MEMRA_CPU_EXPERT_FREEZE_CACHE").as_deref() == Ok("1");
4450        let caller_warms_before_freeze = std::env::var("MEMRA_CPU_EXPERT_FREEZE_WARMUP_TOKENS")
4451            .ok()
4452            .and_then(|value| value.parse::<usize>().ok())
4453            .is_some_and(|tokens| tokens > 0);
4454        if cpu_hybrid && freeze_cpu_residency && !caller_warms_before_freeze {
4455            e.freeze_moe_cache();
4456        }
4457        let cache_frozen = use_cache && e.moe_cache_frozen();
4458        let cache_dispatch = use_cache && (!cache_frozen || cpu_hybrid);
4459
4460        // 1. ROUTER: one selector is shared with expert-grouped prefill so changing dispatch
4461        // cannot change logits, selected expert ids, or routing weights.
4462        let logits = Self::moe_router_logits(e, m, z, t, cfg)?;
4463        if let Some(sig) = cfg.sigmoid_router() {
4464            Self::trace_sigmoid_router_logits(e, il, t, n_expert, n_used, &logits, m, sig)?;
4465        }
4466
4467        // LAUNCH-STRUCTURE STAGE 3 (2026-07-05, MEMRA_MOE_DEV default ON, =0 rollback): ZERO-DtoH
4468        // device-dispatch when this layer's expert blocks are ALL cache-resident. The fused
4469        // router's sel/w stay ON DEVICE; the expert weight pointers come from the per-layer
4470        // device table of fixed slot addresses; gate/up/silu + down/fma run as the same TWO
4471        // launches per token as gdec. Removes the per-layer router DtoH + stream sync — the
4472        // per-token host stall that dominated the 35B decode wall after stages 1+2.
4473        // BIT-IDENTITY: the router kernel is selection-exact vs the host oracle (kernel-check
4474        // tie gate) and the _dev matvec twins reproduce the gdec kernels' exact FP chains; the
4475        // only difference is where sel/w/pointers are READ from (device instead of params).
4476        // Residency: one-shot PREWARM force-admits the layer while free slots cover it
4477        // (MEMRA_MOE_PREWARM=0 -> organic residency, dev path fires when the SLRU fills).
4478        // Any non-resident layer falls through to host routing + the gdec/sequential path.
4479        // FITS-VRAM RESIDENT EXPERTS (2026-07-06): the layer's expert slabs are device-resident
4480        // (load-time decision) — fire the zero-DtoH dev path unconditionally with the prebuilt
4481        // pointer row. No cache, no dispatch, no residency check: the llama full-offload regime
4482        // (it measured 169.55 vs the cache path's 28.5 on the local 35B — the residency-gate
4483        // all-or-nothing fallback was the 6x). BIT-IDENTITY: same _dev kernels, same math; only
4484        // the pointer table's provenance differs (slab base+stride vs SLRU slot addresses).
4485        // MoE PREFILL PAIR-BATCH (2026-07-06, the 16x pp hole): t>1 on resident experts — ONE
4486        // launch per proj covers ALL (token,expert) pairs (grid.y=pair, warp-per-row), replacing
4487        // the per-expert loop (256 experts x 3-4 launches x tiny m_e). Scatter is slot-ordered
4488        // per token (the sequential-axpy bit-identity class). Requires q8-supported qtypes +
4489        // resident slabs. MEMRA_MOE_PAIRS=0 rollback.
4490        // t >= PRIME_MIN_T only (2026-07-06 exactness fix, verify-probe proof): spec VERIFY
4491        // batches (t = 2..K+2) previously rode these pairs kernels while T=1 decode rode the
4492        // dev_q8 loop — different FP chains, verify-T2 logit maxdiff 2.6e-1 vs eager -> greedy
4493        // flips at tight margins -> 35B real-prompt spec self-consistency FAIL (the 27B
4494        // "verify must be kernel-DISPATCH-identical to decode" lesson, MoE edition). Small-t
4495        // now rides the dev loop below (same kernels per token as decode); pairs serves real
4496        // prefill (t >= 16, where spec never verifies).
4497        // sigmoid-router archs must NOT enter the pairs/dev arms: those dispatch arms still route
4498        // via the fused SOFTMAX device router (moe_router_topk) — silently wrong experts (the
4499        // M3 gate-MISMATCH 74602-vs-92 lesson, 2026-07-07). Sigmoid routing has its own device
4500        // top-k kernel; MEMRA_SIG_ROUTER=0 selects the host sigmoid oracle instead.
4501        // Per-expert macro-scales (compressed-tensors NVFP4 ST class, e.g. unsloth 35B-A3B):
4502        // the device-dispatch kernels (pairs/dev) do NOT fold them — those checkpoints must
4503        // ride the macro-aware sequential/staged paths below or every expert output is off by
4504        // its global scale (~3e4x, measured garbage 2026-07-16). GGUF experts: macros None.
4505        let no_exp_macros = m.gate_exps.macros.is_none()
4506            && m.up_exps.macros.is_none()
4507            && m.down_exps.macros.is_none();
4508        // A clamped-SwiGLU layer (step35 43/44) also cannot ride pairs: moe_pairs_silu_mul's
4509        // epilogue is plain silu(gate)*up with no clamp form, and moe_ffn_pairs takes no `il`
4510        // so it cannot even see the per-layer limit.
4511        if cfg.sigmoid_router().is_none()
4512            && cfg.m3.is_none()
4513            && cfg.hy3.is_none()
4514            && !cfg.swiglu_clamped_at(il as u32)
4515            && no_exp_macros
4516            && t >= PRIME_MIN_T
4517            && m.dev_exps.is_some()
4518            && moe_q8_enabled()
4519            && q8_expert_supported(m.gate_exps.qtype)
4520            && q8_expert_supported(m.up_exps.qtype)
4521            && q8_expert_supported(m.down_exps.qtype)
4522            && std::env::var("MEMRA_MOE_PAIRS")
4523                .map(|v| v != "0")
4524                .unwrap_or(true)
4525            && std::env::var("MEMRA_MOE_STATS").is_err()
4526        {
4527            return Self::moe_ffn_pairs(e, m, z, &logits, t, cfg);
4528        }
4529
4530        // t < PRIME_MIN_T: moe_ffn_dev loops tokens serially (1 launch-pair per token) — the
4531        // decode path AND the spec-verify path (dispatch parity = exactness; see pairs gate
4532        // above). Serial launches are fine at t<=10 (K+2); real prefill never lands here.
4533        // moe_ffn_dev routes via the FUSED SOFTMAX device router (moe_router_topk). Sigmoid
4534        // routing has a separate device top-k kernel, but this dev arm cannot consume its sel/w
4535        // contract, so sigmoid-router arches must NOT enter it: with MOE_CACHE=1 M3 silently
4536        // routed softmax = wrong experts (gate MISMATCH 74602 vs 92, caught 2026-07-07).
4537        // MEMRA_SIG_ROUTER=0 selects the host sigmoid oracle for the supported sigmoid path.
4538        // macro-carrying experts are handled inside the dev path now (epilogue fold + w-scale);
4539        // pairs/gdec/csr keep their macro gates until their kernels grow the fold.
4540        // step35 (2026-08-06) is the third sigmoid-router arch and the deny above was written as
4541        // two arch names, not as the mechanism — so `dev_ok` let it into moe_ffn_dev's
4542        // moe_router_topk (softmax, no exp_probs_b bias, no expert_weights_scale) = silently
4543        // wrong experts, the same failure the pairs gate at :2254 already blocks by predicate.
4544        // Keyed off sigmoid_router() so arch #4 is denied by construction.
4545        // The clamped-SwiGLU layers must also fall through: every moe_ffn_dev kernel's fused
4546        // epilogue is plain silu(gate)*up (moe_gate_up_silu8_dev_*), which has no clamp form.
4547        let dev_ok = uniform_experts
4548            && cfg.sigmoid_router().is_none()
4549            && cfg.m3.is_none()
4550            && cfg.hy3.is_none()
4551            && !cfg.swiglu_clamped_at(il as u32);
4552        // Observation modes must route through the host-visible selection below. Otherwise a fully
4553        // resident layer returns through device dispatch before its trace/stats row is recorded,
4554        // silently biasing calibration toward only non-resident layers on large-VRAM machines.
4555        let observe_routes = std::env::var("MEMRA_MOE_STATS").is_ok()
4556            || std::env::var("MEMRA_MOE_TRACE").is_ok()
4557            || std::env::var("MEMRA_MOE_WEIGHT_TRACE").is_ok()
4558            || std::env::var("MEMRA_MOE_INPUT_TRACE_DIR").is_ok();
4559        if dev_ok
4560            && t < PRIME_MIN_T
4561            && m.dev_exps.is_some()
4562            && n_used <= 8
4563            && moe_dev_enabled()
4564            && !observe_routes
4565        {
4566            return Self::moe_ffn_dev(e, m, z, zq8, &logits, t, cfg, il, max_block);
4567        }
4568        if dev_ok && use_cache && n_used <= 8 && moe_dev_enabled() && !observe_routes {
4569            let row_ok = e.with_moe_cache(max_block, |c, eng| {
4570                if moe_prewarm_enabled() {
4571                    c.prewarm_layer(il, m, eng)?;
4572                }
4573                Ok(c.layer_dev_row(il, n_expert, eng)?.is_some())
4574            })?;
4575            if row_ok {
4576                return Self::moe_ffn_dev(e, m, z, zq8, &logits, t, cfg, il, max_block);
4577            }
4578        }
4579
4580        // Per-token (sel[8], w[8]) — sigmoid host oracle (M3/Hy3), else fused-router/softmax.
4581        let (sel_all, w_all, routed_cpu_input) = if let Some(sig) = cfg.sigmoid_router() {
4582            if cpu_hybrid {
4583                let (sel, w, input) = Self::moe_route_sigmoid_with_input(
4584                    e,
4585                    &logits,
4586                    z,
4587                    t,
4588                    n_expert,
4589                    n_used,
4590                    m.exp_probs_b.as_deref(),
4591                    sig,
4592                    m.active_experts.as_deref(),
4593                )?;
4594                (sel, w, Some(input))
4595            } else {
4596                let (sel, w) =
4597                    Self::moe_route_sigmoid_cfg(e, &logits, t, n_expert, n_used, m, sig)?;
4598                (sel, w, None)
4599            }
4600        } else {
4601            let (sel, w) =
4602                Self::moe_route_cfg(e, &logits, t, n_expert, n_used, m.active_experts.as_deref())?;
4603            (sel, w, None)
4604        };
4605        crate::moesd::record_host_routes(il, n_expert, n_used, &sel_all)?;
4606
4607        // MEMRA_MOE_TRACE=<path>: append one line per (layer, step) with the selected expert ids —
4608        // offline analysis derives the decode working set + step-to-step reuse (the go/no-go
4609        // measurement for resident-expert tiering; see rig5090.jsonl 2026-07-07 pinned-tier row).
4610        Self::trace_moe_routes(il, t, &sel_all, &w_all)?;
4611        Self::trace_moe_input(e, il, t, n_embd, z)?;
4612
4613        // Worker reads ordinarily overlap NVMe with the current expert but leave their H2D copies
4614        // demand-serialized on the compute stream. Hy3 decode already has a safe owner-thread
4615        // boundary here: sigmoid routing read `logits` back to the host, which synchronized all
4616        // earlier-layer compute. Submit the complete selected set first, then reserve only victims
4617        // outside that set and queue its H2D copies on the copy stream. Dispatch inserts an event
4618        // wait for each pending block, so later copies can overlap the earlier expert kernels while
4619        // the selected weights and numerical kernels remain unchanged. Restrict this experiment to
4620        // T=1; batched forwards can have token-local consumers still in flight between selections.
4621        // The CPU backend reads mmap-backed model bytes directly. Positioned-read worker buffers
4622        // are CUDA write-combined staging allocations and are intentionally not CPU-compute inputs;
4623        // do not spend NVMe bandwidth filling them for a layer whose misses go to CPU.
4624        let worker_disk_prefetch =
4625            cache_dispatch && crate::spill_pread::worker_enabled() && !cpu_hybrid;
4626        let promote_worker_h2d =
4627            t == 1 && worker_disk_prefetch && crate::spill_pread::copy_h2d_enabled();
4628        if promote_worker_h2d {
4629            let mut selected_blocks = Vec::with_capacity(n_used * 3);
4630            for &ex in sel_all.iter().take(n_used) {
4631                let ex = ex as u16;
4632                selected_blocks.extend([
4633                    BlockId::new(il, PROJ_GATE, ex),
4634                    BlockId::new(il, PROJ_UP, ex),
4635                    BlockId::new(il, PROJ_DOWN, ex),
4636                ]);
4637            }
4638            for &ex in sel_all.iter().take(n_used) {
4639                Self::moe_prefetch_disk_expert(e, il, ex as usize, m, max_block, &selected_blocks)?;
4640            }
4641            e.with_moe_cache(max_block, |cache, eng| {
4642                cache.promote_worker_reads_at_safe_boundary(
4643                    &selected_blocks,
4644                    &selected_blocks,
4645                    eng,
4646                )?;
4647                Ok(())
4648            })?;
4649        }
4650
4651        // MEMRA_MOE_STATS: per-layer routing stats for the A2 (expert-grouped prefill) baseline —
4652        // per-token expert-id entropy, active-expert coverage, tokens-per-expert group sizes.
4653        if t > 1 && std::env::var("MEMRA_MOE_STATS").is_ok() {
4654            let mut cnt = vec![0u32; n_expert];
4655            for &s in sel_all.iter() {
4656                cnt[s as usize] += 1;
4657            }
4658            let total = sel_all.len() as f64;
4659            let mut h = 0.0f64;
4660            let mut active = 0usize;
4661            for &c in &cnt {
4662                if c > 0 {
4663                    active += 1;
4664                    let p = c as f64 / total;
4665                    h -= p * p.log2();
4666                }
4667            }
4668            let maxc = cnt.iter().copied().max().unwrap_or(0);
4669            println!(
4670                "moe-stats il={} t={} assignments={} active={}/{} entropy={:.3}b (max {:.3}b) mean_tok_per_active={:.2} max_tok_per_expert={}",
4671                il,
4672                t,
4673                sel_all.len(),
4674                active,
4675                n_expert,
4676                h,
4677                (n_expert as f64).log2(),
4678                total / active.max(1) as f64,
4679                maxc
4680            );
4681        }
4682
4683        // LAUNCH-STRUCTURE STAGE 2 (2026-07-05): moe_out memset elision on the gdec path.
4684        // moe_down8_fma_f32 FULLY overwrites its token row (dst[o] = the in-kernel FMA chain that
4685        // starts at 0.0f — numerically the axpy-into-zeroed-row chain), so when the grouped-decode
4686        // path fires the upfront `e.zeros(t*n_embd)` memset is pure launch churn. Allocate uninit
4687        // when gdec CAN fire (any t — decode t=1 AND the spec verify t=K+1 route here per token)
4688        // and lazily zero ONLY the row of a token that falls through to the sequential axpy loop.
4689        // BIT-IDENTITY: unchanged — every row is either fully overwritten (gdec) or
4690        // zeroed-then-accumulated exactly as before (fallback).
4691        // `!swiglu_clamped_at`: the grouped-decode kernels' fused epilogue is plain
4692        // silu(gate)*up (see the m3 note at the call sites below) — a clamped layer must fall
4693        // through to the sequential loop's ffn_act_lim. Hoisted into the fire predicate so the
4694        // uninit/memset invariant at :2418 stays consistent with what actually dispatches.
4695        let gdec_may_fire = uniform_experts
4696            && use_cache
4697            && n_used <= 8
4698            && gdec_enabled()
4699            && !cfg.swiglu_clamped_at(il as u32);
4700        // SLAB-LOCAL RESIDENT ARM (lane/pp-leverb 2026-08-08): the fits-VRAM resident slabs
4701        // (`dev_exps`, sized per PP device by cx-503b) were consumed ONLY by the pairs/dev
4702        // arms — every one of which DENIES sigmoid-router archs (step35/M3/Hy3). So on those
4703        // archs the slabs were uploaded but never read, and every expert went through the
4704        // SLRU (37 GB H2D staging per pp4096 prime on the Step SKU; and post-cx-503b the
4705        // dead slabs additionally STARVE the SLRU, which sizes itself on free-VRAM-after-
4706        // residents). This arm reads the slabs through the SAME kernels the SLRU arm runs —
4707        // gdec's fused pair when eligible, the per-expert qmatvec twins otherwise — with the
4708        // slab base + ex*stride as the pointer provenance (the exact bit-identity class
4709        // `moe_ffn_dev` documents for its resident-vs-SLRU arms; bytes are the same HostExps
4710        // bytes either way). LOCALITY GATE: `d.dev == e.ctx().ordinal()` — a slab on another
4711        // device must NOT be dereferenced (m=1 peer reads are the measured 34-150x class,
4712        // strictly worse than staging); under PP-2 without the prime walker this admits
4713        // stage-0 layers on dev0 and leaves stage-1 layers on the SLRU, which is the honest
4714        // interim shape. MEMRA_MOE_SLAB=0 = rollback/A-B seam (read per call).
4715        let slab_local = m
4716            .dev_exps
4717            .as_ref()
4718            .filter(|d| !d.gu_il && moe_slab_enabled() && d.dev == e.ctx().ordinal());
4719        let slab_bases = slab_local.map(|d| {
4720            use cudarc::driver::DevicePtr;
4721            let s = e.stream();
4722            let (pg, _g0) = d.gate.device_ptr(&s);
4723            let (pu, _g1) = d.up.device_ptr(&s);
4724            let (pd, _g2) = d.down.device_ptr(&s);
4725            (pg as u64, pu as u64, pd as u64)
4726        });
4727        // Fused-pair eligibility mirrors the gdec call sites exactly (plain-SiLU epilogue,
4728        // no macros, m3 clamp excluded, q8-supported qtypes, <=8 experts) — INCLUDING
4729        // `gdec_enabled()`: the pair IS the gdec kernel pair, so MEMRA_MOE_GDEC=0 must
4730        // disable it here too. That seam is also the exactness localizer: with GDEC=0 both
4731        // provenances run the SAME per-expert qmatvec kernels (slab base+stride vs SLRU
4732        // slot) and must be BIT-IDENTICAL — the true provenance-only pair — while GDEC=1
4733        // compares the fused-pair class against the SLRU's hit/miss MIX (gdec for
4734        // all-resident tokens, staged loop for misses), which is a dispatch-class
4735        // comparison, not a provenance one.
4736        let slab_fused_may_fire = slab_bases.is_some()
4737            && n_used <= 8
4738            && gdec_enabled()
4739            && !cfg.swiglu_clamped_at(il as u32)
4740            && cfg.m3.is_none()
4741            && no_exp_macros
4742            && moe_q8;
4743        // moe_out memset elision: BOTH full-row-overwrite arms (gdec + slab fused) allocate
4744        // uninit; a token that falls through to any accumulating loop zeroes its own row.
4745        let mut moe_out = if gdec_may_fire || slab_fused_may_fire {
4746            e.uninit(t * n_embd)?
4747        } else {
4748            e.zeros(t * n_embd)?
4749        };
4750        // The router readback above already established a host boundary. Copy each small-t hidden
4751        // row once so CPU miss experts can start while the owner thread queues resident GPU work.
4752        let cpu_input = if cpu_hybrid {
4753            Some(routed_cpu_input.ok_or("CPU expert routing did not return the MoE input")?)
4754        } else {
4755            None
4756        };
4757
4758        // GPU scratch: one slot per proj, big enough for ONE expert (default stage-every-token path).
4759        // STAGE 2: LAZY — allocated only if the no-cache staging path actually runs (under
4760        // MEMRA_MOE_CACHE they were 3 dead ~1MB alloc_zeros + memset + free per layer per token,
4761        // measured ~123 memsets/token of the decode wall).
4762        let g_len = m.gate_exps.max_expert_bytes(); // 860160 for the uniform 35B gate
4763        let u_len = m.up_exps.max_expert_bytes(); // 860160 for the uniform 35B up
4764        let d_len = m.down_exps.max_expert_bytes(); // 1114112 for the uniform 35B down
4765        let mut scratch_g: Option<CudaSlice<u8>> = None;
4766        let mut scratch_u: Option<CudaSlice<u8>> = None;
4767        let mut scratch_d: Option<CudaSlice<u8>> = None;
4768        // `max_block` (the GLOBAL max expert stride across all layers) is passed in — the cache slots
4769        // are FIXED-ADDRESS and must fit any layer's block (UD/dynamic GGUFs vary quant per layer).
4770
4771        // EDGE-1 §C.2/C.3: the optional pipeline queues the next selected expert's MISS blocks on
4772        // the copy stream before launching the current expert's compute. Pending slots stay invisible
4773        // to cache hits until dispatch inserts the completion-event wait; current gate/up/down ids are
4774        // protected from eviction. This changes scheduling only, never the GGUF bytes or GEMM path.
4775        let page_window = moe_page_prefetch_window();
4776
4777        // 2. PER TOKEN: routed-expert loop. The ONE dispatch change vs Stage-1: a resident slot
4778        //    (cache HIT, no H2D) OR a staged slot (MISS) feeds the SAME unchanged qmatvec_view.
4779        for tok in 0..t {
4780            let sel = &sel_all[tok * n_used..(tok + 1) * n_used];
4781            let w = &w_all[tok * n_used..(tok + 1) * n_used];
4782            let zt = z.slice(tok * n_embd..(tok + 1) * n_embd); // CudaView<f32>
4783            let mut tok_q8: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
4784
4785            // STAGE-2 GROUPED DECODE (2026-07-04, MEMRA_MOE_GDEC default ON, =0 rollback): fold
4786            // this token's whole routed-expert FFN (8x gate/up/silu + 8x down/axpy = 40 launches)
4787            // into TWO launches via expert-pointer indirection over the fixed-address cache slots.
4788            // Fires only when ALL 3*n_used blocks are ALREADY cache-resident (pure-HIT: zero
4789            // memcpy, zero admission, so no slot can move under the collected pointers) — any
4790            // miss falls through to the sequential loop below, which admits as before. In steady
4791            // state on a fully-resident rig every token-layer takes the grouped path.
4792            // BIT-IDENTITY: each in-kernel dot reproduces qmatvec_f32's exact reduction; SiLU is
4793            // silu_mul_f32's exact expression; the down accumulation is a slot-ordered
4794            // __fmaf_rn chain == the sequential axpy_f32 chain (MEMRA_MOE_GDEC_GATE compares).
4795            // cfg.m3: the grouped kernels' fused epilogues are plain SiLU — M3's swigluoai must
4796            // NOT take them until the kernels grow the clamped variant. NVFP4 experts carry
4797            // per-expert macro-scales the fused kernels don't fold — those fall through too.
4798            let no_macros = m.gate_exps.macros.is_none()
4799                && m.up_exps.macros.is_none()
4800                && m.down_exps.macros.is_none();
4801            // SLAB-LOCAL FUSED PAIR (lane/pp-leverb 2026-08-08): the gdec launch pair
4802            // (moe_gate_up_silu8_q8 + moe_down8_fma_q8 — the EXACT kernels, same FP chains)
4803            // with pointers computed from the resident slab base + ex*stride instead of
4804            // collected SLRU slot addresses. No cache lock, no residency predicate — the
4805            // slab holds every expert by construction, so this arm never falls through
4806            // (gdec's P(all 24 resident) ≈ 0.37 coin-flip and the miss path's ~49-launch
4807            // staging both die). Bit-identity class: pointer provenance only, the same
4808            // slab-vs-SLRU equivalence moe_ffn_dev documents. Ordered BEFORE gdec: when a
4809            // slab exists it is strictly better (no lock, no miss).
4810            if slab_fused_may_fire {
4811                let (pg, pu, pd) = slab_bases.unwrap();
4812                let mut gp = [0u64; 8];
4813                let mut up = [0u64; 8];
4814                let mut dp = [0u64; 8];
4815                for (j, &ex) in sel.iter().enumerate() {
4816                    let ex = ex as usize;
4817                    gp[j] = pg + (ex * m.gate_exps.expert_stride) as u64;
4818                    up[j] = pu + (ex * m.up_exps.expert_stride) as u64;
4819                    dp[j] = pd + (ex * m.down_exps.expert_stride) as u64;
4820                }
4821                let mut wv = [0f32; 8];
4822                wv[..n_used].copy_from_slice(w);
4823                if tok_q8.is_none() {
4824                    tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
4825                }
4826                let (zq, zd) = tok_q8.as_ref().unwrap();
4827                let act = e.moe_gate_up_silu8_q8(
4828                    crate::WPtr8(gp),
4829                    crate::WPtr8(up),
4830                    zq,
4831                    zd,
4832                    n_embd,
4833                    n_ff_exp,
4834                    n_used,
4835                    m.gate_exps.qtype,
4836                    m.up_exps.qtype,
4837                    m.gate_exps.row_bytes,
4838                    m.up_exps.row_bytes,
4839                )?;
4840                let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
4841                let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
4842                e.moe_down8_fma_q8(
4843                    crate::WPtr8(dp),
4844                    crate::F32x8(wv),
4845                    &aq2,
4846                    &ad2,
4847                    &mut dst,
4848                    n_ff_exp,
4849                    n_embd,
4850                    n_used,
4851                    m.down_exps.qtype,
4852                    m.down_exps.row_bytes,
4853                )?;
4854                continue;
4855            }
4856            if gdec_may_fire && moe_q8 && cfg.m3.is_none() && no_macros {
4857                if tok_q8.is_none() {
4858                    tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
4859                }
4860                let (zq, zd) = tok_q8.as_ref().unwrap();
4861                if Self::moe_gdec_token_q8(
4862                    e,
4863                    m,
4864                    il,
4865                    max_block,
4866                    zq,
4867                    zd,
4868                    sel,
4869                    w,
4870                    &mut moe_out,
4871                    tok,
4872                    n_embd,
4873                    n_ff_exp,
4874                    n_used,
4875                )? {
4876                    continue;
4877                }
4878            } else if gdec_may_fire
4879                && cfg.m3.is_none()
4880                && no_macros
4881                && Self::moe_gdec_token(
4882                    e,
4883                    m,
4884                    il,
4885                    max_block,
4886                    &zt,
4887                    sel,
4888                    w,
4889                    &mut moe_out,
4890                    tok,
4891                    n_embd,
4892                    n_ff_exp,
4893                    n_used,
4894                )?
4895            {
4896                continue;
4897            }
4898
4899            // STAGE 2 memset-elision invariant: moe_out was allocated UNINIT when gdec or the
4900            // slab pair could fire. This token fell through to a sequential axpy loop, which
4901            // ACCUMULATES — zero its row first (row-sized memset; other rows are owned by the
4902            // full-overwrite arms). slab_fused_may_fire never actually reaches here (its arm
4903            // has no fallible predicate), included for the allocation invariant's symmetry.
4904            if gdec_may_fire || slab_fused_may_fire {
4905                let mut row = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
4906                e.memset_zeros_view(&mut row)?;
4907            }
4908
4909            // Fiddler-style split at whole-expert granularity. Partial residency is deliberately a
4910            // CPU assignment: transferring the missing projection would reintroduce the PCIe/NVMe
4911            // stall this path exists to remove, while mixing projections would require another
4912            // activation round-trip. Weight addresses remain valid until this worker is joined at
4913            // the bottom of the token scope.
4914            let mut cpu_mask = vec![false; sel.len()];
4915            let cpu_worker = if let Some(host_input) = cpu_input.as_ref() {
4916                let gpu_resident = if use_cache {
4917                    e.with_moe_cache(max_block, |cache, _| {
4918                        Ok(sel
4919                            .iter()
4920                            .map(|&expert| {
4921                                let expert = expert as u16;
4922                                [PROJ_GATE, PROJ_UP, PROJ_DOWN]
4923                                    .into_iter()
4924                                    .filter(|&projection| {
4925                                        cache
4926                                            .resident(BlockId::new(il, projection, expert))
4927                                            .is_some()
4928                                    })
4929                                    .count()
4930                            })
4931                            .collect::<Vec<_>>())
4932                    })?
4933                } else {
4934                    vec![0; sel.len()]
4935                };
4936                let mut cpu_selected = Vec::new();
4937                for (index, (&expert, &route_weight)) in sel.iter().zip(w).enumerate() {
4938                    if gpu_resident[index] != 3 {
4939                        cpu_mask[index] = true;
4940                        crate::cpu_experts::record_incomplete_gpu_residency(gpu_resident[index]);
4941                        let expert = expert as usize;
4942                        cpu_selected.push((expert, route_weight));
4943                    }
4944                }
4945                if crate::cpu_experts::predictor_enabled() {
4946                    // Fire-and-forget lookahead: the predictor worker scores layers il+1..
4947                    // from this layer's MoE input and prefetches predicted-and-missing
4948                    // experts into the companion RAM cache. Never blocks this thread.
4949                    let row = &host_input[tok * n_embd..(tok + 1) * n_embd];
4950                    crate::cpu_experts::predictor_submit(il, row);
4951                }
4952                if cpu_selected.is_empty() {
4953                    None
4954                } else {
4955                    let row = &host_input[tok * n_embd..(tok + 1) * n_embd];
4956                    let job = crate::cpu_experts::prepare_job(m, il, &cpu_selected, row)
4957                        .map_err(std::io::Error::other)?;
4958                    Some(crate::cpu_experts::submit(job).map_err(std::io::Error::other)?)
4959                }
4960            } else {
4961                None
4962            };
4963
4964            let worker_window = worker_disk_prefetch
4965                .then(worker_prefetch_window)
4966                .unwrap_or(0);
4967            for (j, &ex) in sel.iter().enumerate() {
4968                if cpu_mask[j] {
4969                    continue;
4970                }
4971                let ex = ex as usize;
4972                // PER-EXPERT SLAB READ (lane/pp-leverb 2026-08-08): the layers the fused
4973                // pair above excludes — step35's CLAMPED layers 43/44 (ffn_act_lim has no
4974                // fused form) and macro-carrying artifacts — still have their bytes in the
4975                // local resident slab. Same kernels as the SLRU arms (`qmatvec_expert_q8` /
4976                // `qmatvec_view`), same ffn_act_lim/macro folds; provenance = slab base +
4977                // ex*stride. No dispatch lock, no admission, no prefetch — nothing to miss.
4978                if let Some(d) = slab_local {
4979                    let gl = m.gate_exps.expert_layout(ex);
4980                    let ul = m.up_exps.expert_layout(ex);
4981                    let dl = m.down_exps.expert_layout(ex);
4982                    let (g0, u0, d0) = (
4983                        ex * m.gate_exps.expert_stride,
4984                        ex * m.up_exps.expert_stride,
4985                        ex * m.down_exps.expert_stride,
4986                    );
4987                    let (gate, up) = if moe_q8 {
4988                        if tok_q8.is_none() {
4989                            tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
4990                        }
4991                        let (zq, zd) = tok_q8.as_ref().unwrap();
4992                        (
4993                            e.qmatvec_expert_q8(
4994                                &d.gate,
4995                                g0..g0 + gl.len,
4996                                zq,
4997                                zd,
4998                                1,
4999                                m.gate_exps.in_f,
5000                                m.gate_exps.out_f,
5001                                gl.qtype,
5002                                gl.row_bytes,
5003                            )?,
5004                            e.qmatvec_expert_q8(
5005                                &d.up,
5006                                u0..u0 + ul.len,
5007                                zq,
5008                                zd,
5009                                1,
5010                                m.up_exps.in_f,
5011                                m.up_exps.out_f,
5012                                ul.qtype,
5013                                ul.row_bytes,
5014                            )?,
5015                        )
5016                    } else {
5017                        (
5018                            e.qmatvec_view(
5019                                &d.gate,
5020                                g0..g0 + gl.len,
5021                                &zt,
5022                                1,
5023                                m.gate_exps.in_f,
5024                                m.gate_exps.out_f,
5025                                gl.qtype,
5026                                gl.row_bytes,
5027                            )?,
5028                            e.qmatvec_view(
5029                                &d.up,
5030                                u0..u0 + ul.len,
5031                                &zt,
5032                                1,
5033                                m.up_exps.in_f,
5034                                m.up_exps.out_f,
5035                                ul.qtype,
5036                                ul.row_bytes,
5037                            )?,
5038                        )
5039                    };
5040                    let mut act = e.uninit(n_ff_exp)?;
5041                    Self::ffn_act_lim(
5042                        e,
5043                        cfg,
5044                        &gate,
5045                        &up,
5046                        m.gate_exps.macro_scale(ex),
5047                        m.up_exps.macro_scale(ex),
5048                        lim_exp,
5049                        &mut act,
5050                        n_ff_exp,
5051                    )?;
5052                    let y = if moe_q8 {
5053                        let (aq2, ad2) = e.quantize_q8_1(&act, 1, n_ff_exp)?;
5054                        e.qmatvec_expert_q8(
5055                            &d.down,
5056                            d0..d0 + dl.len,
5057                            &aq2,
5058                            &ad2,
5059                            1,
5060                            m.down_exps.in_f,
5061                            m.down_exps.out_f,
5062                            dl.qtype,
5063                            dl.row_bytes,
5064                        )?
5065                    } else {
5066                        let actv = act.slice(0..n_ff_exp);
5067                        e.qmatvec_view(
5068                            &d.down,
5069                            d0..d0 + dl.len,
5070                            &actv,
5071                            1,
5072                            m.down_exps.in_f,
5073                            m.down_exps.out_f,
5074                            dl.qtype,
5075                            dl.row_bytes,
5076                        )?
5077                    };
5078                    let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
5079                    e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
5080                    continue;
5081                }
5082                for next in page_prefetch_positions(j, sel.len(), page_window) {
5083                    Self::moe_prefetch_host_expert(sel[next] as usize, m);
5084                }
5085                let keep = [
5086                    crate::moe_cache::BlockId::new(il, crate::moe_cache::PROJ_GATE, ex as u16),
5087                    crate::moe_cache::BlockId::new(il, crate::moe_cache::PROJ_UP, ex as u16),
5088                    crate::moe_cache::BlockId::new(il, crate::moe_cache::PROJ_DOWN, ex as u16),
5089                ];
5090                if worker_disk_prefetch && worker_window > 0 {
5091                    for next in worker_prefetch_positions(j, sel.len(), worker_window) {
5092                        Self::moe_prefetch_disk_expert(
5093                            e,
5094                            il,
5095                            sel[next] as usize,
5096                            m,
5097                            max_block,
5098                            &keep,
5099                        )?;
5100                    }
5101                } else if cache_dispatch
5102                    && !cpu_hybrid
5103                    && moe_prefetch_enabled()
5104                    && j + 1 < sel.len()
5105                {
5106                    let next = sel[j + 1] as usize;
5107                    Self::moe_prefetch_expert(e, il, next, m, max_block, &keep)?;
5108                }
5109                let [gate_q8, up_q8, down_q8] = [moe_q8; 3];
5110                if cache_dispatch && (gate_q8 || up_q8 || down_q8) {
5111                    // dp4a EXPERT PATH (MEMRA_MOE_Q8): quantize z-row once per token. Mixed expert
5112                    // layouts stay on the metadata-aware f32 path.
5113                    if (gate_q8 || up_q8) && tok_q8.is_none() {
5114                        tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
5115                    }
5116                    let gate = if gate_q8 {
5117                        let (zq, zd) = tok_q8.as_ref().unwrap();
5118                        Self::moe_cached_gemm_q8(e, il, PROJ_GATE, ex, m, max_block, zq, zd)?
5119                    } else {
5120                        Self::moe_cached_gemm(e, il, PROJ_GATE, ex, m, max_block, &zt)?
5121                    };
5122                    let up = if up_q8 {
5123                        let (zq, zd) = tok_q8.as_ref().unwrap();
5124                        Self::moe_cached_gemm_q8(e, il, PROJ_UP, ex, m, max_block, zq, zd)?
5125                    } else {
5126                        Self::moe_cached_gemm(e, il, PROJ_UP, ex, m, max_block, &zt)?
5127                    };
5128                    let mut act = e.uninit(n_ff_exp)?;
5129                    Self::ffn_act_lim(
5130                        e,
5131                        cfg,
5132                        &gate,
5133                        &up,
5134                        m.gate_exps.macro_scale(ex),
5135                        m.up_exps.macro_scale(ex),
5136                        lim_exp,
5137                        &mut act,
5138                        n_ff_exp,
5139                    )?;
5140                    let y = if down_q8 {
5141                        let (aq2, ad2) = e.quantize_q8_1(&act, 1, n_ff_exp)?;
5142                        Self::moe_cached_gemm_q8(e, il, PROJ_DOWN, ex, m, max_block, &aq2, &ad2)?
5143                    } else {
5144                        let actv = act.slice(0..n_ff_exp);
5145                        Self::moe_cached_gemm(e, il, PROJ_DOWN, ex, m, max_block, &actv)?
5146                    };
5147                    let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
5148                    // down-proj macro folds into the accumulate weight (1.0 for non-macro archs).
5149                    e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
5150                } else if cache_dispatch {
5151                    // SLRU residency cache: per-projection, dispatch the block (HIT => resident slot,
5152                    // MISS => staged slot) then run the SAME unchanged qmatvec_view from that slot.
5153                    // The bytes the kernel reads are byte-for-byte the same GGUF block (§B.3); the
5154                    // only difference between HIT and MISS is whether the memcpy_htod ran.
5155                    let gate = Self::moe_cached_gemm(e, il, PROJ_GATE, ex, m, max_block, &zt)?;
5156                    let up = Self::moe_cached_gemm(e, il, PROJ_UP, ex, m, max_block, &zt)?;
5157                    let mut act = e.uninit(n_ff_exp)?; // activation fully overwrites
5158                    Self::ffn_act_lim(
5159                        e,
5160                        cfg,
5161                        &gate,
5162                        &up,
5163                        m.gate_exps.macro_scale(ex),
5164                        m.up_exps.macro_scale(ex),
5165                        lim_exp,
5166                        &mut act,
5167                        n_ff_exp,
5168                    )?;
5169                    let actv = act.slice(0..n_ff_exp);
5170                    let y = Self::moe_cached_gemm(e, il, PROJ_DOWN, ex, m, max_block, &actv)?;
5171                    let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
5172                    // down-proj macro folds into the accumulate weight (post-matmul linear scale).
5173                    e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
5174                } else if cache_frozen {
5175                    // A later prompt prime must not change the CPU/GPU assignment frozen after the
5176                    // first prime. Reuse every fixed resident projection directly and stage only a
5177                    // true miss through the ordinary scratch slot. This preserves the established
5178                    // f32-dequant numeric path while avoiding a full-bank reread on every prime.
5179                    let gate = Self::moe_frozen_gemm(
5180                        e,
5181                        il,
5182                        PROJ_GATE,
5183                        ex,
5184                        m,
5185                        max_block,
5186                        &zt,
5187                        &mut scratch_g,
5188                        g_len,
5189                    )?;
5190                    let up = Self::moe_frozen_gemm(
5191                        e,
5192                        il,
5193                        PROJ_UP,
5194                        ex,
5195                        m,
5196                        max_block,
5197                        &zt,
5198                        &mut scratch_u,
5199                        u_len,
5200                    )?;
5201                    let mut act = e.uninit(n_ff_exp)?;
5202                    Self::ffn_act_lim(
5203                        e,
5204                        cfg,
5205                        &gate,
5206                        &up,
5207                        m.gate_exps.macro_scale(ex),
5208                        m.up_exps.macro_scale(ex),
5209                        lim_exp,
5210                        &mut act,
5211                        n_ff_exp,
5212                    )?;
5213                    let actv = act.slice(0..n_ff_exp);
5214                    let y = Self::moe_frozen_gemm(
5215                        e,
5216                        il,
5217                        PROJ_DOWN,
5218                        ex,
5219                        m,
5220                        max_block,
5221                        &actv,
5222                        &mut scratch_d,
5223                        d_len,
5224                    )?;
5225                    let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
5226                    e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
5227                } else {
5228                    // Stage-1: stage gate/up/down for expert `ex` into the scratch slots, then GEMM.
5229                    // Lazy scratch: first no-cache expert allocates the 3 slots (uninit — stage_expert
5230                    // fully overwrites the byte range the GEMM reads).
5231                    if scratch_g.is_none() {
5232                        scratch_g = Some(e.alloc_u8_uninit(g_len)?);
5233                        scratch_u = Some(e.alloc_u8_uninit(u_len)?);
5234                        scratch_d = Some(e.alloc_u8_uninit(d_len)?);
5235                    }
5236                    let (sg, su, sd) = (
5237                        scratch_g.as_mut().unwrap(),
5238                        scratch_u.as_mut().unwrap(),
5239                        scratch_d.as_mut().unwrap(),
5240                    );
5241                    let gl = m.gate_exps.expert_layout(ex);
5242                    let ul = m.up_exps.expert_layout(ex);
5243                    let dl = m.down_exps.expert_layout(ex);
5244                    e.stage_expert(m.gate_exps.expert_bytes(ex), sg, 0)?;
5245                    let gate = e.qmatvec_view(
5246                        sg,
5247                        0..gl.len,
5248                        &zt,
5249                        1,
5250                        m.gate_exps.in_f,
5251                        m.gate_exps.out_f,
5252                        gl.qtype,
5253                        gl.row_bytes,
5254                    )?;
5255
5256                    e.stage_expert(m.up_exps.expert_bytes(ex), su, 0)?;
5257                    let up = e.qmatvec_view(
5258                        su,
5259                        0..ul.len,
5260                        &zt,
5261                        1,
5262                        m.up_exps.in_f,
5263                        m.up_exps.out_f,
5264                        ul.qtype,
5265                        ul.row_bytes,
5266                    )?;
5267
5268                    let mut act = e.uninit(n_ff_exp)?; // activation fully overwrites
5269                    Self::ffn_act_lim(
5270                        e,
5271                        cfg,
5272                        &gate,
5273                        &up,
5274                        m.gate_exps.macro_scale(ex),
5275                        m.up_exps.macro_scale(ex),
5276                        lim_exp,
5277                        &mut act,
5278                        n_ff_exp,
5279                    )?;
5280
5281                    e.stage_expert(m.down_exps.expert_bytes(ex), sd, 0)?;
5282                    let actv = act.slice(0..n_ff_exp);
5283                    let y = e.qmatvec_view(
5284                        sd,
5285                        0..dl.len,
5286                        &actv,
5287                        1,
5288                        m.down_exps.in_f,
5289                        m.down_exps.out_f,
5290                        dl.qtype,
5291                        dl.row_bytes,
5292                    )?;
5293
5294                    let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
5295                    e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
5296                }
5297            }
5298            if let Some(worker) = cpu_worker {
5299                let cpu_output = worker.wait().map_err(std::io::Error::other)?;
5300                let cpu_output = e.htod(&cpu_output)?;
5301                let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
5302                e.axpy_into(&cpu_output, 1.0, &mut dst, n_embd)?;
5303            }
5304            if cpu_hybrid && !cache_frozen && cpu_expert_profile_admit_enabled() {
5305                for (j, &ex) in sel.iter().enumerate() {
5306                    if cpu_mask[j] {
5307                        Self::moe_profile_admit_expert(e, il, ex as usize, m, max_block)?;
5308                    }
5309                }
5310            }
5311        }
5312
5313        // 3. SHARED EXPERT (ALWAYS-ON, no routing) on the SAME z — qwen35moe only. OLMoE and most
5314        //    vanilla MoE have NO shared expert (the shexp tensors are absent / `None`); skip it then.
5315        // gate_inp_shexp is OPTIONAL: qwen35moe gates the shared expert (sigmoid(gate_inp) x sh);
5316        // MiniMax-M3 (DeepSeek-V3 class) has NO shexp gate — the shared expert adds directly.
5317        if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
5318            (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
5319        {
5320            let n_ff_sh = gate_shexp.out_features(); // 512
5321            // Q8 TRUNK-FUSION (decode t=1): gate_shexp+up_shexp are Q8_0 same-shape on the 35B —
5322            // ONE fused2 launch (also folds the two per-matmul re-quantizes of z into one).
5323            // Bit-identical per (tensor,row); falls back to the two matmul calls when ineligible.
5324            // Small-t (spec verify 2..15) rides matmul_decode_exact so shexp FP chains match the
5325            // t==1 decode chain per column (cuBLASLt n-dependence + dp4a-vs-mmvq class); real
5326            // prefill keeps the batched matmul. Activation routes through ffn_act (SiLU for
5327            // softmax archs, clamped swigluoai for M3 — identical to silu_mul when cfg.m3 is None).
5328            let verify_t = t > 1 && t < PRIME_MIN_T;
5329            let (sg_gate, sg_up) = if t == 1 {
5330                match e.matmul_q8_fused2_x(gate_shexp, up_shexp, z)? {
5331                    Some(pair) => pair,
5332                    None => (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?),
5333                }
5334            } else if verify_t {
5335                (
5336                    e.matmul_decode_exact(gate_shexp, z, t)?,
5337                    e.matmul_decode_exact(up_shexp, z, t)?,
5338                )
5339            } else {
5340                (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?) // [T, 512] each
5341            };
5342            let mut sa = e.uninit(t * n_ff_sh)?; // activation fully overwrites
5343            Self::ffn_act_lim(
5344                e,
5345                cfg,
5346                &sg_gate,
5347                &sg_up,
5348                1.0,
5349                1.0,
5350                lim_shexp,
5351                &mut sa,
5352                t * n_ff_sh,
5353            )?;
5354            let sh = if verify_t {
5355                e.matmul_decode_exact(down_shexp, &sa, t)?
5356            } else {
5357                e.matmul(down_shexp, &sa, t)?
5358            }; // [T, n_embd]
5359
5360            // shexp gate: qwen35moe sigmoid-gates via ffn_gate_inp_shexp (1-D ne=[n_embd] ->
5361            // out_f=1); M3 has no gate tensor -> weight 1.0. Decode + verify ride the fused
5362            // sigmoid-dot kernel (one fold order for both chains; kills the per-layer
5363            // cuBLASLt m=1 splitK GEMM — 40x/step, ~10% of the H100 q35 decode step).
5364            // SERVE ISOLATION (lane/concat-prime-exact, 2026-08-02): PREFILL rides it too.
5365            // This out_f=1 cuBLASLt GEMV is the SECOND m-dependent op in the trunk (probed:
5366            // rows [0,19) move by 1.07e-4 between m=74 and m=75 while sigmoid_dot_rows is
5367            // BIT-IDENTICAL — allw-shexpgate-o35b.log). The gate multiplies the shared
5368            // expert's contribution into every token's residual, so under cross-request
5369            // concat prefill a session's hidden state depended on its co-arrivals' token
5370            // count. MEMRA_ROUTER_PREFILL_EXACT=0 reverts to the batched cuBLASLt linear.
5371            let g = match &m.gate_inp_shexp {
5372                Some(gate_inp_shexp) => {
5373                    if t < PRIME_MIN_T || crate::router_prefill_exact_on() {
5374                        e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
5375                    } else {
5376                        let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
5377                        let mut g = e.uninit(t)?; // sigmoid fully overwrites
5378                        e.sigmoid(&gs, &mut g, t)?;
5379                        g
5380                    }
5381                }
5382                None => e.htod(&vec![1.0f32; t])?,
5383            };
5384            // moe_out[r, :] += sh[r, :] * g[r]   (per-token scalar gate; g=1 ungated)
5385            e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
5386        }
5387
5388        Ok(moe_out)
5389    }
5390
5391    /// Stage-1 (no-cache) per-DECODE-TOKEN H2D bytes: every routed block re-staged every layer every
5392    /// token = sum over MoE layers of n_used * (gate+up+down expert_stride). The §D.4 PCIe baseline.
5393    pub fn stage1_h2d_per_token(&self) -> u64 {
5394        use crate::hybrid::Ffn;
5395        let n_used = self
5396            .cfg
5397            .moe
5398            .as_ref()
5399            .map(|m| m.expert_used_count as u64)
5400            .unwrap_or(0);
5401        let mut bytes = 0u64;
5402        for l in self.layers.iter() {
5403            if let Ffn::Moe(m) = &l.ffn {
5404                bytes += n_used
5405                    * (m.gate_exps.max_expert_bytes()
5406                        + m.up_exps.max_expert_bytes()
5407                        + m.down_exps.max_expert_bytes()) as u64;
5408            }
5409        }
5410        bytes
5411    }
5412
5413    /// Largest expert block (bytes) across ALL MoE layers + the MTP head — the fixed cache slot size.
5414    /// UD/dynamic GGUFs quant different layers differently, so `expert_stride` varies per layer; the
5415    /// residency cache slots are fixed-address and must fit any block, so size to this global max.
5416    pub(crate) fn max_moe_block(&self) -> usize {
5417        use crate::hybrid::Ffn;
5418        let mut mx = 0usize;
5419        let mut scan = |ffn: &Ffn| {
5420            if let Ffn::Moe(m) = ffn {
5421                mx = mx
5422                    .max(m.gate_exps.max_expert_bytes())
5423                    .max(m.up_exps.max_expert_bytes())
5424                    .max(m.down_exps.max_expert_bytes());
5425            }
5426        };
5427        for l in self.layers.iter() {
5428            scan(&l.ffn);
5429        }
5430        if let Some(mtp) = self.mtp.as_ref() {
5431            scan(&mtp.ffn);
5432        }
5433        mx
5434    }
5435
5436    /// Exact retained projection lengths for cache class sizing. Pruned ids keep router positions
5437    /// but have no bytes and therefore consume no residency slot.
5438    pub(crate) fn moe_cache_block_sizes(&self) -> Vec<usize> {
5439        use crate::hybrid::Ffn;
5440        let mut sizes = Vec::new();
5441        let mut scan = |ffn: &Ffn| {
5442            let Ffn::Moe(m) = ffn else { return };
5443            for ex in 0..m.gate_exps.n_expert {
5444                if m.active_experts.as_ref().is_some_and(|active| !active[ex]) {
5445                    continue;
5446                }
5447                for exps in [&m.gate_exps, &m.up_exps, &m.down_exps] {
5448                    let len = exps.expert_layout(ex).len;
5449                    if len > 0 {
5450                        sizes.push(len);
5451                    }
5452                }
5453            }
5454        };
5455        for layer in &self.layers {
5456            scan(&layer.ffn);
5457        }
5458        if let Some(mtp) = &self.mtp {
5459            scan(&mtp.ffn);
5460        }
5461        sizes
5462    }
5463
5464    /// Persist the frozen residency set so a later process can restage it directly and skip
5465    /// the profiling warmup. Plain text: a versioned header binding slot geometry, then one
5466    /// `layer proj ex` triple per line. A mismatched or stale profile is rejected at load
5467    /// (header check) or degrades to fewer restaged blocks (per-id checks); either way the
5468    /// post-freeze argmax gate still validates the serving assignment.
5469    pub fn save_cpu_expert_residency_profile(
5470        &self,
5471        e: &Engine,
5472        path: &std::path::Path,
5473    ) -> Result<(), Box<dyn std::error::Error>> {
5474        let Some(ids) = e.export_moe_residency() else {
5475            return Err("no MoE residency cache to persist".into());
5476        };
5477        let mut body = format!(
5478            "memra-freeze-profile v1 max_block={} blocks={}\n",
5479            self.max_moe_block(),
5480            ids.len()
5481        );
5482        for (layer, proj, ex) in &ids {
5483            body.push_str(&format!("{layer} {proj} {ex}\n"));
5484        }
5485        let tmp = path.with_extension("tmp");
5486        std::fs::write(&tmp, body)?;
5487        std::fs::rename(&tmp, path)?;
5488        println!(
5489            "[moe-cache] freeze profile saved: {} blocks -> {}",
5490            ids.len(),
5491            path.display()
5492        );
5493        Ok(())
5494    }
5495
5496    /// Restage a saved freeze profile and freeze immediately, skipping the profiling warmup.
5497    /// Returns false (leaving the cache untouched for a normal warmup) when the profile is
5498    /// missing or its header does not match this model's slot geometry.
5499    pub fn restore_cpu_expert_residency_profile(
5500        &self,
5501        e: &Engine,
5502        path: &std::path::Path,
5503    ) -> Result<bool, Box<dyn std::error::Error>> {
5504        use crate::hybrid::Ffn;
5505        use crate::moe_cache::BlockId;
5506        let Ok(content) = std::fs::read_to_string(path) else {
5507            return Ok(false);
5508        };
5509        let mut lines = content.lines();
5510        let Some(header) = lines.next() else {
5511            return Ok(false);
5512        };
5513        let expected = format!("memra-freeze-profile v1 max_block={}", self.max_moe_block());
5514        if !header.starts_with(&expected) {
5515            println!(
5516                "[moe-cache] freeze profile ignored (geometry mismatch): {}",
5517                path.display()
5518            );
5519            return Ok(false);
5520        }
5521        let mut by_layer: std::collections::HashMap<u16, Vec<BlockId>> =
5522            std::collections::HashMap::new();
5523        for line in lines {
5524            let mut fields = line.split_whitespace();
5525            let (Some(layer), Some(proj), Some(ex)) = (fields.next(), fields.next(), fields.next())
5526            else {
5527                continue;
5528            };
5529            let (Ok(layer), Ok(proj), Ok(ex)) =
5530                (layer.parse::<u16>(), proj.parse::<u8>(), ex.parse::<u16>())
5531            else {
5532                continue;
5533            };
5534            by_layer
5535                .entry(layer)
5536                .or_default()
5537                .push(BlockId::new(layer, proj, ex));
5538        }
5539        let requested: usize = by_layer.values().map(Vec::len).sum();
5540        if requested == 0 {
5541            return Ok(false);
5542        }
5543        let max_block = self.max_moe_block();
5544        let mut restaged = 0usize;
5545        let mut stage_layer =
5546            |layer_index: u16, ffn: &Ffn| -> Result<(), Box<dyn std::error::Error>> {
5547                let Ffn::Moe(m) = ffn else { return Ok(()) };
5548                let Some(ids) = by_layer.get(&layer_index) else {
5549                    return Ok(());
5550                };
5551                e.with_moe_cache(max_block, |cache, eng| {
5552                    for id in ids {
5553                        if cache.restage_block(*id, m, eng)? {
5554                            restaged += 1;
5555                        }
5556                    }
5557                    Ok(())
5558                })
5559            };
5560        for (index, layer) in self.layers.iter().enumerate() {
5561            stage_layer(index as u16, &layer.ffn)?;
5562        }
5563        if let Some(mtp) = self.mtp.as_ref() {
5564            stage_layer(u16::MAX, &mtp.ffn)?;
5565        }
5566        e.freeze_moe_cache();
5567        println!(
5568            "[moe-cache] freeze profile restored: {restaged}/{requested} blocks restaged from {}",
5569            path.display()
5570        );
5571        Ok(true)
5572    }
5573
5574    /// Freeze the heterogeneous CPU/GPU split after the caller's discarded profile warmup.
5575    pub fn freeze_cpu_expert_residency(
5576        &self,
5577        e: &Engine,
5578    ) -> Result<(), Box<dyn std::error::Error>> {
5579        e.freeze_moe_cache();
5580        Ok(())
5581    }
5582
5583    /// FFN activation dispatch: swigluoai (clamped, alpha/limit) when cfg.m3 says so, else the
5584    /// standard SiLU*up. One seam so every FFN site (dense, routed expert, shared expert) follows
5585    /// the model's activation exactly.
5586    ///
5587    /// NO-`il` FORM: cannot apply step35's PER-LAYER SwiGLU clamp. Only call it from a site whose
5588    /// layer provably has no live limit (dense-FFN layers, MTP blocks) — `ffn_act_lim` is the
5589    /// form for anything that can land on a clamped layer.
5590    pub fn ffn_act(
5591        e: &Engine,
5592        cfg: &ModelConfig,
5593        gate: &CudaSlice<f32>,
5594        up: &CudaSlice<f32>,
5595        act: &mut CudaSlice<f32>,
5596        n: usize,
5597    ) -> Result<(), Box<dyn std::error::Error>> {
5598        Self::ffn_act_scaled(e, cfg, gate, up, 1.0, 1.0, act, n)
5599    }
5600
5601    /// ffn_act with per-tensor post-matmul macro-scales folded in (gs/us == 1.0 -> identical
5602    /// float ops to ffn_act; used by the ModelOpt NVFP4 expert path where each expert tensor
5603    /// carries a `weight_scale_2`). Same no-`il` contract as `ffn_act`.
5604    #[allow(clippy::too_many_arguments)]
5605    pub(crate) fn ffn_act_scaled(
5606        e: &Engine,
5607        cfg: &ModelConfig,
5608        gate: &CudaSlice<f32>,
5609        up: &CudaSlice<f32>,
5610        gs: f32,
5611        us: f32,
5612        act: &mut CudaSlice<f32>,
5613        n: usize,
5614    ) -> Result<(), Box<dyn std::error::Error>> {
5615        Self::ffn_act_lim(e, cfg, gate, up, gs, us, None, act, n)
5616    }
5617
5618    /// ffn_act_scaled + step35's PER-LAYER clamped SwiGLU. `limit`:
5619    ///   * `None`   -> the unclamped dispatch (every arch except step35's layers 43-44).
5620    ///   * `Some(l)`-> `min(silu(gate*gs), l) * clamp(up*us, +-l)` (llama-graph.cpp:2146/1751,
5621    ///                 non-DEEPSEEK4 branch). Callers source it from `cfg.clamp_exp_at(il)`
5622    ///                 (routed experts) or `cfg.clamp_shexp_at(il)` (shared expert) — the two
5623    ///                 arrays are SEPARATE and a layer can have one without the other.
5624    /// The `> 1e-6` eps gate lives in `clamp_exp_at`/`clamp_shexp_at`, so a `Some` here is
5625    /// already known live.
5626    #[allow(clippy::too_many_arguments)]
5627    pub(crate) fn ffn_act_lim(
5628        e: &Engine,
5629        cfg: &ModelConfig,
5630        gate: &CudaSlice<f32>,
5631        up: &CudaSlice<f32>,
5632        gs: f32,
5633        us: f32,
5634        limit: Option<f32>,
5635        act: &mut CudaSlice<f32>,
5636        n: usize,
5637    ) -> Result<(), Box<dyn std::error::Error>> {
5638        if let Some(m3) = cfg.m3.as_ref() {
5639            debug_assert!(
5640                limit.is_none(),
5641                "m3 swigluoai and step35 clamp are different archs"
5642            );
5643            return e.swigluoai_mul_scaled(
5644                gate,
5645                up,
5646                gs,
5647                us,
5648                m3.swiglu_alpha,
5649                m3.swiglu_limit,
5650                act,
5651                n,
5652            );
5653        }
5654        if let Some(l) = limit {
5655            return e.swiglu_clamped_mul_scaled(gate, up, gs, us, l, act, n);
5656        }
5657        if gs == 1.0 && us == 1.0 {
5658            return e.silu_mul(gate, up, act, n);
5659        }
5660        e.silu_mul_scaled(gate, up, gs, us, act, n)
5661    }
5662
5663    /// Routing for the whole batch: returns (sel [T*n_used] expert ids, w [T*n_used] renorm weights),
5664    /// token-major. Default = the Stage-1 host path (dtoh logits, softmax-256, stable DESC top-k,
5665    /// renorm). MEMRA_FUSED_ROUTER = the device kernel (§A) which reproduces the same numerics; we
5666    /// still dtoh the tiny [T,n_used] sel/w buffers (64 B/token vs 1 KB/token) — the host loop
5667    /// indexes HostExps.bytes on the CPU to choose the DMA source (§A.2 output staging).
5668    fn moe_route(
5669        e: &Engine,
5670        logits: &CudaSlice<f32>,
5671        t: usize,
5672        n_expert: usize,
5673        n_used: usize,
5674    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
5675        Self::moe_route_cfg(e, logits, t, n_expert, n_used, None)
5676    }
5677
5678    /// DeepSeek-V3-class sigmoid routing (Step-3.7, MiniMax-M3, Hy3, GLM-DSA). Reference:
5679    /// M3/Hy3 modeling code — scores = sigmoid(logits); selection over scores + expert bias
5680    /// (M3 `e_score_correction_bias` / Hy3 `expert_bias`, both surfaced as `exp_probs_b`);
5681    /// weights = un-biased scores of the selected experts, sum-normalized when `route_norm`,
5682    /// x scaling factor (M3 routed_scaling_factor 2.0 / Hy3 router_scaling_factor 2.826).
5683    /// Default uses one device kernel plus the same pinned sel/w readback contract as the fused
5684    /// softmax router. `MEMRA_SIG_ROUTER=0` restores the full-logit DtoH host oracle.
5685    #[allow(clippy::too_many_arguments)]
5686    fn moe_route_sigmoid_cfg(
5687        e: &Engine,
5688        logits: &CudaSlice<f32>,
5689        t: usize,
5690        n_expert: usize,
5691        n_used: usize,
5692        m: &MoeWeights,
5693        (sf, route_norm): (f32, bool),
5694    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
5695        if sigmoid_router_enabled() {
5696            return e.moe_router_sigmoid_topk_host(
5697                logits,
5698                t,
5699                n_expert,
5700                n_used,
5701                m.active_count(),
5702                &m.exp_probs_b_dev,
5703                &m.active_experts_dev,
5704                sf,
5705                route_norm,
5706            );
5707        }
5708        let lg = e.dtoh(logits)?;
5709        Self::moe_route_sigmoid_host(
5710            &lg,
5711            t,
5712            n_expert,
5713            n_used,
5714            m.exp_probs_b.as_deref(),
5715            sf,
5716            route_norm,
5717            m.active_experts.as_deref(),
5718        )
5719    }
5720
5721    /// Softmax router for qwen35moe/OLMoE. An active-mask overlay retains the host oracle because
5722    /// the existing softmax device kernel has no mask input.
5723    fn moe_route_cfg(
5724        e: &Engine,
5725        logits: &CudaSlice<f32>,
5726        t: usize,
5727        n_expert: usize,
5728        n_used: usize,
5729        active: Option<&[bool]>,
5730    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
5731        // LAUNCH-STRUCTURE STAGE 1 (2026-07-05): fused router DEFAULT ON (MEMRA_FUSED_ROUTER=0
5732        // rollback) via the single-sync pinned readback — softmax arch only.
5733        if active.is_none() && !matches!(std::env::var("MEMRA_FUSED_ROUTER").as_deref(), Ok("0")) {
5734            return e.moe_router_topk_host(logits, t, n_expert, n_used);
5735        }
5736        // Host oracle (the §D bit-identity reference).
5737        let lg = e.dtoh(logits)?; // [T*n_expert] host
5738        let mut sel = vec![0u32; t * n_used];
5739        let mut w_out = vec![0f32; t * n_used];
5740        for tok in 0..t {
5741            let row = &lg[tok * n_expert..(tok + 1) * n_expert];
5742            // softmax over ALL n_expert (stable: subtract max)
5743            let maxl = row
5744                .iter()
5745                .enumerate()
5746                .filter(|(i, _)| active.is_none_or(|mask| mask[*i]))
5747                .map(|(_, &x)| x)
5748                .fold(f32::NEG_INFINITY, f32::max);
5749            let mut probs = vec![0f32; n_expert];
5750            let mut den = 0f32;
5751            for i in 0..n_expert {
5752                if active.is_some_and(|mask| !mask[i]) {
5753                    continue;
5754                }
5755                let x = (row[i] - maxl).exp();
5756                probs[i] = x;
5757                den += x;
5758            }
5759            for p in probs.iter_mut() {
5760                *p /= den;
5761            }
5762            // stable DESC sort: prob DESC, ascending-index tiebreak.
5763            let mut idx: Vec<usize> = (0..n_expert)
5764                .filter(|&i| active.is_none_or(|mask| mask[i]))
5765                .collect();
5766            idx.sort_by(|&a, &b| probs[b].total_cmp(&probs[a]).then(a.cmp(&b)));
5767            let sl = &idx[..n_used];
5768            let mut wv: Vec<f32> = sl.iter().map(|&i| probs[i]).collect();
5769            let mut ws: f32 = wv.iter().sum();
5770            ws = ws.max(6.103515625e-5_f32); // F16 smallest normal, clamp BEFORE divide
5771            for x in wv.iter_mut() {
5772                *x /= ws;
5773            }
5774            for j in 0..n_used {
5775                sel[tok * n_used + j] = sl[j] as u32;
5776                w_out[tok * n_used + j] = wv[j];
5777            }
5778        }
5779        Ok((sel, w_out))
5780    }
5781
5782    #[allow(clippy::too_many_arguments)]
5783    fn moe_route_sigmoid_with_input(
5784        e: &Engine,
5785        logits: &CudaSlice<f32>,
5786        input: &CudaSlice<f32>,
5787        t: usize,
5788        n_expert: usize,
5789        n_used: usize,
5790        bias: Option<&[f32]>,
5791        (sf, route_norm): (f32, bool),
5792        active: Option<&[bool]>,
5793    ) -> Result<(Vec<u32>, Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
5794        let (lg, input) = e.dtoh_pair(logits, input)?;
5795        let (sel, w) =
5796            Self::moe_route_sigmoid_host(&lg, t, n_expert, n_used, bias, sf, route_norm, active)?;
5797        Ok((sel, w, input))
5798    }
5799
5800    /// Start the prediction-guided prefetch worker (MEMRA_MOE_PREFETCH=depth). Call after
5801    /// residency freeze: the worker filters against a static snapshot of the frozen HBM set.
5802    /// Builds a fully-owned per-layer table (host router copies via one-time DtoH, bias,
5803    /// active mask, prebuilt projection descriptors) so no model reference escapes.
5804    pub fn start_moe_prefetch_predictor(
5805        &self,
5806        e: &Engine,
5807        cfg: &ModelConfig,
5808    ) -> Result<(), Box<dyn std::error::Error>> {
5809        use crate::hybrid::Ffn;
5810        let Some(sig) = cfg.sigmoid_router() else {
5811            return Err("prefetch predictor requires a sigmoid-router arch".into());
5812        };
5813        let resident: std::collections::HashSet<(u16, u8, u16)> = e
5814            .export_moe_residency()
5815            .ok_or("prefetch predictor needs the frozen MoE residency cache")?
5816            .into_iter()
5817            .collect();
5818        let mut layers = Vec::new();
5819        for (index, layer) in self.layers.iter().enumerate() {
5820            let Ffn::Moe(m) = &layer.ffn else { continue };
5821            let crate::model::GpuTensor::Float { data, .. } = &m.gate_inp else {
5822                continue;
5823            };
5824            let router = e.dtoh(data)?;
5825            let n_expert = m.gate_exps.n_expert;
5826            let n_embd = m.gate_exps.in_f;
5827            if router.len() != n_embd * n_expert {
5828                continue;
5829            }
5830            let build = |exps: &crate::model::HostExps| {
5831                (0..n_expert)
5832                    .map(|expert| crate::cpu_experts::predictor_projection(exps, expert))
5833                    .collect::<Vec<_>>()
5834            };
5835            layers.push((
5836                index as u16,
5837                crate::cpu_experts::PredictLayerInit {
5838                    router,
5839                    bias: m.exp_probs_b.clone(),
5840                    active: m.active_experts.clone(),
5841                    n_embd,
5842                    n_used: cfg
5843                        .moe
5844                        .as_ref()
5845                        .map(|moe| moe.expert_used_count as usize)
5846                        .ok_or("prefetch predictor requires MoE config")?,
5847                    sig,
5848                    weights_n_expert: n_expert,
5849                    gate: build(&m.gate_exps),
5850                    up: build(&m.up_exps),
5851                    down: build(&m.down_exps),
5852                },
5853            ));
5854        }
5855        crate::cpu_experts::start_prefetch_predictor(layers, resident).map_err(|error| error.into())
5856    }
5857
5858    /// Sigmoid-routing oracle shared by the prefetch predictor and `kernel-check`: identical
5859    /// selection math to the rollback runtime, applied to host-computed logits.
5860    #[allow(clippy::too_many_arguments)]
5861    pub fn moe_route_sigmoid_host_public(
5862        logits: &[f32],
5863        t: usize,
5864        n_expert: usize,
5865        n_used: usize,
5866        bias: Option<&[f32]>,
5867        sf: f32,
5868        route_norm: bool,
5869        active: Option<&[bool]>,
5870    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
5871        Self::moe_route_sigmoid_host(logits, t, n_expert, n_used, bias, sf, route_norm, active)
5872    }
5873
5874    #[allow(clippy::too_many_arguments)]
5875    fn moe_route_sigmoid_host(
5876        lg: &[f32],
5877        t: usize,
5878        n_expert: usize,
5879        n_used: usize,
5880        bias: Option<&[f32]>,
5881        sf: f32,
5882        route_norm: bool,
5883        active: Option<&[bool]>,
5884    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
5885        let active_count = active
5886            .map(|mask| mask.iter().filter(|&&enabled| enabled).count())
5887            .unwrap_or(n_expert);
5888        crate::sigrouter_contract::validate_active_count(n_used, active_count)?;
5889        if lg.len() != t * n_expert {
5890            return Err(format!(
5891                "sigmoid router logits length mismatch: got {}, expected {}",
5892                lg.len(),
5893                t * n_expert,
5894            )
5895            .into());
5896        }
5897        let mut sel = vec![0u32; t * n_used];
5898        let mut w_out = vec![0f32; t * n_used];
5899        for tok in 0..t {
5900            let row = &lg[tok * n_expert..(tok + 1) * n_expert];
5901            let scores: Vec<f32> = row.iter().map(|&x| 1.0 / (1.0 + (-x).exp())).collect();
5902            // selection score = sigmoid + bias; weight = plain sigmoid.
5903            let selsc: Vec<f32> = match bias {
5904                Some(b) => scores.iter().zip(b).map(|(s, bb)| s + bb).collect(),
5905                None => scores.clone(),
5906            };
5907            let mut idx: Vec<usize> = (0..n_expert)
5908                .filter(|&i| active.is_none_or(|mask| mask[i]))
5909                .collect();
5910            idx.sort_by(|&a, &b| selsc[b].total_cmp(&selsc[a]).then(a.cmp(&b)));
5911            let sl = &idx[..n_used];
5912            let mut wv: Vec<f32> = sl.iter().map(|&i| scores[i]).collect();
5913            if route_norm {
5914                let ws: f32 = wv.iter().sum::<f32>().max(1e-20);
5915                for x in wv.iter_mut() {
5916                    *x = *x / ws * sf;
5917                }
5918            } else {
5919                for x in wv.iter_mut() {
5920                    *x *= sf;
5921                }
5922            }
5923            for j in 0..n_used {
5924                sel[tok * n_used + j] = sl[j] as u32;
5925                w_out[tok * n_used + j] = wv[j];
5926            }
5927        }
5928        Ok((sel, w_out))
5929    }
5930
5931    /// Step-3.7 resident-slab dispatch: sigmoid top-k outputs remain on device and feed the
5932    /// existing slot-ordered q8 expert kernels directly. Mixed layouts, spill, remote slabs,
5933    /// macro-scaled experts, and observation modes are denied by the caller.
5934    #[allow(clippy::too_many_arguments)]
5935    fn moe_ffn_sigmoid_dev(
5936        e: &Engine,
5937        m: &MoeWeights,
5938        z: &CudaSlice<f32>,
5939        zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
5940        logits: &CudaSlice<f32>,
5941        t: usize,
5942        cfg: &ModelConfig,
5943        il: u16,
5944        (scaling_factor, route_norm): (f32, bool),
5945    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5946        let moe = cfg.moe.as_ref().unwrap();
5947        let n_embd = cfg.n_embd as usize;
5948        let n_expert = moe.expert_count as usize;
5949        let n_used = moe.expert_used_count as usize;
5950        let n_ff_exp = moe.expert_ff_length as usize;
5951        let dev = m.dev_exps.as_ref().unwrap();
5952        debug_assert!(cfg.step35.is_some());
5953        debug_assert_eq!(dev.dev, e.ctx().ordinal());
5954        debug_assert!(m.has_uniform_expert_layout());
5955        debug_assert!(!m.has_macros);
5956
5957        let (sel_d, w_d) = e.moe_router_sigmoid_topk(
5958            logits,
5959            t,
5960            n_expert,
5961            n_used,
5962            m.active_count(),
5963            &m.exp_probs_b_dev,
5964            &m.active_experts_dev,
5965            scaling_factor,
5966            route_norm,
5967        )?;
5968        crate::moesd::record_device_routes(e, il, n_expert, n_used, &sel_d)?;
5969        let (gate_row_bytes, up_row_bytes) = if dev.gu_il {
5970            let combined = m.gate_exps.row_bytes + m.up_exps.row_bytes;
5971            (combined, combined)
5972        } else {
5973            (m.gate_exps.row_bytes, m.up_exps.row_bytes)
5974        };
5975        let (zq, zd) = match (t, zq8) {
5976            (1, Some((q, d))) => (q.clone(), d.clone()),
5977            _ => e.quantize_q8_1(z, t, n_embd)?,
5978        };
5979        let n_pairs = t * n_used;
5980        let mut moe_out = if cfg.clamp_exp_at(il as u32).is_some() {
5981            // The final Step layers retain the established separate gate/up -> clamp -> down
5982            // arithmetic. Pair rows are derived from token position; selected expert ids and
5983            // routing weights remain the device router's buffers throughout.
5984            let pair_tok: Vec<i32> = (0..n_pairs).map(|pair| (pair / n_used) as i32).collect();
5985            let pair_tok_d = e.htod_i32(&pair_tok)?;
5986            let gate = e.moe_pairs_matvec_q8(
5987                &dev.ptr_row,
5988                0,
5989                &pair_tok_d,
5990                &sel_d,
5991                &zq,
5992                &zd,
5993                n_embd,
5994                n_ff_exp,
5995                n_expert,
5996                n_pairs,
5997                m.gate_exps.qtype,
5998                gate_row_bytes,
5999            )?;
6000            let up = e.moe_pairs_matvec_q8(
6001                &dev.ptr_row,
6002                1,
6003                &pair_tok_d,
6004                &sel_d,
6005                &zq,
6006                &zd,
6007                n_embd,
6008                n_ff_exp,
6009                n_expert,
6010                n_pairs,
6011                m.up_exps.qtype,
6012                up_row_bytes,
6013            )?;
6014            let mut act = e.uninit(n_pairs * n_ff_exp)?;
6015            Self::ffn_act_lim(
6016                e,
6017                cfg,
6018                &gate,
6019                &up,
6020                1.0,
6021                1.0,
6022                cfg.clamp_exp_at(il as u32),
6023                &mut act,
6024                n_pairs * n_ff_exp,
6025            )?;
6026            let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
6027            let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
6028            let pair_self_d = e.htod_i32(&pair_self)?;
6029            let down = e.moe_pairs_matvec_q8(
6030                &dev.ptr_row,
6031                2,
6032                &pair_self_d,
6033                &sel_d,
6034                &aq2,
6035                &ad2,
6036                n_ff_exp,
6037                n_embd,
6038                n_expert,
6039                n_pairs,
6040                m.down_exps.qtype,
6041                m.down_exps.row_bytes,
6042            )?;
6043            let tok_off: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
6044            let tok_ids: Vec<i32> = (0..n_pairs as i32).collect();
6045            let tok_off_d = e.htod_i32(&tok_off)?;
6046            let tok_ids_d = e.htod_i32(&tok_ids)?;
6047            let mut output = e.uninit(t * n_embd)?;
6048            e.moe_pairs_scatter(&down, &w_d, &tok_off_d, &tok_ids_d, &mut output, t, n_embd)?;
6049            output
6050        } else {
6051            let act = e.moe_gate_up_silu8_dev_q8_rows(
6052                &dev.ptr_row,
6053                &sel_d,
6054                &zq,
6055                &zd,
6056                t,
6057                n_embd,
6058                n_ff_exp,
6059                n_used,
6060                n_expert,
6061                m.gate_exps.qtype,
6062                m.up_exps.qtype,
6063                gate_row_bytes,
6064                up_row_bytes,
6065                &m.dev_macros,
6066            )?;
6067            let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
6068            let mut output = e.uninit(t * n_embd)?;
6069            e.moe_down8_fma_dev_q8_rows_g(
6070                &dev.ptr_row,
6071                &sel_d,
6072                &w_d,
6073                &aq2,
6074                &ad2,
6075                &mut output,
6076                t,
6077                n_ff_exp,
6078                n_embd,
6079                n_used,
6080                n_expert,
6081                m.down_exps.qtype,
6082                m.down_exps.row_bytes,
6083            )?;
6084            output
6085        };
6086
6087        if std::env::var("MEMRA_SIG_ROUTER_DISPATCH_TRACE").as_deref() == Ok("1") {
6088            eprintln!(
6089                "[sigrouter-dev] layer={il} tokens={t} experts={n_expert} used={n_used} clamp={} gu_il={}",
6090                cfg.clamp_exp_at(il as u32).is_some(),
6091                dev.gu_il,
6092            );
6093        }
6094        Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut moe_out)?;
6095        Ok(moe_out)
6096    }
6097
6098    /// LAUNCH-STRUCTURE STAGE 3: the ZERO-DtoH fully-resident MoE FFN. Caller guarantees the
6099    /// layer's device pointer row exists (checked under the cache lock). Router top-k runs on
6100    /// device; sel/w are consumed by the `_dev` matvec twins directly; NOTHING crosses PCIe.
6101    /// Same numerics as the fused-router + gdec chain (kernel-level bit-identity, see the
6102    /// MoE PREFILL PAIR-BATCH: host routing (sel/w like the sequential path), then 5 launches
6103    /// TOTAL per layer (quantize z, gate-pairs, up-pairs, silu, act-quantize, down-pairs,
6104    /// scatter) regardless of T or expert count. Bit-identity class: per (pair,row) dot =
6105    /// qmatvec_expert_q8 order; per-token accumulation slot-ordered (scatter kernel).
6106    fn moe_ffn_pairs(
6107        e: &Engine,
6108        m: &MoeWeights,
6109        z: &CudaSlice<f32>,
6110        logits: &CudaSlice<f32>,
6111        t: usize,
6112        cfg: &ModelConfig,
6113    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6114        let moe = cfg.moe.as_ref().unwrap();
6115        let n_embd = cfg.n_embd as usize;
6116        let n_expert = moe.expert_count as usize;
6117        let n_used = moe.expert_used_count as usize;
6118        let n_ff_exp = moe.expert_ff_length as usize;
6119        // This arm has no `il` and every kernel in it fuses PLAIN silu(gate)*up, so a clamped
6120        // model must never reach it. The caller's gate at `moe_ffn_sequential_zq8` denies per
6121        // layer via `swiglu_clamped_at(il)`; assert the whole-model form here so a future caller
6122        // that forgets the gate fails loudly in debug instead of returning wrong logits.
6123        debug_assert!(
6124            !cfg.swiglu_clamped_anywhere(),
6125            "moe_ffn_pairs has no per-layer clamp: fused epilogues are plain SiLU"
6126        );
6127        let dev = m.dev_exps.as_ref().unwrap();
6128        // WALL-GAP ARC: interleaved gate/up slab strides (see moe_ffn_dev).
6129        let (rbg_d, rbu_d) = if dev.gu_il {
6130            let sxx = m.gate_exps.row_bytes + m.up_exps.row_bytes;
6131            (sxx, sxx)
6132        } else {
6133            (m.gate_exps.row_bytes, m.up_exps.row_bytes)
6134        };
6135
6136        let (sel_all, w_all) = Self::moe_route(e, logits, t, n_expert, n_used)?;
6137        let n_pairs = t * n_used;
6138        // pair arrays: pair p = (token p/n_used, slot p%n_used) — ALREADY slot-ordered per token,
6139        // so the CSR is trivial: tok_pair_off[tok] = tok*n_used, ids identity.
6140        let pair_tok: Vec<i32> = (0..n_pairs).map(|p| (p / n_used) as i32).collect();
6141        let pair_ex: Vec<i32> = sel_all.iter().map(|&x| x as i32).collect();
6142        let pair_w: Vec<f32> = w_all.clone();
6143        let tok_off: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
6144        let tok_ids: Vec<i32> = (0..n_pairs as i32).collect();
6145        let pt = e.htod_i32(&pair_tok)?;
6146        let px = e.htod_i32(&pair_ex)?;
6147        let pw = e.htod(&pair_w)?;
6148        let toff = e.htod_i32(&tok_off)?;
6149        let tids = e.htod_i32(&tok_ids)?;
6150
6151        // z quantized ONCE for all tokens; gate/up pair matvecs; silu; act quantize; down; scatter.
6152        // EXPERT-MAJOR CSR (rung 2): pairs grouped by expert -> the kernel reuses each weight
6153        // row across the expert's token group (llama-MMQ's core win). Host grouping is O(pairs).
6154        let mut by_ex: Vec<Vec<i32>> = vec![Vec::new(); n_expert];
6155        for p in 0..n_pairs {
6156            by_ex[pair_ex[p] as usize].push(p as i32);
6157        }
6158        let mut ex_ids: Vec<i32> = Vec::new();
6159        let mut ex_off: Vec<i32> = vec![0];
6160        let mut ex_pairs: Vec<i32> = Vec::with_capacity(n_pairs);
6161        for (ex, list) in by_ex.iter().enumerate() {
6162            if list.is_empty() {
6163                continue;
6164            }
6165            ex_ids.push(ex as i32);
6166            ex_pairs.extend_from_slice(list);
6167            ex_off.push(ex_pairs.len() as i32);
6168        }
6169        let n_active = ex_ids.len();
6170        let exi = e.htod_i32(&ex_ids)?;
6171        let exo = e.htod_i32(&ex_off)?;
6172        let exp_d = e.htod_i32(&ex_pairs)?;
6173        let _ = &px; // pair-major twin keeps it; em path uses CSR
6174
6175        // INT8-MMA EXPERT MMQ (MEMRA_MOE_MMA=1, opt-in): the m16n8k16.s8 tensor-core analog of the
6176        // _dec dp4a kernel (cu/mmq_iq_experts.cu). Same CSR grouping; per-expert matvec runs as a
6177        // 128x128-tile int8 MMA GEMM over the expert's token group. Weight IQ nibbles decode to int8
6178        // at tile-load + per-32 float scale; activation is q8_1_mmq (D4, same quant class as dp4a).
6179        // FP-ORDER differs from dp4a (MMA reduction) — logits SHIFT, gated on argmax/spec/closeness,
6180        // NOT byte-identity (like the W4A8 path). Requires IQ3_S/IQ4_XS + in_f % 256 == 0.
6181        // t >= 16 (GEMM_M-class rule): the MMA tile needs token volume (crossover ~200 tok/expert;
6182        // microbench: dp4a wins at tiny groups). ALSO an exactness requirement — spec verify
6183        // batches (t=2..K+2) must ride the dp4a path whose FP order matches the T=1 decode chain,
6184        // else K=1 self-consistency FAILs (caught 2026-07-06: MMA at T=2 flipped a verify argmax).
6185        // DEFAULT ON (2026-07-06, third flip — this time with the real culprit fixed): the
6186        // "MMA prime breaks spec" failure was the ROUTER's cuBLASLt n-dependence (d994271),
6187        // not MMA's own FP order — both this and the k-quant arms were innocent suspects whose
6188        // margin shifts surfaced the router bug. With the router decode-exact at verify t, the
6189        // full battery is green with MMA on (spec p1/p2/p3 PASS, raw K=1..8 PASS, argmax MATCH,
6190        // pp6257 2862 = 2.1x dec). t>=16 floor still required: verify batches must ride dp4a
6191        // (dispatch parity with the T=1 decode chain). MEMRA_MOE_MMA=0 rollback;
6192        // MEMRA_MOE_MMA_T overrides the floor (bisect seam).
6193        static MMA_T: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
6194        let mma_t = *MMA_T.get_or_init(|| {
6195            std::env::var("MEMRA_MOE_MMA_T")
6196                .ok()
6197                .and_then(|v| v.parse().ok())
6198                .unwrap_or(16)
6199        });
6200        let use_mma = std::env::var("MEMRA_MOE_MMA")
6201            .map(|v| v != "0")
6202            .unwrap_or(true)
6203            && t >= mma_t
6204            && q8_expert_dec_supported(m.gate_exps.qtype)
6205            && q8_expert_dec_supported(m.up_exps.qtype)
6206            && q8_expert_dec_supported(m.down_exps.qtype)
6207            && n_embd % 256 == 0
6208            && n_ff_exp % 256 == 0;
6209        // GROUPED f16 LANE admission (MEMRA_MOE_F16G, rounds 46-49): its own door, no longer a
6210        // subset of the MMQ arm — q35's k-quant stragglers (Q6_K/Q4_K down, one Q3_K gate/up
6211        // layer) fail q8_expert_dec_supported but dequant fine to f16, so f16g must be able to
6212        // take a layer the MMQ arm would reject. Same t >= mma_t floor as MMA: decode and
6213        // spec-verify batches must ride the dp4a path whose FP order matches the T=1 chain.
6214        // MODE 2 (sm_120a naked default, lane/f16g-default-rearb 2026-08-02): every layer
6215        // whose three projections pass f16g_proj_ok rides the sk visitor with direct tile
6216        // loaders — with IQ4_XS/IQ3_S direct coverage the sk arm beats the int8-MMA MMQ
6217        // tiles on the IQ-bank models too (q35 board-2048 +33.9%, KAT pp512 +46.7%).
6218        // AUTO-KQUANT (mode 3, MEMRA_MOE_F16G=3, lane/q4k-expert-prefill): admit f16g ONLY
6219        // where the MMA arm can't take the layer's QTYPES — the k-quant expert class whose
6220        // baseline is the per-pair _em fallback (Ornith-35B Q4_K: board-2048 3.14x). Its
6221        // MMA-capable carve-out (IQ3_S/IQ4_XS/Q4_0 banks on MMQ) was priced pre-IQ-direct;
6222        // it survives as the rollback seam. Keyed on qtype capability, NOT use_mma, so
6223        // MEMRA_MOE_MMA=0 stays a pure dp4a rollback seam.
6224        let mma_capable = q8_expert_dec_supported(m.gate_exps.qtype)
6225            && q8_expert_dec_supported(m.up_exps.qtype)
6226            && q8_expert_dec_supported(m.down_exps.qtype)
6227            && n_embd % 256 == 0
6228            && n_ff_exp % 256 == 0;
6229        let f16g_mode = crate::moe_f16g_mode();
6230        let f16g = f16g_mode != 0
6231            && t >= mma_t
6232            && (f16g_mode != 3 || !mma_capable)
6233            && f16g_proj_ok(m.gate_exps.qtype, n_embd)
6234            && f16g_proj_ok(m.up_exps.qtype, n_embd)
6235            && f16g_proj_ok(m.down_exps.qtype, n_ff_exp);
6236        if use_mma || f16g {
6237            // GROUPED f16 LANE (MEMRA_MOE_F16G, rounds 46-49, experimental door): dequant
6238            // the active experts ONCE per projection to f16 and run one grouped f16 GEMM over
6239            // the CSR groups. f16-mirror numeric class — argmax/spec gated before promotion.
6240            // NOTE: the pair-gather kernel needs pair p ordered EXPERT-MAJOR (ex_pairs order);
6241            // the y rows come back in that same CSR order, so gate/up/down all stay pair-major
6242            // in ex_pairs order — but moe_pairs_silu_mul and the scatter consume PAIR-ID order.
6243            // We therefore gather activations per ex_pairs and scatter y back through ex_pairs.
6244            let y_down = if f16g {
6245                // CSR order end-to-end: gather z rows by the pair's TOKEN (pair p's token is
6246                // p / n_used — the trivial CSR above), silu in CSR order (elementwise), one
6247                // permute at the very end back to pair-id order for the scatter.
6248                let csr_tok: Vec<i32> = ex_pairs.iter().map(|&p| p / n_used as i32).collect();
6249                let csr_tok_d = e.htod_i32(&csr_tok)?;
6250                let (z_f16, z_s) = e.moe_f16g_act(z, Some(&csr_tok_d), n_embd, n_pairs)?;
6251                let g_csr = e.moe_f16_grouped(
6252                    &dev.ptr_row,
6253                    0,
6254                    n_expert,
6255                    &exi,
6256                    &ex_off,
6257                    &exo,
6258                    &z_f16,
6259                    &z_s,
6260                    n_embd,
6261                    n_ff_exp,
6262                    n_active,
6263                    n_pairs,
6264                    m.gate_exps.qtype,
6265                    rbg_d,
6266                )?;
6267                let u_csr = e.moe_f16_grouped(
6268                    &dev.ptr_row,
6269                    1,
6270                    n_expert,
6271                    &exi,
6272                    &ex_off,
6273                    &exo,
6274                    &z_f16,
6275                    &z_s,
6276                    n_embd,
6277                    n_ff_exp,
6278                    n_active,
6279                    n_pairs,
6280                    m.up_exps.qtype,
6281                    rbu_d,
6282                )?;
6283                let act_csr = e.moe_pairs_silu_mul(&g_csr, &u_csr, n_pairs * n_ff_exp)?;
6284                let (a_f16, a_s) = e.moe_f16g_act(&act_csr, None, n_ff_exp, n_pairs)?;
6285                let d_csr = e.moe_f16_grouped(
6286                    &dev.ptr_row,
6287                    2,
6288                    n_expert,
6289                    &exi,
6290                    &ex_off,
6291                    &exo,
6292                    &a_f16,
6293                    &a_s,
6294                    n_ff_exp,
6295                    n_embd,
6296                    n_active,
6297                    n_pairs,
6298                    m.down_exps.qtype,
6299                    m.down_exps.row_bytes,
6300                )?;
6301                e.rows_permute(&d_csr, &exp_d, n_pairs, n_embd)?
6302            } else {
6303                // gate/up: activation = z, token-major over t tokens; pair_tok gathers the routed row.
6304                let z_scr = e.mmq_iq_quantize_act(z, n_embd, t)?;
6305                let gate = e.mmq_iq_experts(
6306                    &dev.ptr_row,
6307                    0,
6308                    n_expert,
6309                    &exi,
6310                    &exo,
6311                    &exp_d,
6312                    &pt,
6313                    &z_scr,
6314                    n_embd,
6315                    n_ff_exp,
6316                    n_active,
6317                    n_pairs,
6318                    t,
6319                    m.gate_exps.qtype,
6320                    rbg_d,
6321                )?;
6322                let up = e.mmq_iq_experts(
6323                    &dev.ptr_row,
6324                    1,
6325                    n_expert,
6326                    &exi,
6327                    &exo,
6328                    &exp_d,
6329                    &pt,
6330                    &z_scr,
6331                    n_embd,
6332                    n_ff_exp,
6333                    n_active,
6334                    n_pairs,
6335                    t,
6336                    m.up_exps.qtype,
6337                    rbu_d,
6338                )?;
6339                // down: activation = silu(gate)*up, pair-major [n_pairs, n_ff_exp]; pair_tok =
6340                // identity. FUSED ACT-EPILOGUE (default on): one launch computes the activation in
6341                // registers and writes ONLY the quantized scratch — the two-pass chain
6342                // (moe_pairs_silu_mul writes act f32, mmq_iq_quantize_act re-reads it) is the
6343                // MEMRA_MOE_FUSE_ACTQ=0 rollback. Scratch bytes are BYTE-IDENTICAL (kernel-check).
6344                let a_scr = if crate::moe_fuse_actq_on() {
6345                    e.mmq_iq_fused_act_quant(&gate, &up, n_ff_exp, n_pairs, 0)?
6346                } else {
6347                    let act = e.moe_pairs_silu_mul(&gate, &up, n_pairs * n_ff_exp)?;
6348                    e.mmq_iq_quantize_act(&act, n_ff_exp, n_pairs)?
6349                };
6350                let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
6351                let pself = e.htod_i32(&pair_self)?;
6352                e.mmq_iq_experts(
6353                    &dev.ptr_row,
6354                    2,
6355                    n_expert,
6356                    &exi,
6357                    &exo,
6358                    &exp_d,
6359                    &pself,
6360                    &a_scr,
6361                    n_ff_exp,
6362                    n_embd,
6363                    n_active,
6364                    n_pairs,
6365                    n_pairs,
6366                    m.down_exps.qtype,
6367                    m.down_exps.row_bytes,
6368                )?
6369            };
6370            let mut moe_out = e.uninit(t * n_embd)?;
6371            e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
6372            if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
6373                (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
6374            {
6375                let n_ff_sh = gate_shexp.out_features();
6376                let sg_gate = e.matmul(gate_shexp, z, t)?;
6377                let sg_up = e.matmul(up_shexp, z, t)?;
6378                let mut sa = e.uninit(t * n_ff_sh)?;
6379                Self::ffn_act(e, cfg, &sg_gate, &sg_up, &mut sa, t * n_ff_sh)?;
6380                let sh = e.matmul(down_shexp, &sa, t)?;
6381                // shexp gate: qwen35moe sigmoid-gates; M3 has no gate tensor -> weight 1.0.
6382                // SERVE ISOLATION (lane/concat-prime-exact): the fused sigmoid-dot is the
6383                // m-INVARIANT form — see the sequential arm's note. This is the PAIRS arm,
6384                // i.e. the one real prefill actually takes on a resident-expert MoE model,
6385                // so the concat-prime isolation fix has to land here as well.
6386                let g = match &m.gate_inp_shexp {
6387                    Some(gate_inp_shexp) if crate::router_prefill_exact_on() => {
6388                        e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
6389                    }
6390                    Some(gate_inp_shexp) => {
6391                        let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
6392                        let mut g = e.uninit(t)?;
6393                        e.sigmoid(&gs, &mut g, t)?;
6394                        g
6395                    }
6396                    None => e.htod(&vec![1.0f32; t])?,
6397                };
6398                e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
6399            }
6400            return Ok(moe_out);
6401        }
6402
6403        // DECODE-ONCE MMQ (rung 3, MEMRA_MOE_DEC=1 default-on): dequant each weight group once per
6404        // (row,group) then dp4a across the expert's tokens. _em re-decoded per token (NEUTRAL).
6405        let dec = std::env::var("MEMRA_MOE_DEC")
6406            .map(|v| v != "0")
6407            .unwrap_or(true);
6408        let matvec = |proj,
6409                      exi: &_,
6410                      exo: &_,
6411                      exp_d: &_,
6412                      pt: &_,
6413                      aq: &_,
6414                      ad: &_,
6415                      inf,
6416                      outf,
6417                      qtype,
6418                      rb|
6419         -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6420            // _dec's decode-once extractors are IQ-only; k-quant expert layers take the _em dot path.
6421            let dec = dec && q8_expert_dec_supported(qtype);
6422            if dec {
6423                e.moe_pairs_matvec_q8_dec(
6424                    &dev.ptr_row,
6425                    proj,
6426                    exi,
6427                    exo,
6428                    exp_d,
6429                    pt,
6430                    aq,
6431                    ad,
6432                    inf,
6433                    outf,
6434                    n_expert,
6435                    n_active,
6436                    n_pairs,
6437                    qtype,
6438                    rb,
6439                )
6440            } else {
6441                e.moe_pairs_matvec_q8_em(
6442                    &dev.ptr_row,
6443                    proj,
6444                    exi,
6445                    exo,
6446                    exp_d,
6447                    pt,
6448                    aq,
6449                    ad,
6450                    inf,
6451                    outf,
6452                    n_expert,
6453                    n_active,
6454                    n_pairs,
6455                    qtype,
6456                    rb,
6457                )
6458            }
6459        };
6460        let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
6461        let gate = matvec(
6462            0,
6463            &exi,
6464            &exo,
6465            &exp_d,
6466            &pt,
6467            &zq,
6468            &zd,
6469            n_embd,
6470            n_ff_exp,
6471            m.gate_exps.qtype,
6472            rbg_d,
6473        )?;
6474        let up = matvec(
6475            1,
6476            &exi,
6477            &exo,
6478            &exp_d,
6479            &pt,
6480            &zq,
6481            &zd,
6482            n_embd,
6483            n_ff_exp,
6484            m.up_exps.qtype,
6485            rbu_d,
6486        )?;
6487        let act = e.moe_pairs_silu_mul(&gate, &up, n_pairs * n_ff_exp)?;
6488        let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
6489        // down consumes PAIR-major activation rows: pair_tok = identity.
6490        let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
6491        let pself = e.htod_i32(&pair_self)?;
6492        let y_down = matvec(
6493            2,
6494            &exi,
6495            &exo,
6496            &exp_d,
6497            &pself,
6498            &aq2,
6499            &ad2,
6500            n_ff_exp,
6501            n_embd,
6502            m.down_exps.qtype,
6503            m.down_exps.row_bytes,
6504        )?;
6505        let mut moe_out = e.uninit(t * n_embd)?; // scatter fully overwrites per (token,col)
6506        e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
6507
6508        // SHARED EXPERT epilogue — same as the other paths.
6509        // gate_inp_shexp is OPTIONAL: qwen35moe gates the shared expert (sigmoid(gate_inp) x sh);
6510        // MiniMax-M3 (DeepSeek-V3 class) has NO shexp gate — the shared expert adds directly.
6511        if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
6512            (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
6513        {
6514            let n_ff_sh = gate_shexp.out_features();
6515            // These decode-exact forms are required by the new Step resident arm. Keep the
6516            // established grouped shared-expert program for every other architecture: widening
6517            // this to Gemma changed its speculative acceptance despite green argmax gates.
6518            let step_exact = cfg.step35.is_some();
6519            let verify_t = step_exact && t > 1 && t < PRIME_MIN_T;
6520            let (sg_gate, sg_up) = if step_exact && t == 1 {
6521                match e.matmul_q8_fused2_x(gate_shexp, up_shexp, z)? {
6522                    Some(pair) => pair,
6523                    None => (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?),
6524                }
6525            } else if verify_t {
6526                let mut fused = None;
6527                if crate::spec::spec_fused_t()
6528                    && (2..=4).contains(&t)
6529                    && e.uses_q8_1_fast(gate_shexp)
6530                    && e.uses_q8_1_fast(up_shexp)
6531                {
6532                    let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
6533                    fused = e.matmul_q8_fused2_t(gate_shexp, up_shexp, &zq, &zd, t)?;
6534                }
6535                match fused {
6536                    Some(pair) => pair,
6537                    None => (
6538                        e.matmul_decode_exact(gate_shexp, z, t)?,
6539                        e.matmul_decode_exact(up_shexp, z, t)?,
6540                    ),
6541                }
6542            } else {
6543                (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?)
6544            };
6545            let mut sa = e.uninit(t * n_ff_sh)?;
6546            e.silu_mul(&sg_gate, &sg_up, &mut sa, t * n_ff_sh)?;
6547            let sh = if verify_t {
6548                e.matmul_decode_exact(down_shexp, &sa, t)?
6549            } else {
6550                e.matmul(down_shexp, &sa, t)?
6551            };
6552            // shexp gate: qwen35moe sigmoid-gates; M3 has no gate tensor -> weight 1.0.
6553            // m-INVARIANT fused sigmoid-dot under router_prefill_exact (serve isolation,
6554            // lane/concat-prime-exact) — every shexp-gate arm shares the same form so a
6555            // dispatch choice cannot change bits.
6556            let g = match &m.gate_inp_shexp {
6557                Some(gate_inp_shexp) if crate::router_prefill_exact_on() => {
6558                    e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
6559                }
6560                Some(gate_inp_shexp) => {
6561                    let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
6562                    let mut g = e.uninit(t)?;
6563                    e.sigmoid(&gs, &mut g, t)?;
6564                    g
6565                }
6566                None => e.htod(&vec![1.0f32; t])?,
6567            };
6568            e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
6569        }
6570        Ok(moe_out)
6571    }
6572
6573    /// kernel headers); the shared-expert epilogue is byte-identical to moe_ffn_sequential's.
6574    #[allow(clippy::too_many_arguments)]
6575    #[allow(clippy::too_many_arguments)]
6576    fn moe_ffn_dev(
6577        e: &Engine,
6578        m: &MoeWeights,
6579        z: &CudaSlice<f32>,
6580        zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
6581        logits: &CudaSlice<f32>,
6582        t: usize,
6583        cfg: &ModelConfig,
6584        il: u16,
6585        max_block: usize,
6586    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6587        let moe = cfg.moe.as_ref().unwrap();
6588        let n_embd = cfg.n_embd as usize;
6589        let n_expert = moe.expert_count as usize;
6590        let n_used = moe.expert_used_count as usize;
6591        let n_ff_exp = moe.expert_ff_length as usize;
6592        // moe_router_topk is SOFTMAX-only (no exp_probs_b bias, no expert_weights_scale) and every
6593        // gate_up kernel here fuses PLAIN silu(gate)*up. `dev_ok` denies sigmoid-router archs and
6594        // clamped layers; assert both so a future caller that skips the gate fails loudly.
6595        debug_assert!(
6596            cfg.sigmoid_router().is_none(),
6597            "moe_ffn_dev routes SOFTMAX: a sigmoid-router arch would pick wrong experts"
6598        );
6599        debug_assert!(
6600            !cfg.swiglu_clamped_at(il as u32),
6601            "moe_ffn_dev's fused epilogue is plain SiLU: no clamped form"
6602        );
6603
6604        // device top-k: sel [t, n_used] i32, w [t, n_used] f32 — stays on device.
6605        let (sel_d, mut w_d) = e.moe_router_topk(logits, t, n_expert, n_used)?;
6606        // Down-projection macro fold (compressed-tensors NVFP4 artifacts): one tiny launch,
6607        // skipped entirely for macro-free experts (every k-quant GGUF).
6608        if m.has_macros {
6609            e.moe_w_scale_by_expert(&mut w_d, &sel_d, &m.dev_macros, n_expert, t * n_used)?;
6610        }
6611
6612        // moe_out rows are FULLY overwritten by moe_down8_fma_dev — uninit (stage-2 rule).
6613        let mut moe_out = e.uninit(t * n_embd)?;
6614
6615        // RESIDENT-EXPERTS arm: the pointer row comes from the load-time slab (no cache, no
6616        // lock). Same kernels/loop as the SLRU arm below — only the row's provenance differs.
6617        if let Some(dev) = m.dev_exps.as_ref() {
6618            // WALL-GAP ARC: interleaved gate/up slab (MEMRA_MOE_GU_IL) -> both projections use
6619            // the combined stride; up's base is offset in the ptr table. Down unchanged.
6620            let (rbg_d, rbu_d) = if dev.gu_il {
6621                let sxx = m.gate_exps.row_bytes + m.up_exps.row_bytes;
6622                (sxx, sxx)
6623            } else {
6624                (m.gate_exps.row_bytes, m.up_exps.row_bytes)
6625            };
6626            let q8 = moe_q8_enabled()
6627                && q8_expert_supported(m.gate_exps.qtype)
6628                && q8_expert_supported(m.up_exps.qtype)
6629                && q8_expert_supported(m.down_exps.qtype);
6630            // SMALL-M ROWS ARM (MEMRA_SPEC_M2, lane/spec-m2): batch the verify token loop —
6631            // ONE batched z-quantize + ONE gate_up rows launch + ONE act quantize + ONE down
6632            // rows launch (4 launches/layer, was 4t). BIT-IDENTICAL per token to the serial
6633            // loop below (rows twins = the _v/w8h2v per-token programs on a grid.z token axis;
6634            // quantize_q8_1 is per-32-block row-independent). Gated to the AUTO kernel modes —
6635            // a custom MEMRA_MOE_DEVQ8_GU/DOWN diagnostic run keeps the serial loop so the
6636            // dispatched kernel stays exactly the env-selected one — and to the w8h2v shape
6637            // (n_ff_exp==512, n_used<=8), the same contract the AUTO down dispatch keys on.
6638            let rows_arm = q8
6639                && t > 1
6640                && crate::spec::spec_m2()
6641                && n_ff_exp == 512
6642                && n_used <= 8
6643                && std::env::var("MEMRA_MOE_DEVQ8_GU")
6644                    .map(|v| v.is_empty() || v == "v")
6645                    .unwrap_or(true)
6646                && std::env::var("MEMRA_MOE_DEVQ8_DOWN")
6647                    .map(|v| v.is_empty() || v == "w8h2v")
6648                    .unwrap_or(true);
6649            // CSR EXPERT-DEDUP gate_up (verify-cost target #1, DEFAULT ON 2026-07-10): the
6650            // owner-scan kernel serves every (token, slot) pair of an expert from ONE block,
6651            // deduping the 38-40% duplicated weight-stream+decode the overlap probe measured
6652            // (gate_up 55.0 -> 39.7us/launch; +1.3-2.1% spec e2e p2, +0.6-1.7% p3, all K).
6653            // Bit-identical to the _rows twins (explicit-intrinsic accumulate — the ULP/fmad
6654            // lesson; =2 byte-compare verified zero diffs). down stays on _rows (CSR down
6655            // measured 23.5 -> 37.5us: 16-group rows can't amortize the serial pair loop).
6656            // MEMRA_MOE_CSR=0 rollback; =2 runs BOTH paths and byte-compares (debug).
6657            let csr_mode = std::env::var("MEMRA_MOE_CSR")
6658                .ok()
6659                .and_then(|v| v.parse::<i32>().ok())
6660                .unwrap_or(1);
6661            let csr_qt = |qt: i32| qt == crate::QT_IQ4_XS || qt == crate::QT_IQ3_S;
6662            let csr_arm = rows_arm
6663                && csr_mode > 0
6664                && t <= 10
6665                && csr_qt(m.gate_exps.qtype)
6666                && csr_qt(m.up_exps.qtype)
6667                && csr_qt(m.down_exps.qtype);
6668            if csr_arm {
6669                if csr_mode == 2 {
6670                    static ENGAGED: std::sync::Once = std::sync::Once::new();
6671                    ENGAGED.call_once(|| eprintln!("[memra] moe CSR byte-compare mode ON (t={t})"));
6672                }
6673                let n_pairs = t * n_used;
6674                let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
6675                let act = e.moe_gate_up_silu8_dev_q8_csr(
6676                    &dev.ptr_row,
6677                    &sel_d,
6678                    &zq,
6679                    &zd,
6680                    n_pairs,
6681                    n_embd,
6682                    n_ff_exp,
6683                    n_used,
6684                    n_expert,
6685                    m.gate_exps.qtype,
6686                    m.up_exps.qtype,
6687                    rbg_d,
6688                    rbu_d,
6689                )?;
6690                let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
6691                // down stays on the _rows twin — BOTH CSR down variants measured negative
6692                // (v1 serial pairs 23.5->37.5us; v2 warp-parallel+SMEM -14% e2e, K=8 -37%):
6693                // 16-group rows have too little decode to amortize any dedup structure.
6694                e.moe_down8_fma_dev_q8_rows(
6695                    &dev.ptr_row,
6696                    &sel_d,
6697                    &w_d,
6698                    &aq2,
6699                    &ad2,
6700                    &mut moe_out,
6701                    t,
6702                    n_ff_exp,
6703                    n_embd,
6704                    n_used,
6705                    n_expert,
6706                    m.down_exps.qtype,
6707                    m.down_exps.row_bytes,
6708                )?;
6709                if csr_mode == 2 {
6710                    // DEBUG BYTE-COMPARE: run the _rows twins on the same inputs, diff bits.
6711                    let act_r = e.moe_gate_up_silu8_dev_q8_rows(
6712                        &dev.ptr_row,
6713                        &sel_d,
6714                        &zq,
6715                        &zd,
6716                        t,
6717                        n_embd,
6718                        n_ff_exp,
6719                        n_used,
6720                        n_expert,
6721                        m.gate_exps.qtype,
6722                        m.up_exps.qtype,
6723                        rbg_d,
6724                        rbu_d,
6725                        &m.dev_macros,
6726                    )?;
6727                    let mut out_r = e.uninit(t * n_embd)?;
6728                    let (aq2r, ad2r) = e.quantize_q8_1(&act_r, n_pairs, n_ff_exp)?;
6729                    e.moe_down8_fma_dev_q8_rows(
6730                        &dev.ptr_row,
6731                        &sel_d,
6732                        &w_d,
6733                        &aq2r,
6734                        &ad2r,
6735                        &mut out_r,
6736                        t,
6737                        n_ff_exp,
6738                        n_embd,
6739                        n_used,
6740                        n_expert,
6741                        m.down_exps.qtype,
6742                        m.down_exps.row_bytes,
6743                    )?;
6744                    let (a1, a2) = (e.dtoh(&act)?, e.dtoh(&act_r)?);
6745                    let (o1, o2) = (e.dtoh(&moe_out)?, e.dtoh(&out_r)?);
6746                    let ba = a1
6747                        .iter()
6748                        .zip(&a2)
6749                        .filter(|(x, y)| x.to_bits() != y.to_bits())
6750                        .count();
6751                    let bo = o1
6752                        .iter()
6753                        .zip(&o2)
6754                        .filter(|(x, y)| x.to_bits() != y.to_bits())
6755                        .count();
6756                    if ba + bo > 0 {
6757                        eprintln!(
6758                            "[csr-check] il={il} t={t} ACT diffs={ba}/{} OUT diffs={bo}/{}",
6759                            a1.len(),
6760                            o1.len()
6761                        );
6762                        // First 4 differing ACT elements: (pair, o, csr, rows) + that pair's expert
6763                        let sel_h = e.dtoh_i32(&sel_d)?;
6764                        let mut shown = 0;
6765                        for (i, (x, y)) in a1.iter().zip(&a2).enumerate() {
6766                            if x.to_bits() != y.to_bits() && shown < 4 {
6767                                let (p, o) = (i / n_ff_exp, i % n_ff_exp);
6768                                let ex = sel_h[p];
6769                                let npx = sel_h.iter().filter(|&&v| v == ex).count();
6770                                eprintln!(
6771                                    "  ACT p={p} ex={ex} np={npx} o={o} csr={x:e} rows={y:e}"
6772                                );
6773                                shown += 1;
6774                            }
6775                        }
6776                        std::process::exit(3);
6777                    }
6778                }
6779            } else if rows_arm {
6780                // TEMP PROBE (MEMRA_MOE_OVERLAP=1): cross-token expert-activation overlap at
6781                // verify — sizes the CSR dedup win (unique experts vs t*n_used pairs).
6782                if std::env::var("MEMRA_MOE_OVERLAP").as_deref() == Ok("1") {
6783                    use std::sync::atomic::{AtomicU64, Ordering};
6784                    static PAIRS: AtomicU64 = AtomicU64::new(0);
6785                    static UNIQ: AtomicU64 = AtomicU64::new(0);
6786                    static CALLS: AtomicU64 = AtomicU64::new(0);
6787                    let sel_h = e.dtoh_i32(&sel_d)?;
6788                    let mut u: Vec<i32> = sel_h.clone();
6789                    u.sort_unstable();
6790                    u.dedup();
6791                    PAIRS.fetch_add(sel_h.len() as u64, Ordering::Relaxed);
6792                    UNIQ.fetch_add(u.len() as u64, Ordering::Relaxed);
6793                    let c = CALLS.fetch_add(1, Ordering::Relaxed) + 1;
6794                    if c % 480 == 0 {
6795                        let p = PAIRS.load(Ordering::Relaxed);
6796                        let q = UNIQ.load(Ordering::Relaxed);
6797                        eprintln!(
6798                            "[overlap] calls={c} pairs={p} unique={q} ratio={:.3} (t={t})",
6799                            q as f64 / p as f64
6800                        );
6801                    }
6802                }
6803                let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
6804                let act = e.moe_gate_up_silu8_dev_q8_rows(
6805                    &dev.ptr_row,
6806                    &sel_d,
6807                    &zq,
6808                    &zd,
6809                    t,
6810                    n_embd,
6811                    n_ff_exp,
6812                    n_used,
6813                    n_expert,
6814                    m.gate_exps.qtype,
6815                    m.up_exps.qtype,
6816                    rbg_d,
6817                    rbu_d,
6818                    &m.dev_macros,
6819                )?;
6820                let (aq2, ad2) = e.quantize_q8_1(&act, t * n_used, n_ff_exp)?;
6821                e.moe_down8_fma_dev_q8_rows(
6822                    &dev.ptr_row,
6823                    &sel_d,
6824                    &w_d,
6825                    &aq2,
6826                    &ad2,
6827                    &mut moe_out,
6828                    t,
6829                    n_ff_exp,
6830                    n_embd,
6831                    n_used,
6832                    n_expert,
6833                    m.down_exps.qtype,
6834                    m.down_exps.row_bytes,
6835                )?;
6836            } else {
6837                for tok in 0..t {
6838                    let zt = z.slice(tok * n_embd..(tok + 1) * n_embd);
6839                    let selt = sel_d.slice(tok * n_used..(tok + 1) * n_used);
6840                    let wt = w_d.slice(tok * n_used..(tok + 1) * n_used);
6841                    let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
6842                    if q8 {
6843                        let (zq, zd) = match (t, zq8) {
6844                            (1, Some((q, d))) => (q.clone(), d.clone()),
6845                            _ => e.quantize_q8_1_view(&zt, 1, n_embd)?,
6846                        };
6847                        let act = e.moe_gate_up_silu8_dev_q8(
6848                            &dev.ptr_row,
6849                            &selt,
6850                            &zq,
6851                            &zd,
6852                            n_embd,
6853                            n_ff_exp,
6854                            n_used,
6855                            n_expert,
6856                            m.gate_exps.qtype,
6857                            m.up_exps.qtype,
6858                            rbg_d,
6859                            rbu_d,
6860                            &m.dev_macros,
6861                        )?;
6862                        let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
6863                        e.moe_down8_fma_dev_q8(
6864                            &dev.ptr_row,
6865                            &selt,
6866                            &wt,
6867                            &aq2,
6868                            &ad2,
6869                            &mut dst,
6870                            n_ff_exp,
6871                            n_embd,
6872                            n_used,
6873                            n_expert,
6874                            m.down_exps.qtype,
6875                            m.down_exps.row_bytes,
6876                        )?;
6877                    } else {
6878                        let act = e.moe_gate_up_silu8_dev(
6879                            &dev.ptr_row,
6880                            &selt,
6881                            &zt,
6882                            n_embd,
6883                            n_ff_exp,
6884                            n_used,
6885                            n_expert,
6886                            m.gate_exps.qtype,
6887                            m.up_exps.qtype,
6888                            rbg_d,
6889                            rbu_d,
6890                            &m.dev_macros,
6891                        )?;
6892                        e.moe_down8_fma_dev(
6893                            &dev.ptr_row,
6894                            &selt,
6895                            &wt,
6896                            &act,
6897                            &mut dst,
6898                            n_ff_exp,
6899                            n_embd,
6900                            n_used,
6901                            n_expert,
6902                            m.down_exps.qtype,
6903                            m.down_exps.row_bytes,
6904                        )?;
6905                    }
6906                }
6907            }
6908        } else {
6909            // Launch under the cache lock: the row borrow lives as long as the closure, and the
6910            // lock covers only launch ISSUE (µs), same policy as moe_cached_gemm.
6911            // Q8 ARM PARITY (2026-07-06): the SLRU arm ran the f32-dequant _dev kernels only —
6912            // 80us/launch vs the q8 twins' 15us on the SAME shapes (fixed-build profile: 228
6913            // f32 launches = 36ms of the 64-tok window). Same q8 gate + kernels as the resident
6914            // arm above; MEMRA_MOE_Q8=0 restores the byte-identical f32 path.
6915            let q8 = moe_q8_enabled()
6916                && q8_expert_supported(m.gate_exps.qtype)
6917                && q8_expert_supported(m.up_exps.qtype)
6918                && q8_expert_supported(m.down_exps.qtype);
6919            e.with_moe_cache(max_block, |c, eng| {
6920                let row = c
6921                    .layer_dev_row(il, n_expert, eng)?
6922                    .ok_or("moe_ffn_dev: layer row vanished under the lock")?;
6923                for tok in 0..t {
6924                    let zt = z.slice(tok * n_embd..(tok + 1) * n_embd);
6925                    let selt = sel_d.slice(tok * n_used..(tok + 1) * n_used);
6926                    let wt = w_d.slice(tok * n_used..(tok + 1) * n_used);
6927                    let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
6928                    if q8 {
6929                        let (zq, zd) = match (t, zq8) {
6930                            (1, Some((q, d))) => (q.clone(), d.clone()),
6931                            _ => eng.quantize_q8_1_view(&zt, 1, n_embd)?,
6932                        };
6933                        let act = eng.moe_gate_up_silu8_dev_q8(
6934                            row,
6935                            &selt,
6936                            &zq,
6937                            &zd,
6938                            n_embd,
6939                            n_ff_exp,
6940                            n_used,
6941                            n_expert,
6942                            m.gate_exps.qtype,
6943                            m.up_exps.qtype,
6944                            m.gate_exps.row_bytes,
6945                            m.up_exps.row_bytes,
6946                            &m.dev_macros,
6947                        )?;
6948                        let (aq2, ad2) = eng.quantize_q8_1(&act, n_used, n_ff_exp)?;
6949                        eng.moe_down8_fma_dev_q8(
6950                            row,
6951                            &selt,
6952                            &wt,
6953                            &aq2,
6954                            &ad2,
6955                            &mut dst,
6956                            n_ff_exp,
6957                            n_embd,
6958                            n_used,
6959                            n_expert,
6960                            m.down_exps.qtype,
6961                            m.down_exps.row_bytes,
6962                        )?;
6963                    } else {
6964                        let act = eng.moe_gate_up_silu8_dev(
6965                            row,
6966                            &selt,
6967                            &zt,
6968                            n_embd,
6969                            n_ff_exp,
6970                            n_used,
6971                            n_expert,
6972                            m.gate_exps.qtype,
6973                            m.up_exps.qtype,
6974                            m.gate_exps.row_bytes,
6975                            m.up_exps.row_bytes,
6976                            &m.dev_macros,
6977                        )?;
6978                        eng.moe_down8_fma_dev(
6979                            row,
6980                            &selt,
6981                            &wt,
6982                            &act,
6983                            &mut dst,
6984                            n_ff_exp,
6985                            n_embd,
6986                            n_used,
6987                            n_expert,
6988                            m.down_exps.qtype,
6989                            m.down_exps.row_bytes,
6990                        )?;
6991                    }
6992                }
6993                // instrumentation parity with the host paths (3 blocks/expert-slot, all hits).
6994                c.hits += (t * 3 * n_used) as u64;
6995                Ok(())
6996            })?;
6997        }
6998
6999        // SHARED EXPERT epilogue — byte-identical to moe_ffn_sequential step 3 (incl. its Q8
7000        // TRUNK-FUSION arm: fused2 is bit-identical to the two matmul calls per (tensor,row)).
7001        // gate_inp_shexp is OPTIONAL: qwen35moe gates the shared expert (sigmoid(gate_inp) x sh);
7002        // MiniMax-M3 (DeepSeek-V3 class) has NO shexp gate — the shared expert adds directly.
7003        if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
7004            (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
7005        {
7006            let n_ff_sh = gate_shexp.out_features();
7007            // verify-t (2..15) decode-exact arm: this fn now serves the spec verify batches
7008            // (pairs gate moved to t>=PRIME_MIN_T), so the shexp chain must match t==1 per col.
7009            let verify_t = t > 1 && t < PRIME_MIN_T;
7010            let (sg_gate, sg_up) = if t == 1 {
7011                match e.matmul_q8_fused2_x(gate_shexp, up_shexp, z)? {
7012                    Some(pair) => pair,
7013                    None => (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?),
7014                }
7015            } else if verify_t {
7016                // VERIFY-TIER TRUNK FUSION (MEMRA_SPEC_FUSED_T, t=2-4): the shexp gate+up pair
7017                // rides one shared quantize + one fused2 batched launch instead of two
7018                // decode-exact calls. Bit-identical per (tensor,token,row) — see spec_fused_t.
7019                let mut fused = None;
7020                if crate::spec::spec_fused_t()
7021                    && (2..=4).contains(&t)
7022                    && e.uses_q8_1_fast(gate_shexp)
7023                    && e.uses_q8_1_fast(up_shexp)
7024                {
7025                    let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
7026                    fused = e.matmul_q8_fused2_t(gate_shexp, up_shexp, &zq, &zd, t)?;
7027                }
7028                match fused {
7029                    Some(pair) => pair,
7030                    None => (
7031                        e.matmul_decode_exact(gate_shexp, z, t)?,
7032                        e.matmul_decode_exact(up_shexp, z, t)?,
7033                    ),
7034                }
7035            } else {
7036                (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?)
7037            };
7038            let mut sa = e.uninit(t * n_ff_sh)?; // silu_mul fully overwrites
7039            e.silu_mul(&sg_gate, &sg_up, &mut sa, t * n_ff_sh)?;
7040            let sh = if verify_t {
7041                e.matmul_decode_exact(down_shexp, &sa, t)?
7042            } else {
7043                e.matmul(down_shexp, &sa, t)?
7044            };
7045            // shexp gate: qwen35moe sigmoid-gates; M3 has no gate tensor -> weight 1.0.
7046            // Same fused sigmoid-dot as moe_ffn_sequential step 3 (byte-identity contract
7047            // between the two arms; prefill keeps the batched cuBLASLt linear).
7048            let g = match &m.gate_inp_shexp {
7049                Some(gate_inp_shexp) => {
7050                    // router_prefill_exact_on(): the fused sigmoid-dot is m-INVARIANT and
7051                    // serves EVERY t (serve isolation, lane/concat-prime-exact).
7052                    if t < PRIME_MIN_T || crate::router_prefill_exact_on() {
7053                        e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
7054                    } else {
7055                        let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
7056                        let mut g = e.uninit(t)?;
7057                        e.sigmoid(&gs, &mut g, t)?;
7058                        g
7059                    }
7060                }
7061                None => e.htod(&vec![1.0f32; t])?,
7062            };
7063            e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
7064        }
7065
7066        Ok(moe_out)
7067    }
7068
7069    /// STAGE-2 GROUPED DECODE (2026-07-04): run ONE token's whole routed-expert FFN in TWO
7070    /// launches when every one of its 3*n_used blocks is ALREADY cache-resident. Returns
7071    /// Ok(true) if the grouped path ran (caller skips the sequential loop for this token);
7072    /// Ok(false) on ANY miss (caller falls through — the sequential loop stages/admits as
7073    /// before, so the NEXT occurrence takes the grouped path). Pointer safety: cache slots are
7074    /// fixed-address for the engine's lifetime and the pure-HIT path performs no admission, so
7075    /// the collected raw pointers cannot move between collection and launch (single-threaded
7076    /// decode; the lock is held only for collection, launches are stream-ordered after any
7077    /// prior same-stream staging writes).
7078    #[allow(clippy::too_many_arguments)]
7079    /// q8 twin of moe_gdec_token (dp4a arc): same residency check + 2-launch shape; the fused
7080    /// kernels consume the pre-quantized z-row and re-quantize act per slot batch.
7081    #[allow(clippy::too_many_arguments)]
7082    fn moe_gdec_token_q8(
7083        e: &Engine,
7084        m: &MoeWeights,
7085        il: u16,
7086        max_block: usize,
7087        zq: &CudaSlice<i8>,
7088        zd: &CudaSlice<f32>,
7089        sel: &[u32],
7090        w: &[f32],
7091        moe_out: &mut CudaSlice<f32>,
7092        tok: usize,
7093        n_embd: usize,
7094        n_ff_exp: usize,
7095        n_used: usize,
7096    ) -> Result<bool, Box<dyn std::error::Error>> {
7097        use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
7098        use cudarc::driver::DevicePtr;
7099        let ptrs = e.with_moe_cache(max_block, |c, eng| {
7100            let mut g = [0u64; 8];
7101            let mut u = [0u64; 8];
7102            let mut d = [0u64; 8];
7103            for (j, &ex) in sel.iter().enumerate() {
7104                let ex = ex as u16;
7105                let (Some(sg), Some(su), Some(sd)) = (
7106                    c.resident(BlockId::new(il, PROJ_GATE, ex)),
7107                    c.resident(BlockId::new(il, PROJ_UP, ex)),
7108                    c.resident(BlockId::new(il, PROJ_DOWN, ex)),
7109                ) else {
7110                    return Ok(None);
7111                };
7112                let __s = eng.stream();
7113                let (pg, _e0) = c.slot(sg).device_ptr(&__s);
7114                let (pu, _e1) = c.slot(su).device_ptr(&__s);
7115                let (pd, _e2) = c.slot(sd).device_ptr(&__s);
7116                g[j] = pg as u64;
7117                u[j] = pu as u64;
7118                d[j] = pd as u64;
7119            }
7120            if cpu_expert_profile_admit_enabled() && !c.is_frozen() {
7121                for &ex in sel {
7122                    let ex = ex as u16;
7123                    for proj in [PROJ_GATE, PROJ_UP, PROJ_DOWN] {
7124                        c.note_profile_hit(BlockId::new(il, proj, ex));
7125                    }
7126                }
7127            }
7128            c.hits += (3 * n_used) as u64;
7129            Ok(Some((g, u, d)))
7130        })?;
7131        let Some((g, u, d)) = ptrs else {
7132            return Ok(false);
7133        };
7134        let mut wv = [0f32; 8];
7135        wv[..n_used].copy_from_slice(w);
7136        let act = e.moe_gate_up_silu8_q8(
7137            crate::WPtr8(g),
7138            crate::WPtr8(u),
7139            zq,
7140            zd,
7141            n_embd,
7142            n_ff_exp,
7143            n_used,
7144            m.gate_exps.qtype,
7145            m.up_exps.qtype,
7146            m.gate_exps.row_bytes,
7147            m.up_exps.row_bytes,
7148        )?;
7149        // per-slot act quantize: [n_used, n_ff] rows in one quantize launch.
7150        let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
7151        let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
7152        e.moe_down8_fma_q8(
7153            crate::WPtr8(d),
7154            crate::F32x8(wv),
7155            &aq2,
7156            &ad2,
7157            &mut dst,
7158            n_ff_exp,
7159            n_embd,
7160            n_used,
7161            m.down_exps.qtype,
7162            m.down_exps.row_bytes,
7163        )?;
7164        Ok(true)
7165    }
7166
7167    fn moe_gdec_token(
7168        e: &Engine,
7169        m: &MoeWeights,
7170        il: u16,
7171        max_block: usize,
7172        zt: &cudarc::driver::CudaView<f32>,
7173        sel: &[u32],
7174        w: &[f32],
7175        moe_out: &mut CudaSlice<f32>,
7176        tok: usize,
7177        n_embd: usize,
7178        n_ff_exp: usize,
7179        n_used: usize,
7180    ) -> Result<bool, Box<dyn std::error::Error>> {
7181        use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
7182        use cudarc::driver::DevicePtr;
7183        // One lock hold: residency-check all 3*n_used blocks, collect raw slot pointers.
7184        let ptrs = e.with_moe_cache(max_block, |c, eng| {
7185            let mut g = [0u64; 8];
7186            let mut u = [0u64; 8];
7187            let mut d = [0u64; 8];
7188            for (j, &ex) in sel.iter().enumerate() {
7189                let ex = ex as u16;
7190                let (Some(sg), Some(su), Some(sd)) = (
7191                    c.resident(BlockId::new(il, PROJ_GATE, ex)),
7192                    c.resident(BlockId::new(il, PROJ_UP, ex)),
7193                    c.resident(BlockId::new(il, PROJ_DOWN, ex)),
7194                ) else {
7195                    return Ok(None);
7196                };
7197                let __s = eng.stream();
7198                let (pg, _e0) = c.slot(sg).device_ptr(&__s);
7199                let (pu, _e1) = c.slot(su).device_ptr(&__s);
7200                let (pd, _e2) = c.slot(sd).device_ptr(&__s);
7201                g[j] = pg as u64;
7202                u[j] = pu as u64;
7203                d[j] = pd as u64;
7204            }
7205            if cpu_expert_profile_admit_enabled() && !c.is_frozen() {
7206                for &ex in sel {
7207                    let ex = ex as u16;
7208                    for proj in [PROJ_GATE, PROJ_UP, PROJ_DOWN] {
7209                        c.note_profile_hit(BlockId::new(il, proj, ex));
7210                    }
7211                }
7212            }
7213            c.hits += (3 * n_used) as u64; // instrumentation parity with dispatch()
7214            Ok(Some((g, u, d)))
7215        })?;
7216        let Some((g, u, d)) = ptrs else {
7217            return Ok(false);
7218        };
7219        let mut wv = [0f32; 8];
7220        wv[..n_used].copy_from_slice(w);
7221        // 2 launches: (gate+up+silu) x8, then (down + slot-ordered FMA accumulate) x8.
7222        let act = e.moe_gate_up_silu8(
7223            crate::WPtr8(g),
7224            crate::WPtr8(u),
7225            zt,
7226            n_embd,
7227            n_ff_exp,
7228            n_used,
7229            m.gate_exps.qtype,
7230            m.up_exps.qtype,
7231            m.gate_exps.row_bytes,
7232            m.up_exps.row_bytes,
7233        )?;
7234        let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
7235        e.moe_down8_fma_into(
7236            crate::WPtr8(d),
7237            crate::F32x8(wv),
7238            &act,
7239            &mut dst,
7240            n_ff_exp,
7241            n_embd,
7242            n_used,
7243            m.down_exps.qtype,
7244            m.down_exps.row_bytes,
7245        )?;
7246        Ok(true)
7247    }
7248
7249    /// EDGE-1 §B.3: dispatch one expert projection through the SLRU cache, then run the SAME
7250    /// `qmatvec_view` from whichever slot it landed in (resident HIT or staged MISS). `x` is the
7251    /// sliced activation row. `proj` selects the gate/up/down HostExps tensor. Returns y = W_expert @ x.
7252    /// q8 twin of moe_cached_gemm: same dispatch/slot mechanics, dp4a expert kernel.
7253    fn moe_cached_gemm_q8(
7254        e: &Engine,
7255        il: u16,
7256        proj: u8,
7257        ex: usize,
7258        m: &MoeWeights,
7259        max_block: usize,
7260        aq: &CudaSlice<i8>,
7261        ad: &CudaSlice<f32>,
7262    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7263        use crate::moe_cache::{BlockId, DispatchSlot, PROJ_GATE, PROJ_UP};
7264        let exps = match proj {
7265            PROJ_GATE => &m.gate_exps,
7266            PROJ_UP => &m.up_exps,
7267            _ => &m.down_exps,
7268        };
7269        let layout = exps.expert_layout(ex);
7270        let id = BlockId::new(il, proj, ex as u16);
7271        let source = exps.expert_source(ex);
7272        e.with_moe_cache(max_block, |c, eng| {
7273            let slot = c.dispatch_source(id, source, eng)?;
7274            let DispatchSlot::Resident(sl) = slot;
7275            let buf = c.slot(sl);
7276            eng.qmatvec_expert_q8(
7277                buf,
7278                0..layout.len,
7279                aq,
7280                ad,
7281                1,
7282                exps.in_f,
7283                exps.out_f,
7284                layout.qtype,
7285                layout.row_bytes,
7286            )
7287        })
7288    }
7289
7290    fn moe_cached_gemm(
7291        e: &Engine,
7292        il: u16,
7293        proj: u8,
7294        ex: usize,
7295        m: &MoeWeights,
7296        max_block: usize,
7297        x: &cudarc::driver::CudaView<f32>,
7298    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7299        use crate::moe_cache::{BlockId, DispatchSlot, PROJ_GATE, PROJ_UP};
7300        let exps = match proj {
7301            PROJ_GATE => &m.gate_exps,
7302            PROJ_UP => &m.up_exps,
7303            _ => &m.down_exps,
7304        };
7305        let layout = exps.expert_layout(ex);
7306        let id = BlockId::new(il, proj, ex as u16);
7307        let source = exps.expert_source(ex);
7308        // dispatch under the lock (lookup/admit/memcpy-issue), then resolve the slot and GEMM.
7309        e.with_moe_cache(max_block, |c, eng| {
7310            let slot = c.dispatch_source(id, source, eng)?;
7311            // resolve the device buffer for this slot; the GEMM is enqueued on the compute stream
7312            // (the same stream the memcpy was issued on, so ordering holds without extra sync).
7313            let DispatchSlot::Resident(sl) = slot;
7314            let buf = c.slot(sl);
7315            eng.qmatvec_view(
7316                buf,
7317                0..layout.len,
7318                x,
7319                1,
7320                exps.in_f,
7321                exps.out_f,
7322                layout.qtype,
7323                layout.row_bytes,
7324            )
7325        })
7326    }
7327
7328    /// Populate the warmup cache for an expert whose current-token output intentionally ran on
7329    /// CPU. No GEMM is launched and callers invoke this only after the CPU result has completed,
7330    /// so the current forward's backend assignment and output remain unchanged.
7331    fn moe_profile_admit_expert(
7332        e: &Engine,
7333        il: u16,
7334        ex: usize,
7335        m: &MoeWeights,
7336        max_block: usize,
7337    ) -> Result<(), Box<dyn std::error::Error>> {
7338        use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
7339        e.with_moe_cache(max_block, |cache, eng| {
7340            for (proj, exps) in [
7341                (PROJ_GATE, &m.gate_exps),
7342                (PROJ_UP, &m.up_exps),
7343                (PROJ_DOWN, &m.down_exps),
7344            ] {
7345                let id = BlockId::new(il, proj, ex as u16);
7346                let _ = cache.dispatch_source(id, exps.expert_source(ex), eng)?;
7347            }
7348            Ok(())
7349        })
7350    }
7351
7352    /// Read a projection from the immutable residency set when present; otherwise use one
7353    /// transient slot without admitting or evicting anything. Used only by post-freeze prefill.
7354    #[allow(clippy::too_many_arguments)]
7355    fn moe_frozen_gemm(
7356        e: &Engine,
7357        il: u16,
7358        proj: u8,
7359        ex: usize,
7360        m: &MoeWeights,
7361        max_block: usize,
7362        x: &cudarc::driver::CudaView<f32>,
7363        scratch: &mut Option<CudaSlice<u8>>,
7364        scratch_len: usize,
7365    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7366        use crate::moe_cache::{BlockId, PROJ_GATE, PROJ_UP};
7367        let exps = match proj {
7368            PROJ_GATE => &m.gate_exps,
7369            PROJ_UP => &m.up_exps,
7370            _ => &m.down_exps,
7371        };
7372        let layout = exps.expert_layout(ex);
7373        let id = BlockId::new(il, proj, ex as u16);
7374        if let Some(output) = e.with_moe_cache(max_block, |cache, eng| {
7375            let Some(slot) = cache.resident(id) else {
7376                return Ok(None);
7377            };
7378            let buf = cache.slot(slot);
7379            Ok(Some(eng.qmatvec_view(
7380                buf,
7381                0..layout.len,
7382                x,
7383                1,
7384                exps.in_f,
7385                exps.out_f,
7386                layout.qtype,
7387                layout.row_bytes,
7388            )?))
7389        })? {
7390            return Ok(output);
7391        }
7392        if scratch.is_none() {
7393            *scratch = Some(e.alloc_u8_uninit(scratch_len)?);
7394        }
7395        let scratch = scratch.as_mut().unwrap();
7396        e.stage_expert(exps.expert_bytes(ex), scratch, 0)?;
7397        e.qmatvec_view(
7398            scratch,
7399            0..layout.len,
7400            x,
7401            1,
7402            exps.in_f,
7403            exps.out_f,
7404            layout.qtype,
7405            layout.row_bytes,
7406        )
7407    }
7408
7409    fn moe_prefetch_expert(
7410        e: &Engine,
7411        il: u16,
7412        ex: usize,
7413        m: &MoeWeights,
7414        max_block: usize,
7415        keep: &[crate::moe_cache::BlockId],
7416    ) -> Result<(), Box<dyn std::error::Error>> {
7417        use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
7418        e.with_moe_cache(max_block, |c, eng| {
7419            for (proj, exps) in [
7420                (PROJ_GATE, &m.gate_exps),
7421                (PROJ_UP, &m.up_exps),
7422                (PROJ_DOWN, &m.down_exps),
7423            ] {
7424                let id = BlockId::new(il, proj, ex as u16);
7425                let _ = c.prefetch_source(id, exps.expert_source(ex), keep, eng)?;
7426            }
7427            Ok(())
7428        })
7429    }
7430
7431    /// Worker-mode disk lookahead for grouped prefill. Memory sources are deliberately skipped so
7432    /// selecting `worker` changes only storage scheduling here; all CUDA work stays in dispatch.
7433    fn moe_prefetch_disk_expert(
7434        e: &Engine,
7435        il: u16,
7436        ex: usize,
7437        m: &MoeWeights,
7438        max_block: usize,
7439        keep: &[crate::moe_cache::BlockId],
7440    ) -> Result<(), Box<dyn std::error::Error>> {
7441        use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
7442        e.with_moe_cache(max_block, |c, eng| {
7443            for (proj, exps) in [
7444                (PROJ_GATE, &m.gate_exps),
7445                (PROJ_UP, &m.up_exps),
7446                (PROJ_DOWN, &m.down_exps),
7447            ] {
7448                let source = exps.expert_source(ex);
7449                if let crate::model::ExpertSource::Disk { .. } = &source {
7450                    let id = BlockId::new(il, proj, ex as u16);
7451                    let _ = c.prefetch_source(id, source, keep, eng)?;
7452                }
7453            }
7454            Ok(())
7455        })
7456    }
7457
7458    #[inline]
7459    fn moe_prefetch_host_expert(ex: usize, m: &MoeWeights) {
7460        let _ = m.gate_exps.prefetch_expert_pages(ex);
7461        let _ = m.up_exps.prefetch_expert_pages(ex);
7462        let _ = m.down_exps.prefetch_expert_pages(ex);
7463    }
7464}
7465
7466// ================================================================================================
7467// A2: EXPERT-GROUPED MoE PREFILL (MEMRA_MOE_GROUPED=1). Resident-case prototype.
7468//
7469// Instead of the per-token loop (T * 8 experts * 3 projections = 12024 individual m=1 matvecs),
7470// this groups tokens by expert and runs ONE matmul per active expert per projection at m=m_e.
7471// On a 501-token prefill with ~170 active experts, that's ~510 matmuls (vs 12024).
7472//
7473// EXACTNESS: per-token accumulation across its 8 experts is reordered (grouped processes experts
7474// in expert-id order, not the router's top-k order). To preserve bit-identity with the sequential
7475// loop, we use an 8-SLOT scheme: expert outputs are scattered into slots keyed by the token's
7476// top-k position (0..7), then reduced in that fixed order. This makes the f32 addition order
7477// identical to the per-token loop regardless of expert processing order.
7478//
7479// Memory: T * 8 * n_embd * 4 = 501 * 8 * 2048 * 4 = ~32 MB (slot buffer). Fine on 96GB.
7480// ================================================================================================
7481
7482impl HybridModel {
7483    /// Resident-slab fast path for host-routed grouped prefill. Unclamped layers batch the
7484    /// sequential fused q8 program over the token axis; clamped layers use the separate
7485    /// expert-major q8 chain so `ffn_act_lim` remains authoritative.
7486    #[allow(clippy::too_many_arguments)]
7487    fn moe_ffn_grouped_resident_q8(
7488        e: &Engine,
7489        m: &MoeWeights,
7490        z: &CudaSlice<f32>,
7491        t: usize,
7492        cfg: &ModelConfig,
7493        il: u16,
7494        sel_all: &[u32],
7495        w_all: &[f32],
7496        table: &CudaSlice<u64>,
7497        gu_il: bool,
7498    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7499        let moe = cfg.moe.as_ref().unwrap();
7500        let n_embd = cfg.n_embd as usize;
7501        let n_expert = moe.expert_count as usize;
7502        let n_used = moe.expert_used_count as usize;
7503        let n_ff_exp = moe.expert_ff_length as usize;
7504        let n_pairs = t * n_used;
7505        debug_assert_eq!(sel_all.len(), n_pairs);
7506        debug_assert_eq!(w_all.len(), n_pairs);
7507        debug_assert!(
7508            m.gate_exps.macros.is_none()
7509                && m.up_exps.macros.is_none()
7510                && m.down_exps.macros.is_none(),
7511            "resident grouped q8 does not fold per-expert macro scales",
7512        );
7513
7514        // The rows twins run the resident sequential program verbatim on grid.z = token:
7515        // fused gate/up/SiLU per slot, batched activation quantization, then the original
7516        // slot-ordered down/FMA chain. Routing remains the sigmoid selector above (device by
7517        // default; MEMRA_SIG_ROUTER=0 is the host oracle); these kernels consume sel/w only and
7518        // never enter the softmax router.
7519        if !cfg.swiglu_clamped_at(il as u32) {
7520            let sel: Vec<i32> = sel_all.iter().map(|&expert| expert as i32).collect();
7521            let sel_d = e.htod_i32(&sel)?;
7522            let w_d = e.htod(w_all)?;
7523            let (gate_row_bytes, up_row_bytes) = if gu_il {
7524                let combined = m.gate_exps.row_bytes + m.up_exps.row_bytes;
7525                (combined, combined)
7526            } else {
7527                (m.gate_exps.row_bytes, m.up_exps.row_bytes)
7528            };
7529            let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
7530            let act = e.moe_gate_up_silu8_dev_q8_rows(
7531                table,
7532                &sel_d,
7533                &zq,
7534                &zd,
7535                t,
7536                n_embd,
7537                n_ff_exp,
7538                n_used,
7539                n_expert,
7540                m.gate_exps.qtype,
7541                m.up_exps.qtype,
7542                gate_row_bytes,
7543                up_row_bytes,
7544                &m.dev_macros,
7545            )?;
7546            let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
7547            let mut moe_out = e.uninit(t * n_embd)?;
7548            e.moe_down8_fma_dev_q8_rows_g(
7549                table,
7550                &sel_d,
7551                &w_d,
7552                &aq2,
7553                &ad2,
7554                &mut moe_out,
7555                t,
7556                n_ff_exp,
7557                n_embd,
7558                n_used,
7559                n_expert,
7560                m.down_exps.qtype,
7561                m.down_exps.row_bytes,
7562            )?;
7563
7564            if std::env::var("MEMRA_MOE_STATS").is_ok() {
7565                let mut counts = vec![0usize; n_expert];
7566                for &expert in sel_all {
7567                    counts[expert as usize] += 1;
7568                }
7569                let mut sizes: Vec<usize> =
7570                    counts.into_iter().filter(|&count| count != 0).collect();
7571                sizes.sort_unstable();
7572                let mean = sizes.iter().sum::<usize>() as f64 / sizes.len().max(1) as f64;
7573                println!(
7574                    "moe-grouped il={il} t={t} dispatch=resident-q8-rows active={}/{} \
7575                     m_e: min={} median={} mean={mean:.1} max={}",
7576                    sizes.len(),
7577                    n_expert,
7578                    sizes.first().copied().unwrap_or(0),
7579                    sizes.get(sizes.len() / 2).copied().unwrap_or(0),
7580                    sizes.last().copied().unwrap_or(0),
7581                );
7582            }
7583            return Ok(moe_out);
7584        }
7585
7586        // Clamped layers keep gate/up, activation, and down as separate stages. The pair-major
7587        // matvec body is qmatvec_expert_q8 verbatim, while one launch covers every routed pair.
7588        // Pair ids stay in router slot order so scatter preserves the sequential FMA chain.
7589        let pair_tok: Vec<i32> = (0..n_pairs).map(|pair| (pair / n_used) as i32).collect();
7590        let pair_ex: Vec<i32> = sel_all.iter().map(|&expert| expert as i32).collect();
7591        let tok_off: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
7592        let tok_ids: Vec<i32> = (0..n_pairs as i32).collect();
7593
7594        let mut by_expert: Vec<Vec<i32>> = vec![Vec::new(); n_expert];
7595        for (pair, &expert) in pair_ex.iter().enumerate() {
7596            by_expert[expert as usize].push(pair as i32);
7597        }
7598
7599        let pair_tok_d = e.htod_i32(&pair_tok)?;
7600        let pair_ex_d = e.htod_i32(&pair_ex)?;
7601        let pair_w_d = e.htod(w_all)?;
7602        let tok_off_d = e.htod_i32(&tok_off)?;
7603        let tok_ids_d = e.htod_i32(&tok_ids)?;
7604
7605        let matvec = |proj: i32,
7606                      pair_rows: &CudaSlice<i32>,
7607                      aq: &CudaSlice<i8>,
7608                      ad: &CudaSlice<f32>,
7609                      in_f: usize,
7610                      out_f: usize,
7611                      qtype: i32,
7612                      row_bytes: usize|
7613         -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7614            e.moe_pairs_matvec_q8(
7615                table, proj, pair_rows, &pair_ex_d, aq, ad, in_f, out_f, n_expert, n_pairs, qtype,
7616                row_bytes,
7617            )
7618        };
7619
7620        let (gate_row_bytes, up_row_bytes) = if gu_il {
7621            let combined = m.gate_exps.row_bytes + m.up_exps.row_bytes;
7622            (combined, combined)
7623        } else {
7624            (m.gate_exps.row_bytes, m.up_exps.row_bytes)
7625        };
7626        let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
7627        let gate = matvec(
7628            0,
7629            &pair_tok_d,
7630            &zq,
7631            &zd,
7632            n_embd,
7633            n_ff_exp,
7634            m.gate_exps.qtype,
7635            gate_row_bytes,
7636        )?;
7637        let up = matvec(
7638            1,
7639            &pair_tok_d,
7640            &zq,
7641            &zd,
7642            n_embd,
7643            n_ff_exp,
7644            m.up_exps.qtype,
7645            up_row_bytes,
7646        )?;
7647        let mut act = e.uninit(n_pairs * n_ff_exp)?;
7648        Self::ffn_act_lim(
7649            e,
7650            cfg,
7651            &gate,
7652            &up,
7653            1.0,
7654            1.0,
7655            cfg.clamp_exp_at(il as u32),
7656            &mut act,
7657            n_pairs * n_ff_exp,
7658        )?;
7659        let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
7660        let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
7661        let pair_self_d = e.htod_i32(&pair_self)?;
7662        let down = matvec(
7663            2,
7664            &pair_self_d,
7665            &aq2,
7666            &ad2,
7667            n_ff_exp,
7668            n_embd,
7669            m.down_exps.qtype,
7670            m.down_exps.row_bytes,
7671        )?;
7672        let mut moe_out = e.uninit(t * n_embd)?;
7673        e.moe_pairs_scatter(
7674            &down,
7675            &pair_w_d,
7676            &tok_off_d,
7677            &tok_ids_d,
7678            &mut moe_out,
7679            t,
7680            n_embd,
7681        )?;
7682
7683        if std::env::var("MEMRA_MOE_STATS").is_ok() {
7684            let mut sizes: Vec<usize> = by_expert
7685                .iter()
7686                .filter_map(|pairs| (!pairs.is_empty()).then_some(pairs.len()))
7687                .collect();
7688            sizes.sort_unstable();
7689            let mean = sizes.iter().sum::<usize>() as f64 / sizes.len().max(1) as f64;
7690            println!(
7691                "moe-grouped il={il} t={t} dispatch=resident-q8-clamped-pairs active={}/{} \
7692                 m_e: min={} median={} mean={mean:.1} max={}",
7693                sizes.len(),
7694                n_expert,
7695                sizes.first().copied().unwrap_or(0),
7696                sizes.get(sizes.len() / 2).copied().unwrap_or(0),
7697                sizes.last().copied().unwrap_or(0),
7698            );
7699        }
7700        Ok(moe_out)
7701    }
7702
7703    fn moe_ffn_grouped_add_shared(
7704        e: &Engine,
7705        m: &MoeWeights,
7706        z: &CudaSlice<f32>,
7707        t: usize,
7708        cfg: &ModelConfig,
7709        il: u16,
7710        moe_out: &mut CudaSlice<f32>,
7711    ) -> Result<(), Box<dyn std::error::Error>> {
7712        let n_embd = cfg.n_embd as usize;
7713        if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
7714            (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
7715        {
7716            let n_ff_sh = gate_shexp.out_features();
7717            let sg_gate = e.matmul(gate_shexp, z, t)?;
7718            let sg_up = e.matmul(up_shexp, z, t)?;
7719            let mut sa = e.uninit(t * n_ff_sh)?;
7720            Self::ffn_act_lim(
7721                e,
7722                cfg,
7723                &sg_gate,
7724                &sg_up,
7725                1.0,
7726                1.0,
7727                cfg.clamp_shexp_at(il as u32),
7728                &mut sa,
7729                t * n_ff_sh,
7730            )?;
7731            let sh = e.matmul(down_shexp, &sa, t)?;
7732            let gate = match &m.gate_inp_shexp {
7733                Some(gate_inp_shexp) => {
7734                    if t < PRIME_MIN_T || crate::router_prefill_exact_on() {
7735                        e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
7736                    } else {
7737                        let raw = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
7738                        let mut gate = e.uninit(t)?;
7739                        e.sigmoid(&raw, &mut gate, t)?;
7740                        gate
7741                    }
7742                }
7743                None => e.htod(&vec![1.0f32; t])?,
7744            };
7745            e.add_scaled_rows(&sh, &gate, moe_out, n_embd, t)?;
7746        }
7747        Ok(())
7748    }
7749
7750    /// A2 expert-grouped MoE FFN (prefill path, MEMRA_MOE_GROUPED=1). Same semantics as moe_ffn:
7751    /// z [T, n_embd] -> moe_out [T, n_embd]. BIT-IDENTICAL to moe_ffn when using the slot scheme.
7752    pub(crate) fn moe_ffn_grouped(
7753        e: &Engine,
7754        m: &MoeWeights,
7755        z: &CudaSlice<f32>,
7756        t: usize,
7757        cfg: &ModelConfig,
7758        il: u16,
7759        max_block: usize,
7760    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7761        let moe = cfg.moe.as_ref().unwrap();
7762        let n_embd = cfg.n_embd as usize;
7763        let n_expert = moe.expert_count as usize;
7764        let n_used = moe.expert_used_count as usize;
7765        let n_ff_exp = moe.expert_ff_length as usize;
7766        // step35 per-layer SwiGLU clamp; None on every other arch / unclamped layer.
7767        let lim_exp = cfg.clamp_exp_at(il as u32);
7768
7769        // 1. ROUTER: exactly the same m-invariant sigmoid selector as the sequential path
7770        // (device by default; MEMRA_SIG_ROUTER=0 is the host oracle). The grouped dispatch never
7771        // enters the softmax-only pairs/dev router.
7772        let logits = Self::moe_router_logits(e, m, z, t, cfg)?;
7773        if let Some(sig) = cfg.sigmoid_router() {
7774            Self::trace_sigmoid_router_logits(e, il, t, n_expert, n_used, &logits, m, sig)?;
7775        }
7776        let (sel_all, w_all) = if let Some(sig) = cfg.sigmoid_router() {
7777            Self::moe_route_sigmoid_cfg(e, &logits, t, n_expert, n_used, m, sig)?
7778        } else {
7779            Self::moe_route_cfg(e, &logits, t, n_expert, n_used, m.active_experts.as_deref())?
7780        };
7781        crate::moesd::record_host_routes(il, n_expert, n_used, &sel_all)?;
7782        Self::trace_moe_routes(il, t, &sel_all, &w_all)?;
7783        Self::trace_moe_input(e, il, t, n_embd, z)?;
7784
7785        // Fits-VRAM Step35 under the PP stage split: use the expert-major q8 arithmetic directly
7786        // from the owning device's uniform resident slab. This is not `moe_ffn_pairs`: routing
7787        // already happened above and the activation is clamp-aware. Mixed layouts, remote slabs,
7788        // macro-scaled experts, q8-disabled configs, and spill fall through to metadata-aware A2.
7789        let no_exp_macros = m.gate_exps.macros.is_none()
7790            && m.up_exps.macros.is_none()
7791            && m.down_exps.macros.is_none();
7792        let resident_q8 = m.dev_exps.as_ref().filter(|dev| {
7793            m.has_uniform_expert_layout()
7794                && no_exp_macros
7795                && moe_q8_enabled()
7796                && q8_expert_supported(m.gate_exps.qtype)
7797                && q8_expert_supported(m.up_exps.qtype)
7798                && q8_expert_supported(m.down_exps.qtype)
7799                && moe_slab_enabled()
7800                && dev.dev == e.ctx().ordinal()
7801        });
7802        if let Some(dev) = resident_q8 {
7803            let mut moe_out = Self::moe_ffn_grouped_resident_q8(
7804                e,
7805                m,
7806                z,
7807                t,
7808                cfg,
7809                il,
7810                &sel_all,
7811                &w_all,
7812                &dev.ptr_row,
7813                dev.gu_il,
7814            )?;
7815            Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut moe_out)?;
7816            return Ok(moe_out);
7817        }
7818
7819        // 2. BUILD PER-EXPERT TOKEN LISTS (host-side grouping).
7820        // For each expert e, we need: which tokens use it, their positions in z, their top-k
7821        // slot index (for bit-identical accumulation), and their weights.
7822        struct ExpertGroup {
7823            tok_indices: Vec<i32>,  // indices into z rows (0..T-1)
7824            slot_indices: Vec<i32>, // top-k slot (0..n_used-1) for that token-expert pair
7825            weights: Vec<f32>,      // renormalized weight for that token-expert pair
7826        }
7827        let mut groups: Vec<ExpertGroup> = (0..n_expert)
7828            .map(|_| ExpertGroup {
7829                tok_indices: Vec::new(),
7830                slot_indices: Vec::new(),
7831                weights: Vec::new(),
7832            })
7833            .collect();
7834
7835        for tok in 0..t {
7836            for j in 0..n_used {
7837                let ex = sel_all[tok * n_used + j] as usize;
7838                let w = w_all[tok * n_used + j];
7839                groups[ex].tok_indices.push(tok as i32);
7840                groups[ex].slot_indices.push(j as i32);
7841                groups[ex].weights.push(w);
7842            }
7843        }
7844
7845        // 3. ALLOCATE SLOT BUFFER: [T, n_used, n_embd] f32, zero-initialized.
7846        // Each token's 8 expert contributions land in their respective slots.
7847        let mut slot_buf = e.zeros(t * n_used * n_embd)?;
7848        let mut wbuf = e.zeros(t * n_used)?; // [T, n_used] weight buffer for FMA reduce
7849
7850        // Expert weight dimensions (used in both cache and staging paths).
7851        let g_len = m.gate_exps.max_expert_bytes();
7852        let u_len = m.up_exps.max_expert_bytes();
7853        let d_len = m.down_exps.max_expert_bytes();
7854        let moe_q8 = m.has_uniform_expert_layout()
7855            && moe_q8_enabled()
7856            && q8_expert_supported(m.gate_exps.qtype)
7857            && q8_expert_supported(m.up_exps.qtype)
7858            && q8_expert_supported(m.down_exps.qtype);
7859        // The per-expert slab view is legal only for the ordinary contiguous uniform layout.
7860        // Interleaved GU slabs require the pointer-table fast path above.
7861        let slab_local = m
7862            .dev_exps
7863            .as_ref()
7864            .filter(|dev| !dev.gu_il && moe_slab_enabled() && dev.dev == e.ctx().ordinal());
7865        let use_cache =
7866            slab_local.is_none() && Engine::moe_cache_enabled() && !e.moe_cache_frozen();
7867        // The sequential no-cache/frozen staging oracle is f32. Use q8 only where sequential
7868        // also does: a local resident slab or a live SLRU dispatch.
7869        let grouped_q8 = moe_q8 && (slab_local.is_some() || use_cache);
7870
7871        // GPU scratch for staging (only allocated without a local slab or cache).
7872        let (mut scratch_g, mut scratch_u, mut scratch_d) = if slab_local.is_none() && !use_cache {
7873            (
7874                Some(e.alloc_u8(g_len)?),
7875                Some(e.alloc_u8(u_len)?),
7876                Some(e.alloc_u8(d_len)?),
7877            )
7878        } else {
7879            (None, None, None)
7880        };
7881
7882        // 4. PER ACTIVE EXPERT: gather, compute, scatter.
7883        // Processing ORDER: DESCENDING m_e (biggest token batches first) — the concluded winner
7884        // (rig5090 2026-07-04, the ascending-id arm and its MEMRA_MOE_ORDER seam removed): desc is
7885        // a first-forward win at partial cache capacity — the hot (big-m_e) experts are admitted
7886        // to the SLRU before the small-m_e tail can pollute it, so residency converges in ONE
7887        // forward instead of several: auto-cache T=501 126.9 -> 169.9 tok/s (1.34x), cap512
7888        // 119.6 -> 160.8 (and kills the rep-to-rep bimodal); wash (<2%) at cap64 pure-spill and
7889        // at long prompts where every expert stages regardless. Order is FREE to change without
7890        // breaking the byte-identity gate: the slot scheme pins each token's accumulation order
7891        // regardless of expert processing order (the whole point of the slots).
7892        let mut order: Vec<usize> = (0..n_expert)
7893            .filter(|&ex| !groups[ex].tok_indices.is_empty())
7894            .collect();
7895        order.sort_by(|&a, &b| {
7896            groups[b]
7897                .tok_indices
7898                .len()
7899                .cmp(&groups[a].tok_indices.len())
7900                .then(a.cmp(&b))
7901        });
7902        let mut m_dist: Vec<usize> = Vec::new(); // for stats
7903        let page_window = moe_page_prefetch_window();
7904        let worker_disk_prefetch = use_cache && crate::spill_pread::worker_enabled();
7905        if worker_disk_prefetch {
7906            if let Some(first) = grouped_worker_prefetch_position(order.len(), None) {
7907                Self::moe_prefetch_disk_expert(e, il, order[first], m, max_block, &[])?;
7908            }
7909        }
7910        for (order_pos, &ex) in order.iter().enumerate() {
7911            for next in page_prefetch_positions(order_pos, order.len(), page_window) {
7912                Self::moe_prefetch_host_expert(order[next], m);
7913            }
7914            if worker_disk_prefetch {
7915                if let Some(next) = grouped_worker_prefetch_position(order.len(), Some(order_pos)) {
7916                    use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
7917                    let keep = [
7918                        BlockId::new(il, PROJ_GATE, ex as u16),
7919                        BlockId::new(il, PROJ_UP, ex as u16),
7920                        BlockId::new(il, PROJ_DOWN, ex as u16),
7921                    ];
7922                    Self::moe_prefetch_disk_expert(e, il, order[next], m, max_block, &keep)?;
7923                }
7924            }
7925            let grp = &groups[ex];
7926            let m_e = grp.tok_indices.len();
7927            m_dist.push(m_e);
7928            let gl = m.gate_exps.expert_layout(ex);
7929            let ul = m.up_exps.expert_layout(ex);
7930            let dl = m.down_exps.expert_layout(ex);
7931
7932            // Upload index/weight arrays to device. The down-proj per-expert macro-scale
7933            // (ModelOpt weight_scale_2) folds into the scatter weights — post-matmul linear,
7934            // same fold as the sequential loop's `w[j] * macro_scale(ex)`. 1.0 for GGUF experts.
7935            let tok_idx_d = e.htod_i32(&grp.tok_indices)?;
7936            let slot_idx_d = e.htod_i32(&grp.slot_indices)?;
7937            let dmac = m.down_exps.macro_scale(ex);
7938            let weight_d = if dmac == 1.0 {
7939                e.htod(&grp.weights)?
7940            } else {
7941                let scaled: Vec<f32> = grp.weights.iter().map(|&w| w * dmac).collect();
7942                e.htod(&scaled)?
7943            };
7944
7945            // GATHER: collect m_e activation rows from z into a contiguous buffer.
7946            let mut gathered = e.zeros(m_e * n_embd)?;
7947            e.gather_rows(z, &tok_idx_d, &mut gathered, n_embd, m_e)?;
7948            let gv = gathered.slice(0..m_e * n_embd);
7949
7950            // Compute gate/up/down from the local slab, metadata-aware cache, or staging. The q8
7951            // form keeps each gathered row in the sequential arithmetic class at `m=m_e`.
7952            let y = if let Some(dev) = slab_local {
7953                let gate_start = ex * m.gate_exps.expert_stride;
7954                let up_start = ex * m.up_exps.expert_stride;
7955                let down_start = ex * m.down_exps.expert_stride;
7956                if grouped_q8 {
7957                    let (zq, zd) = e.quantize_q8_1(&gathered, m_e, n_embd)?;
7958                    let gate = e.qmatvec_expert_q8(
7959                        &dev.gate,
7960                        gate_start..gate_start + gl.len,
7961                        &zq,
7962                        &zd,
7963                        m_e,
7964                        m.gate_exps.in_f,
7965                        m.gate_exps.out_f,
7966                        gl.qtype,
7967                        gl.row_bytes,
7968                    )?;
7969                    let up = e.qmatvec_expert_q8(
7970                        &dev.up,
7971                        up_start..up_start + ul.len,
7972                        &zq,
7973                        &zd,
7974                        m_e,
7975                        m.up_exps.in_f,
7976                        m.up_exps.out_f,
7977                        ul.qtype,
7978                        ul.row_bytes,
7979                    )?;
7980                    let mut act = e.uninit(m_e * n_ff_exp)?;
7981                    Self::ffn_act_lim(
7982                        e,
7983                        cfg,
7984                        &gate,
7985                        &up,
7986                        m.gate_exps.macro_scale(ex),
7987                        m.up_exps.macro_scale(ex),
7988                        lim_exp,
7989                        &mut act,
7990                        m_e * n_ff_exp,
7991                    )?;
7992                    let (aq2, ad2) = e.quantize_q8_1(&act, m_e, n_ff_exp)?;
7993                    e.qmatvec_expert_q8(
7994                        &dev.down,
7995                        down_start..down_start + dl.len,
7996                        &aq2,
7997                        &ad2,
7998                        m_e,
7999                        m.down_exps.in_f,
8000                        m.down_exps.out_f,
8001                        dl.qtype,
8002                        dl.row_bytes,
8003                    )?
8004                } else {
8005                    let gate = e.qmatvec_view(
8006                        &dev.gate,
8007                        gate_start..gate_start + gl.len,
8008                        &gv,
8009                        m_e,
8010                        m.gate_exps.in_f,
8011                        m.gate_exps.out_f,
8012                        gl.qtype,
8013                        gl.row_bytes,
8014                    )?;
8015                    let up = e.qmatvec_view(
8016                        &dev.up,
8017                        up_start..up_start + ul.len,
8018                        &gv,
8019                        m_e,
8020                        m.up_exps.in_f,
8021                        m.up_exps.out_f,
8022                        ul.qtype,
8023                        ul.row_bytes,
8024                    )?;
8025                    let mut act = e.uninit(m_e * n_ff_exp)?;
8026                    Self::ffn_act_lim(
8027                        e,
8028                        cfg,
8029                        &gate,
8030                        &up,
8031                        m.gate_exps.macro_scale(ex),
8032                        m.up_exps.macro_scale(ex),
8033                        lim_exp,
8034                        &mut act,
8035                        m_e * n_ff_exp,
8036                    )?;
8037                    let actv = act.slice(0..m_e * n_ff_exp);
8038                    e.qmatvec_view(
8039                        &dev.down,
8040                        down_start..down_start + dl.len,
8041                        &actv,
8042                        m_e,
8043                        m.down_exps.in_f,
8044                        m.down_exps.out_f,
8045                        dl.qtype,
8046                        dl.row_bytes,
8047                    )?
8048                }
8049            } else if use_cache {
8050                use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
8051                if grouped_q8 {
8052                    let (zq, zd) = e.quantize_q8_1(&gathered, m_e, n_embd)?;
8053                    let gate = e.with_moe_cache(max_block, |cache, eng| {
8054                        let id = BlockId::new(il, PROJ_GATE, ex as u16);
8055                        let slot = cache.dispatch_source(id, m.gate_exps.expert_source(ex), eng)?;
8056                        eng.qmatvec_expert_q8(
8057                            cache.buf(slot),
8058                            0..gl.len,
8059                            &zq,
8060                            &zd,
8061                            m_e,
8062                            m.gate_exps.in_f,
8063                            m.gate_exps.out_f,
8064                            gl.qtype,
8065                            gl.row_bytes,
8066                        )
8067                    })?;
8068                    let up = e.with_moe_cache(max_block, |cache, eng| {
8069                        let id = BlockId::new(il, PROJ_UP, ex as u16);
8070                        let slot = cache.dispatch_source(id, m.up_exps.expert_source(ex), eng)?;
8071                        eng.qmatvec_expert_q8(
8072                            cache.buf(slot),
8073                            0..ul.len,
8074                            &zq,
8075                            &zd,
8076                            m_e,
8077                            m.up_exps.in_f,
8078                            m.up_exps.out_f,
8079                            ul.qtype,
8080                            ul.row_bytes,
8081                        )
8082                    })?;
8083                    let mut act = e.uninit(m_e * n_ff_exp)?;
8084                    Self::ffn_act_lim(
8085                        e,
8086                        cfg,
8087                        &gate,
8088                        &up,
8089                        m.gate_exps.macro_scale(ex),
8090                        m.up_exps.macro_scale(ex),
8091                        lim_exp,
8092                        &mut act,
8093                        m_e * n_ff_exp,
8094                    )?;
8095                    let (aq2, ad2) = e.quantize_q8_1(&act, m_e, n_ff_exp)?;
8096                    e.with_moe_cache(max_block, |cache, eng| {
8097                        let id = BlockId::new(il, PROJ_DOWN, ex as u16);
8098                        let slot = cache.dispatch_source(id, m.down_exps.expert_source(ex), eng)?;
8099                        eng.qmatvec_expert_q8(
8100                            cache.buf(slot),
8101                            0..dl.len,
8102                            &aq2,
8103                            &ad2,
8104                            m_e,
8105                            m.down_exps.in_f,
8106                            m.down_exps.out_f,
8107                            dl.qtype,
8108                            dl.row_bytes,
8109                        )
8110                    })?
8111                } else {
8112                    let gate = e.with_moe_cache(max_block, |cache, eng| {
8113                        let id = BlockId::new(il, PROJ_GATE, ex as u16);
8114                        let slot = cache.dispatch_source(id, m.gate_exps.expert_source(ex), eng)?;
8115                        eng.qmatvec_view(
8116                            cache.buf(slot),
8117                            0..gl.len,
8118                            &gv,
8119                            m_e,
8120                            m.gate_exps.in_f,
8121                            m.gate_exps.out_f,
8122                            gl.qtype,
8123                            gl.row_bytes,
8124                        )
8125                    })?;
8126                    let up = e.with_moe_cache(max_block, |cache, eng| {
8127                        let id = BlockId::new(il, PROJ_UP, ex as u16);
8128                        let slot = cache.dispatch_source(id, m.up_exps.expert_source(ex), eng)?;
8129                        eng.qmatvec_view(
8130                            cache.buf(slot),
8131                            0..ul.len,
8132                            &gv,
8133                            m_e,
8134                            m.up_exps.in_f,
8135                            m.up_exps.out_f,
8136                            ul.qtype,
8137                            ul.row_bytes,
8138                        )
8139                    })?;
8140                    let mut act = e.uninit(m_e * n_ff_exp)?;
8141                    Self::ffn_act_lim(
8142                        e,
8143                        cfg,
8144                        &gate,
8145                        &up,
8146                        m.gate_exps.macro_scale(ex),
8147                        m.up_exps.macro_scale(ex),
8148                        lim_exp,
8149                        &mut act,
8150                        m_e * n_ff_exp,
8151                    )?;
8152                    let actv = act.slice(0..m_e * n_ff_exp);
8153                    e.with_moe_cache(max_block, |cache, eng| {
8154                        let id = BlockId::new(il, PROJ_DOWN, ex as u16);
8155                        let slot = cache.dispatch_source(id, m.down_exps.expert_source(ex), eng)?;
8156                        eng.qmatvec_view(
8157                            cache.buf(slot),
8158                            0..dl.len,
8159                            &actv,
8160                            m_e,
8161                            m.down_exps.in_f,
8162                            m.down_exps.out_f,
8163                            dl.qtype,
8164                            dl.row_bytes,
8165                        )
8166                    })?
8167                }
8168            } else {
8169                let sg = scratch_g.as_mut().unwrap();
8170                let su = scratch_u.as_mut().unwrap();
8171                let sd = scratch_d.as_mut().unwrap();
8172                e.stage_expert(m.gate_exps.expert_bytes(ex), sg, 0)?;
8173                e.stage_expert(m.up_exps.expert_bytes(ex), su, 0)?;
8174                e.stage_expert(m.down_exps.expert_bytes(ex), sd, 0)?;
8175                if grouped_q8 {
8176                    let (zq, zd) = e.quantize_q8_1(&gathered, m_e, n_embd)?;
8177                    let gate = e.qmatvec_expert_q8(
8178                        sg,
8179                        0..gl.len,
8180                        &zq,
8181                        &zd,
8182                        m_e,
8183                        m.gate_exps.in_f,
8184                        m.gate_exps.out_f,
8185                        gl.qtype,
8186                        gl.row_bytes,
8187                    )?;
8188                    let up = e.qmatvec_expert_q8(
8189                        su,
8190                        0..ul.len,
8191                        &zq,
8192                        &zd,
8193                        m_e,
8194                        m.up_exps.in_f,
8195                        m.up_exps.out_f,
8196                        ul.qtype,
8197                        ul.row_bytes,
8198                    )?;
8199                    let mut act = e.uninit(m_e * n_ff_exp)?;
8200                    Self::ffn_act_lim(
8201                        e,
8202                        cfg,
8203                        &gate,
8204                        &up,
8205                        m.gate_exps.macro_scale(ex),
8206                        m.up_exps.macro_scale(ex),
8207                        lim_exp,
8208                        &mut act,
8209                        m_e * n_ff_exp,
8210                    )?;
8211                    let (aq2, ad2) = e.quantize_q8_1(&act, m_e, n_ff_exp)?;
8212                    e.qmatvec_expert_q8(
8213                        sd,
8214                        0..dl.len,
8215                        &aq2,
8216                        &ad2,
8217                        m_e,
8218                        m.down_exps.in_f,
8219                        m.down_exps.out_f,
8220                        dl.qtype,
8221                        dl.row_bytes,
8222                    )?
8223                } else {
8224                    let gate = e.qmatvec_view(
8225                        sg,
8226                        0..gl.len,
8227                        &gv,
8228                        m_e,
8229                        m.gate_exps.in_f,
8230                        m.gate_exps.out_f,
8231                        gl.qtype,
8232                        gl.row_bytes,
8233                    )?;
8234                    let up = e.qmatvec_view(
8235                        su,
8236                        0..ul.len,
8237                        &gv,
8238                        m_e,
8239                        m.up_exps.in_f,
8240                        m.up_exps.out_f,
8241                        ul.qtype,
8242                        ul.row_bytes,
8243                    )?;
8244                    let mut act = e.uninit(m_e * n_ff_exp)?;
8245                    Self::ffn_act_lim(
8246                        e,
8247                        cfg,
8248                        &gate,
8249                        &up,
8250                        m.gate_exps.macro_scale(ex),
8251                        m.up_exps.macro_scale(ex),
8252                        lim_exp,
8253                        &mut act,
8254                        m_e * n_ff_exp,
8255                    )?;
8256                    let actv = act.slice(0..m_e * n_ff_exp);
8257                    e.qmatvec_view(
8258                        sd,
8259                        0..dl.len,
8260                        &actv,
8261                        m_e,
8262                        m.down_exps.in_f,
8263                        m.down_exps.out_f,
8264                        dl.qtype,
8265                        dl.row_bytes,
8266                    )?
8267                }
8268            };
8269
8270            // SCATTER into slot buffer: each row goes to slot_buf[tok, slot, :].
8271            e.scatter_slot(
8272                &y,
8273                &tok_idx_d,
8274                &slot_idx_d,
8275                &weight_d,
8276                &mut slot_buf,
8277                &mut wbuf,
8278                n_embd,
8279                n_used,
8280                m_e,
8281            )?;
8282        }
8283
8284        // 5. REDUCE SLOTS: sum the 8 slots per token into the final moe_out.
8285        let mut moe_out = e.zeros(t * n_embd)?;
8286        e.reduce_slots(&slot_buf, &wbuf, &mut moe_out, n_embd, n_used, t)?;
8287
8288        // STATS: print m-distribution when MEMRA_MOE_STATS is set.
8289        if std::env::var("MEMRA_MOE_STATS").is_ok() && !m_dist.is_empty() {
8290            m_dist.sort_unstable();
8291            let active = m_dist.len();
8292            let mean = m_dist.iter().sum::<usize>() as f64 / active as f64;
8293            let median = m_dist[active / 2];
8294            let max_m = *m_dist.last().unwrap();
8295            let min_m = m_dist[0];
8296            let above16 = m_dist.iter().filter(|&&x| x >= 16).count();
8297            println!(
8298                "moe-grouped il={il} t={t} active={active}/{n_expert} \
8299                      m_e: min={min_m} median={median} mean={mean:.1} max={max_m} \
8300                      above_gemm_threshold(>=16)={above16}/{active}"
8301            );
8302        }
8303
8304        Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut moe_out)?;
8305        Ok(moe_out)
8306    }
8307
8308    /// Lane-3 M2: cross-stream MoE for lockstep decode. Routes all m stream rows in one
8309    /// batch, executes fully-HBM-resident experts through the grouped gather/GEMM/scatter
8310    /// machinery at m_e>1 (weight reads amortized across streams), and assigns any expert
8311    /// with a missing projection to that row's CPU companion call (whole-expert granularity,
8312    /// same rule as the sequential frozen path). Slot-pinned accumulation keeps each row's
8313    /// expert-sum order identical to the sequential path.
8314    pub(crate) fn moe_ffn_lockstep(
8315        &self,
8316        e: &Engine,
8317        m: &MoeWeights,
8318        zbatch: &CudaSlice<f32>,
8319        mrows: usize,
8320        il: u16,
8321        max_block: usize,
8322    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8323        use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
8324        let cfg = &self.cfg;
8325        let moe = cfg.moe.as_ref().unwrap();
8326        let n_embd = cfg.n_embd as usize;
8327        let n_expert = moe.expert_count as usize;
8328        let n_used = moe.expert_used_count as usize;
8329        let n_ff_exp = moe.expert_ff_length as usize;
8330        // step35 per-layer SwiGLU clamp; None on every other arch / unclamped layer.
8331        let lim_exp = cfg.clamp_exp_at(il as u32);
8332        let lim_shexp = cfg.clamp_shexp_at(il as u32);
8333
8334        let logits = e.matmul(&m.gate_inp, zbatch, mrows)?;
8335        if let Some(sig) = cfg.sigmoid_router() {
8336            Self::trace_sigmoid_router_logits(e, il, mrows, n_expert, n_used, &logits, m, sig)?;
8337        }
8338        let (sel_all, w_all) = if let Some(sig) = cfg.sigmoid_router() {
8339            Self::moe_route_sigmoid_cfg(e, &logits, mrows, n_expert, n_used, m, sig)?
8340        } else {
8341            Self::moe_route_cfg(
8342                e,
8343                &logits,
8344                mrows,
8345                n_expert,
8346                n_used,
8347                m.active_experts.as_deref(),
8348            )?
8349        };
8350        Self::trace_moe_routes(il, mrows, &sel_all, &w_all)?;
8351
8352        // Residency split at whole-expert granularity against the (frozen) cache.
8353        let resident_expert: Vec<bool> = e.with_moe_cache(max_block, |c, _| {
8354            Ok((0..n_expert)
8355                .map(|ex| {
8356                    [PROJ_GATE, PROJ_UP, PROJ_DOWN]
8357                        .into_iter()
8358                        .all(|p| c.resident(BlockId::new(il, p, ex as u16)).is_some())
8359                })
8360                .collect())
8361        })?;
8362
8363        struct Group {
8364            rows: Vec<i32>,
8365            slots: Vec<i32>,
8366            weights: Vec<f32>,
8367        }
8368        let mut groups: std::collections::HashMap<usize, Group> = Default::default();
8369        let mut cpu_rows: Vec<Vec<(usize, f32)>> = vec![Vec::new(); mrows];
8370        let mut cpu_by_expert: std::collections::HashMap<usize, Vec<(usize, f32)>> =
8371            Default::default();
8372        for row in 0..mrows {
8373            for j in 0..n_used {
8374                let ex = sel_all[row * n_used + j] as usize;
8375                let w = w_all[row * n_used + j];
8376                if resident_expert[ex] {
8377                    let group = groups.entry(ex).or_insert_with(|| Group {
8378                        rows: Vec::new(),
8379                        slots: Vec::new(),
8380                        weights: Vec::new(),
8381                    });
8382                    group.rows.push(row as i32);
8383                    group.slots.push(j as i32);
8384                    group.weights.push(w);
8385                } else {
8386                    crate::cpu_experts::record_incomplete_gpu_residency(0);
8387                    cpu_rows[row].push((ex, w));
8388                    cpu_by_expert.entry(ex).or_default().push((row, w));
8389                }
8390            }
8391        }
8392
8393        // CPU tickets first: reads/compute overlap the GPU grouped work below. Experts routed
8394        // by >=2 streams go through the multi-row ABI (weight decode amortized across rows);
8395        // each row's remaining experts stay one ordinary per-row call. Contribution FP-sum
8396        // order per row differs from the sequential single-call chunk — part of the
8397        // documented lockstep numeric class.
8398        let host_rows = e.dtoh(zbatch)?;
8399        let rows_ok = crate::cpu_experts::rows_supported();
8400        enum CpuPart {
8401            Single { row: usize },
8402            Rows { rows: Vec<usize> },
8403        }
8404        let mut tickets: Vec<(CpuPart, crate::cpu_experts::CpuExpertTicket)> = Vec::new();
8405        let mut rows_served: std::collections::HashSet<(usize, usize)> = Default::default();
8406        if rows_ok {
8407            let mut shared: Vec<(usize, Vec<(usize, f32)>)> = cpu_by_expert
8408                .into_iter()
8409                .filter(|(_, rows)| rows.len() >= 2)
8410                .collect();
8411            shared.sort_by_key(|(ex, _)| *ex);
8412            for (ex, mut row_weights) in shared {
8413                row_weights.sort_by_key(|(row, _)| *row);
8414                let inputs: Vec<(&[f32], f32)> = row_weights
8415                    .iter()
8416                    .map(|&(row, w)| (&host_rows[row * n_embd..(row + 1) * n_embd], w))
8417                    .collect();
8418                let job = crate::cpu_experts::prepare_rows_job(m, ex, &inputs)
8419                    .map_err(std::io::Error::other)?;
8420                for &(row, _) in &row_weights {
8421                    rows_served.insert((row, ex));
8422                }
8423                tickets.push((
8424                    CpuPart::Rows {
8425                        rows: row_weights.iter().map(|&(row, _)| row).collect(),
8426                    },
8427                    crate::cpu_experts::submit_rows(job).map_err(std::io::Error::other)?,
8428                ));
8429            }
8430        }
8431        for (row, selected) in cpu_rows.iter().enumerate() {
8432            let leftover: Vec<(usize, f32)> = selected
8433                .iter()
8434                .copied()
8435                .filter(|&(ex, _)| !rows_served.contains(&(row, ex)))
8436                .collect();
8437            if leftover.is_empty() {
8438                continue;
8439            }
8440            let host_row = &host_rows[row * n_embd..(row + 1) * n_embd];
8441            let job = crate::cpu_experts::prepare_job(m, il, &leftover, host_row)
8442                .map_err(std::io::Error::other)?;
8443            tickets.push((
8444                CpuPart::Single { row },
8445                crate::cpu_experts::submit(job).map_err(std::io::Error::other)?,
8446            ));
8447        }
8448
8449        let mut slot_buf = e.zeros(mrows * n_used * n_embd)?;
8450        let mut wbuf = e.zeros(mrows * n_used)?;
8451        let mut order: Vec<usize> = groups.keys().copied().collect();
8452        order.sort_by(|&a, &b| {
8453            groups[&b]
8454                .rows
8455                .len()
8456                .cmp(&groups[&a].rows.len())
8457                .then(a.cmp(&b))
8458        });
8459        for &ex in &order {
8460            let group = &groups[&ex];
8461            let m_e = group.rows.len();
8462            let gl = m.gate_exps.expert_layout(ex);
8463            let ul = m.up_exps.expert_layout(ex);
8464            let dl = m.down_exps.expert_layout(ex);
8465            let row_idx_d = e.htod_i32(&group.rows)?;
8466            let slot_idx_d = e.htod_i32(&group.slots)?;
8467            let dmac = m.down_exps.macro_scale(ex);
8468            let weight_d = if dmac == 1.0 {
8469                e.htod(&group.weights)?
8470            } else {
8471                let scaled: Vec<f32> = group.weights.iter().map(|&w| w * dmac).collect();
8472                e.htod(&scaled)?
8473            };
8474            let mut gathered = e.zeros(m_e * n_embd)?;
8475            e.gather_rows(zbatch, &row_idx_d, &mut gathered, n_embd, m_e)?;
8476            let gv = gathered.slice(0..m_e * n_embd);
8477            let gate = e.with_moe_cache(max_block, |c, eng| {
8478                let slot = c
8479                    .resident(BlockId::new(il, PROJ_GATE, ex as u16))
8480                    .ok_or("lockstep resident expert vanished (cache not frozen?)")?;
8481                eng.qmatvec_view(
8482                    c.buf(crate::moe_cache::DispatchSlot::Resident(slot)),
8483                    0..gl.len,
8484                    &gv,
8485                    m_e,
8486                    m.gate_exps.in_f,
8487                    m.gate_exps.out_f,
8488                    gl.qtype,
8489                    gl.row_bytes,
8490                )
8491            })?;
8492            let up = e.with_moe_cache(max_block, |c, eng| {
8493                let slot = c
8494                    .resident(BlockId::new(il, PROJ_UP, ex as u16))
8495                    .ok_or("lockstep resident expert vanished (cache not frozen?)")?;
8496                eng.qmatvec_view(
8497                    c.buf(crate::moe_cache::DispatchSlot::Resident(slot)),
8498                    0..ul.len,
8499                    &gv,
8500                    m_e,
8501                    m.up_exps.in_f,
8502                    m.up_exps.out_f,
8503                    ul.qtype,
8504                    ul.row_bytes,
8505                )
8506            })?;
8507            let mut act = e.zeros(m_e * n_ff_exp)?;
8508            Self::ffn_act_lim(
8509                e,
8510                cfg,
8511                &gate,
8512                &up,
8513                m.gate_exps.macro_scale(ex),
8514                m.up_exps.macro_scale(ex),
8515                lim_exp,
8516                &mut act,
8517                m_e * n_ff_exp,
8518            )?;
8519            let actv = act.slice(0..m_e * n_ff_exp);
8520            let y = e.with_moe_cache(max_block, |c, eng| {
8521                let slot = c
8522                    .resident(BlockId::new(il, PROJ_DOWN, ex as u16))
8523                    .ok_or("lockstep resident expert vanished (cache not frozen?)")?;
8524                eng.qmatvec_view(
8525                    c.buf(crate::moe_cache::DispatchSlot::Resident(slot)),
8526                    0..dl.len,
8527                    &actv,
8528                    m_e,
8529                    m.down_exps.in_f,
8530                    m.down_exps.out_f,
8531                    dl.qtype,
8532                    dl.row_bytes,
8533                )
8534            })?;
8535            e.scatter_slot(
8536                &y,
8537                &row_idx_d,
8538                &slot_idx_d,
8539                &weight_d,
8540                &mut slot_buf,
8541                &mut wbuf,
8542                n_embd,
8543                n_used,
8544                m_e,
8545            )?;
8546        }
8547        let mut moe_out = e.zeros(mrows * n_embd)?;
8548        e.reduce_slots(&slot_buf, &wbuf, &mut moe_out, n_embd, n_used, mrows)?;
8549
8550        // CPU contributions join BEFORE the shared expert (the sequential path's placement).
8551        let mut row_sums: Vec<Option<Vec<f32>>> = vec![None; mrows];
8552        for (part, ticket) in tickets {
8553            let cpu_output = ticket.wait().map_err(std::io::Error::other)?;
8554            let mut add_row = |row: usize, chunk: &[f32]| {
8555                let sum = row_sums[row].get_or_insert_with(|| vec![0.0f32; n_embd]);
8556                for (accumulator, value) in sum.iter_mut().zip(chunk) {
8557                    *accumulator += value;
8558                }
8559            };
8560            match part {
8561                CpuPart::Single { row } => add_row(row, &cpu_output),
8562                CpuPart::Rows { rows } => {
8563                    for (slot, row) in rows.into_iter().enumerate() {
8564                        add_row(row, &cpu_output[slot * n_embd..(slot + 1) * n_embd]);
8565                    }
8566                }
8567            }
8568        }
8569        for (row, sum) in row_sums.into_iter().enumerate() {
8570            let Some(sum) = sum else { continue };
8571            let cpu_output = e.htod(&sum)?;
8572            let mut dst = moe_out.slice_mut(row * n_embd..(row + 1) * n_embd);
8573            e.axpy_into(&cpu_output, 1.0, &mut dst, n_embd)?;
8574        }
8575
8576        if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
8577            (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
8578        {
8579            let n_ff_sh = gate_shexp.out_features();
8580            let sg_gate = e.matmul(gate_shexp, zbatch, mrows)?;
8581            let sg_up = e.matmul(up_shexp, zbatch, mrows)?;
8582            let mut sa = e.zeros(mrows * n_ff_sh)?;
8583            Self::ffn_act_lim(
8584                e,
8585                cfg,
8586                &sg_gate,
8587                &sg_up,
8588                1.0,
8589                1.0,
8590                lim_shexp,
8591                &mut sa,
8592                mrows * n_ff_sh,
8593            )?;
8594            let sh = e.matmul(down_shexp, &sa, mrows)?;
8595            // lockstep rows ARE decode tokens: fused sigmoid-dot per row so batched serving
8596            // decode matches the single-sequence decode chain bit-for-bit.
8597            let g = match &m.gate_inp_shexp {
8598                Some(gate_inp_shexp) => {
8599                    e.sigmoid_dot_rows(zbatch, gate_inp_shexp.float_data(), n_embd, mrows)?
8600                }
8601                None => e.htod(&vec![1.0f32; mrows])?,
8602            };
8603            e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, mrows)?;
8604        }
8605
8606        Ok(moe_out)
8607    }
8608}
8609
8610// ============================ gemma4 (R8 verified wiring) ==================================
8611// Node-for-node vs llama.cpp src/models/gemma4.cpp:180-405 (HANDOVER "R8 VERIFIED WIRING").
8612// v0 bring-up: full attention everywhere (exact for prompts < sliding_window 1024 — R6 masking
8613// later), sdpa_naive (hd 512 has no FA stamp), sequential host-staged MoE (the perf arms grow
8614// gemma variants after the correctness gate).
8615impl HybridModel {
8616    /// Per-layer attention geometry (R5): (head_dim, n_kv, n_head, rope_base, scale, is_swa).
8617    pub(crate) fn gemma4_geom(&self, il: usize) -> (usize, usize, usize, f32, f32, bool) {
8618        let g = self.cfg.gemma4.as_ref().unwrap();
8619        let swa = g.swa_pattern[il];
8620        let hd = if swa {
8621            g.key_length_swa
8622        } else {
8623            g.key_length_global
8624        } as usize;
8625        // attention scale = 1.0 (llama gemma4.cpp:11 "Gemma4 uses self.scaling = 1.0" — q/k are
8626        // per-head rms-normed; NOT the 1/sqrt(hd) default. Bring-up bug: 1/sqrt(hd) left token-0
8627        // rows exact (softmax over one element) while every later position drifted).
8628        (
8629            hd,
8630            g.head_count_kv[il] as usize,
8631            self.cfg.n_head as usize,
8632            if swa {
8633                g.rope_base_swa
8634            } else {
8635                g.rope_base_global
8636            },
8637            1.0,
8638            swa,
8639        )
8640    }
8641
8642    /// Suppress-token mask over t logits rows (tokenizer.ggml.suppress_tokens; no-op when the
8643    /// model ships none). NOT monotonic like softcap — must run before every argmax/sample, so
8644    /// every gemma4 logits tail (forward/prime/decode/dc/verify/e4b) calls this on device ld.
8645    pub(crate) fn gemma4_suppress(
8646        &self,
8647        e: &Engine,
8648        ld: &mut CudaSlice<f32>,
8649        t: usize,
8650    ) -> Result<(), Box<dyn std::error::Error>> {
8651        if let Some((ids, n)) = self.gemma4_aux.as_ref().and_then(|a| a.suppress_d.as_ref()) {
8652            // suppress_d is primary-owned by the audited invariant (serving picks the last/head
8653            // stage as primary, and this tail runs only after the last stage). The assert turns
8654            // that argued invariant into a checked one: any topology violating primary==head
8655            // trips here in debug instead of silently peer-reading a device-0 buffer.
8656            #[cfg(debug_assertions)]
8657            crate::debug_assert_tensor_stream_device(
8658                ids,
8659                &e.stream(),
8660                "gemma4_suppress.suppress_d",
8661            );
8662            e.mask_ids_rows(ld, ids, *n, self.output.out_features(), t)?;
8663        }
8664        Ok(())
8665    }
8666
8667    /// gemma4 attention (R5 geometry, R7 weightless V-norm on the RAW K projection, R9 dual rope).
8668    /// `cache`: Some => PRIME mode — append the T post-rope K / normed V rows into the quantized
8669    /// KV cache (same per-row quantize math as the decode append) and advance len. Fresh-prompt
8670    /// only (v0): attends within `tokens` via the f32 sdpa.
8671    #[allow(clippy::too_many_arguments)]
8672    fn gemma4_attn_prime(
8673        &self,
8674        e: &Engine,
8675        fa: &crate::hybrid::FullAttnLayer,
8676        il: usize,
8677        h: &CudaSlice<f32>,
8678        pos_d: &CudaSlice<i32>,
8679        t: usize,
8680        cache: Option<&mut Cache>,
8681        island: Option<&CudaSlice<i32>>,
8682    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8683        let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
8684        let eps = self.cfg.rms_eps;
8685        let aux = self.gemma4_aux.as_ref().unwrap();
8686        let ones = aux.ones(e);
8687        #[cfg(debug_assertions)]
8688        crate::debug_assert_tensor_stream_device(ones, &e.stream(), "gemma4_attn_prime.ones");
8689
8690        // quantize-once window: q/k/v share `h` — the MMQ D4 activation quantizes once
8691        // (h stays borrowed across the triple, so the cache key can't go stale).
8692        e.mmq_act_begin();
8693        let q0 = e.matmul(&fa.wq, h, t)?; // [t, nh*hd]
8694        if il == 0 && std::env::var("MEMRA_G4_PRIME_TRACE").as_deref() == Ok("1") {
8695            let v = e.dtoh(&q0)?;
8696            let nan = v.iter().filter(|x| x.is_nan()).count();
8697            let amax = v.iter().fold(0f32, |a, x| a.max(x.abs()));
8698            eprintln!(
8699                "[g4-prime-trace] L0 q0: nan={nan}/{} amax={amax:.3}",
8700                v.len()
8701            );
8702        }
8703        let k0 = e.matmul(&fa.wk, h, t)?; // [t, nkv*hd]
8704        // globals ship no v_proj (wv := wk at load) — V input is the SAME projection output;
8705        // reuse k0 instead of re-running the identical matmul (K=V dedup, 5 layers).
8706        let v0 = if swa {
8707            e.matmul(&fa.wv, h, t)?
8708        } else {
8709            e.clone_dtod(&k0)?
8710        };
8711        if il == 0 && std::env::var("MEMRA_G4_PRIME_TRACE").as_deref() == Ok("1") {
8712            for (tag, buf) in [("k0", &k0), ("v0", &v0)] {
8713                let v = e.dtoh(buf)?;
8714                let nan = v.iter().filter(|x| x.is_nan()).count();
8715                let amax = v.iter().fold(0f32, |a, x| a.max(x.abs()));
8716                eprintln!(
8717                    "[g4-prime-trace] L0 {tag}: nan={nan}/{} amax={amax:.3}",
8718                    v.len()
8719                );
8720            }
8721        }
8722
8723        let mut q = e.uninit(t * nh * hd)?;
8724        let mut k = e.uninit(t * nkv * hd)?;
8725        // R7: V = weightless rms_norm of the raw projection; NEVER roped.
8726        let mut v = e.uninit(t * nkv * hd)?;
8727        // 31B glue lane: producers emit the bf16 FA operands (norm emits vb; rope emits qb/kb
8728        // post-rope) — kills 3 f32->bf16 converts + re-reads per layer. Bit-identical operands
8729        // (same __float2bfloat16); MEMRA_FA_EMIT=0 reverts to the convert-in-FA path.
8730        static EMIT: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8731        // Island primes take the mask-capable naive kernel below; keep the operands f32
8732        // (the bf16 FA emit path has no island consumer). Text-only keeps emit unchanged.
8733        let emit = island.is_none()
8734            && t >= 16
8735            && crate::Engine::qkvnorm_w_on_prefill(nh * t + 2 * nkv * t, hd)
8736            && *EMIT.get_or_init(|| {
8737                std::env::var("MEMRA_FA_EMIT")
8738                    .map(|s| s != "0")
8739                    .unwrap_or(true)
8740            });
8741        let mut qb = e.alloc_uninit::<u8>(if emit { t * nh * hd * 2 } else { 1 })?;
8742        let mut kb = e.alloc_uninit::<u8>(if emit { t * nkv * hd * 2 } else { 1 })?;
8743        let mut vb = e.alloc_uninit::<u8>(if emit { t * nkv * hd * 2 } else { 1 })?;
8744        // f16-P/V door: emit V as f16 straight from the norm when this layer's FA consumer
8745        // reads f16 — kills the per-layer bf16->f16 re-encode (1.4GB/pass on 31B SWA).
8746        let v_f16 = emit
8747            && crate::fa_f16pv_on()
8748            && match hd {
8749                512 => true,
8750                256 => swa && crate::faw_hp_on() && nh % 2 == 0 && (nh / nkv) % 2 == 0,
8751                _ => false,
8752            };
8753        if emit {
8754            e.rms_norm_qkv_w4b(
8755                &q0,
8756                &k0,
8757                &v0,
8758                fa.q_norm.float_data(),
8759                fa.k_norm.float_data(),
8760                ones,
8761                &mut q,
8762                &mut k,
8763                &mut v,
8764                &mut vb,
8765                hd,
8766                nh * t,
8767                nkv * t,
8768                eps,
8769                v_f16,
8770            )?;
8771        } else {
8772            e.rms_norm_qkv(
8773                &q0,
8774                &k0,
8775                &v0,
8776                fa.q_norm.float_data(),
8777                fa.k_norm.float_data(),
8778                ones,
8779                &mut q,
8780                &mut k,
8781                &mut v,
8782                hd,
8783                nh * t,
8784                nkv * t,
8785                eps,
8786            )?;
8787        }
8788
8789        let ff = if swa {
8790            None
8791        } else {
8792            Some(
8793                aux.rope_freqs(e)
8794                    .expect("gemma4 global rope needs rope_freqs.weight"),
8795            )
8796        };
8797        #[cfg(debug_assertions)]
8798        if let Some(ff) = ff {
8799            crate::debug_assert_tensor_stream_device(
8800                ff,
8801                &e.stream(),
8802                "gemma4_attn_prime.rope_freqs",
8803            );
8804        }
8805        if emit {
8806            e.rope_neox2_bf16e(
8807                &mut q, &mut k, &mut qb, &mut kb, pos_d, hd, hd, nh, nkv, t, base, 1.0, ff,
8808            )?;
8809        } else {
8810            e.rope_neox2(&mut q, &mut k, pos_d, hd, hd, nh, nkv, t, base, 1.0, ff)?;
8811        }
8812
8813        if let Some(cache) = cache {
8814            let kvl = cache.kv[il].as_mut().unwrap();
8815            assert_eq!(kvl.len, 0, "gemma4 prime is fresh-prompt only (v0)");
8816            e.append_kv_quantized_rows(
8817                &k,
8818                &v,
8819                &mut kvl.k,
8820                &mut kvl.v,
8821                kvl.len,
8822                t,
8823                kvl.kv_dim_k,
8824                kvl.kv_dim_v,
8825                kvl.k_tok_bytes,
8826                kvl.v_tok_bytes,
8827                (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on()),
8828            )?;
8829            kvl.len += t;
8830        }
8831        let mut attn = e.zeros(t * nh * hd)?;
8832        // R6: SWA layers mask keys older than sliding_window once the prompt exceeds it
8833        // (windowed naive twin; fa windowed stamps later). Under the window, full attention
8834        // is exact — SWA rides fa_prefill (hd-256 stamp), the hd-512 globals stay naive.
8835        let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
8836        if let Some(span) = island {
8837            // Masked-prefill arm: every layer routes through the island-aware naive
8838            // kernel (correctness-first, same posture as the vision tower v1). The
8839            // window argument keeps the R6 shortcut: 0 while the prompt fits the
8840            // window, the real window beyond it.
8841            let w = if swa && t > win { win } else { 0 };
8842            e.sdpa_naive_island(&q, &k, &v, &mut attn, span, hd, nh, nkv, t, t, scale, w)?;
8843        } else if swa && t > win {
8844            if hd == 256 && std::env::var("MEMRA_NOFA").is_err() {
8845                if emit {
8846                    e.fa_prefill_w_pre(
8847                        &qb, &kb, &vb, &mut attn, hd, nh, nkv, t, t, scale, true, win, v_f16,
8848                    )?;
8849                } else {
8850                    e.fa_prefill_w(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
8851                }
8852            } else {
8853                e.sdpa_naive_w(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
8854            }
8855        } else if hd == 256 && std::env::var("MEMRA_NOFA").is_err() {
8856            e.fa_prefill(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
8857        } else if hd == 512 && std::env::var("MEMRA_NOFA").is_err() {
8858            if emit {
8859                e.fa_prefill_hd512_pre(
8860                    &qb, &kb, &vb, &mut attn, hd, nh, nkv, t, t, scale, true, v_f16,
8861                )?;
8862            } else {
8863                e.fa_prefill_hd512(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
8864            }
8865        } else {
8866            e.sdpa_naive(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
8867        }
8868        Ok(e.matmul(&fa.wo, &attn, t)?)
8869    }
8870
8871    /// Back-compat wrapper (pure prefill, no cache).
8872    fn gemma4_attn(
8873        &self,
8874        e: &Engine,
8875        fa: &crate::hybrid::FullAttnLayer,
8876        il: usize,
8877        h: &CudaSlice<f32>,
8878        pos_d: &CudaSlice<i32>,
8879        t: usize,
8880    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8881        self.gemma4_attn_prime(e, fa, il, h, pos_d, t, None, None)
8882    }
8883
8884    /// gemma4 MoE with the expert input PRE-QUANTIZED (the q8z tail fusion). Caller guarantees
8885    /// the fast-arm conditions (resident dev slabs + dp4a qtypes) — decode t=1 and verify rows
8886    /// arms only, per-token kernel chains identical to the f32-input path (same quantize bytes:
8887    /// the q8z epilogue is quantize_q8_1 verbatim).
8888    fn gemma4_moe_q8(
8889        &self,
8890        e: &Engine,
8891        m: &crate::hybrid::MoeWeights,
8892        bits: &crate::hybrid::Gemma4MoeBits,
8893        mq: &(CudaSlice<i8>, CudaSlice<f32>),
8894        router_in: &CudaSlice<f32>,
8895        t: usize,
8896    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8897        let cfg = &self.cfg;
8898        let moe = cfg.moe.as_ref().unwrap();
8899        let n_embd = cfg.n_embd as usize;
8900        let n_expert = moe.expert_count as usize;
8901        let n_used = moe.expert_used_count as usize;
8902        let n_ff_exp = moe.expert_ff_length as usize;
8903        // router stays the two-launch pair: THREE fuse variants measured worse (serial-dot
8904        // -50%, 1024-thread warp-parallel -12% on topk sync overhead — jsonl 2026-07-11);
8905        // the pair's 12us is kernel time, not launch gaps.
8906        let logits = if crate::router_kernel_on() {
8907            e.router_gemv(m.gate_inp.float_data(), router_in, n_embd, n_expert, t)?
8908        } else {
8909            e.matmul(&m.gate_inp, router_in, t)?
8910        };
8911        let dev = m.dev_exps.as_ref().unwrap();
8912        let (sel_d, w_d) =
8913            e.moe_router_topk_scaled(&logits, t, n_expert, n_used, &bits.per_expert_scale_d)?;
8914        let (zq, zd) = mq;
8915        if t == 1 {
8916            let selv = sel_d.slice(0..n_used);
8917            let wv = w_d.slice(0..n_used);
8918            let act = e.moe_gate_up_gelu8_dev_q8(
8919                &dev.ptr_row,
8920                &selv,
8921                zq,
8922                zd,
8923                n_embd,
8924                n_ff_exp,
8925                n_used,
8926                n_expert,
8927                m.gate_exps.qtype,
8928                m.up_exps.qtype,
8929                m.gate_exps.row_bytes,
8930                m.up_exps.row_bytes,
8931            )?;
8932            let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
8933            let mut moe_out = e.uninit(n_embd)?;
8934            e.moe_down8_fma_dev_q8(
8935                &dev.ptr_row,
8936                &selv,
8937                &wv,
8938                &aq2,
8939                &ad2,
8940                &mut moe_out.slice_mut(0..n_embd),
8941                n_ff_exp,
8942                n_embd,
8943                n_used,
8944                n_expert,
8945                m.down_exps.qtype,
8946                m.down_exps.row_bytes,
8947            )?;
8948            return Ok(moe_out);
8949        }
8950        let csr = t <= 10 && std::env::var("MEMRA_GEMMA_CSR").as_deref() != Ok("0");
8951        let act = if csr {
8952            e.moe_gate_up_gelu8_dev_q8_csr(
8953                &dev.ptr_row,
8954                &sel_d,
8955                zq,
8956                zd,
8957                t * n_used,
8958                n_embd,
8959                n_ff_exp,
8960                n_used,
8961                n_expert,
8962                m.gate_exps.qtype,
8963                m.up_exps.qtype,
8964                m.gate_exps.row_bytes,
8965                m.up_exps.row_bytes,
8966            )?
8967        } else {
8968            e.moe_gate_up_gelu8_dev_q8_rows(
8969                &dev.ptr_row,
8970                &sel_d,
8971                zq,
8972                zd,
8973                t,
8974                n_embd,
8975                n_ff_exp,
8976                n_used,
8977                n_expert,
8978                m.gate_exps.qtype,
8979                m.up_exps.qtype,
8980                m.gate_exps.row_bytes,
8981                m.up_exps.row_bytes,
8982            )?
8983        };
8984        let (aq2, ad2) = e.quantize_q8_1(&act, t * n_used, n_ff_exp)?;
8985        let mut moe_out = e.uninit(t * n_embd)?;
8986        // down stays rows_g: the CSR dedup twin measured NEGATIVE at nsb=22 too (189.4 vs
8987        // 207.0 depth spec, bitwise-exact — jsonl 2026-07-10; qwen nsb=16 same verdict).
8988        e.moe_down8_fma_dev_q8_rows_g(
8989            &dev.ptr_row,
8990            &sel_d,
8991            &w_d,
8992            &aq2,
8993            &ad2,
8994            &mut moe_out,
8995            t,
8996            n_ff_exp,
8997            n_embd,
8998            n_used,
8999            n_expert,
9000            m.down_exps.qtype,
9001            m.down_exps.row_bytes,
9002        )?;
9003        Ok(moe_out)
9004    }
9005
9006    /// gemma4 MoE (R2 router prologue input supplied by caller, R3 per-expert output scale).
9007    /// Sequential host-staged v0 — softmax gating + renorm (moe_route, the qwen recipe), GELU
9008    /// experts, scale folded into the accumulate weight (post-matmul linear scale, exact fold).
9009    fn gemma4_moe(
9010        &self,
9011        e: &Engine,
9012        m: &crate::hybrid::MoeWeights,
9013        bits: &crate::hybrid::Gemma4MoeBits,
9014        moe_in: &CudaSlice<f32>,
9015        router_in: &CudaSlice<f32>,
9016        t: usize,
9017    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9018        let cfg = &self.cfg;
9019        let moe = cfg.moe.as_ref().unwrap();
9020        let n_embd = cfg.n_embd as usize;
9021        let n_expert = moe.expert_count as usize;
9022        let n_used = moe.expert_used_count as usize;
9023        let n_ff_exp = moe.expert_ff_length as usize;
9024
9025        // Router: the in-house GEMV for ALL small t (decode AND verify ride the same per-column
9026        // kernel — the cuBLASLt n-dependence flipped top-k at verify t on the 27B, d994271);
9027        // batched matmul only at real prefill.
9028        let logits = if t < PRIME_MIN_T && crate::router_kernel_on() {
9029            e.router_gemv(m.gate_inp.float_data(), router_in, n_embd, n_expert, t)?
9030        } else {
9031            e.matmul(&m.gate_inp, router_in, t)?
9032        };
9033
9034        // FAST SMALL-T ARM (decode t=1 AND spec verify t=2..15): device softmax-topk router,
9035        // then PER TOKEN the same fused gate_up GELU + down8 FMA launch pair over the resident
9036        // dev slabs — verify rides the EXACT decode kernel chain per token (dispatch-parity
9037        // law; the qwen "verify must be kernel-dispatch-identical to decode" lesson).
9038        if t < PRIME_MIN_T
9039            && m.dev_exps.as_ref().is_some_and(|d| !d.gu_il)
9040            && expert_dp4a_supported(m.gate_exps.qtype)
9041            && expert_dp4a_supported(m.up_exps.qtype)
9042            && expert_dp4a_supported(m.down_exps.qtype)
9043            && std::env::var("MEMRA_GEMMA_MOE_FAST").as_deref() != Ok("0")
9044        {
9045            let dev = m.dev_exps.as_ref().unwrap();
9046            let (sel_d, w_d) =
9047                e.moe_router_topk_scaled(&logits, t, n_expert, n_used, &bits.per_expert_scale_d)?;
9048            if t == 1 {
9049                let (zq, zd) = e.quantize_q8_1(moe_in, 1, n_embd)?;
9050                let selv = sel_d.slice(0..n_used);
9051                let wv = w_d.slice(0..n_used);
9052                let act = e.moe_gate_up_gelu8_dev_q8(
9053                    &dev.ptr_row,
9054                    &selv,
9055                    &zq,
9056                    &zd,
9057                    n_embd,
9058                    n_ff_exp,
9059                    n_used,
9060                    n_expert,
9061                    m.gate_exps.qtype,
9062                    m.up_exps.qtype,
9063                    m.gate_exps.row_bytes,
9064                    m.up_exps.row_bytes,
9065                )?;
9066                let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
9067                let mut moe_out = e.uninit(n_embd)?;
9068                e.moe_down8_fma_dev_q8(
9069                    &dev.ptr_row,
9070                    &selv,
9071                    &wv,
9072                    &aq2,
9073                    &ad2,
9074                    &mut moe_out.slice_mut(0..n_embd),
9075                    n_ff_exp,
9076                    n_embd,
9077                    n_used,
9078                    n_expert,
9079                    m.down_exps.qtype,
9080                    m.down_exps.row_bytes,
9081                )?;
9082                return Ok(moe_out);
9083            }
9084            // VERIFY ROWS TWINS (t=2..15): ONE launch pair for all tokens; per (token,row,slot)
9085            // bodies are the t=1 kernels VERBATIM (bit-identical to the per-token loop).
9086            // gate_up rides the CSR owner-scan dedup when t <= 10 (each duplicated expert's
9087            // weight stream decoded once; the qwen CSR class). MEMRA_GEMMA_CSR=0 -> rows.
9088            let (zq, zd) = e.quantize_q8_1(moe_in, t, n_embd)?;
9089            let csr = t <= 10 && std::env::var("MEMRA_GEMMA_CSR").as_deref() != Ok("0");
9090            let act = if csr {
9091                e.moe_gate_up_gelu8_dev_q8_csr(
9092                    &dev.ptr_row,
9093                    &sel_d,
9094                    &zq,
9095                    &zd,
9096                    t * n_used,
9097                    n_embd,
9098                    n_ff_exp,
9099                    n_used,
9100                    n_expert,
9101                    m.gate_exps.qtype,
9102                    m.up_exps.qtype,
9103                    m.gate_exps.row_bytes,
9104                    m.up_exps.row_bytes,
9105                )?
9106            } else {
9107                e.moe_gate_up_gelu8_dev_q8_rows(
9108                    &dev.ptr_row,
9109                    &sel_d,
9110                    &zq,
9111                    &zd,
9112                    t,
9113                    n_embd,
9114                    n_ff_exp,
9115                    n_used,
9116                    n_expert,
9117                    m.gate_exps.qtype,
9118                    m.up_exps.qtype,
9119                    m.gate_exps.row_bytes,
9120                    m.up_exps.row_bytes,
9121                )?
9122            };
9123            let (aq2, ad2) = e.quantize_q8_1(&act, t * n_used, n_ff_exp)?;
9124            let mut moe_out = e.uninit(t * n_embd)?;
9125            e.moe_down8_fma_dev_q8_rows_g(
9126                &dev.ptr_row,
9127                &sel_d,
9128                &w_d,
9129                &aq2,
9130                &ad2,
9131                &mut moe_out,
9132                t,
9133                n_ff_exp,
9134                n_embd,
9135                n_used,
9136                n_expert,
9137                m.down_exps.qtype,
9138                m.down_exps.row_bytes,
9139            )?;
9140            return Ok(moe_out);
9141        }
9142
9143        let (sel_all, mut w_all) = Self::moe_route(e, &logits, t, n_expert, n_used)?;
9144        for (i, &sx) in sel_all.iter().enumerate() {
9145            w_all[i] *= bits.per_expert_scale[sx as usize];
9146        }
9147
9148        // PREFILL PAIRS ARM (t >= 16): expert-major CSR over the resident slabs — ONE launch per
9149        // projection covers ALL (token,expert) pairs (the qwen pairs recipe, _em dot = expert_dot_g
9150        // which has the Q4_0 body; GELU pairs epilogue; R3 scale folded into pair_w).
9151        if t >= PRIME_MIN_T
9152            && m.dev_exps.as_ref().is_some_and(|d| !d.gu_il)
9153            && expert_dp4a_supported(m.gate_exps.qtype)
9154            && expert_dp4a_supported(m.up_exps.qtype)
9155            && expert_dp4a_supported(m.down_exps.qtype)
9156            && std::env::var("MEMRA_GEMMA_MOE_PAIRS").as_deref() != Ok("0")
9157        {
9158            let dev = m.dev_exps.as_ref().unwrap();
9159            let n_pairs = t * n_used;
9160            let pair_ex: Vec<i32> = sel_all.iter().map(|&x| x as i32).collect();
9161            let pair_tok: Vec<i32> = (0..n_pairs).map(|p| (p / n_used) as i32).collect();
9162            let tok_off: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
9163            let tok_ids: Vec<i32> = (0..n_pairs as i32).collect();
9164            let pt = e.htod_i32(&pair_tok)?;
9165            let pw = e.htod(&w_all)?;
9166            let toff = e.htod_i32(&tok_off)?;
9167            let tids = e.htod_i32(&tok_ids)?;
9168            let mut by_ex: Vec<Vec<i32>> = vec![Vec::new(); n_expert];
9169            for p in 0..n_pairs {
9170                by_ex[pair_ex[p] as usize].push(p as i32);
9171            }
9172            let mut ex_ids: Vec<i32> = Vec::new();
9173            let mut ex_off: Vec<i32> = vec![0];
9174            let mut ex_pairs: Vec<i32> = Vec::with_capacity(n_pairs);
9175            for (ex, list) in by_ex.iter().enumerate() {
9176                if list.is_empty() {
9177                    continue;
9178                }
9179                ex_ids.push(ex as i32);
9180                ex_pairs.extend_from_slice(list);
9181                ex_off.push(ex_pairs.len() as i32);
9182            }
9183            let n_active = ex_ids.len();
9184            let exi = e.htod_i32(&ex_ids)?;
9185            let exo = e.htod_i32(&ex_off)?;
9186            let exp_d = e.htod_i32(&ex_pairs)?;
9187            // GROUPED f16 LANE (MEMRA_MOE_F16G=1, round 46 arc 2): dequant active experts to
9188            // f16 once per projection + one grouped f16 GEMM over the CSR groups; CSR order
9189            // end-to-end (gelu is elementwise), one row permute before the scatter. The
9190            // ragged down k (704) needs no padding here — cublas takes any k.
9191            // f16-mirror numeric class, argmax/spec gated. PER-MODEL default OFF (round 50):
9192            // the gelu class regressed g26 board-2048 prefill -8.3% under the round-49
9193            // Hopper default — see moe_f16g_gemma_on.
9194            if crate::moe_f16g_gemma_on()
9195                && f16g_proj_ok(m.gate_exps.qtype, n_embd)
9196                && f16g_proj_ok(m.up_exps.qtype, n_embd)
9197                && f16g_proj_ok(m.down_exps.qtype, n_ff_exp)
9198            {
9199                let csr_tok: Vec<i32> = ex_pairs.iter().map(|&p| p / n_used as i32).collect();
9200                let csr_tok_d = e.htod_i32(&csr_tok)?;
9201                let (z_f16, z_s) = e.moe_f16g_act(moe_in, Some(&csr_tok_d), n_embd, n_pairs)?;
9202                let g_csr = e.moe_f16_grouped(
9203                    &dev.ptr_row,
9204                    0,
9205                    n_expert,
9206                    &exi,
9207                    &ex_off,
9208                    &exo,
9209                    &z_f16,
9210                    &z_s,
9211                    n_embd,
9212                    n_ff_exp,
9213                    n_active,
9214                    n_pairs,
9215                    m.gate_exps.qtype,
9216                    m.gate_exps.row_bytes,
9217                )?;
9218                let u_csr = e.moe_f16_grouped(
9219                    &dev.ptr_row,
9220                    1,
9221                    n_expert,
9222                    &exi,
9223                    &ex_off,
9224                    &exo,
9225                    &z_f16,
9226                    &z_s,
9227                    n_embd,
9228                    n_ff_exp,
9229                    n_active,
9230                    n_pairs,
9231                    m.up_exps.qtype,
9232                    m.up_exps.row_bytes,
9233                )?;
9234                let act_csr = e.moe_pairs_gelu_mul(&g_csr, &u_csr, n_pairs * n_ff_exp)?;
9235                let (a_f16, a_s) = e.moe_f16g_act(&act_csr, None, n_ff_exp, n_pairs)?;
9236                let d_csr = e.moe_f16_grouped(
9237                    &dev.ptr_row,
9238                    2,
9239                    n_expert,
9240                    &exi,
9241                    &ex_off,
9242                    &exo,
9243                    &a_f16,
9244                    &a_s,
9245                    n_ff_exp,
9246                    n_embd,
9247                    n_active,
9248                    n_pairs,
9249                    m.down_exps.qtype,
9250                    m.down_exps.row_bytes,
9251                )?;
9252                let y_down = e.rows_permute(&d_csr, &exp_d, n_pairs, n_embd)?;
9253                let mut moe_out = e.uninit(t * n_embd)?;
9254                e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
9255                if std::env::var("MEMRA_F16G_DEBUG").is_ok() {
9256                    let scan = |v: &[f32]| v.iter().filter(|x| !x.is_finite()).count();
9257                    let (yd, mo) = (e.dtoh(&y_down)?, e.dtoh(&moe_out)?);
9258                    eprintln!(
9259                        "[f16g-debug] post-permute bad={} post-scatter bad={}",
9260                        scan(&yd),
9261                        scan(&mo)
9262                    );
9263                }
9264                return Ok(moe_out);
9265            }
9266            // gate/up: int8-MMA expert GEMM (in_f = n_embd 2816 = 11x256 tiles ok); down keeps
9267            // the decode-once dp4a (in_f 704 fails the 256-superblock tiling).
9268            let mma =
9269                n_embd % 256 == 0 && std::env::var("MEMRA_GEMMA_MOE_MMA").as_deref() != Ok("0");
9270            let (gate, up) = if mma {
9271                let z_scr = e.mmq_iq_quantize_act(moe_in, n_embd, t)?;
9272                (
9273                    e.mmq_iq_experts(
9274                        &dev.ptr_row,
9275                        0,
9276                        n_expert,
9277                        &exi,
9278                        &exo,
9279                        &exp_d,
9280                        &pt,
9281                        &z_scr,
9282                        n_embd,
9283                        n_ff_exp,
9284                        n_active,
9285                        n_pairs,
9286                        t,
9287                        m.gate_exps.qtype,
9288                        m.gate_exps.row_bytes,
9289                    )?,
9290                    e.mmq_iq_experts(
9291                        &dev.ptr_row,
9292                        1,
9293                        n_expert,
9294                        &exi,
9295                        &exo,
9296                        &exp_d,
9297                        &pt,
9298                        &z_scr,
9299                        n_embd,
9300                        n_ff_exp,
9301                        n_active,
9302                        n_pairs,
9303                        t,
9304                        m.up_exps.qtype,
9305                        m.up_exps.row_bytes,
9306                    )?,
9307                )
9308            } else {
9309                let (zq, zd) = e.quantize_q8_1(moe_in, t, n_embd)?;
9310                (
9311                    e.moe_pairs_matvec_q8_dec(
9312                        &dev.ptr_row,
9313                        0,
9314                        &exi,
9315                        &exo,
9316                        &exp_d,
9317                        &pt,
9318                        &zq,
9319                        &zd,
9320                        n_embd,
9321                        n_ff_exp,
9322                        n_expert,
9323                        n_active,
9324                        n_pairs,
9325                        m.gate_exps.qtype,
9326                        m.gate_exps.row_bytes,
9327                    )?,
9328                    e.moe_pairs_matvec_q8_dec(
9329                        &dev.ptr_row,
9330                        1,
9331                        &exi,
9332                        &exo,
9333                        &exp_d,
9334                        &pt,
9335                        &zq,
9336                        &zd,
9337                        n_embd,
9338                        n_ff_exp,
9339                        n_expert,
9340                        n_active,
9341                        n_pairs,
9342                        m.up_exps.qtype,
9343                        m.up_exps.row_bytes,
9344                    )?,
9345                )
9346            };
9347            let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
9348            let pself = e.htod_i32(&pair_self)?;
9349            // DOWN through the int8-MMA expert GEMM (2026-07-31, g26 prefill lever): the
9350            // ragged k (n_ff_exp=704 on the 26B) rides a PADDED k-walk — in_f rounds up
9351            // to the 256-val superblock (768) while the act quantizer's zero padding
9352            // makes every padded-k product exactly zero (weight overread bytes multiply
9353            // zero int8 act values; the dev slab carries 144B tail slack for the OOB).
9354            // The old dp4a matvec was 11.3ms/call at m=T (the 0.07x prefill wall).
9355            // MEMRA_GEMMA_MOE_MMA=0 reverts down together with gate/up.
9356            // FUSED ACT-EPILOGUE (default on): gelu_tanh(gate)*up + D4 quantize in one
9357            // launch — no f32 act buffer (the fused kernel zero-pads the ragged tail
9358            // exactly like the two-pass quantizer). MEMRA_MOE_FUSE_ACTQ=0 rollback.
9359            // Scratch bytes are BYTE-IDENTICAL (kernel-check gated).
9360            let y_down = if mma {
9361                let in_pad = n_ff_exp.div_ceil(256) * 256;
9362                let a_scr = if crate::moe_fuse_actq_on() {
9363                    e.mmq_iq_fused_act_quant(&gate, &up, n_ff_exp, n_pairs, 1)?
9364                } else {
9365                    let act = e.moe_pairs_gelu_mul(&gate, &up, n_pairs * n_ff_exp)?;
9366                    e.mmq_iq_quantize_act(&act, n_ff_exp, n_pairs)?
9367                };
9368                e.mmq_iq_experts(
9369                    &dev.ptr_row,
9370                    2,
9371                    n_expert,
9372                    &exi,
9373                    &exo,
9374                    &exp_d,
9375                    &pself,
9376                    &a_scr,
9377                    in_pad,
9378                    n_embd,
9379                    n_active,
9380                    n_pairs,
9381                    n_pairs,
9382                    m.down_exps.qtype,
9383                    m.down_exps.row_bytes,
9384                )?
9385            } else {
9386                let act = e.moe_pairs_gelu_mul(&gate, &up, n_pairs * n_ff_exp)?;
9387                let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
9388                e.moe_pairs_matvec_q8_dec(
9389                    &dev.ptr_row,
9390                    2,
9391                    &exi,
9392                    &exo,
9393                    &exp_d,
9394                    &pself,
9395                    &aq2,
9396                    &ad2,
9397                    n_ff_exp,
9398                    n_embd,
9399                    n_expert,
9400                    n_active,
9401                    n_pairs,
9402                    m.down_exps.qtype,
9403                    m.down_exps.row_bytes,
9404                )?
9405            };
9406            let mut moe_out = e.uninit(t * n_embd)?;
9407            e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
9408            return Ok(moe_out);
9409        }
9410
9411        let g_len = m.gate_exps.expert_stride;
9412        let u_len = m.up_exps.expert_stride;
9413        let d_len = m.down_exps.expert_stride;
9414        // Resident dev slabs (fits-VRAM regime): read each expert straight from the device slab
9415        // at ex*stride — zero H2D, SAME qmatvec_view kernel/bytes as the staged path. Staging is
9416        // the spill fallback.
9417        let dev = m.dev_exps.as_ref().filter(|d| !d.gu_il);
9418        let (mut sg, mut su, mut sd) = if dev.is_some() {
9419            (None, None, None)
9420        } else {
9421            (
9422                Some(e.alloc_u8_uninit(g_len)?),
9423                Some(e.alloc_u8_uninit(u_len)?),
9424                Some(e.alloc_u8_uninit(d_len)?),
9425            )
9426        };
9427        let mut moe_out = e.zeros(t * n_embd)?;
9428        for tok in 0..t {
9429            let sel = &sel_all[tok * n_used..(tok + 1) * n_used];
9430            let w = &w_all[tok * n_used..(tok + 1) * n_used];
9431            let zt = moe_in.slice(tok * n_embd..(tok + 1) * n_embd);
9432            for (j, &ex) in sel.iter().enumerate() {
9433                let ex = ex as usize;
9434                let gate = match dev {
9435                    Some(d) => e.qmatvec_view(
9436                        &d.gate,
9437                        ex * g_len..(ex + 1) * g_len,
9438                        &zt,
9439                        1,
9440                        m.gate_exps.in_f,
9441                        m.gate_exps.out_f,
9442                        m.gate_exps.qtype,
9443                        m.gate_exps.row_bytes,
9444                    )?,
9445                    None => {
9446                        let sg = sg.as_mut().unwrap();
9447                        e.stage_expert(m.gate_exps.expert_bytes(ex), sg, 0)?;
9448                        e.qmatvec_view(
9449                            sg,
9450                            0..g_len,
9451                            &zt,
9452                            1,
9453                            m.gate_exps.in_f,
9454                            m.gate_exps.out_f,
9455                            m.gate_exps.qtype,
9456                            m.gate_exps.row_bytes,
9457                        )?
9458                    }
9459                };
9460                let up = match dev {
9461                    Some(d) => e.qmatvec_view(
9462                        &d.up,
9463                        ex * u_len..(ex + 1) * u_len,
9464                        &zt,
9465                        1,
9466                        m.up_exps.in_f,
9467                        m.up_exps.out_f,
9468                        m.up_exps.qtype,
9469                        m.up_exps.row_bytes,
9470                    )?,
9471                    None => {
9472                        let su = su.as_mut().unwrap();
9473                        e.stage_expert(m.up_exps.expert_bytes(ex), su, 0)?;
9474                        e.qmatvec_view(
9475                            su,
9476                            0..u_len,
9477                            &zt,
9478                            1,
9479                            m.up_exps.in_f,
9480                            m.up_exps.out_f,
9481                            m.up_exps.qtype,
9482                            m.up_exps.row_bytes,
9483                        )?
9484                    }
9485                };
9486                let mut act = e.uninit(n_ff_exp)?;
9487                e.gelu_tanh_mul(&gate, &up, &mut act, n_ff_exp)?;
9488                let actv = act.slice(0..n_ff_exp);
9489                let y = match dev {
9490                    Some(d) => e.qmatvec_view(
9491                        &d.down,
9492                        ex * d_len..(ex + 1) * d_len,
9493                        &actv,
9494                        1,
9495                        m.down_exps.in_f,
9496                        m.down_exps.out_f,
9497                        m.down_exps.qtype,
9498                        m.down_exps.row_bytes,
9499                    )?,
9500                    None => {
9501                        let sd = sd.as_mut().unwrap();
9502                        e.stage_expert(m.down_exps.expert_bytes(ex), sd, 0)?;
9503                        e.qmatvec_view(
9504                            sd,
9505                            0..d_len,
9506                            &actv,
9507                            1,
9508                            m.down_exps.in_f,
9509                            m.down_exps.out_f,
9510                            m.down_exps.qtype,
9511                            m.down_exps.row_bytes,
9512                        )?
9513                    }
9514                };
9515                let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
9516                e.axpy_into(&y, w[j], &mut dst, n_embd)?;
9517            }
9518        }
9519        Ok(moe_out)
9520    }
9521
9522    /// One gemma4 trunk layer (R8): x -> x_next.
9523    fn gemma4_layer(
9524        &self,
9525        e: &Engine,
9526        il: usize,
9527        layer: &crate::hybrid::HybridLayer,
9528        x: &CudaSlice<f32>,
9529        pos_d: &CudaSlice<i32>,
9530        t: usize,
9531    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9532        let n_embd = self.cfg.n_embd as usize;
9533        let eps = self.cfg.rms_eps;
9534
9535        let mut h = e.zeros(t * n_embd)?;
9536        e.rms_norm(x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
9537        let Mixer::Full(fa) = &layer.mixer else {
9538            panic!("gemma4 layer {il} not full-attn")
9539        };
9540        let o = self.gemma4_attn(e, fa, il, &h, pos_d, t)?;
9541        // gemma order: post_attention_norm applies to the ATTENTION OUTPUT, then the residual.
9542        let mut cur = e.zeros(t * n_embd)?;
9543        e.rms_norm(
9544            &o,
9545            layer.post_attn_norm.float_data(),
9546            &mut cur,
9547            n_embd,
9548            t,
9549            eps,
9550        )?;
9551        self.gemma4_layer_tail_add(e, layer, &cur, x, t)
9552    }
9553
9554    /// Everything after the attention output in a gemma4 layer: the residual add (cur + x ->
9555    /// attn_out) FUSED with the three attn_out norms, then shared FFN + router + MoE + combine +
9556    /// layer scale — shared verbatim by the prefill, decode and verify paths.
9557    fn gemma4_layer_tail_add(
9558        &self,
9559        e: &Engine,
9560        layer: &crate::hybrid::HybridLayer,
9561        cur: &CudaSlice<f32>,
9562        x: &CudaSlice<f32>,
9563        t: usize,
9564    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9565        Ok(self.gemma4_layer_tail_add_n(e, layer, cur, x, t, None)?.0)
9566    }
9567
9568    /// tail_add with the NEXT layer's attn_norm fused into the closing add+scale (one launch
9569    /// produces both x_next and h_next — the cross-layer fusion). `next_norm` None = last layer.
9570    fn gemma4_layer_tail_add_n(
9571        &self,
9572        e: &Engine,
9573        layer: &crate::hybrid::HybridLayer,
9574        cur: &CudaSlice<f32>,
9575        x: &CudaSlice<f32>,
9576        t: usize,
9577        next_norm: Option<&CudaSlice<f32>>,
9578    ) -> Result<(CudaSlice<f32>, Option<CudaSlice<f32>>), Box<dyn std::error::Error>> {
9579        let n_embd = self.cfg.n_embd as usize;
9580        let bits = layer.gemma4.as_ref().unwrap();
9581        let (sn, attn_out) = self.gemma4_layer_tail_core(e, layer, cur, x, t)?;
9582        let mut xn = e.uninit(t * n_embd)?;
9583        match next_norm {
9584            Some(w) => {
9585                let mut hn = e.uninit(t * n_embd)?;
9586                e.add_scale_rms_norm(
9587                    &sn,
9588                    &attn_out,
9589                    bits.layer_scale,
9590                    w,
9591                    &mut xn,
9592                    &mut hn,
9593                    n_embd,
9594                    t,
9595                    self.cfg.rms_eps,
9596                )?;
9597                Ok((xn, Some(hn)))
9598            }
9599            None => {
9600                e.add_scale(&sn, &attn_out, bits.layer_scale, &mut xn, t * n_embd)?;
9601                Ok((xn, None))
9602            }
9603        }
9604    }
9605
9606    /// Tail core: attn_out = cur+x (fused with the 3 norms), both FFN branches, the combine
9607    /// norm — returns (sn, attn_out) for the closing add+scale variants.
9608    fn gemma4_layer_tail_core(
9609        &self,
9610        e: &Engine,
9611        layer: &crate::hybrid::HybridLayer,
9612        cur: &CudaSlice<f32>,
9613        x: &CudaSlice<f32>,
9614        t: usize,
9615    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9616        self.gemma4_layer_tail_core_pn(e, layer, cur, x, t, None, false)
9617    }
9618
9619    /// tail_core with an optional PRE-NORM fold (glue-fusion lane): `pre_norm = Some(wa)`
9620    /// means `cur` is the RAW attention output and the dense entry runs
9621    /// rms(cur, wa) + residual-add + ffn_norm as ONE launch (E4B's post-attn norm fold).
9622    /// `defer_post_norm`: dense-arm exit returns RAW f0 (ffn_down output) instead of
9623    /// sn = rms(f0, post_ffw) — the caller fuses the post-norm into its residual emit
9624    /// (rms_pre_add_q8_1, E4B glue wave 5). MoE arm ignores it.
9625    fn gemma4_layer_tail_core_pn(
9626        &self,
9627        e: &Engine,
9628        layer: &crate::hybrid::HybridLayer,
9629        cur: &CudaSlice<f32>,
9630        x: &CudaSlice<f32>,
9631        t: usize,
9632        pre_norm: Option<&CudaSlice<f32>>,
9633        defer_post_norm: bool,
9634    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9635        let n_embd = self.cfg.n_embd as usize;
9636        let eps = self.cfg.rms_eps;
9637        let bits = layer.gemma4.as_ref().unwrap();
9638
9639        // DENSE gemma4 variants (31B/E4B — no MoE, no parallel branch): attn_out = cur + x
9640        // fused with the single ffn_norm; GELU_PAR ffn; post_ffw_norm.
9641        let Some(mbits) = bits.moe_bits.as_ref() else {
9642            let crate::hybrid::Ffn::Dense {
9643                ffn_gate,
9644                ffn_up,
9645                ffn_down,
9646            } = &layer.ffn
9647            else {
9648                panic!("gemma4 dense layer without Dense ffn")
9649            };
9650            let mut attn_out = e.uninit(t * n_embd)?;
9651            let mut zsh = e.uninit(t * n_embd)?;
9652            // wave-2: with the pre-norm fold active the entry ALSO emits zsh q8_1 — the
9653            // t=1 fused2 gate/up consume the pair with no standalone quantize launch.
9654            let mut zpair: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
9655            match pre_norm {
9656                Some(wa) if t == 1 => {
9657                    zpair = Some(e.rms_pre_add_rms_norm_q8z(
9658                        cur,
9659                        wa,
9660                        x,
9661                        bits.ffn_norm.float_data(),
9662                        &mut attn_out,
9663                        &mut zsh,
9664                        n_embd,
9665                        t,
9666                        eps,
9667                    )?);
9668                }
9669                Some(wa) => e.rms_pre_add_rms_norm(
9670                    cur,
9671                    wa,
9672                    x,
9673                    bits.ffn_norm.float_data(),
9674                    &mut attn_out,
9675                    &mut zsh,
9676                    n_embd,
9677                    t,
9678                    eps,
9679                )?,
9680                None => e.add_rms_norm(
9681                    cur,
9682                    x,
9683                    bits.ffn_norm.float_data(),
9684                    &mut attn_out,
9685                    &mut zsh,
9686                    n_embd,
9687                    t,
9688                    eps,
9689                )?,
9690            }
9691            let n_ff = ffn_gate.out_features();
9692            // FFN persistent slab (counter-barrier form) FALSIFIED here 2026-07-14
9693            // (falsification #7, jsonl row): PDL glue already hides the launch boundaries
9694            // it fused; barrier + worst-segment occupancy net −0.3% (31B depth) / −2.3%
9695            // (E4B spec). Down's act dependency is all-to-all, so sentinel sync cannot
9696            // rescue segment C — the megakernel front is closed for the dense tail.
9697            let (gate, up) = if t == 1 {
9698                let (zq, zd) = match zpair {
9699                    Some(p) => p,
9700                    None => e.quantize_q8_1(&zsh, 1, n_embd)?,
9701                };
9702                match e.matmul_q4_fused2(ffn_gate, ffn_up, &zq, &zd)? {
9703                    Some(p) => p,
9704                    // NVFP4mix dense trunk: gate/up are both NVFP4 (down stays Q8_0).
9705                    None => match e.matmul_nvfp4_fused2(ffn_gate, ffn_up, &zq, &zd, 1)? {
9706                        Some(p) => p,
9707                        None => (
9708                            e.matmul_pre(ffn_gate, &zq, &zd, &zsh, 1)?,
9709                            e.matmul_pre(ffn_up, &zq, &zd, &zsh, 1)?,
9710                        ),
9711                    },
9712                }
9713            } else {
9714                // BATCHED FUSED2 (DEFAULT ON 2026-07-13, MEMRA_F2B=0 seam): one segmented
9715                // launch for the verify's gate+up — the up segment's blocks fill SMs as
9716                // the gate segment drains (the launch-tail mechanism behind the b-tier
9717                // plateau; first positive after six falsified in-kernel variants).
9718                static F2B: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9719                let f2b = *F2B.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0"));
9720                let fused = if f2b {
9721                    let (zq, zd) = e.quantize_q8_1(&zsh, t, n_embd)?;
9722                    e.matmul_q4_fused2_batched(ffn_gate, ffn_up, &zq, &zd, t)?
9723                } else {
9724                    None
9725                };
9726                match fused {
9727                    Some(p) => p,
9728                    None => {
9729                        // quantize-once window: gate/up share `zsh` (borrowed across the pair).
9730                        e.mmq_act_begin();
9731                        (e.matmul(ffn_gate, &zsh, t)?, e.matmul(ffn_up, &zsh, t)?)
9732                    }
9733                }
9734            };
9735            let mut act = e.uninit(t * n_ff)?;
9736            // act quantize folds into the GELU epilogue (bit-identical q8_1 rounding);
9737            // ffn_down rides matmul_pre — one quantize launch fewer per layer.
9738            let f0 = if e.uses_q8_1_fast(ffn_down) {
9739                let upv = e.view(&up, t * n_ff);
9740                let up_all = upv.slice(0..t * n_ff);
9741                let (aq, ad) = e.gelu_tanh_mul_q8_1(&gate, &up_all, &mut act, n_ff, t)?;
9742                e.matmul_pre(ffn_down, &aq, &ad, &act, t)?
9743            } else {
9744                e.gelu_tanh_mul(&gate, &up, &mut act, t * n_ff)?;
9745                e.matmul(ffn_down, &act, t)?
9746            };
9747            if defer_post_norm {
9748                return Ok((f0, attn_out));
9749            }
9750            let mut sn = e.uninit(t * n_embd)?;
9751            e.rms_norm(
9752                &f0,
9753                bits.post_ffw_norm.float_data(),
9754                &mut sn,
9755                n_embd,
9756                t,
9757                eps,
9758            )?;
9759            return Ok((sn, attn_out));
9760        };
9761
9762        assert!(pre_norm.is_none(), "pre-norm fold is dense-entry only");
9763        // MoE variant (26B): attn_out = cur + x fused with the three attn_out norms
9764        // (ffn_norm + router-scale + pre_ffw_norm_2): ONE launch, chains verbatim. At small t
9765        // (decode + verify) the zsh and moe_in outputs are EMITTED q8_1 (both consumers are
9766        // quantized matmuls — two quantize launches + two f32 round-trips fold away).
9767        let mut attn_out = e.uninit(t * n_embd)?;
9768        let mut router_in = e.uninit(t * n_embd)?;
9769        let fast_moe = match &layer.ffn {
9770            crate::hybrid::Ffn::Moe(m) => {
9771                m.dev_exps.as_ref().is_some_and(|d| !d.gu_il)
9772                    && expert_dp4a_supported(m.gate_exps.qtype)
9773                    && expert_dp4a_supported(m.up_exps.qtype)
9774                    && expert_dp4a_supported(m.down_exps.qtype)
9775                    && std::env::var("MEMRA_GEMMA_MOE_FAST").as_deref() != Ok("0")
9776            }
9777            _ => false,
9778        };
9779        let q8z = t < PRIME_MIN_T && fast_moe;
9780        let (zsh_f32, zsh_q8, moe_q8) = if q8z {
9781            let (z0, m2) = e.add_rms_norm3_q8z(
9782                cur,
9783                x,
9784                bits.ffn_norm.float_data(),
9785                &mbits.router_scale_pre,
9786                mbits.pre_ffw_norm_2.float_data(),
9787                &mut attn_out,
9788                &mut router_in,
9789                n_embd,
9790                t,
9791                eps,
9792            )?;
9793            (None, Some(z0), Some(m2))
9794        } else {
9795            let mut zsh = e.uninit(t * n_embd)?;
9796            let mut moe_in = e.uninit(t * n_embd)?;
9797            e.add_rms_norm3(
9798                cur,
9799                x,
9800                bits.ffn_norm.float_data(),
9801                &mbits.router_scale_pre,
9802                mbits.pre_ffw_norm_2.float_data(),
9803                &mut attn_out,
9804                &mut zsh,
9805                &mut router_in,
9806                &mut moe_in,
9807                n_embd,
9808                t,
9809                eps,
9810            )?;
9811            (Some((zsh, moe_in)), None, None)
9812        };
9813        let attn_out2 = attn_out;
9814        #[allow(unused_variables)]
9815        let attn_out = &attn_out2;
9816        let n_ff = mbits.shared_gate.out_features();
9817        let (gate, up) = if let Some((zq, zd)) = zsh_q8.as_ref() {
9818            if t == 1 {
9819                match e.matmul_q4_fused2(&mbits.shared_gate, &mbits.shared_up, zq, zd)? {
9820                    Some(p) => p,
9821                    None => match e.matmul_nvfp4_fused2(
9822                        &mbits.shared_gate,
9823                        &mbits.shared_up,
9824                        zq,
9825                        zd,
9826                        1,
9827                    )? {
9828                        Some(p) => p,
9829                        None => {
9830                            let h0 = e.zeros(0)?;
9831                            (
9832                                e.matmul_pre(&mbits.shared_gate, zq, zd, &h0, 1)?,
9833                                e.matmul_pre(&mbits.shared_up, zq, zd, &h0, 1)?,
9834                            )
9835                        }
9836                    },
9837                }
9838            } else {
9839                // verify t 2..15: the batched mmvq twins consume the pre-quantized pair.
9840                let h0 = e.zeros(0)?;
9841                (
9842                    e.matmul_pre(&mbits.shared_gate, zq, zd, &h0, t)?,
9843                    e.matmul_pre(&mbits.shared_up, zq, zd, &h0, t)?,
9844                )
9845            }
9846        } else {
9847            let (zsh, _) = zsh_f32.as_ref().unwrap();
9848            (
9849                e.matmul(&mbits.shared_gate, zsh, t)?,
9850                e.matmul(&mbits.shared_up, zsh, t)?,
9851            )
9852        };
9853        let mut act = e.uninit(t * n_ff)?;
9854        e.gelu_tanh_mul(&gate, &up, &mut act, t * n_ff)?;
9855        let mlp0 = e.matmul(&mbits.shared_down, &act, t)?;
9856        let crate::hybrid::Ffn::Moe(m) = &layer.ffn else {
9857            panic!("gemma4 layer not MoE")
9858        };
9859        let moe0 = match (&moe_q8, &zsh_f32) {
9860            (Some(mq), _) => self.gemma4_moe_q8(e, m, mbits, mq, &router_in, t)?,
9861            (None, Some((_, moe_in))) => self.gemma4_moe(e, m, mbits, moe_in, &router_in, t)?,
9862            _ => unreachable!(),
9863        };
9864        // post_ffw_norm_1(mlp0) + post_ffw_norm_2(moe0): one fused launch, per-row verbatim.
9865        let mut mlp = e.uninit(t * n_embd)?;
9866        let mut moe = e.uninit(t * n_embd)?;
9867        e.rms_norm2x(
9868            &mlp0,
9869            &moe0,
9870            mbits.post_ffw_norm_1.float_data(),
9871            mbits.post_ffw_norm_2.float_data(),
9872            &mut mlp,
9873            &mut moe,
9874            n_embd,
9875            t,
9876            eps,
9877        )?;
9878
9879        // combine: rms_norm(mlp + moe, post_ffw_norm) + attn_out, then the layer output scalar.
9880        // add+norm fused (add_rms_norm == add then rms_norm, kernel-check-pinned identity).
9881        let mut sum = e.uninit(t * n_embd)?;
9882        let mut sn = e.uninit(t * n_embd)?;
9883        e.add_rms_norm(
9884            &mlp,
9885            &moe,
9886            bits.post_ffw_norm.float_data(),
9887            &mut sum,
9888            &mut sn,
9889            n_embd,
9890            t,
9891            eps,
9892        )?;
9893        Ok((sn, attn_out2))
9894    }
9895
9896    /// tail_add with the next attn_norm emitted PRE-QUANTIZED q8_1 (decode/verify loops).
9897    /// pn-fold front of `gemma4_layer_tail_add_nq` (MEMRA_G4_PNFOLD, GAP-DIAGNOSIS
9898    /// verdict 7): takes the RAW attention output `o` and, on the dense trunk, folds
9899    /// post_attn_norm into the tail entry and post_ffw_norm into the residual exit —
9900    /// the E4B glue chain (13594) backported to the decode/verify trio. Kills the
9901    /// standalone rms_norm(o) + rms_norm(f0) (+ffn quantize) launches per layer.
9902    /// BITS-CHANGING vs the two-launch chain; every caller rides this ONE front so
9903    /// decode == verify == graph parity holds by construction at either seam value.
9904    /// Non-dense (26B MoE) and seam-off fall to the unfused chain unchanged.
9905    pub(crate) fn gemma4_layer_tail_add_nq_pn(
9906        &self,
9907        e: &Engine,
9908        layer: &crate::hybrid::HybridLayer,
9909        o: &CudaSlice<f32>,
9910        x: &CudaSlice<f32>,
9911        t: usize,
9912        next_norm: Option<&CudaSlice<f32>>,
9913    ) -> Result<(CudaSlice<f32>, Option<(CudaSlice<i8>, CudaSlice<f32>)>), Box<dyn std::error::Error>>
9914    {
9915        let n_embd = self.cfg.n_embd as usize;
9916        let eps = self.cfg.rms_eps;
9917        let bits = layer.gemma4.as_ref().unwrap();
9918        if Engine::g4_pnfold_on() && matches!(layer.ffn, crate::hybrid::Ffn::Dense { .. }) {
9919            let (f0, attn_out) = self.gemma4_layer_tail_core_pn(
9920                e,
9921                layer,
9922                o,
9923                x,
9924                t,
9925                Some(layer.post_attn_norm.float_data()),
9926                true,
9927            )?;
9928            let mut xn = e.uninit(t * n_embd)?;
9929            return match next_norm {
9930                Some(w) => {
9931                    let pair = e.rms_pre_add_scale_rms_norm_q8_1(
9932                        &f0,
9933                        bits.post_ffw_norm.float_data(),
9934                        &attn_out,
9935                        bits.layer_scale,
9936                        w,
9937                        &mut xn,
9938                        n_embd,
9939                        t,
9940                        eps,
9941                    )?;
9942                    Ok((xn, Some(pair)))
9943                }
9944                None => {
9945                    let mut sn = e.uninit(t * n_embd)?;
9946                    e.rms_norm(
9947                        &f0,
9948                        bits.post_ffw_norm.float_data(),
9949                        &mut sn,
9950                        n_embd,
9951                        t,
9952                        eps,
9953                    )?;
9954                    e.add_scale(&sn, &attn_out, bits.layer_scale, &mut xn, t * n_embd)?;
9955                    Ok((xn, None))
9956                }
9957            };
9958        }
9959        let mut cur = e.uninit(t * n_embd)?;
9960        e.rms_norm(
9961            o,
9962            layer.post_attn_norm.float_data(),
9963            &mut cur,
9964            n_embd,
9965            t,
9966            eps,
9967        )?;
9968        self.gemma4_layer_tail_add_nq(e, layer, &cur, x, t, next_norm)
9969    }
9970
9971    pub(crate) fn gemma4_layer_tail_add_nq(
9972        &self,
9973        e: &Engine,
9974        layer: &crate::hybrid::HybridLayer,
9975        cur: &CudaSlice<f32>,
9976        x: &CudaSlice<f32>,
9977        t: usize,
9978        next_norm: Option<&CudaSlice<f32>>,
9979    ) -> Result<(CudaSlice<f32>, Option<(CudaSlice<i8>, CudaSlice<f32>)>), Box<dyn std::error::Error>>
9980    {
9981        let n_embd = self.cfg.n_embd as usize;
9982        let bits = layer.gemma4.as_ref().unwrap();
9983        let (sn, attn_out) = self.gemma4_layer_tail_core(e, layer, cur, x, t)?;
9984        let mut xn = e.uninit(t * n_embd)?;
9985        match next_norm {
9986            Some(w) => {
9987                let pair = e.add_scale_rms_norm_q8_1(
9988                    &sn,
9989                    &attn_out,
9990                    bits.layer_scale,
9991                    w,
9992                    &mut xn,
9993                    n_embd,
9994                    t,
9995                    self.cfg.rms_eps,
9996                )?;
9997                Ok((xn, Some(pair)))
9998            }
9999            None => {
10000                e.add_scale(&sn, &attn_out, bits.layer_scale, &mut xn, t * n_embd)?;
10001                Ok((xn, None))
10002            }
10003        }
10004    }
10005
10006    /// gemma4 prefill: `last_only` = forward_last semantics (lm_head on the final row only).
10007    /// R4: final logits softcapped 30*tanh(l/30) on host (monotonic — argmax unaffected).
10008    fn gemma4_forward(
10009        &self,
10010        e: &Engine,
10011        tokens: &[u32],
10012        last_only: bool,
10013    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
10014        // E4B routes to its own forward regardless of the caller's entry point (forward /
10015        // forward_last / prime paths all funnel here for gemma4).
10016        if self.is_gemma4_e4b() {
10017            return self.gemma4_e4b_forward(e, tokens, last_only);
10018        }
10019        let n_embd = self.cfg.n_embd as usize;
10020        let t = tokens.len();
10021        let pos: Vec<i32> = (0..t as i32).collect();
10022        let pos_d = e.htod_i32(&pos)?;
10023
10024        let mut x = self.embed(e, tokens)?;
10025        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
10026        // MEMRA_GEMMA_PROBE=1: per-layer trunk stats (host rms + max of the LAST token row) —
10027        // the bring-up bisect vs llama-eval-callback node stats.
10028        let probe = std::env::var("MEMRA_GEMMA_PROBE").is_ok();
10029        let stat =
10030            |e: &Engine, x: &CudaSlice<f32>, tag: &str| -> Result<(), Box<dyn std::error::Error>> {
10031                let h = e.dtoh(x)?;
10032                let bad = h.iter().filter(|v| !v.is_finite()).count();
10033                let mx = h
10034                    .iter()
10035                    .filter(|v| v.is_finite())
10036                    .fold(0.0f32, |m, v| m.max(v.abs()));
10037                eprintln!(
10038                    "[gemma-probe] {tag}: tok0_first3={:?} bad={bad} max={mx:.3e}",
10039                    &h[..3]
10040                );
10041                Ok(())
10042            };
10043        if probe {
10044            stat(e, &x, "embed")?;
10045        }
10046        for (il, layer) in self.layers.iter().enumerate() {
10047            x = self.gemma4_layer(e, il, layer, &x, &pos_d, t)?;
10048            if probe {
10049                stat(e, &x, &format!("L{il}"))?;
10050            }
10051        }
10052        let mut hn = e.zeros(t * n_embd)?;
10053        e.rms_norm(
10054            &x,
10055            self.output_norm.float_data(),
10056            &mut hn,
10057            n_embd,
10058            t,
10059            self.cfg.rms_eps,
10060        )?;
10061        let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
10062        let n_vocab = self.output.out_features();
10063        let logits = if last_only {
10064            let hv = e.view(&hn, t * n_embd);
10065            let last_row = hv.slice((t - 1) * n_embd..t * n_embd);
10066            let mut hlast = e.zeros(n_embd)?;
10067            e.copy_view_into(&mut hlast, 0, &last_row, n_embd)?;
10068            let mut ld = e.matmul(&self.output, &hlast, 1)?;
10069            e.softcap(&mut ld, cap, n_vocab)?;
10070            self.gemma4_suppress(e, &mut ld, 1)?;
10071            e.dtoh(&ld)?
10072        } else {
10073            let mut ld = e.matmul(&self.output, &hn, t)?;
10074            e.softcap(&mut ld, cap, t * n_vocab)?;
10075            self.gemma4_suppress(e, &mut ld, t)?;
10076            e.dtoh(&ld)?
10077        };
10078        Ok(logits)
10079    }
10080
10081    /// gemma4 BATCHED PROMPT PRIME v0 (fresh cache only): the prefill graph over the whole
10082    /// prompt with each layer's post-rope K / weightless-normed V appended into the quantized
10083    /// KV cache (decode-append row math). Returns (last-row softcapped logits, h_seed = last
10084    /// pre-output_norm hidden, hiddens = full pre-output_norm stack [T, n_embd]).
10085    pub(crate) fn gemma4_prime(
10086        &self,
10087        e: &Engine,
10088        tokens: &[u32],
10089        cache: &mut Cache,
10090        overlay: Option<&crate::vision::EmbedOverlay>,
10091    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10092        // Err, not assert (2026-08-07, lane/gemma4-serve-gaps): a served gemma4 prompt longer
10093        // than the worker's prefill tick used to chunk here, and chunk 2 (pos > 0) killed the
10094        // whole worker process on this line. The worker now primes gemma4 monolithically and
10095        // routes continuation suffixes tokenwise; this is the per-request backstop.
10096        if cache.pos != 0 {
10097            return Err(
10098                "gemma4 prime v0 is fresh-prompt only (no continuation/chunked prime) \
10099                        — prime the full prompt in one call or decode tokenwise"
10100                    .into(),
10101            );
10102        }
10103        let n_embd = self.cfg.n_embd as usize;
10104        let eps = self.cfg.rms_eps;
10105        let t = tokens.len();
10106        let pos: Vec<i32> = (0..t as i32).collect();
10107        let pos_d = e.htod_i32(&pos)?;
10108        let mut x = self.embed(e, tokens)?;
10109        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
10110        // Masked-prefill arm (lane/gemma-vision): splice the tower rows AFTER the
10111        // sqrt(n_embd) text scale — the reference scales token batches only
10112        // (llama.cpp gemma4.cpp:182 `ubatch.token ? sqrtf(n_embd) : 1.0f`), so image
10113        // embeddings enter the trunk raw. Each span becomes an attention ISLAND:
10114        // bidirectional within itself, causal+SWA everywhere else, matching the
10115        // reference's llama_set_causal_attn(false) image batch exactly.
10116        let island: Option<CudaSlice<i32>> = match overlay {
10117            Some(ov) => {
10118                let mut span_id = vec![-1i32; t];
10119                for (i, &(pos, row_off, n_rows)) in ov.spans.iter().enumerate() {
10120                    if pos + n_rows > t {
10121                        return Err(format!(
10122                            "gemma4 overlay span {i} [{pos}, {}) exceeds the prompt ({t})",
10123                            pos + n_rows
10124                        )
10125                        .into());
10126                    }
10127                    let view = ov.rows.slice(row_off * n_embd..(row_off + n_rows) * n_embd);
10128                    e.copy_view_into(&mut x, pos * n_embd, &view, n_rows * n_embd)?;
10129                    for s in span_id.iter_mut().skip(pos).take(n_rows) {
10130                        *s = i as i32;
10131                    }
10132                }
10133                // MEMRA_GV_FORCE_CAUSAL=1: WRONG-ARM probe seam — splice the rows but
10134                // keep the plain causal mask. Exists only so the decisive probe can show
10135                // the island mask itself changes the answer; never on in serving.
10136                if std::env::var("MEMRA_GV_FORCE_CAUSAL").as_deref() == Ok("1") {
10137                    None
10138                } else {
10139                    Some(e.htod_i32(&span_id)?)
10140                }
10141            }
10142            None => None,
10143        };
10144        for (il, layer) in self.layers.iter().enumerate() {
10145            let mut h = e.zeros(t * n_embd)?;
10146            e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
10147            let Mixer::Full(fa) = &layer.mixer else {
10148                panic!("gemma4 layer not full-attn")
10149            };
10150            let trace = il == 0 && std::env::var("MEMRA_G4_PRIME_TRACE").as_deref() == Ok("1");
10151            if trace {
10152                let v = e.dtoh(&h)?;
10153                let nan = v.iter().filter(|x| x.is_nan()).count();
10154                eprintln!("[g4-prime-trace] L0 post-attn_norm: nan={nan}/{}", v.len());
10155            }
10156            let o =
10157                self.gemma4_attn_prime(e, fa, il, &h, &pos_d, t, Some(cache), island.as_ref())?;
10158            if trace {
10159                let v = e.dtoh(&o)?;
10160                let nan = v.iter().filter(|x| x.is_nan()).count();
10161                eprintln!("[g4-prime-trace] L0 post-attn: nan={nan}/{}", v.len());
10162            }
10163            let mut cur = e.zeros(t * n_embd)?;
10164            e.rms_norm(
10165                &o,
10166                layer.post_attn_norm.float_data(),
10167                &mut cur,
10168                n_embd,
10169                t,
10170                eps,
10171            )?;
10172            x = self.gemma4_layer_tail_add(e, layer, &cur, &x, t)?;
10173            self.dflash_tap(e, cache, il, &x, t)?;
10174            // MEMRA_G4_PRIME_TRACE=1: per-layer NaN hunt (recipe-probe diagnostic; off = zero cost)
10175            if std::env::var("MEMRA_G4_PRIME_TRACE").as_deref() == Ok("1") {
10176                let h = e.dtoh(&x)?;
10177                let nan = h.iter().filter(|v| v.is_nan()).count();
10178                let amax = h.iter().fold(0f32, |a, v| a.max(v.abs()));
10179                eprintln!(
10180                    "[g4-prime-trace] layer {il}: nan={nan}/{} amax={amax:.3}",
10181                    h.len()
10182                );
10183                if nan > 0 {
10184                    return Err(format!("g4-prime-trace: first NaN at layer {il}").into());
10185                }
10186            }
10187        }
10188        cache.pos += t;
10189        let hiddens = e.clone_dtod(&x)?;
10190        let xv = e.view(&x, t * n_embd);
10191        let last_row = xv.slice((t - 1) * n_embd..t * n_embd);
10192        let mut h_seed = e.zeros(n_embd)?;
10193        e.copy_view_into(&mut h_seed, 0, &last_row, n_embd)?;
10194        let mut hn = e.uninit(n_embd)?;
10195        e.rms_norm(
10196            &h_seed,
10197            self.output_norm.float_data(),
10198            &mut hn,
10199            n_embd,
10200            1,
10201            eps,
10202        )?;
10203        let mut ld = e.matmul(&self.output, &hn, 1)?;
10204        let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
10205        e.softcap(&mut ld, cap, self.output.out_features())?;
10206        self.gemma4_suppress(e, &mut ld, 1)?;
10207        let logits = e.dtoh(&ld)?;
10208        Ok((logits, h_seed, hiddens))
10209    }
10210
10211    /// gemma4 T=1 decode attention: per-layer geometry, quantized-KV append + fa_decode
10212    /// (vec kernels at hd 256, generic scalar at the globals' hd 512), weightless V-norm,
10213    /// dual rope, scale 1.0. Takes the attn-normed input PRE-QUANTIZED (the cross-layer
10214    /// fused norm emits q8 directly — the f32 h never materializes).
10215    fn gemma4_decode_attn(
10216        &self,
10217        e: &Engine,
10218        fa: &crate::hybrid::FullAttnLayer,
10219        il: usize,
10220        hq: &CudaSlice<i8>,
10221        hdq: &CudaSlice<f32>,
10222        pos_d: &CudaSlice<i32>,
10223        cache: &mut Cache,
10224    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10225        let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
10226        let eps = self.cfg.rms_eps;
10227        let aux = self.gemma4_aux.as_ref().unwrap();
10228        let ones = aux.ones(e);
10229        #[cfg(debug_assertions)]
10230        crate::debug_assert_tensor_stream_device(ones, &e.stream(), "gemma4_decode_attn.ones");
10231        let (hq, hdq) = (hq, hdq);
10232        let h0 = e.zeros(0)?;
10233        let h = &h0;
10234        let (q0, k0, v0) = if swa {
10235            match e.matmul_q4_fused3(&fa.wq, &fa.wk, &fa.wv, &hq, &hdq)? {
10236                Some(t3) => t3,
10237                // NVFP4mix trio is MIXED-type (wv stays Q8_0), so fused3 can never
10238                // match — fuse the uniform (q,k) pair and take v as its own single.
10239                None => match e.matmul_nvfp4_fused2(&fa.wq, &fa.wk, &hq, &hdq, 1)? {
10240                    Some((q0, k0)) => {
10241                        let v0 = e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?;
10242                        (q0, k0, v0)
10243                    }
10244                    None => (
10245                        e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
10246                        e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
10247                        e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?,
10248                    ),
10249                },
10250            }
10251        } else {
10252            let (q0, k0) = match e.matmul_q4_fused2(&fa.wq, &fa.wk, &hq, &hdq)? {
10253                Some(p) => p,
10254                None => match e.matmul_nvfp4_fused2(&fa.wq, &fa.wk, &hq, &hdq, 1)? {
10255                    Some(p) => p,
10256                    None => (
10257                        e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
10258                        e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
10259                    ),
10260                },
10261            };
10262            let v0 = e.clone_dtod(&k0)?;
10263            (q0, k0, v0)
10264        };
10265        let mut q = e.uninit(nh * hd)?;
10266        let mut k = e.uninit(nkv * hd)?;
10267        let mut v = e.uninit(nkv * hd)?;
10268        // E4B wave-3 fold, m=1 completion (2026-07-23): the rows arms took this fold at
10269        // 550fcfa5; the decode-step trio kept 2 launches/layer it doesn't need.
10270        let ff = if swa {
10271            None
10272        } else {
10273            Some(
10274                aux.rope_freqs(e)
10275                    .expect("gemma4 global rope needs rope_freqs.weight"),
10276            )
10277        };
10278        #[cfg(debug_assertions)]
10279        if let Some(ff) = ff {
10280            crate::debug_assert_tensor_stream_device(
10281                ff,
10282                &e.stream(),
10283                "gemma4_decode_attn.rope_freqs",
10284            );
10285        }
10286        let kvl = cache.kv[il].as_mut().unwrap();
10287        let kv_fp8 = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
10288        if crate::Engine::qkv_append_on() {
10289            // append fold, eager completion (zoo-fusion arc): the dc arms took this fold
10290            // 2026-07-23; the eager step kept the 2-launch pair it doesn't need. Host-len
10291            // twin of the dc fold — bit-identical bodies, one launch per layer.
10292            e.rms_norm_qkv_rope_append(
10293                &q0,
10294                &k0,
10295                &v0,
10296                fa.q_norm.float_data(),
10297                fa.k_norm.float_data(),
10298                ones,
10299                &mut q,
10300                &mut k,
10301                &mut v,
10302                hd,
10303                nh,
10304                nkv,
10305                pos_d,
10306                nh,
10307                nkv,
10308                base,
10309                1.0,
10310                ff,
10311                eps,
10312                &mut kvl.k,
10313                &mut kvl.v,
10314                kvl.len,
10315                kvl.k_tok_bytes,
10316                kvl.v_tok_bytes,
10317                kv_fp8,
10318            )?;
10319        } else {
10320            e.rms_norm_qkv_rope(
10321                &q0,
10322                &k0,
10323                &v0,
10324                fa.q_norm.float_data(),
10325                fa.k_norm.float_data(),
10326                ones,
10327                &mut q,
10328                &mut k,
10329                &mut v,
10330                hd,
10331                nh,
10332                nkv,
10333                pos_d,
10334                nh,
10335                nkv,
10336                base,
10337                1.0,
10338                ff,
10339                eps,
10340            )?;
10341            e.append_kv_quantized(
10342                &k,
10343                &v,
10344                &mut kvl.k,
10345                &mut kvl.v,
10346                kvl.len,
10347                kvl.kv_dim_k,
10348                kvl.kv_dim_v,
10349                kvl.k_tok_bytes,
10350                kvl.v_tok_bytes,
10351                kv_fp8,
10352            )?;
10353        }
10354        kvl.len += 1;
10355        // R6 decode: SWA layers attend only the last `sliding_window` keys — a token-aligned
10356        // VIEW OFFSET into the quantized cache (keys carry absolute rope; the mask is purely
10357        // positional). Globals attend the full history.
10358        let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
10359        let mut attn = e.uninit(nh * hd)?;
10360        // global layers: SAME hd512 rows twin as verify with t=1 (parity law).
10361        if !swa
10362            && hd == 512
10363            && kvl.len >= crate::fa512_min_tkv()
10364            && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
10365        {
10366            let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
10367            let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
10368            // device-len: eager syncs the counter (async arg-store), dc keeps it live.
10369            let base = kvl.len as i32;
10370            e.i32_set_k(&mut kvl.len_d, base)?;
10371            e.fa_decode_rows(
10372                &q,
10373                &kp,
10374                &vp,
10375                &mut attn,
10376                hd,
10377                nh,
10378                nkv,
10379                kvl.len - 1,
10380                1,
10381                scale,
10382                kvl.k_tok_bytes,
10383                kvl.v_tok_bytes,
10384                Some((&kvl.len_d, -1)),
10385                false,
10386                false,
10387                None,
10388            )?;
10389            return Ok(e.matmul(&fa.wo, &attn, 1)?);
10390        }
10391        // windowed regime: SAME rows_w kernel as verify with t=1 (parity law — see verify_attn).
10392        if swa
10393            && kvl.len > win
10394            && hd == 256
10395            && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
10396        {
10397            let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
10398            let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
10399            let base = kvl.len as i32;
10400            e.i32_set_k(&mut kvl.len_d, base)?;
10401            e.fa_decode_rows_w(
10402                &q,
10403                &kp,
10404                &vp,
10405                &mut attn,
10406                hd,
10407                nh,
10408                nkv,
10409                &kvl.len_d,
10410                -1,
10411                1,
10412                scale,
10413                win,
10414                kvl.k_tok_bytes,
10415                kvl.v_tok_bytes,
10416                None,
10417            )?;
10418            return Ok(e.matmul(&fa.wo, &attn, 1)?);
10419        }
10420        let (off_tok, t_kv) = if swa && kvl.len > win {
10421            (kvl.len - win, win)
10422        } else {
10423            (0, kvl.len)
10424        };
10425        let k_view = e.view_u8_range(
10426            &kvl.k,
10427            off_tok * kvl.k_tok_bytes,
10428            (off_tok + t_kv) * kvl.k_tok_bytes,
10429        );
10430        let v_view = e.view_u8_range(
10431            &kvl.v,
10432            off_tok * kvl.v_tok_bytes,
10433            (off_tok + t_kv) * kvl.v_tok_bytes,
10434        );
10435        e.fa_decode_kvmod(
10436            &q,
10437            &k_view,
10438            &v_view,
10439            &mut attn,
10440            hd,
10441            nh,
10442            nkv,
10443            t_kv,
10444            scale,
10445            kvl.k_tok_bytes,
10446            kvl.v_tok_bytes,
10447            swa && crate::Engine::wkv_on(),
10448        )?;
10449        Ok(e.matmul(&fa.wo, &attn, 1)?)
10450    }
10451
10452    /// gemma4 DEVICE-COUNTER decode step (graph arc): token id + rope pos + KV lengths live in
10453    /// device counters; ZERO varying host kernel args. `cap_bucket_max` = Some(bucket) for graph
10454    /// capture (host mirrors untouched, n_splits from bucket, full-buffer KV views) / None for
10455    /// the eager-dc gate path (host mirrors advanced, live geometry — bit-identical target =
10456    /// gemma4_decode_step_h's token stream). V1 scope: t_kv <= sliding_window (no window views
10457    /// in-graph; the driver gates).
10458    #[allow(clippy::too_many_arguments)]
10459    pub fn gemma4_decode_step_dc(
10460        &self,
10461        e: &Engine,
10462        token_d: &CudaSlice<u32>,
10463        pos_d: &mut CudaSlice<i32>,
10464        embd_gpu: &CudaSlice<u8>,
10465        embd_qt: i32,
10466        embd_rb: usize,
10467        cache: &mut Cache,
10468        n_vocab: usize,
10469        cap_bucket_max: Option<(usize, usize)>,
10470    ) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
10471        let mut tok_out = e.stream().alloc_zeros::<u32>(1)?;
10472        self.gemma4_decode_step_dc_into(
10473            e,
10474            token_d,
10475            pos_d,
10476            embd_gpu,
10477            embd_qt,
10478            embd_rb,
10479            cache,
10480            n_vocab,
10481            cap_bucket_max,
10482            &mut tok_out,
10483        )?;
10484        Ok(tok_out)
10485    }
10486
10487    /// CAPTURE body: argmax lands in the PERSISTENT `tok_out` (same buffer = same address on
10488    /// every replay; pass `token_d` itself for the self-feeding graph loop).
10489    #[allow(clippy::too_many_arguments)]
10490    pub fn gemma4_decode_step_dc_into(
10491        &self,
10492        e: &Engine,
10493        token_d: &CudaSlice<u32>,
10494        pos_d: &mut CudaSlice<i32>,
10495        embd_gpu: &CudaSlice<u8>,
10496        embd_qt: i32,
10497        embd_rb: usize,
10498        cache: &mut Cache,
10499        n_vocab: usize,
10500        cap_bucket_max: Option<(usize, usize)>,
10501        tok_out: &mut CudaSlice<u32>,
10502    ) -> Result<(), Box<dyn std::error::Error>> {
10503        let n_embd = self.cfg.n_embd as usize;
10504        let eps = self.cfg.rms_eps;
10505        let mut x = e.embed_gather_device(embd_gpu, token_d, n_embd, embd_qt, embd_rb)?;
10506        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
10507        let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
10508        let n_layers = self.layers.len();
10509        for (il, layer) in self.layers.iter().enumerate() {
10510            let (hq, hdq) = match h_carry.take() {
10511                Some(p) => p,
10512                None => {
10513                    e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, 1, eps)?
10514                }
10515            };
10516            let Mixer::Full(fa) = &layer.mixer else {
10517                panic!("gemma4 layer {il} not full-attn")
10518            };
10519            let o =
10520                self.gemma4_decode_attn_dc(e, fa, il, &hq, &hdq, pos_d, cache, cap_bucket_max)?;
10521            let next_norm = if il + 1 < n_layers {
10522                Some(self.layers[il + 1].attn_norm.float_data())
10523            } else {
10524                None
10525            };
10526            let (xn, hn) = self.gemma4_layer_tail_add_nq_pn(e, layer, &o, &x, 1, next_norm)?;
10527            x = xn;
10528            h_carry = hn;
10529        }
10530        let mut hn = e.uninit(n_embd)?;
10531        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
10532        let mut logits = e.matmul(&self.output, &hn, 1)?;
10533        self.gemma4_suppress(e, &mut logits, 1)?; // cap skipped (monotonic); the mask is not
10534        e.argmax_token_device_into(&logits, tok_out, n_vocab)?;
10535        e.inc_seqlen(pos_d)?;
10536        if cap_bucket_max.is_none() {
10537            cache.pos += 1;
10538        }
10539        Ok(())
10540    }
10541
10542    /// Persistent transient slots for the ALLOC-FREE captured dc step (the graph door):
10543    /// every buffer the step produces per token lives here, allocated ONCE pre-capture, so
10544    /// the captured graph carries zero cuMemAllocAsync/Free nodes (the 226us/launch tax,
10545    /// osrt 2026-07-23). Sized for the model's max per-layer shapes.
10546
10547    /// Build the slot set (call OUTSIDE any capture).
10548    pub fn g4_dc_slots(&self, e: &Engine) -> Result<G4DcSlots, Box<dyn std::error::Error>> {
10549        let n_embd = self.cfg.n_embd as usize;
10550        let n_vocab = self.output.out_features();
10551        let n_layers = self.layers.len();
10552        let (mut qmax, mut kvmax, mut ffmax) = (0usize, 0usize, 0usize);
10553        for il in 0..n_layers {
10554            let (hd, nkv, nh, _b, _s, _w) = self.gemma4_geom(il);
10555            qmax = qmax.max(nh * hd);
10556            kvmax = kvmax.max(nkv * hd);
10557            if let crate::hybrid::Ffn::Dense { ffn_gate, .. } = &self.layers[il].ffn {
10558                ffmax = ffmax.max(ffn_gate.out_features());
10559            }
10560        }
10561        Ok(G4DcSlots {
10562            x: e.uninit(n_embd)?,
10563            xn: e.uninit(n_embd)?,
10564            cur: e.uninit(n_embd)?,
10565            hq: e.alloc_i8_uninit(n_embd)?,
10566            hd_: e.uninit(n_embd / 32)?,
10567            q0: e.uninit(qmax)?,
10568            k0: e.uninit(kvmax)?,
10569            v0: e.uninit(kvmax)?,
10570            q: e.uninit(qmax)?,
10571            k: e.uninit(kvmax)?,
10572            v: e.uninit(kvmax)?,
10573            attn: e.uninit(qmax)?,
10574            o: e.uninit(n_embd)?,
10575            attn_out: e.uninit(n_embd)?,
10576            zsh: e.uninit(n_embd)?,
10577            // zq/zd feed the wo matvec (nh*hd rows — 4096 on hd512 globals > n_embd),
10578            // the ffn entry (n_embd) and the lm_head (n_embd): size for the max.
10579            zq: e.alloc_i8_uninit(n_embd.max(qmax))?,
10580            zd: e.uninit(n_embd.max(qmax) / 32)?,
10581            gate: e.uninit(ffmax)?,
10582            up: e.uninit(ffmax)?,
10583            act: e.uninit(ffmax)?,
10584            actq: e.alloc_i8_uninit(ffmax)?,
10585            actd: e.uninit(ffmax / 32)?,
10586            f0: e.uninit(n_embd)?,
10587            sn: e.uninit(n_embd)?,
10588            hn: e.uninit(n_embd)?,
10589            logits: e.uninit(n_vocab)?,
10590        })
10591    }
10592
10593    /// m=1 pre-quantized matvec into a slot — mirrors matmul_pre's m=1 mmvq route exactly
10594    /// (rp4-mirror bytes; mmvq_supports guaranteed for gemma4 q4_0/q6_K).
10595    fn g4_matvec_m1_into(
10596        &self,
10597        e: &Engine,
10598        w: &crate::model::GpuTensor,
10599        aq: &CudaSlice<i8>,
10600        ad: &CudaSlice<f32>,
10601        y: &mut CudaSlice<f32>,
10602    ) -> Result<(), Box<dyn std::error::Error>> {
10603        use crate::model::GpuTensor;
10604        let (bytes, qtype, row_bytes, scale, rp) = match w {
10605            GpuTensor::Quant {
10606                bytes,
10607                qtype,
10608                row_bytes,
10609                scale,
10610                rp,
10611                ..
10612            } => (bytes, *qtype, *row_bytes, *scale, *rp),
10613            _ => return Err("g4_matvec_m1_into: non-quant tensor".into()),
10614        };
10615        let (mbytes, mrp) = match w {
10616            GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
10617            _ => (bytes, rp),
10618        };
10619        e.qmatvec_mmvq_into(
10620            mbytes,
10621            aq,
10622            ad,
10623            1,
10624            w.in_features(),
10625            w.out_features(),
10626            qtype,
10627            row_bytes,
10628            scale,
10629            mrp,
10630            y,
10631        )
10632    }
10633
10634    /// ALLOC-FREE dc step (capture body): kernel-for-kernel mirror of
10635    /// `gemma4_decode_step_dc_into` at t=1 with every transient slot-fed. Dense gemma4 only
10636    /// (12B/31B; uniform q4_0 trunk guarantees the fused2/3 arms).
10637    #[allow(clippy::too_many_arguments)]
10638    pub fn gemma4_decode_step_dc_slotted(
10639        &self,
10640        e: &Engine,
10641        token_d: &CudaSlice<u32>,
10642        pos_d: &mut CudaSlice<i32>,
10643        embd_gpu: &CudaSlice<u8>,
10644        embd_qt: i32,
10645        embd_rb: usize,
10646        cache: &mut Cache,
10647        n_vocab: usize,
10648        cap_bucket_max: Option<(usize, usize)>,
10649        sl: &mut G4DcSlots,
10650        tok_out: &mut CudaSlice<u32>,
10651        ring: Option<(&mut CudaSlice<u32>, usize)>,
10652    ) -> Result<(), Box<dyn std::error::Error>> {
10653        let n_embd = self.cfg.n_embd as usize;
10654        let eps = self.cfg.rms_eps;
10655        e.embed_gather_device_into(embd_gpu, token_d, &mut sl.x, n_embd, embd_qt, embd_rb)?;
10656        e.scale_inplace(&mut sl.x, (n_embd as f32).sqrt(), n_embd)?;
10657        let n_layers = self.layers.len();
10658        let mut has_carry = false;
10659        for il in 0..n_layers {
10660            if !has_carry {
10661                e.rms_norm_q8_1_into(
10662                    &sl.x,
10663                    self.layers[il].attn_norm.float_data(),
10664                    n_embd,
10665                    1,
10666                    eps,
10667                    &mut sl.hq,
10668                    &mut sl.hd_,
10669                )?;
10670            }
10671            has_carry = true;
10672            let layer = &self.layers[il];
10673            let Mixer::Full(fa) = &layer.mixer else {
10674                panic!("gemma4 layer {il} not full-attn")
10675            };
10676            self.gemma4_decode_attn_dc_slotted(e, fa, il, pos_d, cache, cap_bucket_max, sl)?;
10677            // pn-fold: the tail's entry launch consumes RAW sl.o (post_attn folded there);
10678            // the standalone norm only survives on the unfused seam arm.
10679            if !Engine::g4_pnfold_on() {
10680                e.rms_norm(
10681                    &sl.o,
10682                    layer.post_attn_norm.float_data(),
10683                    &mut sl.cur,
10684                    n_embd,
10685                    1,
10686                    eps,
10687                )?;
10688            }
10689            let next_norm = if il + 1 < n_layers {
10690                Some(self.layers[il + 1].attn_norm.float_data())
10691            } else {
10692                None
10693            };
10694            self.gemma4_layer_tail_slotted(e, layer, next_norm, sl)?;
10695            std::mem::swap(&mut sl.x, &mut sl.xn);
10696        }
10697        e.rms_norm(
10698            &sl.x,
10699            self.output_norm.float_data(),
10700            &mut sl.hn,
10701            n_embd,
10702            1,
10703            eps,
10704        )?;
10705        e.quantize_q8_1_into(&sl.hn, 1, n_embd, &mut sl.zq, &mut sl.zd)?;
10706        // lm_head via the same m=1 mmvq route as matmul (q6_K on the gemma4 family).
10707        {
10708            let (zq, zd) = (&sl.zq, &sl.zd);
10709            let zq = unsafe { &*(zq as *const CudaSlice<i8>) };
10710            let zd = unsafe { &*(zd as *const CudaSlice<f32>) };
10711            self.g4_matvec_m1_into(e, &self.output, zq, zd, &mut sl.logits)?;
10712        }
10713        self.gemma4_suppress(e, &mut sl.logits, 1)?;
10714        e.argmax_token_device_into(&sl.logits, tok_out, n_vocab)?;
10715        if let Some((ring, base)) = ring {
10716            // in-graph token parking: slot = (pos - base) % ring.len(); the door drains
10717            // with ONE sync per chunk instead of a per-token dtoh (the launch-serialization
10718            // fix — llama's 885us/launch overlaps GPU work because nothing waits per token).
10719            e.plain_tok_ring(tok_out, pos_d, base, ring)?;
10720        }
10721        e.inc_seqlen(pos_d)?;
10722        if cap_bucket_max.is_none() {
10723            cache.pos += 1;
10724        }
10725        Ok(())
10726    }
10727
10728    /// Slot-fed dc attention: mirrors gemma4_decode_attn_dc's CAPTURE arm kernel-for-kernel.
10729    #[allow(clippy::too_many_arguments)]
10730    fn gemma4_decode_attn_dc_slotted(
10731        &self,
10732        e: &Engine,
10733        fa: &crate::hybrid::FullAttnLayer,
10734        il: usize,
10735        pos_d: &CudaSlice<i32>,
10736        cache: &mut Cache,
10737        cap_bucket_max: Option<(usize, usize)>,
10738        sl: &mut G4DcSlots,
10739    ) -> Result<(), Box<dyn std::error::Error>> {
10740        let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
10741        let eps = self.cfg.rms_eps;
10742        let aux = self.gemma4_aux.as_ref().unwrap();
10743        let ones = aux.ones(e);
10744        #[cfg(debug_assertions)]
10745        crate::debug_assert_tensor_stream_device(
10746            ones,
10747            &e.stream(),
10748            "gemma4_decode_attn_dc_slotted.ones",
10749        );
10750        {
10751            let hq = unsafe { &*(&sl.hq as *const CudaSlice<i8>) };
10752            let hdq = unsafe { &*(&sl.hd_ as *const CudaSlice<f32>) };
10753            if swa {
10754                if !e.matmul_q4_fused3_into(
10755                    &fa.wq, &fa.wk, &fa.wv, hq, hdq, &mut sl.q0, &mut sl.k0, &mut sl.v0,
10756                )? {
10757                    // NVFP4mix trio is MIXED-type (wv stays Q8_0): fuse the uniform
10758                    // (q,k) pair, v through the generic m1 slot matvec — the same two
10759                    // bodies the eager arm runs, slot-fed for the alloc-free capture.
10760                    if e.matmul_nvfp4_fused2_into(&fa.wq, &fa.wk, hq, hdq, &mut sl.q0, &mut sl.k0)?
10761                    {
10762                        self.g4_matvec_m1_into(e, &fa.wv, hq, hdq, &mut sl.v0)?;
10763                    } else {
10764                        return Err("slotted step: fused3 unavailable (non-uniform trunk)".into());
10765                    }
10766                }
10767            } else {
10768                if !e.matmul_q4_fused2_into(&fa.wq, &fa.wk, hq, hdq, &mut sl.q0, &mut sl.k0)?
10769                    && !e
10770                        .matmul_nvfp4_fused2_into(&fa.wq, &fa.wk, hq, hdq, &mut sl.q0, &mut sl.k0)?
10771                {
10772                    return Err("slotted step: fused2 unavailable".into());
10773                }
10774                let k0r = unsafe { &*(&sl.k0 as *const CudaSlice<f32>) };
10775                e.copy_into(&mut sl.v0, 0, k0r, nkv * hd)?;
10776            }
10777        }
10778        // E4B wave-3 fold, m=1 completion (2026-07-23) — MUST mirror the dc_into arm
10779        // kernel-for-kernel (graph stream-identity gate).
10780        let ff = if swa {
10781            None
10782        } else {
10783            Some(
10784                aux.rope_freqs(e)
10785                    .expect("gemma4 global rope needs rope_freqs.weight"),
10786            )
10787        };
10788        #[cfg(debug_assertions)]
10789        if let Some(ff) = ff {
10790            crate::debug_assert_tensor_stream_device(
10791                ff,
10792                &e.stream(),
10793                "gemma4_decode_attn_dc_slotted.rope_freqs",
10794            );
10795        }
10796        let kvl = cache.kv[il].as_mut().unwrap();
10797        let kv_fp8 = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
10798        if crate::Engine::qkv_append_on() {
10799            // append fold (2026-07-23): mirrors dc_into.
10800            e.rms_norm_qkv_rope_append_dc(
10801                &sl.q0,
10802                &sl.k0,
10803                &sl.v0,
10804                fa.q_norm.float_data(),
10805                fa.k_norm.float_data(),
10806                ones,
10807                &mut sl.q,
10808                &mut sl.k,
10809                &mut sl.v,
10810                hd,
10811                nh,
10812                nkv,
10813                pos_d,
10814                nh,
10815                nkv,
10816                base,
10817                1.0,
10818                ff,
10819                eps,
10820                &mut kvl.k,
10821                &mut kvl.v,
10822                &kvl.len_d,
10823                kvl.k_tok_bytes,
10824                kvl.v_tok_bytes,
10825                kv_fp8,
10826            )?;
10827        } else {
10828            e.rms_norm_qkv_rope(
10829                &sl.q0,
10830                &sl.k0,
10831                &sl.v0,
10832                fa.q_norm.float_data(),
10833                fa.k_norm.float_data(),
10834                ones,
10835                &mut sl.q,
10836                &mut sl.k,
10837                &mut sl.v,
10838                hd,
10839                nh,
10840                nkv,
10841                pos_d,
10842                nh,
10843                nkv,
10844                base,
10845                1.0,
10846                ff,
10847                eps,
10848            )?;
10849            e.append_kv_quantized_dc(
10850                &sl.k,
10851                &sl.v,
10852                &mut kvl.k,
10853                &mut kvl.v,
10854                &kvl.len_d,
10855                kvl.kv_dim_k,
10856                kvl.kv_dim_v,
10857                kvl.k_tok_bytes,
10858                kvl.v_tok_bytes,
10859                kv_fp8,
10860            )?;
10861        }
10862        e.inc_seqlen(&mut kvl.len_d)?;
10863        let (b_swa, b_glob) = cap_bucket_max.expect("slotted step is capture-only");
10864        let k_view = e.view_u8(&kvl.k, kvl.k.len());
10865        let v_view = e.view_u8(&kvl.v, kvl.v.len());
10866        let rows_on = std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0");
10867        let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
10868        // combine-q8 emit (wave-5b m=1 port): the rows arms quantize inside the combine —
10869        // the standalone quantize launch runs only on the non-rows fallback. MUST mirror
10870        // the dc_into arm branch-for-branch (stream gate).
10871        let mut fa_q8 = false;
10872        if !swa && hd == 512 && b_glob >= crate::fa512_min_tkv() && rows_on {
10873            e.fa_decode_rows(
10874                &sl.q,
10875                &k_view,
10876                &v_view,
10877                &mut sl.attn,
10878                hd,
10879                nh,
10880                nkv,
10881                b_glob - 1,
10882                1,
10883                scale,
10884                kvl.k_tok_bytes,
10885                kvl.v_tok_bytes,
10886                Some((&kvl.len_d, -1)),
10887                false,
10888                false,
10889                Some((&mut sl.zq, &mut sl.zd)),
10890            )?;
10891            fa_q8 = true;
10892        } else if swa && b_swa > win && hd == 256 && rows_on {
10893            e.fa_decode_rows_w(
10894                &sl.q,
10895                &k_view,
10896                &v_view,
10897                &mut sl.attn,
10898                hd,
10899                nh,
10900                nkv,
10901                &kvl.len_d,
10902                -1,
10903                1,
10904                scale,
10905                win,
10906                kvl.k_tok_bytes,
10907                kvl.v_tok_bytes,
10908                Some((&mut sl.zq, &mut sl.zd)),
10909            )?;
10910            fa_q8 = true;
10911        } else {
10912            let b = if swa { b_swa } else { b_glob };
10913            e.fa_decode_dc(
10914                &sl.q,
10915                &k_view,
10916                &v_view,
10917                &mut sl.attn,
10918                hd,
10919                nh,
10920                nkv,
10921                &kvl.len_d,
10922                b,
10923                scale,
10924                kvl.k_tok_bytes,
10925                kvl.v_tok_bytes,
10926                swa && crate::Engine::wkv_on(),
10927            )?;
10928        }
10929        if !fa_q8 {
10930            let aq = unsafe { &*(&sl.attn as *const CudaSlice<f32>) };
10931            e.quantize_q8_1_into(aq, 1, nh * hd, &mut sl.zq, &mut sl.zd)?;
10932        }
10933        {
10934            let zq = unsafe { &*(&sl.zq as *const CudaSlice<i8>) };
10935            let zd = unsafe { &*(&sl.zd as *const CudaSlice<f32>) };
10936            self.g4_matvec_m1_into(e, &fa.wo, zq, zd, &mut sl.o)?;
10937        }
10938        Ok(())
10939    }
10940
10941    /// Slot-fed dense layer tail: mirrors gemma4_layer_tail_core (t=1, no pre-norm fold) +
10942    /// tail_add_nq kernel-for-kernel; the next layer's (hq, hd_) carry lands in the slots.
10943    fn gemma4_layer_tail_slotted(
10944        &self,
10945        e: &Engine,
10946        layer: &crate::hybrid::HybridLayer,
10947        next_norm: Option<&CudaSlice<f32>>,
10948        sl: &mut G4DcSlots,
10949    ) -> Result<(), Box<dyn std::error::Error>> {
10950        let n_embd = self.cfg.n_embd as usize;
10951        let eps = self.cfg.rms_eps;
10952        let bits = layer.gemma4.as_ref().unwrap();
10953        let crate::hybrid::Ffn::Dense {
10954            ffn_gate,
10955            ffn_up,
10956            ffn_down,
10957        } = &layer.ffn
10958        else {
10959            return Err("slotted tail: dense ffn only".into());
10960        };
10961        let pnfold = Engine::g4_pnfold_on();
10962        if pnfold {
10963            // pn-fold entry (E4B glue, slot-fed): rms(o, post_attn) + residual + ffn_norm
10964            // + q8 emit in ONE launch — replaces rms_norm + add_rms_norm + quantize.
10965            let or = unsafe { &*(&sl.o as *const CudaSlice<f32>) };
10966            let xr = unsafe { &*(&sl.x as *const CudaSlice<f32>) };
10967            e.rms_pre_add_rms_norm_q8z_into(
10968                or,
10969                layer.post_attn_norm.float_data(),
10970                xr,
10971                bits.ffn_norm.float_data(),
10972                &mut sl.attn_out,
10973                &mut sl.zsh,
10974                n_embd,
10975                1,
10976                eps,
10977                &mut sl.zq,
10978                &mut sl.zd,
10979            )?;
10980        } else {
10981            e.add_rms_norm(
10982                &sl.cur,
10983                &sl.x,
10984                bits.ffn_norm.float_data(),
10985                &mut sl.attn_out,
10986                &mut sl.zsh,
10987                n_embd,
10988                1,
10989                eps,
10990            )?;
10991        }
10992        let n_ff = ffn_gate.out_features();
10993        if !pnfold {
10994            let zshr = unsafe { &*(&sl.zsh as *const CudaSlice<f32>) };
10995            e.quantize_q8_1_into(zshr, 1, n_embd, &mut sl.zq, &mut sl.zd)?;
10996        }
10997        {
10998            let zq = unsafe { &*(&sl.zq as *const CudaSlice<i8>) };
10999            let zd = unsafe { &*(&sl.zd as *const CudaSlice<f32>) };
11000            if !e.matmul_q4_fused2_into(ffn_gate, ffn_up, zq, zd, &mut sl.gate, &mut sl.up)?
11001                && !e.matmul_nvfp4_fused2_into(
11002                    ffn_gate,
11003                    ffn_up,
11004                    zq,
11005                    zd,
11006                    &mut sl.gate,
11007                    &mut sl.up,
11008                )?
11009            {
11010                return Err("slotted tail: ffn fused2 unavailable".into());
11011            }
11012        }
11013        debug_assert!(e.uses_q8_1_fast(ffn_down));
11014        {
11015            let upr = unsafe { &*(&sl.up as *const CudaSlice<f32>) };
11016            let upv = e.view(upr, n_ff);
11017            let up_all = upv.slice(0..n_ff);
11018            let gr = unsafe { &*(&sl.gate as *const CudaSlice<f32>) };
11019            e.gelu_tanh_mul_q8_1_into(
11020                gr,
11021                &up_all,
11022                &mut sl.act,
11023                n_ff,
11024                1,
11025                &mut sl.actq,
11026                &mut sl.actd,
11027            )?;
11028        }
11029        {
11030            let aq = unsafe { &*(&sl.actq as *const CudaSlice<i8>) };
11031            let ad = unsafe { &*(&sl.actd as *const CudaSlice<f32>) };
11032            self.g4_matvec_m1_into(e, ffn_down, aq, ad, &mut sl.f0)?;
11033        }
11034        if pnfold {
11035            // pn-fold exit: rms(f0, post_ffw) + scaled residual + next attn_norm + q8
11036            // emit in ONE launch — replaces rms_norm + add_scale_rms_norm_q8_1.
11037            if let Some(w) = next_norm {
11038                let f0r = unsafe { &*(&sl.f0 as *const CudaSlice<f32>) };
11039                let aor = unsafe { &*(&sl.attn_out as *const CudaSlice<f32>) };
11040                e.rms_pre_add_scale_rms_norm_q8_1_into(
11041                    f0r,
11042                    bits.post_ffw_norm.float_data(),
11043                    aor,
11044                    bits.layer_scale,
11045                    w,
11046                    &mut sl.xn,
11047                    n_embd,
11048                    1,
11049                    eps,
11050                    &mut sl.hq,
11051                    &mut sl.hd_,
11052                )?;
11053                return Ok(());
11054            }
11055        }
11056        e.rms_norm(
11057            &sl.f0,
11058            bits.post_ffw_norm.float_data(),
11059            &mut sl.sn,
11060            n_embd,
11061            1,
11062            eps,
11063        )?;
11064        match next_norm {
11065            Some(w) => {
11066                e.add_scale_rms_norm_q8_1_into(
11067                    &sl.sn,
11068                    &sl.attn_out,
11069                    bits.layer_scale,
11070                    w,
11071                    &mut sl.xn,
11072                    n_embd,
11073                    1,
11074                    eps,
11075                    &mut sl.hq,
11076                    &mut sl.hd_,
11077                )?;
11078            }
11079            None => {
11080                e.add_scale(&sl.sn, &sl.attn_out, bits.layer_scale, &mut sl.xn, n_embd)?;
11081            }
11082        }
11083        Ok(())
11084    }
11085
11086    /// dc attention: same math as gemma4_decode_attn, KV slot/lengths from device counters.
11087    #[allow(clippy::too_many_arguments)]
11088    fn gemma4_decode_attn_dc(
11089        &self,
11090        e: &Engine,
11091        fa: &crate::hybrid::FullAttnLayer,
11092        il: usize,
11093        hq: &CudaSlice<i8>,
11094        hdq: &CudaSlice<f32>,
11095        pos_d: &CudaSlice<i32>,
11096        cache: &mut Cache,
11097        cap_bucket_max: Option<(usize, usize)>,
11098    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11099        let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
11100        let eps = self.cfg.rms_eps;
11101        let aux = self.gemma4_aux.as_ref().unwrap();
11102        let ones = aux.ones(e);
11103        #[cfg(debug_assertions)]
11104        crate::debug_assert_tensor_stream_device(ones, &e.stream(), "gemma4_decode_attn_dc.ones");
11105        let (q0, k0, v0) = if swa {
11106            match e.matmul_q4_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hdq)? {
11107                Some(t3) => t3,
11108                // NVFP4mix: wv is Q8_0 (mixed trio) — fuse the (q,k) pair, v single.
11109                None => match e.matmul_nvfp4_fused2(&fa.wq, &fa.wk, hq, hdq, 1)? {
11110                    Some((q0, k0)) => {
11111                        let h0 = e.zeros(0)?;
11112                        let v0 = e.matmul_pre(&fa.wv, hq, hdq, &h0, 1)?;
11113                        (q0, k0, v0)
11114                    }
11115                    None => {
11116                        let h0 = e.zeros(0)?;
11117                        (
11118                            e.matmul_pre(&fa.wq, hq, hdq, &h0, 1)?,
11119                            e.matmul_pre(&fa.wk, hq, hdq, &h0, 1)?,
11120                            e.matmul_pre(&fa.wv, hq, hdq, &h0, 1)?,
11121                        )
11122                    }
11123                },
11124            }
11125        } else {
11126            let (q0, k0) = match e.matmul_q4_fused2(&fa.wq, &fa.wk, hq, hdq)? {
11127                Some(p) => p,
11128                None => match e.matmul_nvfp4_fused2(&fa.wq, &fa.wk, hq, hdq, 1)? {
11129                    Some(p) => p,
11130                    None => {
11131                        let h0 = e.zeros(0)?;
11132                        (
11133                            e.matmul_pre(&fa.wq, hq, hdq, &h0, 1)?,
11134                            e.matmul_pre(&fa.wk, hq, hdq, &h0, 1)?,
11135                        )
11136                    }
11137                },
11138            };
11139            let v0 = e.clone_dtod(&k0)?;
11140            (q0, k0, v0)
11141        };
11142        let mut q = e.uninit(nh * hd)?;
11143        let mut k = e.uninit(nkv * hd)?;
11144        let mut v = e.uninit(nkv * hd)?;
11145        // E4B wave-3 fold, m=1 completion (2026-07-23): mirrored by the slotted arm.
11146        let ff = if swa {
11147            None
11148        } else {
11149            Some(
11150                aux.rope_freqs(e)
11151                    .expect("gemma4 global rope needs rope_freqs.weight"),
11152            )
11153        };
11154        #[cfg(debug_assertions)]
11155        if let Some(ff) = ff {
11156            crate::debug_assert_tensor_stream_device(
11157                ff,
11158                &e.stream(),
11159                "gemma4_decode_attn_dc.rope_freqs",
11160            );
11161        }
11162        let kvl = cache.kv[il].as_mut().unwrap();
11163        let kv_fp8 = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
11164        if crate::Engine::qkv_append_on() {
11165            // append fold (2026-07-23): norm+rope+cache-append in ONE launch.
11166            e.rms_norm_qkv_rope_append_dc(
11167                &q0,
11168                &k0,
11169                &v0,
11170                fa.q_norm.float_data(),
11171                fa.k_norm.float_data(),
11172                ones,
11173                &mut q,
11174                &mut k,
11175                &mut v,
11176                hd,
11177                nh,
11178                nkv,
11179                pos_d,
11180                nh,
11181                nkv,
11182                base,
11183                1.0,
11184                ff,
11185                eps,
11186                &mut kvl.k,
11187                &mut kvl.v,
11188                &kvl.len_d,
11189                kvl.k_tok_bytes,
11190                kvl.v_tok_bytes,
11191                kv_fp8,
11192            )?;
11193        } else {
11194            e.rms_norm_qkv_rope(
11195                &q0,
11196                &k0,
11197                &v0,
11198                fa.q_norm.float_data(),
11199                fa.k_norm.float_data(),
11200                ones,
11201                &mut q,
11202                &mut k,
11203                &mut v,
11204                hd,
11205                nh,
11206                nkv,
11207                pos_d,
11208                nh,
11209                nkv,
11210                base,
11211                1.0,
11212                ff,
11213                eps,
11214            )?;
11215            e.append_kv_quantized_dc(
11216                &k,
11217                &v,
11218                &mut kvl.k,
11219                &mut kvl.v,
11220                &kvl.len_d,
11221                kvl.kv_dim_k,
11222                kvl.kv_dim_v,
11223                kvl.k_tok_bytes,
11224                kvl.v_tok_bytes,
11225                kv_fp8,
11226            )?;
11227        }
11228        e.inc_seqlen(&mut kvl.len_d)?;
11229        let mut attn = e.uninit(nh * hd)?;
11230        // combine-q8 emit carrier (wave-5b m=1 port): rows arms fill this; the tail then
11231        // rides g4_matvec_m1_into instead of matmul's internal quantize.
11232        let mut fa_q8: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
11233        // Weight prefetch NOT wired here (26B/31B/qwen probes 2026-07-13: 26B flat
11234        // 196.6/196.0 vs 196.4/196.1 — MoE ffn dilutes wo; 31B −0.2% — dense decode sits
11235        // at the DRAM wall, no idle window to front-load into). E4B keeps the arm
11236        // (gemma4_e4b_attn, +0.65% valid window).
11237        match cap_bucket_max {
11238            None => {
11239                // dc-EAGER: host knows the length — R6 window views EXACTLY like the eager
11240                // decode (SWA layers attend the last `sliding_window` keys); the device
11241                // counters carry only the append slot + the graph seam.
11242                kvl.len += 1;
11243                let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
11244                if !swa
11245                    && hd == 512
11246                    && kvl.len >= crate::fa512_min_tkv()
11247                    && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
11248                {
11249                    // global layers: SAME hd512 rows twin as verify, t=1 (parity law).
11250                    // dc: len_d is live (inc_seqlen) — device-len rides it, plus=-1.
11251                    let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
11252                    let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
11253                    let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
11254                    e.fa_decode_rows(
11255                        &q,
11256                        &kp,
11257                        &vp,
11258                        &mut attn,
11259                        hd,
11260                        nh,
11261                        nkv,
11262                        kvl.len - 1,
11263                        1,
11264                        scale,
11265                        kvl.k_tok_bytes,
11266                        kvl.v_tok_bytes,
11267                        Some((&kvl.len_d, -1)),
11268                        false,
11269                        false,
11270                        Some((&mut aq8, &mut ad8)),
11271                    )?;
11272                    fa_q8 = Some((aq8, ad8));
11273                } else if swa
11274                    && kvl.len > win
11275                    && hd == 256
11276                    && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
11277                {
11278                    // windowed regime: SAME rows_w kernel as verify, t=1 (parity law).
11279                    let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
11280                    let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
11281                    let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
11282                    e.fa_decode_rows_w(
11283                        &q,
11284                        &kp,
11285                        &vp,
11286                        &mut attn,
11287                        hd,
11288                        nh,
11289                        nkv,
11290                        &kvl.len_d,
11291                        -1,
11292                        1,
11293                        scale,
11294                        win,
11295                        kvl.k_tok_bytes,
11296                        kvl.v_tok_bytes,
11297                        Some((&mut aq8, &mut ad8)),
11298                    )?;
11299                    fa_q8 = Some((aq8, ad8));
11300                } else {
11301                    let (off_tok, t_kv) = if swa && kvl.len > win {
11302                        (kvl.len - win, win)
11303                    } else {
11304                        (0, kvl.len)
11305                    };
11306                    let k_view = e.view_u8_range(
11307                        &kvl.k,
11308                        off_tok * kvl.k_tok_bytes,
11309                        (off_tok + t_kv) * kvl.k_tok_bytes,
11310                    );
11311                    let v_view = e.view_u8_range(
11312                        &kvl.v,
11313                        off_tok * kvl.v_tok_bytes,
11314                        (off_tok + t_kv) * kvl.v_tok_bytes,
11315                    );
11316                    e.fa_decode_kvmod(
11317                        &q,
11318                        &k_view,
11319                        &v_view,
11320                        &mut attn,
11321                        hd,
11322                        nh,
11323                        nkv,
11324                        t_kv,
11325                        scale,
11326                        kvl.k_tok_bytes,
11327                        kvl.v_tok_bytes,
11328                        swa && crate::Engine::wkv_on(),
11329                    )?;
11330                }
11331            }
11332            Some((b_swa, b_glob)) => {
11333                // capture: full-buffer views + device length. PARITY LAW port (graph arc step
11334                // 3): the SAME regime branches as dc-eager, pinned per arm — b_swa is the
11335                // exact-key t_kv for the dc family (n_splits constant per bucket), b_glob is
11336                // the RUNG max for the rows family (kernels derive per-replay splits from
11337                // kvl.len_d; grid/partials sized for the rung end, excess splits exit).
11338                let k_view = e.view_u8(&kvl.k, kvl.k.len());
11339                let v_view = e.view_u8(&kvl.v, kvl.v.len());
11340                let rows_on = std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0");
11341                let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
11342                if !swa && hd == 512 && b_glob >= crate::fa512_min_tkv() && rows_on {
11343                    let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
11344                    e.fa_decode_rows(
11345                        &q,
11346                        &k_view,
11347                        &v_view,
11348                        &mut attn,
11349                        hd,
11350                        nh,
11351                        nkv,
11352                        b_glob - 1,
11353                        1,
11354                        scale,
11355                        kvl.k_tok_bytes,
11356                        kvl.v_tok_bytes,
11357                        Some((&kvl.len_d, -1)),
11358                        false,
11359                        false,
11360                        Some((&mut aq8, &mut ad8)),
11361                    )?;
11362                    fa_q8 = Some((aq8, ad8));
11363                } else if swa && b_swa > win && hd == 256 && rows_on {
11364                    let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
11365                    e.fa_decode_rows_w(
11366                        &q,
11367                        &k_view,
11368                        &v_view,
11369                        &mut attn,
11370                        hd,
11371                        nh,
11372                        nkv,
11373                        &kvl.len_d,
11374                        -1,
11375                        1,
11376                        scale,
11377                        win,
11378                        kvl.k_tok_bytes,
11379                        kvl.v_tok_bytes,
11380                        Some((&mut aq8, &mut ad8)),
11381                    )?;
11382                    fa_q8 = Some((aq8, ad8));
11383                } else {
11384                    let b = if swa { b_swa } else { b_glob };
11385                    e.fa_decode_dc(
11386                        &q,
11387                        &k_view,
11388                        &v_view,
11389                        &mut attn,
11390                        hd,
11391                        nh,
11392                        nkv,
11393                        &kvl.len_d,
11394                        b,
11395                        scale,
11396                        kvl.k_tok_bytes,
11397                        kvl.v_tok_bytes,
11398                        swa && crate::Engine::wkv_on(),
11399                    )?;
11400                }
11401            }
11402        }
11403        // combine-q8 emit (wave-5b m=1 port): the rows arms produced the wo activation pair
11404        // in-combine — ride the slotted arm's exact matvec route (parity by construction).
11405        if let Some((aq8, ad8)) = fa_q8 {
11406            let mut y = e.uninit(fa.wo.out_features())?;
11407            self.g4_matvec_m1_into(e, &fa.wo, &aq8, &ad8, &mut y)?;
11408            return Ok(y);
11409        }
11410        Ok(e.matmul(&fa.wo, &attn, 1)?)
11411    }
11412
11413    /// gemma4 GRAPH-REPLAY greedy loop: per (swa-key, global-key) fa bucket, capture ONE full
11414    /// dc step (self-feeding: argmax writes token_d in-graph) and replay it — one graph launch
11415    /// per token, one 4B dtoh. V1 scope: whole generation under the sliding window (no window
11416    /// views in-graph); caller gates and falls back to the dc-eager loop.
11417    pub fn gemma4_generate_graph(
11418        &self,
11419        e: &Engine,
11420        prompt_pos: usize,
11421        first_token: u32,
11422        cache: &mut Cache,
11423        max_new: usize,
11424        eos: &[u32],
11425        mut on_token: impl FnMut(u32) -> bool,
11426    ) -> Result<(Vec<u32>, crate::decode::StopReason), Box<dyn std::error::Error>> {
11427        if self.is_gemma4_e4b() {
11428            return Err(
11429                "E4B graph serving is unwired (HANDOVER-E4B.md) — dc-eager is the serving arm"
11430                    .into(),
11431            );
11432        }
11433        use crate::decode::StopReason;
11434        let n_vocab = self.output.out_features();
11435        let n_embd = self.cfg.n_embd as usize;
11436        let embd_gpu = self
11437            .embd_gpu
11438            .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload"));
11439        let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
11440        for kvl in cache.kv.iter_mut().flatten() {
11441            e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
11442        }
11443        let mut token_d = e.stream().clone_htod(&[first_token])?;
11444        let mut pos_d = e.htod_i32(&[prompt_pos as i32])?;
11445        let g4 = self.cfg.gemma4.as_ref().unwrap();
11446        let (hd_s, hd_g) = (g4.key_length_swa as usize, g4.key_length_global as usize);
11447        // per-layer nkv: swa vs global counts from the pattern (uniform within class).
11448        let nkv_s = g4
11449            .head_count_kv
11450            .iter()
11451            .zip(g4.swa_pattern.iter())
11452            .find(|p| *p.1)
11453            .map(|p| *p.0 as usize)
11454            .unwrap_or(8);
11455        let nkv_g = g4
11456            .head_count_kv
11457            .iter()
11458            .zip(g4.swa_pattern.iter())
11459            .find(|p| !*p.1)
11460            .map(|p| *p.0 as usize)
11461            .unwrap_or(2);
11462        let mut graphs: std::collections::HashMap<
11463            ((bool, usize), (bool, usize), bool, bool),
11464            (
11465                cudarc::driver::CudaGraph,
11466                Vec<Box<dyn std::any::Any + Send>>,
11467            ),
11468        > = Default::default();
11469        // ALLOC-FREE capture: persistent transient slots — the captured graph carries zero
11470        // mem nodes (the 226us/launch tax). Slots must outlive every cached graph.
11471        let mut slots = self.g4_dc_slots(e)?;
11472        // Chunked replay ring: tokens park on-device; ONE drain sync per chunk. ring_base is
11473        // baked at the door entry (the modulo keeps every capture valid indefinitely).
11474        const RING: usize = 64;
11475        // DRAIN pinned at 1 (2026-07-23): relaunching the SAME graph exec before its prior
11476        // launch completes is ILLEGAL (chunk=4 -> ILLEGAL_ADDRESS; chunk=1 clean). Pipelining
11477        // needs alternating exec instances — and the measured payoff was ~0 (the ~200us
11478        // cuGraphLaunch host cost already overlaps its own launch's GPU work; llama pays
11479        // 885us/launch the same way). MEMRA_GRAPH_DRAIN raises it only for experiments.
11480        const DRAIN: usize = 1;
11481        let mut ring = e.stream().alloc_zeros::<u32>(RING)?;
11482        let ring_base = prompt_pos;
11483        let mut out = Vec::with_capacity(max_new);
11484        let mut reason = StopReason::MaxNew;
11485        let mut next = first_token;
11486        let mut captures = 0usize;
11487        for _ in 0..max_new {
11488            out.push(next);
11489            if eos.contains(&next) {
11490                reason = StopReason::Eos;
11491                break;
11492            }
11493            if !on_token(next) {
11494                reason = StopReason::Callback;
11495                break;
11496            }
11497            let t_kv = cache.pos + 1;
11498            // Bucket key per ARM (graph arc step 3):
11499            //  - swa component: dc family under the window (exact (fa_vec, n_splits) key —
11500            //    n_splits must be constant per bucket); rows_w above it (window-constant, so
11501            //    the component collapses to a single marker).
11502            //  - global component: dc family under the fa512 floor (exact key); rows_dpl16
11503            //    at/above it — the kernel derives splits from len_d per replay, so buckets
11504            //    are power-of-2 RUNGS (one capture per doubling; grid sized for the rung end).
11505            let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
11506            let f512 = crate::fa512_min_tkv();
11507            let key_s = if t_kv > win {
11508                (true, usize::MAX)
11509            } else {
11510                e.fa_bucket_key(t_kv, hd_s, nkv_s, crate::Engine::wkv_on())
11511            };
11512            let (key_g, rung_end) = if t_kv >= f512 {
11513                // strict upper bound: the rung must ROLL at exact powers (t_kv==1024 starts
11514                // the [1024,2048) bucket) — sizing covers every replayed T_kv < end.
11515                let end = (t_kv + 1).next_power_of_two().max(f512 * 2);
11516                ((true, end), end)
11517            } else {
11518                (e.fa_bucket_key(t_kv, hd_g, nkv_g, false), t_kv)
11519            };
11520            let key = (key_s, key_g, t_kv >= f512, t_kv > win);
11521            if !graphs.contains_key(&key) {
11522                let bucket_max = (t_kv, rung_end);
11523                // snapshot device+host state (the 3 capture-warmup runs leave no residue).
11524                let snap = cache.snapshot(e)?;
11525                let pos_save = e.dtoh_i32_one(&pos_d)?;
11526                let len_save: Vec<Option<i32>> = cache
11527                    .kv
11528                    .iter()
11529                    .map(|k| k.as_ref().map(|kvl| e.dtoh_i32_one(&kvl.len_d).unwrap()))
11530                    .collect();
11531                let tok_save = e.dtoh_u32_one(&token_d)?;
11532                // RETAINED capture (2026-07-23): the plain capture_graph left pool-transient
11533                // clones as dead COPY NODES replayed every launch — the E4B 0.74ms/token
11534                // regression class, and this door's measured -8.8%. The keeper pins warmup
11535                // transients so the captured graph holds kernel nodes only.
11536                let graph = {
11537                    let tok_ref = &mut token_d;
11538                    let pos_ref = &mut pos_d;
11539                    let cache_ref = &mut *cache;
11540                    let slots_ref = &mut slots;
11541                    let ring_ref = &mut ring;
11542                    e.capture_graph_retained_flags(
11543                        cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY,
11544                        |e| {
11545                        // self-feeding: the argmax writes token_d itself.
11546                        let tok_in = unsafe { &*(tok_ref as *const CudaSlice<u32>) };
11547                        let sl = unsafe { &mut *(slots_ref as *mut G4DcSlots) };
11548                        let rg = unsafe { &mut *(ring_ref as *mut CudaSlice<u32>) };
11549                        self.gemma4_decode_step_dc_slotted(e, tok_in, pos_ref, embd_gpu, qt, rb,
11550                                                           cache_ref, n_vocab, Some(bucket_max),
11551                                                           sl, tok_ref, Some((rg, ring_base)))
11552                    })?
11553                };
11554                cache.rollback(e, &snap, 0)?;
11555                e.set_i32_one(&mut pos_d, pos_save)?;
11556                for (il, ls) in len_save.iter().enumerate() {
11557                    if let (Some(kvl), Some(v)) = (cache.kv[il].as_mut(), ls) {
11558                        e.set_i32_one(&mut kvl.len_d, *v)?;
11559                    }
11560                }
11561                e.set_u32_one(&mut token_d, tok_save)?;
11562                if std::env::var("MEMRA_GRAPH_CENSUS").as_deref() == Ok("1") {
11563                    if let Ok(c) = crate::graph_update::node_census(&graph.0) {
11564                        eprintln!("[graph-census] {c:?}");
11565                    }
11566                }
11567                graphs.insert(key, graph);
11568                captures += 1;
11569            }
11570            // CHUNKED REPLAY: enqueue up to DRAIN launches back-to-back (the ~200us host
11571            // cost of each cuGraphLaunch overlaps the PREVIOUS launch's ~11ms GPU work),
11572            // then ONE sync + ring drain. The chunk must stay inside this bucket key and
11573            // the budget; capture warmups already emitted their tokens through the ring.
11574            let mut chunk = 1usize;
11575            let drain_cap: usize = std::env::var("MEMRA_GRAPH_DRAIN")
11576                .ok()
11577                .and_then(|v| v.parse().ok())
11578                .unwrap_or(DRAIN);
11579            while chunk < drain_cap && out.len() + chunk < max_new {
11580                let t_next = cache.pos + 1 + chunk;
11581                let key_s2 = if t_next > win {
11582                    (true, usize::MAX)
11583                } else {
11584                    e.fa_bucket_key(t_next, hd_s, nkv_s, crate::Engine::wkv_on())
11585                };
11586                let key_g2 = if t_next >= f512 {
11587                    (true, (t_next + 1).next_power_of_two().max(f512 * 2))
11588                } else {
11589                    e.fa_bucket_key(t_next, hd_g, nkv_g, false)
11590                };
11591                if (key_s2, key_g2, t_next >= f512, t_next > win) != key {
11592                    break;
11593                }
11594                chunk += 1;
11595            }
11596            let g = &graphs.get(&key).unwrap().0;
11597            for _ in 0..chunk {
11598                g.launch()?;
11599            }
11600            e.stream().synchronize()?;
11601            let ringh = e.dtoh_u32(&ring)?;
11602            for j in 0..chunk {
11603                let pos_j = cache.pos + j;
11604                let tok_j = ringh[(pos_j - ring_base) % RING];
11605                cache.pos += 0; // advanced below in one shot
11606                if j + 1 == chunk {
11607                    next = tok_j;
11608                } else {
11609                    out.push(tok_j);
11610                    if eos.contains(&tok_j) || !on_token(tok_j) {
11611                        reason = if eos.contains(&tok_j) {
11612                            StopReason::Eos
11613                        } else {
11614                            StopReason::Callback
11615                        };
11616                        // roll device/host state back to the stop point.
11617                        let keep = cache.pos + j + 1;
11618                        e.set_i32_one(&mut pos_d, keep as i32)?;
11619                        for kvl in cache.kv.iter_mut().filter_map(|k| k.as_mut()) {
11620                            e.set_i32_one(&mut kvl.len_d, keep as i32)?;
11621                            kvl.len = keep;
11622                        }
11623                        cache.pos = keep;
11624                        if std::env::var("MEMRA_GRAPH_STATS").is_ok() {
11625                            eprintln!("[gemma-graph] captures={captures} buckets={}", graphs.len());
11626                        }
11627                        return Ok((out, reason));
11628                    }
11629                }
11630            }
11631            cache.pos += chunk;
11632            for kvl in cache.kv.iter_mut().filter_map(|k| k.as_mut()) {
11633                kvl.len += chunk;
11634            }
11635        }
11636        if std::env::var("MEMRA_GRAPH_STATS").is_ok() {
11637            eprintln!("[gemma-graph] captures={captures} buckets={}", graphs.len());
11638        }
11639        Ok((out, reason))
11640    }
11641
11642    /// gemma4 VERIFY step (spec decode): t tokens batched through the trunk at positions
11643    /// pos0..pos0+t-1, K/V rows appended to the quantized cache (caller rolls back rejected
11644    /// rows via kvl.len), per-token causal windowed attend (fa_decode per token — the same
11645    /// kernel family as T=1 decode at each token's t_kv). Returns [t, n_vocab] softcapped
11646    /// logits (host) + advances cache.pos by t.
11647    pub(crate) fn gemma4_decode_step_t(
11648        &self,
11649        e: &Engine,
11650        tokens: &[u32],
11651        pos0: usize,
11652        cache: &mut Cache,
11653    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
11654        Ok(self.gemma4_decode_step_t_h(e, tokens, pos0, cache)?.0)
11655    }
11656
11657    /// GREEDY verify: per-row DEVICE argmax (t x 4B host traffic instead of the t x 1MB logits
11658    /// stack — softcap skipped: tanh is monotonic, per-row argmax unaffected). Returns
11659    /// (argmax ids [t], post-output_norm hidden stack [t, n_embd]).
11660    pub(crate) fn gemma4_decode_step_t_am(
11661        &self,
11662        e: &Engine,
11663        tokens: &[u32],
11664        pos0: usize,
11665        cache: &mut Cache,
11666    ) -> Result<(Vec<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11667        let (ld, hn) = self.gemma4_verify_trunk(e, tokens, pos0, cache, None)?;
11668        let t = tokens.len();
11669        let n_vocab = self.output.out_features();
11670        let mut toks = e.stream().alloc_zeros::<u32>(t)?;
11671        for i in 0..t {
11672            e.argmax_token_device_col(&ld, i, n_vocab, &mut toks, i)?;
11673        }
11674        Ok((e.dtoh_u32(&toks)?, hn))
11675    }
11676
11677    /// Device-token verify (async spec round): tokens live in tok_d[0..t]; per-row argmax
11678    /// lands in a DEVICE buffer (no host logits). Returns (vam_d [t] u32 device, hn stack).
11679    pub(crate) fn gemma4_decode_step_t_am_dev(
11680        &self,
11681        e: &Engine,
11682        tok_d: &CudaSlice<u32>,
11683        t: usize,
11684        pos0: usize,
11685        cache: &mut Cache,
11686    ) -> Result<(CudaSlice<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11687        let (ld, hn) = self.gemma4_verify_trunk(e, &vec![0u32; t], pos0, cache, Some(tok_d))?;
11688        let n_vocab = self.output.out_features();
11689        let mut vam = e.stream().alloc_zeros::<u32>(t)?;
11690        for i in 0..t {
11691            e.argmax_token_device_col(&ld, i, n_vocab, &mut vam, i)?;
11692        }
11693        Ok((vam, hn))
11694    }
11695
11696    /// gemma4 verify + the POST-output_norm hidden stack [t, n_embd] (the drafter's h input —
11697    /// llama's h_nextn convention).
11698    pub(crate) fn gemma4_decode_step_t_h(
11699        &self,
11700        e: &Engine,
11701        tokens: &[u32],
11702        pos0: usize,
11703        cache: &mut Cache,
11704    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11705        let (mut ld, hn) = self.gemma4_verify_trunk(e, tokens, pos0, cache, None)?;
11706        let t = tokens.len();
11707        let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
11708        e.softcap(&mut ld, cap, t * self.output.out_features())?;
11709        Ok((e.dtoh(&ld)?, hn))
11710    }
11711
11712    /// Persistent device scratch for the BURST verify stream (one per generation): rope pos
11713    /// rows [cap] + per-row t_kv counters (device-filled from the round counter, zero H2D).
11714    pub(crate) fn verify_stream_scratch(
11715        &self,
11716        e: &Engine,
11717        cap: usize,
11718    ) -> Result<VerifyStreamScratch, Box<dyn std::error::Error>> {
11719        Ok(VerifyStreamScratch {
11720            pos_d: e.htod_i32(&vec![0i32; cap])?,
11721            row_ctrs: (0..cap)
11722                .map(|_| e.htod_i32(&[0]))
11723                .collect::<Result<_, _>>()?,
11724        })
11725    }
11726
11727    /// BURST verify trunk (device-slot twin): tokens from a device buffer, rope positions
11728    /// iota'd from `ctr`, appends/attention at the layers' len_d counters (verify_attn_stream),
11729    /// NO host cache.pos/len advance. Returns (per-row device argmaxes [t], post-norm hidden
11730    /// stack [t, n_embd]) — the burst's accept/seed inputs, zero host readbacks.
11731    /// `scr` = PERSISTENT per-gen scratch (pos rows + per-row t_kv counters): the first cut
11732    /// htod_i32-allocated these per call — 9 pageable H2D copies per round, each a stream
11733    /// sync, exactly the turnaround the burst exists to remove.
11734    pub(crate) fn gemma4_verify_t_am_stream(
11735        &self,
11736        e: &Engine,
11737        tok_d: &CudaSlice<u32>,
11738        t: usize,
11739        ctr: &CudaSlice<i32>,
11740        hint: usize,
11741        cache: &mut Cache,
11742        scr: &mut VerifyStreamScratch,
11743    ) -> Result<(CudaSlice<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11744        let n_embd = self.cfg.n_embd as usize;
11745        let eps = self.cfg.rms_eps;
11746        assert!(t <= scr.row_ctrs.len() && t <= 64);
11747        e.i32_iota_from(ctr, &mut scr.pos_d, t)?;
11748        for i in 0..t {
11749            e.i32_copy_add(ctr, &mut scr.row_ctrs[i], (i + 1) as i32)?;
11750        }
11751        let (pos_d, row_ctrs) = (&scr.pos_d, &scr.row_ctrs);
11752        let embd_gpu = self
11753            .embd_gpu
11754            .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload"));
11755        let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
11756        let mut x = e.embed_gather_device_td(embd_gpu, tok_d, t, n_embd, qt, rb)?;
11757        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
11758        let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
11759        let n_layers = self.layers.len();
11760        for (il, layer) in self.layers.iter().enumerate() {
11761            let (hq, hdq) = match h_carry.take() {
11762                Some(p) => p,
11763                None => {
11764                    e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, t, eps)?
11765                }
11766            };
11767            let Mixer::Full(fa) = &layer.mixer else {
11768                panic!("gemma4 layer {il} not full-attn")
11769            };
11770            let o = self
11771                .gemma4_verify_attn_stream(e, fa, il, &hq, &hdq, pos_d, t, cache, hint, row_ctrs)?;
11772            let next_norm = if il + 1 < n_layers {
11773                Some(self.layers[il + 1].attn_norm.float_data())
11774            } else {
11775                None
11776            };
11777            let (xn, hn) = self.gemma4_layer_tail_add_nq_pn(e, layer, &o, &x, t, next_norm)?;
11778            x = xn;
11779            h_carry = hn;
11780            self.dflash_tap(e, cache, il, &x, t)?;
11781        }
11782        let mut hn = e.uninit(t * n_embd)?;
11783        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
11784        let ld = e.matmul(&self.output, &hn, t)?;
11785        let n_vocab = self.output.out_features();
11786        let mut vam = e.stream().alloc_zeros::<u32>(t)?;
11787        for i in 0..t {
11788            e.argmax_token_device_col(&ld, i, n_vocab, &mut vam, i)?;
11789        }
11790        Ok((vam, hn))
11791    }
11792
11793    /// Verify trunk core: returns (UN-softcapped logits device [t, n_vocab], post-output_norm
11794    /// hidden stack [t, n_embd]); appends KV rows + advances cache.pos.
11795    /// DFlash tap write (dflash lane): copy the post-layer residual rows of `x` into the
11796    /// armed sink at the tap slot for `il` (row-major [t, n_taps*hidden]). Per-row D2D
11797    /// copies — t <= block_size on verify; prime pays t*n_taps once per prompt (dedicated
11798    /// kernel later if it shows in the profile).
11799    fn dflash_tap(
11800        &self,
11801        e: &Engine,
11802        cache: &mut Cache,
11803        il: usize,
11804        x: &CudaSlice<f32>,
11805        t: usize,
11806    ) -> Result<(), Box<dyn std::error::Error>> {
11807        let Some(taps) = cache.dflash_taps.as_mut() else {
11808            return Ok(());
11809        };
11810        let Some(slot) = taps.layer_ids.iter().position(|&l| l == il) else {
11811            return Ok(());
11812        };
11813        let h = taps.hidden;
11814        let n_taps = taps.layer_ids.len();
11815        debug_assert_eq!(taps.t, t);
11816        let xv = e.view(x, t * h);
11817        for r in 0..t {
11818            let row = xv.slice(r * h..(r + 1) * h);
11819            e.copy_view_into(&mut taps.buf, r * n_taps * h + slot * h, &row, h)?;
11820        }
11821        Ok(())
11822    }
11823
11824    fn gemma4_verify_trunk(
11825        &self,
11826        e: &Engine,
11827        tokens: &[u32],
11828        pos0: usize,
11829        cache: &mut Cache,
11830        tok_dev: Option<&CudaSlice<u32>>,
11831    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11832        let n_embd = self.cfg.n_embd as usize;
11833        let eps = self.cfg.rms_eps;
11834        let t = tokens.len();
11835        let pos: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
11836        let pos_d = e.htod_i32(&pos)?;
11837        let mut x = match tok_dev {
11838            Some(td) => {
11839                let embd_gpu = self
11840                    .embd_gpu
11841                    .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload"));
11842                let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
11843                e.embed_gather_device_td(embd_gpu, td, t, n_embd, qt, rb)?
11844            }
11845            None => e.htod(&self.embd.gather(n_embd, tokens))?,
11846        };
11847        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
11848        let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
11849        let n_layers = self.layers.len();
11850        for (il, layer) in self.layers.iter().enumerate() {
11851            let (hq, hdq) = match h_carry.take() {
11852                Some(p) => p,
11853                None => {
11854                    e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, t, eps)?
11855                }
11856            };
11857            let Mixer::Full(fa) = &layer.mixer else {
11858                panic!("gemma4 layer {il} not full-attn")
11859            };
11860            let o = self.gemma4_verify_attn(e, fa, il, &hq, &hdq, &pos_d, t, cache)?;
11861            let next_norm = if il + 1 < n_layers {
11862                Some(self.layers[il + 1].attn_norm.float_data())
11863            } else {
11864                None
11865            };
11866            let (xn, hn) = self.gemma4_layer_tail_add_nq_pn(e, layer, &o, &x, t, next_norm)?;
11867            x = xn;
11868            h_carry = hn;
11869            self.dflash_tap(e, cache, il, &x, t)?;
11870        }
11871        let mut hn = e.uninit(t * n_embd)?;
11872        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
11873        let mut ld = e.matmul(&self.output, &hn, t)?;
11874        self.gemma4_suppress(e, &mut ld, t)?; // before the per-row argmax consumers
11875        cache.pos += t;
11876        Ok((ld, hn))
11877    }
11878
11879    /// Verify attention: project/norm/rope t rows, append them to the cache, then attend each
11880    /// row causally over [win_off_i .. base+i] via the SAME fa_decode dispatch as T=1 decode.
11881    /// BURST verify attention (device-slot twin of `gemma4_verify_attn`): the append lands
11882    /// at the layer's len_d counter (rows_dc), attention bases ride the counter (rows /
11883    /// rows_w with base_dev), and NO host len is read or bumped — `hint` is a host UPPER
11884    /// bound on base_len used only for split sizing and arm gating (the burst loop passes
11885    /// pos0 + burst slack). Host len mirrors re-sync at the burst drain.
11886    #[allow(clippy::too_many_arguments)]
11887    fn gemma4_verify_attn_stream(
11888        &self,
11889        e: &Engine,
11890        fa: &crate::hybrid::FullAttnLayer,
11891        il: usize,
11892        hq: &CudaSlice<i8>,
11893        hdq: &CudaSlice<f32>,
11894        pos_d: &CudaSlice<i32>,
11895        t: usize,
11896        cache: &mut Cache,
11897        hint: usize,
11898        row_ctrs: &[CudaSlice<i32>],
11899    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11900        let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
11901        let eps = self.cfg.rms_eps;
11902        let aux = self.gemma4_aux.as_ref().unwrap();
11903        let ones = aux.ones(e);
11904        #[cfg(debug_assertions)]
11905        crate::debug_assert_tensor_stream_device(
11906            ones,
11907            &e.stream(),
11908            "gemma4_verify_attn_stream.ones",
11909        );
11910        let h0 = e.zeros(0)?;
11911        let h = &h0;
11912        // BATCHED FUSED qkv (MEMRA_F2B=1, megakernel microcosm): swa layers fuse all three,
11913        // globals fuse q,k (v := k clone). Bit-identical per row; segments tail-fill.
11914        static F2B_QKV: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11915        let f2b = *F2B_QKV.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0"));
11916        let fused_qkv = if f2b {
11917            if swa {
11918                e.matmul_q4_fused3_batched(&fa.wq, &fa.wk, &fa.wv, hq, hdq, t)?
11919                    .map(|(a, b, c)| (a, b, Some(c)))
11920            } else {
11921                e.matmul_q4_fused2_batched(&fa.wq, &fa.wk, hq, hdq, t)?
11922                    .map(|(a, b)| (a, b, None))
11923            }
11924        } else {
11925            None
11926        };
11927        let (q0, k0, v0) = match fused_qkv {
11928            Some((a, b, cv)) => {
11929                let v = match cv {
11930                    Some(c) => c,
11931                    None => e.clone_dtod(&b)?,
11932                };
11933                (a, b, v)
11934            }
11935            None => {
11936                let q0 = e.matmul_pre(&fa.wq, hq, hdq, h, t)?;
11937                let k0 = e.matmul_pre(&fa.wk, hq, hdq, h, t)?;
11938                let v0 = if swa {
11939                    e.matmul_pre(&fa.wv, hq, hdq, h, t)?
11940                } else {
11941                    e.clone_dtod(&k0)?
11942                };
11943                (q0, k0, v0)
11944            }
11945        };
11946        let mut q = e.uninit(t * nh * hd)?;
11947        let mut k = e.uninit(t * nkv * hd)?;
11948        let mut v = e.uninit(t * nkv * hd)?;
11949        // E4B wave-3 fold (2026-07-23 decode-dust port): q/k/v norms + q/k rope in ONE
11950        // launch — rope math verbatim on the normed rows; V ones-rms, never roped.
11951        let ff = if swa {
11952            None
11953        } else {
11954            Some(
11955                aux.rope_freqs(e)
11956                    .expect("gemma4 global rope needs rope_freqs.weight"),
11957            )
11958        };
11959        #[cfg(debug_assertions)]
11960        if let Some(ff) = ff {
11961            crate::debug_assert_tensor_stream_device(
11962                ff,
11963                &e.stream(),
11964                "gemma4_verify_attn_stream.rope_freqs",
11965            );
11966        }
11967        e.rms_norm_qkv_rope(
11968            &q0,
11969            &k0,
11970            &v0,
11971            fa.q_norm.float_data(),
11972            fa.k_norm.float_data(),
11973            ones,
11974            &mut q,
11975            &mut k,
11976            &mut v,
11977            hd,
11978            nh * t,
11979            nkv * t,
11980            pos_d,
11981            nh,
11982            nkv,
11983            base,
11984            1.0,
11985            ff,
11986            eps,
11987        )?;
11988        let kvl = cache.kv[il].as_mut().unwrap();
11989        // append at the DEVICE slot; the counter advances by t on-device.
11990        e.append_kv_quantized_rows_dc(
11991            &k,
11992            &v,
11993            &mut kvl.k,
11994            &mut kvl.v,
11995            &kvl.len_d,
11996            t,
11997            kvl.kv_dim_k,
11998            kvl.kv_dim_v,
11999            kvl.k_tok_bytes,
12000            kvl.v_tok_bytes,
12001            (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on()),
12002        )?;
12003        // len_d is NOT advanced here: the burst's device rollback (spec_rollback_stream) is
12004        // the sole len writer after this round's attention (base stays = old len, plus = 0).
12005        let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
12006        let mut attn = e.uninit(t * nh * hd)?;
12007        let k_view = e.view_u8(&kvl.k, kvl.k.len());
12008        let v_view = e.view_u8(&kvl.v, kvl.v.len());
12009        // arm gating on the host HINT (upper bound; burst entry guards hint >= the vec floor
12010        // and a stable window regime — the same rung/regime keys as the draft graph).
12011        if swa && hint + 1 >= win {
12012            // fully-windowed rows: per-row window geometry from the counter (plus = 0: len_d
12013            // still holds the pre-append len; row r's T_kv = ctr + r + 1).
12014            e.fa_decode_rows_w(
12015                &q,
12016                &k_view,
12017                &v_view,
12018                &mut attn,
12019                hd,
12020                nh,
12021                nkv,
12022                &kvl.len_d,
12023                0,
12024                t,
12025                scale,
12026                win,
12027                kvl.k_tok_bytes,
12028                kvl.v_tok_bytes,
12029                None,
12030            )?;
12031        } else if hd == 512 && hint + t < crate::fa512_min_tkv() {
12032            // globals UNDER the fa512 crossover: eager runs the per-row fa_decode_kvmod
12033            // fallback there (parity with t=1 decode) — mirror it with per-row fa_decode_dc
12034            // (bit-correct for any t_kv <= bucket; the drafter dc arc proved the pairing).
12035            // Burst entry gates the horizon onto one side of the crossover, so hint decides
12036            // for every row.
12037            // NO .max(512): fa_decode_dc gates its hd512 dpl16-vs-scalar pick on bucket_max
12038            // (mirroring eager's fa512 floor) — forcing 512 here flipped the arm to dpl16
12039            // while eager ran scalar (il=5 KV drift, the burst's 4/128). SAME LAW capped
12040            // from above (2026-07-13): a regime-pinned hint near the floor (round-graph
12041            // captures use f512-1) pow2-rounds PAST it — clamp under the floor, or the arm
12042            // re-flips to dpl16 (the round-graph 4/64). The scalar unified self-splits, so
12043            // any bucket >= the live length is exact.
12044            let bucket = (hint + t + 2)
12045                .next_power_of_two()
12046                .min(crate::fa512_min_tkv().saturating_sub(1));
12047            let qv = e.view(&q, t * nh * hd);
12048            for i in 0..t {
12049                let q_row = qv.slice(i * nh * hd..(i + 1) * nh * hd);
12050                let mut q_one = e.uninit(nh * hd)?;
12051                e.copy_view_into(&mut q_one, 0, &q_row, nh * hd)?;
12052                let mut a_one = e.uninit(nh * hd)?;
12053                e.fa_decode_dc(
12054                    &q_one,
12055                    &k_view,
12056                    &v_view,
12057                    &mut a_one,
12058                    hd,
12059                    nh,
12060                    nkv,
12061                    &row_ctrs[i],
12062                    bucket,
12063                    scale,
12064                    kvl.k_tok_bytes,
12065                    kvl.v_tok_bytes,
12066                    false,
12067                )?;
12068                e.copy_into(&mut attn, i * nh * hd, &a_one, nh * hd)?;
12069            }
12070        } else if hd == 512 {
12071            // globals past the crossover: dpl16 dc twin via the shared rows wrapper (hint
12072            // sizes splits — upper bound; splits beyond the device len exit in-kernel).
12073            e.fa_decode_rows(
12074                &q,
12075                &k_view,
12076                &v_view,
12077                &mut attn,
12078                hd,
12079                nh,
12080                nkv,
12081                hint,
12082                t,
12083                scale,
12084                kvl.k_tok_bytes,
12085                kvl.v_tok_bytes,
12086                Some((&kvl.len_d, 0)),
12087                false,
12088                false,
12089                None,
12090            )?;
12091        } else {
12092            // hd256 under-window: v4 device-len rows twin.
12093            e.fa_decode_rows_dc(
12094                &q,
12095                &k_view,
12096                &v_view,
12097                &mut attn,
12098                hd,
12099                nh,
12100                nkv,
12101                &kvl.len_d,
12102                hint + t,
12103                t,
12104                scale,
12105                kvl.k_tok_bytes,
12106                kvl.v_tok_bytes,
12107                0,
12108                swa && crate::Engine::wkv_on(),
12109            )?;
12110        }
12111        Ok(e.matmul(&fa.wo, &attn, t)?)
12112    }
12113
12114    fn gemma4_verify_attn(
12115        &self,
12116        e: &Engine,
12117        fa: &crate::hybrid::FullAttnLayer,
12118        il: usize,
12119        hq: &CudaSlice<i8>,
12120        hdq: &CudaSlice<f32>,
12121        pos_d: &CudaSlice<i32>,
12122        t: usize,
12123        cache: &mut Cache,
12124    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12125        let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
12126        let eps = self.cfg.rms_eps;
12127        let aux = self.gemma4_aux.as_ref().unwrap();
12128        let ones = aux.ones(e);
12129        #[cfg(debug_assertions)]
12130        crate::debug_assert_tensor_stream_device(ones, &e.stream(), "gemma4_verify_attn.ones");
12131        let n_embd = self.cfg.n_embd as usize;
12132        let _ = n_embd;
12133
12134        let h0 = e.zeros(0)?;
12135        let h = &h0;
12136        // BATCHED FUSED qkv (MEMRA_F2B=1, megakernel microcosm): swa layers fuse all three,
12137        // globals fuse q,k (v := k clone). Bit-identical per row; segments tail-fill.
12138        static F2B_QKV: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12139        let f2b = *F2B_QKV.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0"));
12140        let fused_qkv = if f2b {
12141            if swa {
12142                e.matmul_q4_fused3_batched(&fa.wq, &fa.wk, &fa.wv, hq, hdq, t)?
12143                    .map(|(a, b, c)| (a, b, Some(c)))
12144            } else {
12145                e.matmul_q4_fused2_batched(&fa.wq, &fa.wk, hq, hdq, t)?
12146                    .map(|(a, b)| (a, b, None))
12147            }
12148        } else {
12149            None
12150        };
12151        let (q0, k0, v0) = match fused_qkv {
12152            Some((a, b, cv)) => {
12153                let v = match cv {
12154                    Some(c) => c,
12155                    None => e.clone_dtod(&b)?,
12156                };
12157                (a, b, v)
12158            }
12159            None => {
12160                let q0 = e.matmul_pre(&fa.wq, hq, hdq, h, t)?;
12161                let k0 = e.matmul_pre(&fa.wk, hq, hdq, h, t)?;
12162                let v0 = if swa {
12163                    e.matmul_pre(&fa.wv, hq, hdq, h, t)?
12164                } else {
12165                    e.clone_dtod(&k0)?
12166                };
12167                (q0, k0, v0)
12168            }
12169        };
12170        let mut q = e.uninit(t * nh * hd)?;
12171        let mut k = e.uninit(t * nkv * hd)?;
12172        let mut v = e.uninit(t * nkv * hd)?;
12173        // E4B wave-3 fold (2026-07-23 decode-dust port): q/k/v norms + q/k rope in ONE
12174        // launch — rope math verbatim on the normed rows; V ones-rms, never roped.
12175        let ff = if swa {
12176            None
12177        } else {
12178            Some(
12179                aux.rope_freqs(e)
12180                    .expect("gemma4 global rope needs rope_freqs.weight"),
12181            )
12182        };
12183        #[cfg(debug_assertions)]
12184        if let Some(ff) = ff {
12185            crate::debug_assert_tensor_stream_device(
12186                ff,
12187                &e.stream(),
12188                "gemma4_verify_attn.rope_freqs",
12189            );
12190        }
12191        e.rms_norm_qkv_rope(
12192            &q0,
12193            &k0,
12194            &v0,
12195            fa.q_norm.float_data(),
12196            fa.k_norm.float_data(),
12197            ones,
12198            &mut q,
12199            &mut k,
12200            &mut v,
12201            hd,
12202            nh * t,
12203            nkv * t,
12204            pos_d,
12205            nh,
12206            nkv,
12207            base,
12208            1.0,
12209            ff,
12210            eps,
12211        )?;
12212        let kvl = cache.kv[il].as_mut().unwrap();
12213        let base_len = kvl.len;
12214        e.append_kv_quantized_rows(
12215            &k,
12216            &v,
12217            &mut kvl.k,
12218            &mut kvl.v,
12219            base_len,
12220            t,
12221            kvl.kv_dim_k,
12222            kvl.kv_dim_v,
12223            kvl.k_tok_bytes,
12224            kvl.v_tok_bytes,
12225            (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on()),
12226        )?;
12227        kvl.len += t;
12228        let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
12229        let mut attn = e.uninit(t * nh * hd)?;
12230        // ROWS twin when no window offset is in play (globals are hd512-ineligible; SWA rows
12231        // under the window need no offset): ONE launch, per-row causal == the per-token loop.
12232        let rows_ok = (hd == 256 && base_len + 1 >= crate::fa_vec_min_tkv())
12233            // gemma globals: hd512 rows twin in the dpl16 vec regime (row 0 gates the batch);
12234            // decode rides the SAME symbol at t=1 (parity law).
12235            || (hd == 512 && !swa && base_len + 1 >= crate::fa512_min_tkv());
12236        if rows_ok && (!swa || base_len + t <= win) {
12237            let k_view = e.view_u8(&kvl.k, (base_len + t) * kvl.k_tok_bytes);
12238            let v_view = e.view_u8(&kvl.v, (base_len + t) * kvl.v_tok_bytes);
12239            if hd == 512 {
12240                // device-len twin: sync the counter to the verify base (async arg-store).
12241                e.i32_set_k(&mut kvl.len_d, base_len as i32)?;
12242                e.fa_decode_rows(
12243                    &q,
12244                    &k_view,
12245                    &v_view,
12246                    &mut attn,
12247                    hd,
12248                    nh,
12249                    nkv,
12250                    base_len,
12251                    t,
12252                    scale,
12253                    kvl.k_tok_bytes,
12254                    kvl.v_tok_bytes,
12255                    Some((&kvl.len_d, 0)),
12256                    false,
12257                    swa && crate::Engine::wkv_on(),
12258                    None,
12259                )?;
12260            } else {
12261                // hd256: the SAME v4_dc symbol the burst verify launches (PARITY LAW — the
12262                // host-base rows_v4 twin compiles apart from v4_dc and the two-symbol split
12263                // drifted the burst's persisted KV at il=8; one symbol, counter arg-stored).
12264                e.i32_set_k(&mut kvl.len_d, base_len as i32)?;
12265                e.fa_decode_rows_dc(
12266                    &q,
12267                    &k_view,
12268                    &v_view,
12269                    &mut attn,
12270                    hd,
12271                    nh,
12272                    nkv,
12273                    &kvl.len_d,
12274                    base_len + t,
12275                    t,
12276                    scale,
12277                    kvl.k_tok_bytes,
12278                    kvl.v_tok_bytes,
12279                    0,
12280                    swa && crate::Engine::wkv_on(),
12281                )?;
12282            }
12283            return Ok(e.matmul(&fa.wo, &attn, t)?);
12284        }
12285        // WINDOWED rows twin (deep ctx, every row fully windowed): one launch, ABSOLUTE-index
12286        // per-row geometry over the prefix view. PARITY LAW (2026-07-10 root-cause): textually
12287        // identical kernels do NOT compile bit-identically (nvcc unrolls fa_decode_vec_q 2x vs
12288        // its rows_w clone — SASS-proven; the unpinned score `+=` chain then rounds apart), so
12289        // decode-vs-verify parity comes from BOTH sides launching THIS SAME rows_w kernel
12290        // (decode passes t=1) — bitwise equal per position by symbol identity, any lane.
12291        // MEMRA_GEMMA_ROWS_W=0 -> per-token loop (decode falls back to fa_decode views too).
12292        if hd == 256
12293            && swa
12294            && base_len + 1 >= win
12295            && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
12296        {
12297            let k_view = e.view_u8(&kvl.k, (base_len + t) * kvl.k_tok_bytes);
12298            let v_view = e.view_u8(&kvl.v, (base_len + t) * kvl.v_tok_bytes);
12299            e.i32_set_k(&mut kvl.len_d, base_len as i32)?;
12300            e.fa_decode_rows_w(
12301                &q,
12302                &k_view,
12303                &v_view,
12304                &mut attn,
12305                hd,
12306                nh,
12307                nkv,
12308                &kvl.len_d,
12309                0,
12310                t,
12311                scale,
12312                win,
12313                kvl.k_tok_bytes,
12314                kvl.v_tok_bytes,
12315                None,
12316            )?;
12317            return Ok(e.matmul(&fa.wo, &attn, t)?);
12318        }
12319        for i in 0..t {
12320            let avail = base_len + i + 1;
12321            let (off_tok, t_kv) = if swa && avail > win {
12322                (avail - win, win)
12323            } else {
12324                (0, avail)
12325            };
12326            let k_view = e.view_u8_range(
12327                &kvl.k,
12328                off_tok * kvl.k_tok_bytes,
12329                (off_tok + t_kv) * kvl.k_tok_bytes,
12330            );
12331            let v_view = e.view_u8_range(
12332                &kvl.v,
12333                off_tok * kvl.v_tok_bytes,
12334                (off_tok + t_kv) * kvl.v_tok_bytes,
12335            );
12336            let qi = e.view(&q, t * nh * hd);
12337            let q_row = qi.slice(i * nh * hd..(i + 1) * nh * hd);
12338            let mut q_one = e.uninit(nh * hd)?;
12339            e.copy_view_into(&mut q_one, 0, &q_row, nh * hd)?;
12340            let mut a_one = e.uninit(nh * hd)?;
12341            // straddle rounds: windowed rows must ride the SAME rows_w kernel decode uses
12342            // (parity law above); global hd512 rows past the fa512 floor likewise ride the
12343            // rows_dpl16 twin; remaining rows keep the gated fa_decode pair.
12344            if swa
12345                && avail > win
12346                && hd == 256
12347                && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
12348            {
12349                let kp = e.view_u8(&kvl.k, avail * kvl.k_tok_bytes);
12350                let vp = e.view_u8(&kvl.v, avail * kvl.v_tok_bytes);
12351                e.i32_set_k(&mut kvl.len_d, (avail - 1) as i32)?;
12352                e.fa_decode_rows_w(
12353                    &q_one,
12354                    &kp,
12355                    &vp,
12356                    &mut a_one,
12357                    hd,
12358                    nh,
12359                    nkv,
12360                    &kvl.len_d,
12361                    0,
12362                    1,
12363                    scale,
12364                    win,
12365                    kvl.k_tok_bytes,
12366                    kvl.v_tok_bytes,
12367                    None,
12368                )?;
12369            } else if !swa
12370                && hd == 512
12371                && avail >= crate::fa512_min_tkv()
12372                && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
12373            {
12374                let kp = e.view_u8(&kvl.k, avail * kvl.k_tok_bytes);
12375                let vp = e.view_u8(&kvl.v, avail * kvl.v_tok_bytes);
12376                e.i32_set_k(&mut kvl.len_d, (avail - 1) as i32)?;
12377                e.fa_decode_rows(
12378                    &q_one,
12379                    &kp,
12380                    &vp,
12381                    &mut a_one,
12382                    hd,
12383                    nh,
12384                    nkv,
12385                    avail - 1,
12386                    1,
12387                    scale,
12388                    kvl.k_tok_bytes,
12389                    kvl.v_tok_bytes,
12390                    Some((&kvl.len_d, 0)),
12391                    false,
12392                    false,
12393                    None,
12394                )?;
12395            } else {
12396                e.fa_decode_kvmod(
12397                    &q_one,
12398                    &k_view,
12399                    &v_view,
12400                    &mut a_one,
12401                    hd,
12402                    nh,
12403                    nkv,
12404                    t_kv,
12405                    scale,
12406                    kvl.k_tok_bytes,
12407                    kvl.v_tok_bytes,
12408                    swa && crate::Engine::wkv_on(),
12409                )?;
12410            }
12411            e.copy_into(&mut attn, i * nh * hd, &a_one, nh * hd)?;
12412        }
12413        Ok(e.matmul(&fa.wo, &attn, t)?)
12414    }
12415
12416    /// gemma4 T=1 decode step: R8 layer graph over the cache; returns (softcapped logits host,
12417    /// h_seed = pre-output_norm hidden). Advances cache.pos.
12418    pub(crate) fn gemma4_decode_step_h(
12419        &self,
12420        e: &Engine,
12421        token: u32,
12422        cache: &mut Cache,
12423    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12424        // M1-PP2 door (crate::pp): 2-stage split with an explicit activation handoff.
12425        // Default OFF — unset env means this branch is never taken. The gemma4 arm stays
12426        // 2-stage in M2 (the N-stage gate model is the generic-arm 9B); N>2 warns + runs
12427        // unsplit rather than guessing a fence.
12428        if let Some(split) = crate::pp::pp2_split(self.layers.len()) {
12429            return self.gemma4_decode_step_h_pp2(e, token, cache, split);
12430        }
12431        if crate::pp::pp_cuts(self.layers.len()).is_some() {
12432            crate::pp::warn_unwired_once("gemma4 eager decode (N>2)");
12433        }
12434        let n_embd = self.cfg.n_embd as usize;
12435        let eps = self.cfg.rms_eps;
12436        let pos_d = e.htod_i32(&[cache.pos as i32])?;
12437        let mut x = e.htod(&self.embd.gather(n_embd, &[token]))?;
12438        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
12439        // cross-layer fusion: each tail's closing add+scale also EMITS the next layer's
12440        // attn-normed input pre-quantized q8_1 (the mixer consumes only quantized matmuls).
12441        let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
12442        let n_layers = self.layers.len();
12443        for (il, layer) in self.layers.iter().enumerate() {
12444            let (hq, hdq) = match h_carry.take() {
12445                Some(p) => p,
12446                None => {
12447                    e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, 1, eps)?
12448                }
12449            };
12450            let Mixer::Full(fa) = &layer.mixer else {
12451                panic!("gemma4 layer {il} not full-attn")
12452            };
12453            let o = self.gemma4_decode_attn(e, fa, il, &hq, &hdq, &pos_d, cache)?;
12454            let next_norm = if il + 1 < n_layers {
12455                Some(self.layers[il + 1].attn_norm.float_data())
12456            } else {
12457                None
12458            };
12459            let (xn, hn) = self.gemma4_layer_tail_add_nq_pn(e, layer, &o, &x, 1, next_norm)?;
12460            x = xn;
12461            h_carry = hn;
12462        }
12463        let mut hn = e.uninit(n_embd)?;
12464        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
12465        let h_seed = e.clone_dtod(&x)?;
12466        let mut ld = e.matmul(&self.output, &hn, 1)?;
12467        let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
12468        e.softcap(&mut ld, cap, self.output.out_features())?; // R4 on device (262k host tanh ~ms/step)
12469        self.gemma4_suppress(e, &mut ld, 1)?;
12470        let logits = e.dtoh(&ld)?;
12471        cache.pos += 1;
12472        Ok((logits, h_seed))
12473    }
12474
12475    /// M1-PP2 stage subgraph (gemma4 arm): layers [lo, hi) of the gemma4 T=1 walk with the
12476    /// pre-quantized next-norm carry LOCAL to the range. Enters with a materialized residual
12477    /// `x` (the range head runs its own `rms_norm_q8_1` against ITS layer's attn_norm — for
12478    /// lo == 0 that is exactly the unsplit loop's il==0 arm); exits with the residual
12479    /// materialized (range tail passes next_norm = None, exactly the unsplit last-layer arm).
12480    /// Bit-identity of the cut relies on the kernel-check-pinned `add_scale_rms_norm_q8_1 ==
12481    /// add_scale then rms_norm_q8_1` identity (`pp2-gate` verifies end-to-end).
12482    fn gemma4_decode_layers(
12483        &self,
12484        e: &Engine,
12485        mut x: CudaSlice<f32>,
12486        lo: usize,
12487        hi: usize,
12488        pos_d: &CudaSlice<i32>,
12489        cache: &mut Cache,
12490    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12491        let n_embd = self.cfg.n_embd as usize;
12492        let eps = self.cfg.rms_eps;
12493        let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
12494        for il in lo..hi {
12495            let layer = &self.layers[il];
12496            let (hq, hdq) = match h_carry.take() {
12497                Some(p) => p,
12498                // range head: il == lo — norm against THIS layer's attn_norm.
12499                None => {
12500                    e.rms_norm_q8_1(&x, self.layers[il].attn_norm.float_data(), n_embd, 1, eps)?
12501                }
12502            };
12503            let Mixer::Full(fa) = &layer.mixer else {
12504                panic!("gemma4 layer {il} not full-attn")
12505            };
12506            let o = self.gemma4_decode_attn(e, fa, il, &hq, &hdq, pos_d, cache)?;
12507            let next_norm = if il + 1 < hi {
12508                Some(self.layers[il + 1].attn_norm.float_data())
12509            } else {
12510                None
12511            };
12512            let (xn, hn) = self.gemma4_layer_tail_add_nq_pn(e, layer, &o, &x, 1, next_norm)?;
12513            x = xn;
12514            h_carry = hn;
12515        }
12516        Ok(x)
12517    }
12518
12519    /// M1-PP2 (increment 2): `gemma4_decode_step_h` as TWO stage subgraphs, each on its
12520    /// own stream (and device, under MEMRA_PP_DEVICES) with the transport-selected
12521    /// boundary handoff — same choreography as the generic arm (decode.rs), same
12522    /// ownership contract (crate::pp). Stage 0 = embed+scale + layers [0, split);
12523    /// stage 1 = layers [split, n) + output_norm + softcapped head.
12524    /// Each stage uploads its own copy of the step's position scalar on its own stream.
12525    /// MEMRA_PP_STREAMS=0 = the increment-1 same-stream seam. Gate: `pp2-gate`.
12526    fn gemma4_decode_step_h_pp2(
12527        &self,
12528        e: &Engine,
12529        token: u32,
12530        cache: &mut Cache,
12531        split: usize,
12532    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12533        if crate::pp::pp2_streams_off() {
12534            return self.gemma4_decode_step_h_pp2_samestream(e, token, cache, split);
12535        }
12536        let rt = crate::pp::Pp2Rt::get(e)?;
12537        let e0 = rt.engine(0, e);
12538        let e1 = rt.engine(1, e);
12539        let n_embd = self.cfg.n_embd as usize;
12540        let eps = self.cfg.rms_eps;
12541        let pos = cache.pos as i32;
12542
12543        // ---- STAGE 0 (its own stream): embed + sqrt(n_embd) scale + layers [0, split) ----
12544        let slot = {
12545            let _st0 = rt.enter(0);
12546            let pos_d = e0.htod_i32(&[pos])?;
12547            #[cfg(debug_assertions)]
12548            crate::debug_assert_tensor_stream_device(
12549                &pos_d,
12550                &e0.stream(),
12551                "gemma4_decode_step_h_pp2.stage0.pos_d",
12552            );
12553            let mut x = e0.htod(&self.embd.gather(n_embd, &[token]))?;
12554            e0.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
12555            let x = self.gemma4_decode_layers(e0, x, 0, split, &pos_d, cache)?;
12556            rt.tx(0, &x, n_embd)?
12557        };
12558
12559        // ---- STAGE 1 (its own stream): RX + layers [split, n) + softcapped head ----
12560        let _st1 = rt.enter(1);
12561        let pos_d = e1.htod_i32(&[pos])?;
12562        #[cfg(debug_assertions)]
12563        crate::debug_assert_tensor_stream_device(
12564            &pos_d,
12565            &e1.stream(),
12566            "gemma4_decode_step_h_pp2.stage1.pos_d",
12567        );
12568        let x = rt.rx(0, slot, n_embd)?;
12569        let x = self.gemma4_decode_layers(e1, x, split, self.layers.len(), &pos_d, cache)?;
12570
12571        let mut hn = e1.uninit(n_embd)?;
12572        e1.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
12573        let h_seed = e1.clone_dtod(&x)?;
12574        let mut ld = e1.matmul(&self.output, &hn, 1)?;
12575        let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
12576        e1.softcap(&mut ld, cap, self.output.out_features())?;
12577        self.gemma4_suppress(e1, &mut ld, 1)?;
12578        let logits = e1.dtoh(&ld)?;
12579        cache.pos += 1;
12580        Ok((logits, h_seed))
12581    }
12582
12583    /// MEMRA_PP_STREAMS=0 rollback seam: the increment-1 gemma4 pp2 body verbatim — both
12584    /// stage subgraphs on the ambient compute stream, boundary = two plain dtod copies.
12585    fn gemma4_decode_step_h_pp2_samestream(
12586        &self,
12587        e: &Engine,
12588        token: u32,
12589        cache: &mut Cache,
12590        split: usize,
12591    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12592        let n_embd = self.cfg.n_embd as usize;
12593        let eps = self.cfg.rms_eps;
12594        let pos_d = e.htod_i32(&[cache.pos as i32])?;
12595
12596        // ---- STAGE 0: embed + sqrt(n_embd) scale + layers [0, split) ----
12597        let mut x = e.htod(&self.embd.gather(n_embd, &[token]))?;
12598        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
12599        let x = self.gemma4_decode_layers(e, x, 0, split, &pos_d, cache)?;
12600
12601        // ---- STAGE BOUNDARY: explicit [n_embd] activation handoff (TX copy, RX copy) ----
12602        let boundary_tx = e.clone_dtod(&x)?;
12603        let boundary_rx = e.clone_dtod(&boundary_tx)?;
12604
12605        // ---- STAGE 1: layers [split, n) + output_norm + softcapped head ----
12606        let x =
12607            self.gemma4_decode_layers(e, boundary_rx, split, self.layers.len(), &pos_d, cache)?;
12608
12609        let mut hn = e.uninit(n_embd)?;
12610        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
12611        let h_seed = e.clone_dtod(&x)?;
12612        let mut ld = e.matmul(&self.output, &hn, 1)?;
12613        let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
12614        e.softcap(&mut ld, cap, self.output.out_features())?;
12615        self.gemma4_suppress(e, &mut ld, 1)?;
12616        let logits = e.dtoh(&ld)?;
12617        cache.pos += 1;
12618        Ok((logits, h_seed))
12619    }
12620}
12621
12622// ============================ step35 (Step-3.7-Flash) ==================================
12623// Node-for-node vs llama.cpp src/models/step35.cpp:216-300. WHY THIS IS A DEDICATED MIXER
12624// FAMILY and not a few branches inside the generic `full_attn*` chain:
12625//
12626//   1. `n_head` IS PER LAYER (`step35.attention.head_count` is an ARRAY: 64 on full-attn
12627//      layers, 96 on SWA). Every generic site reads the `cfg.n_head` scalar, so wq/wo/attn_gate
12628//      shapes and the FA head counts would be wrong on 33 of 45 layers.
12629//   2. `cfg.rope_dim_count` is 128 for this arch, but FULL-attn layers rotate only 64 dims
12630//      (upstream halves `n_rot_full` after the generic loader defaults it). A generic mixer
12631//      would rotate 128 dims on the full layers — silently wrong, plausible-looking logits.
12632//   3. Dual rope base (5e6 full / 1e4 SWA) + `rope_freqs.weight` on FULL layers ONLY.
12633//   4. SWA 3:1 (`sliding_window` 512, pattern [false,true,true,true]).
12634//   5. The head-wise gate is a SEPARATE `attn_gate.weight [n_embd, n_head_l]` (one pre-sigmoid
12635//      scalar per head, broadcast over head_dim) applied to the attention output BEFORE wo —
12636//      not the qwen35 fused-in-wq per-DIM gate that `attn_out_gate()`/`q_gate_split` handle.
12637//      Its input is the POST-attn_norm hidden (`cur`), not the residual.
12638//
12639// Attention scale is the DEFAULT 1/sqrt(n_embd_head_k) (step35.cpp:255) — NOT gemma4's 1.0.
12640impl HybridModel {
12641    /// Step35's declarative geometry row. Missing rows are artifact/config errors; never
12642    /// synthesize a drafter or trunk layer from a neighboring class.
12643    pub(crate) fn step35_geom(&self, il: usize) -> memra_gguf::config::LayerGeometry {
12644        let geometry = self
12645            .cfg
12646            .layer_geometry(il as u32)
12647            .unwrap_or_else(|| panic!("step35 layer {il} has no geometry-table row"));
12648        debug_assert_eq!(
12649            geometry.attention_gate,
12650            memra_gguf::config::AttentionGateKind::SeparateHead
12651        );
12652        geometry
12653    }
12654
12655    /// step35 attention core: everything from the q/k/v projection outputs through the head-wise
12656    /// gate, returning the GATED attention output PRE-`wo` (the caller runs wo, so this composes
12657    /// with both the plain `matmul(&fa.wo, ..)` sites and the f16/into-slab GEMM sites).
12658    ///
12659    /// `hg` = the POST-attn_norm hidden state (upstream's `cur`) — the gate projection's input
12660    /// on single-sequence paths. `gt_pre` supplies the already-projected per-head gate for a
12661    /// cross-request batch; exactly one of `hg` or `gt_pre` must be present.
12662    /// `cache`:
12663    ///   * `Some` => PRIME mode: append this chunk's post-rope K / raw V rows into the resident
12664    ///     quantized cache and attend THROUGH the cache view, exactly like the generic
12665    ///     `full_attn_prime_fa_dispatch` (quantize-then-attend on chunk 0 too, so chunk size
12666    ///     cannot decide where a precision edge falls — the grain-free chunk-invariance
12667    ///     contract, lane/chunkinv-flip).
12668    ///   * `None` => pure prefill (forward/forward_last/t2probe): attend over this batch's f32
12669    ///     q/k/v, no cache side effect.
12670    ///
12671    /// SWA WINDOW, and why prefill needs `sdpa_naive_w_quantized_view`: the view is trimmed to
12672    /// the oldest key ANY query in this chunk can reach (`off = base_len - (win-1)`), which
12673    /// bounds the view at `win-1+t` rows, but the mask is still required — inside the chunk,
12674    /// query `qt` may only see view keys `[qt, qt+win-1]`, so the earlier keys the trimmed view
12675    /// still contains must be masked per query. memra's window convention
12676    /// (`sdpa_naive_w_f32`: mask `t < q_pos-(window-1)`, `q_pos = (T_kv-T)+qt`) is bit-for-bit
12677    /// upstream's `LLAMA_SWA_TYPE_STANDARD` (`llama-hparams.h:359`: mask `p1-p0 >= n_swa`), and
12678    /// step35.cpp:6 sets exactly that swa_type — verified verbatim on both sides.
12679    ///
12680    /// Note the f32 `sdpa_naive_w` floor's shared memory is `t_kv*4` bytes, so the SWA arm needs
12681    /// `t_kv = win-1+t <= 12287` — i.e. it relies on chunked prefill (MEMRA_PRIME_CHUNK, default
12682    /// 4096). A monolithic 32k prime would exceed the 48 KiB dynamic-smem default.
12683    ///
12684    /// SWA ARM SELECTION IS KEYED ON `seq_end`, NOT ON THE CHUNK'S OWN `t_kv` (lane/step35-chunkfix,
12685    /// 2026-08-07). `seq_end` = the ABSOLUTE end position of the whole prime call
12686    /// (`cache.pos + prompt_len` at entry; `t` on the cacheless prefill path) — a property of the
12687    /// REQUEST, identical at every MEMRA_PRIME_CHUNK value. The old predicate `swa && t_kv > win`
12688    /// read the chunk's own extent, which made kernel selection — and therefore the logits, the
12689    /// hidden rows, and the generated text — a function of the chunk size:
12690    ///   a chunk [b,e) has off = max(0, b-(win-1)), t_kv = e-off, so b <= win-1 => off=0 => FA iff
12691    ///   e <= win, while every later chunk has t_kv = t+(win-1) > win. The FA rows were therefore a
12692    ///   contiguous PREFIX [0,P) with P = c*floor(win/c) for c <= win (else 0), and the output
12693    ///   depended only on P. Measured (research/step37-p2-20260806, closed-form receipt
12694    ///   `raw/chunkinv-step35-GAP2-CONFIRMED-20260807.txt`): at T=4883, c=512 and c=64 (P=512) are
12695    ///   byte-identical to each other but DIFFER from the c=4096 default (P=0) by maxdiff 1.813e0
12696    ///   with greedy text diverging at step 6; c=513 (P=0) is bit-identical to the default. A
12697    ///   one-token change in a documented machine-config knob changed the answer.
12698    /// So the pre-fix comment "same cache bytes, same numeric class" was false as written: same
12699    /// bytes, DIFFERENT numeric class — swapping `fa_prefill_view_ws` for the f32 windowed floor on
12700    /// the same rows moves the logits by ~1.8.
12701    ///
12702    /// Keying on `seq_end` makes P identically 0 for every chunk size, so MEMRA_PRIME_CHUNK is a
12703    /// pure memory/transient knob again (research/step35-chunkfix-20260807). It is
12704    /// correct-by-construction rather than a tolerance argument — the windowed kernel computes the
12705    /// same masked attention the FA arm computed (for e <= win the window mask is a no-op under
12706    /// causal, which is exactly why FA was legal there), it is now simply the ONLY arm used once
12707    /// the request passes the window. It also cannot move the shipped default: at chunk=4096 every
12708    /// chunk of a `seq_end > win` prime ALREADY had t_kv > win (chunk 0 has t_kv = min(chunk,
12709    /// seq_end) > win; every later chunk has t_kv >= t+win-1 > win since `PRIME_MIN_T` keeps
12710    /// t >= 16), and a `seq_end <= win` prime keeps every chunk on FA exactly as before — so this
12711    /// is a no-op on both default regimes and only removes the FA prefix at small chunk values.
12712    /// The `t_kv <= 12287` smem ceiling is respected: the chunks whose arm changes are precisely
12713    /// those with t_kv <= win = 512.
12714    #[allow(clippy::too_many_arguments)]
12715    fn step35_attn_pre_wo(
12716        &self,
12717        e: &Engine,
12718        fa: &FullAttnLayer,
12719        mut g3: Vec<CudaSlice<f32>>,
12720        hg: Option<&CudaSlice<f32>>,
12721        gt_pre: Option<&CudaSlice<f32>>,
12722        pos_d: &CudaSlice<i32>,
12723        t: usize,
12724        cache: Option<&mut Cache>,
12725        il: usize,
12726        seq_end: usize,
12727    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12728        let geometry = self.step35_geom(il);
12729        let hd = geometry.head_dim_k as usize;
12730        let nkv = geometry.n_head_kv as usize;
12731        let nh = geometry.n_head as usize;
12732        let rbase = geometry.rope_base;
12733        let scale = geometry.attention_scale();
12734        let swa = geometry.window.is_some();
12735        let eps = self.cfg.rms_eps;
12736        let win = geometry.window.unwrap_or(0) as usize;
12737        let n_rot = geometry.n_rot as usize;
12738
12739        let v = g3.pop().unwrap();
12740        let k0 = g3.pop().unwrap();
12741        let q0 = g3.pop().unwrap();
12742
12743        // q/k RMSNorm over head_dim rows (attn_q_norm/attn_k_norm [128] F32), then the
12744        // per-layer PARTIAL NEOX rope: n_rot = 64 on full layers (dims 64..127 pass through
12745        // unrotated), 128 on SWA. rope_freqs (llama3-style per-dim factors) on FULL only.
12746        let mut q = e.uninit(t * nh * hd)?;
12747        e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, nh * t, eps)?;
12748        let mut k = e.uninit(t * nkv * hd)?;
12749        e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, nkv * t, eps)?;
12750        let ff = if geometry.rope_factors {
12751            self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
12752        } else {
12753            None
12754        };
12755        #[cfg(debug_assertions)]
12756        if let Some(ff) = ff {
12757            crate::debug_assert_tensor_stream_device(
12758                ff,
12759                &e.stream(),
12760                "step35_attn_pre_wo.rope_freqs",
12761            );
12762        }
12763        e.rope_neox2(&mut q, &mut k, pos_d, hd, n_rot, nh, nkv, t, rbase, 1.0, ff)?;
12764
12765        let mut attn = e.uninit(t * nh * hd)?;
12766        match cache {
12767            Some(cache) => {
12768                let base_len = cache.kv[il].as_ref().unwrap().len;
12769                // Read per layer call, never in a measured default.
12770                let legacy_tkv = std::env::var("MEMRA_STEP35_SWA_TKV").as_deref() == Ok("1");
12771                let legacy_calllocal = std::env::var("MEMRA_PRIME_CALLLOCAL").as_deref() == Ok("1");
12772                let off = if swa {
12773                    let raw = base_len.saturating_sub(win - 1);
12774                    if legacy_tkv || legacy_calllocal {
12775                        raw
12776                    } else {
12777                        raw & !31usize
12778                    }
12779                } else {
12780                    0
12781                };
12782                {
12783                    let kvl = cache.kv[il].as_mut().unwrap();
12784                    assert!(kvl.len + t <= cache.max_ctx, "step35 prime: KV overflow");
12785                    let write_row = e.prepare_kv_append(kvl, off, t)?;
12786                    e.append_kv_quantized_rows(
12787                        &k,
12788                        &v,
12789                        &mut kvl.k,
12790                        &mut kvl.v,
12791                        write_row,
12792                        t,
12793                        kvl.kv_dim_k,
12794                        kvl.kv_dim_v,
12795                        kvl.k_tok_bytes,
12796                        kvl.v_tok_bytes,
12797                        crate::Engine::kv_fp8_on(),
12798                    )?;
12799                    kvl.len += t;
12800                    let new_len = kvl.len as i32;
12801                    e.set_i32_one(&mut kvl.len_d, new_len)?;
12802                }
12803                let kvl = cache.kv[il].as_ref().unwrap();
12804                // Both segmentation seams restore the FULL pre-fix arithmetic: their legacy
12805                // arm predicate (`t_kv` below, per-call `seq_end` in prime_cache) AND the
12806                // unaligned view offset here. Both halves are load-bearing for the canaries:
12807                // on the FA default the predicate arms agree bitwise wherever they can differ
12808                // (a t_kv<=win view has no maskable key, so windowed==unwindowed FA bit-for-bit).
12809                // That made predicate-only seams INERT under Lever A — first chunkinv35c, then
12810                // tickinv35c. The raw offset restores the live segmentation-variant mechanism
12811                // on the current FA path: its tile grid starts at the chunk/call boundary.
12812                // SWA: trim the view to the oldest key any query in this chunk can reach —
12813                // ALIGNED DOWN to the FA tile size (BK=32). The raw trim is chunk-DEPENDENT
12814                // (off = base_len-(win-1), and base_len is a chunk boundary), and the FA
12815                // kernel's online-softmax recurrence groups keys into BK tiles relative to
12816                // the VIEW START — so an unaligned off regroups the same absolute keys into
12817                // different tiles at different chunk sizes = different (m,l) rounding =
12818                // chunk-dependent bits. chunkinv35 caught exactly this on the first FA-arm
12819                // battery (first_div == chunk size; raw/leverA-gates-20260807T135541Z.log).
12820                // Aligning off to 32 pins tiles to ABSOLUTE key positions for every chunk
12821                // size; the <=31 extra leading keys are older than EVERY query's window
12822                // (all queries sit at >= base_len, so keys < base_len-(win-1) are masked
12823                // for all of them) and a fully-masked key is an exact-0.0 no-op in both
12824                // kernels (NEG_INF -> p=0.0; l+=0.0 and O+=0.0 are bitwise identity), so
12825                // the floor arm's bits do not move either (gated: G2f, battery 2).
12826                let t_kv = base_len + t - off;
12827                let physical = kvl.physical_rows(off, off + t_kv)?;
12828                let k_view = e.view_u8_range(
12829                    &kvl.k,
12830                    physical.start * kvl.k_tok_bytes,
12831                    physical.end * kvl.k_tok_bytes,
12832                );
12833                let v_view = e.view_u8_range(
12834                    &kvl.v,
12835                    physical.start * kvl.v_tok_bytes,
12836                    physical.end * kvl.v_tok_bytes,
12837                );
12838                // CHUNK-INVARIANT PREDICATE: `seq_end` (the request's absolute end position), not
12839                // this chunk's `t_kv`. See the doc note — keying on t_kv made P = c*floor(win/c)
12840                // rows take FA and the output a function of MEMRA_PRIME_CHUNK.
12841                // MEMRA_STEP35_SWA_TKV=1 is the ROLLBACK SEAM to the pre-fix arithmetic — BOTH
12842                // halves: this predicate AND the unaligned view offset above (`legacy_tkv`).
12843                // It is what gives the step35 chunkinv gate its canary teeth: chunk-VARIANT by
12844                // construction, so the invariance assertion MUST break under it (the seam whose
12845                // absence made the original canary inert — GAP 1 in research/step37-p2-20260806;
12846                // and a predicate-only seam went inert AGAIN under the FA default, battery 2's
12847                // CANARY UNEXPECTEDLY MATCHED — see the offset comment). Read per call, not
12848                // cached (probes flip it in-process). Never on in a measured default run.
12849                let swa_naive = if legacy_tkv {
12850                    t_kv > win
12851                } else {
12852                    seq_end > win
12853                };
12854                if swa && swa_naive {
12855                    // Windowed mask needed (see the doc note). DEFAULT since lane/pp-prefill
12856                    // 2026-08-07: the windowed hd128 FA stamp (`fa_prefill_view_ws_w_hd128`) —
12857                    // the anatomy profile measured the f32 floor at 565 ms/layer on a pp4096
12858                    // (41% of the whole prime) while the unwindowed hd128 FA family did the
12859                    // strictly harder causal-4096 in 3.3 ms. NOTE t_kv can be <= win here
12860                    // (a chunk-0 view on a small chunk); the windowed kernel handles that
12861                    // identically to the unwindowed one modulo the mask, which is the point.
12862                    // NEW NUMERIC CLASS vs the floor (bf16-MMA online softmax vs f32 serial),
12863                    // selected on `seq_end` like every arm here, so the class is uniform for
12864                    // the whole request at every MEMRA_PRIME_CHUNK — chunkinv holds by the
12865                    // same construction as the chunkfix. MEMRA_STEP35_SWA_FA=0 = rollback to
12866                    // the f32 floor (the previous numeric config, kept as the A/B seam).
12867                    if std::env::var("MEMRA_STEP35_SWA_FA").as_deref() == Ok("0") {
12868                        e.sdpa_naive_w_quantized_view(
12869                            &q,
12870                            &k_view,
12871                            &v_view,
12872                            &mut attn,
12873                            hd,
12874                            nh,
12875                            nkv,
12876                            t,
12877                            t_kv,
12878                            scale,
12879                            true,
12880                            win,
12881                            kvl.k_tok_bytes,
12882                            kvl.v_tok_bytes,
12883                        )?;
12884                    } else {
12885                        e.fa_prefill_view_ws_w_hd128(
12886                            &q,
12887                            &k_view,
12888                            &v_view,
12889                            &mut attn,
12890                            hd,
12891                            nh,
12892                            nkv,
12893                            t,
12894                            t_kv,
12895                            scale,
12896                            true,
12897                            win,
12898                            kvl.k_tok_bytes,
12899                            kvl.v_tok_bytes,
12900                        )?;
12901                    }
12902                } else if std::env::var("MEMRA_NOFA").is_ok() {
12903                    e.sdpa_naive_quantized_view(
12904                        &q,
12905                        &k_view,
12906                        &v_view,
12907                        &mut attn,
12908                        hd,
12909                        nh,
12910                        nkv,
12911                        t,
12912                        t_kv,
12913                        scale,
12914                        true,
12915                        kvl.k_tok_bytes,
12916                        kvl.v_tok_bytes,
12917                    )?;
12918                } else {
12919                    // seq_end <= win (or a full-attn layer): no query in the WHOLE request can
12920                    // reach past the window, so the window mask is a no-op under causal and every
12921                    // chunk rides the hd128 dequant-once FA prefill — one arm for the whole
12922                    // request either way, which is what makes the chunk size arithmetic-free.
12923                    e.fa_prefill_view_ws(
12924                        &q,
12925                        &k_view,
12926                        &v_view,
12927                        &mut attn,
12928                        hd,
12929                        nh,
12930                        nkv,
12931                        t,
12932                        t_kv,
12933                        scale,
12934                        true,
12935                        kvl.k_tok_bytes,
12936                        kvl.v_tok_bytes,
12937                        crate::Engine::kv_fp8_on(),
12938                    )?;
12939                }
12940            }
12941            None => {
12942                // Cacheless prefill (forward/forward_last/t2probe) is MONOLITHIC — there is no
12943                // chunk loop on this path, so seq_end == t and the predicate is unchanged. The
12944                // assert pins that: if a chunked cacheless prefill is ever added, it must thread
12945                // seq_end here too or it re-opens the same door.
12946                debug_assert_eq!(
12947                    seq_end, t,
12948                    "step35 cacheless prefill is monolithic (seq_end == t)"
12949                );
12950                if swa && seq_end > win {
12951                    e.sdpa_naive_w(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
12952                } else if std::env::var("MEMRA_NOFA").is_ok() {
12953                    e.sdpa_naive(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
12954                } else {
12955                    e.fa_prefill(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
12956                }
12957            }
12958        }
12959
12960        // HEAD-WISE GATE (step35.cpp:267-285): gate = attn_gate(cur) -> [t, n_head_l];
12961        // attn *= sigmoid(gate) broadcast over head_dim. BEFORE wo.
12962        let gw = fa
12963            .attn_gate
12964            .as_ref()
12965            .ok_or("step35 layer is missing attn_gate.weight (head-wise attention gate)")?;
12966        let gt_owned = if gt_pre.is_none() {
12967            Some(e.matmul(
12968                gw,
12969                hg.ok_or("step35 attention needs hg when gt_pre is absent")?,
12970                t,
12971            )?)
12972        } else {
12973            None
12974        };
12975        let gt = gt_pre.or(gt_owned.as_ref()).unwrap();
12976        let mut ag = e.uninit(t * nh * hd)?;
12977        e.attn_head_gate(&attn, gt, &mut ag, None, hd, nh, t)?;
12978        Ok(ag)
12979    }
12980
12981    /// step35 PREFILL mixer, no cache side effect (the `full_attn` contract: `forward`,
12982    /// `forward_last`, t2probe). Post-`wo`.
12983    pub(crate) fn step35_attn(
12984        &self,
12985        e: &Engine,
12986        fa: &FullAttnLayer,
12987        h: &CudaSlice<f32>,
12988        pos_d: &CudaSlice<i32>,
12989        t: usize,
12990        il: usize,
12991    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12992        let g3 = e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?;
12993        // Monolithic path: the request ends at t (no chunk loop). See step35_attn_pre_wo's note.
12994        let ag = self.step35_attn_pre_wo(e, fa, g3, Some(h), None, pos_d, t, None, il, t)?;
12995        Ok(e.matmul(&fa.wo, &ag, t)?)
12996    }
12997
12998    /// step35 PRIME mixer (the `full_attn_prime` contract: append this chunk's K/V into the
12999    /// resident quantized cache, attend through the cache view). Post-`wo`.
13000    ///
13001    /// `seq_end` = the ABSOLUTE end position of the whole prime request (chunk-invariant); it
13002    /// selects the SWA arm. See `step35_attn_pre_wo`'s doc note for why it cannot be this chunk's
13003    /// own extent.
13004    #[allow(clippy::too_many_arguments)]
13005    pub(crate) fn step35_attn_prime(
13006        &self,
13007        e: &Engine,
13008        fa: &FullAttnLayer,
13009        h: &CudaSlice<f32>,
13010        hx: Option<&CudaSlice<u8>>,
13011        pos_d: &CudaSlice<i32>,
13012        t: usize,
13013        cache: &mut Cache,
13014        il: usize,
13015        seq_end: usize,
13016    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13017        let g3 = match hx {
13018            Some(xh) => e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], h, xh, t)?,
13019            None => e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?,
13020        };
13021        let ag =
13022            self.step35_attn_pre_wo(e, fa, g3, Some(h), None, pos_d, t, Some(cache), il, seq_end)?;
13023        Ok(e.matmul(&fa.wo, &ag, t)?)
13024    }
13025
13026    /// step35 T=1 decode attention (post-`wo`, matching `full_attn_decode_pre`'s contract).
13027    /// `pre_q` = the attn-input norm's q8_1 pair when the caller took the norm-fusion lever
13028    /// (then `h` is a zero-length placeholder and EVERY projection here — including the gate —
13029    /// must be on the q8_1 fast path; `mixer_in_q8_1_fast` enforces that for step35 by also
13030    /// requiring `attn_gate`).
13031    ///
13032    /// SWA decode is a token-aligned VIEW OFFSET into the quantized cache (the gemma4 R6
13033    /// pattern): keys carry absolute rope and the mask is purely positional, so the single
13034    /// query at `len-1` attending the last `win` rows IS the windowed result — no mask kernel.
13035    #[allow(clippy::too_many_arguments)]
13036    pub(crate) fn step35_decode_attn(
13037        &self,
13038        e: &Engine,
13039        fa: &FullAttnLayer,
13040        il: usize,
13041        h: &CudaSlice<f32>,
13042        pre_q: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
13043        pos_d: &CudaSlice<i32>,
13044        cache: &mut Cache,
13045    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13046        let geometry = self.step35_geom(il);
13047        let hd = geometry.head_dim_k as usize;
13048        let nkv = geometry.n_head_kv as usize;
13049        let nh = geometry.n_head as usize;
13050        let rbase = geometry.rope_base;
13051        let scale = geometry.attention_scale();
13052        let swa = geometry.window.is_some();
13053        let eps = self.cfg.rms_eps;
13054        let win = geometry.window.unwrap_or(0) as usize;
13055        let n_rot = geometry.n_rot as usize;
13056        let n_embd = self.cfg.n_embd as usize;
13057        let gw = fa
13058            .attn_gate
13059            .as_ref()
13060            .ok_or("step35 layer is missing attn_gate.weight (head-wise attention gate)")?;
13061
13062        let (q0, k0, v0, gt) = match pre_q {
13063            Some((hq, hdq)) => {
13064                debug_assert!(
13065                    e.uses_q8_1_fast(gw),
13066                    "step35 pre-quantized decode requires attn_gate on the q8_1 fast path \
13067                     (h is a zero-length placeholder here) — see mixer_in_q8_1_fast"
13068                );
13069                let (a, b, c) = match e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hdq)? {
13070                    Some(t3) => t3,
13071                    None => (
13072                        e.matmul_pre(&fa.wq, hq, hdq, h, 1)?,
13073                        e.matmul_pre(&fa.wk, hq, hdq, h, 1)?,
13074                        e.matmul_pre(&fa.wv, hq, hdq, h, 1)?,
13075                    ),
13076                };
13077                let gt = e.matmul_pre(gw, hq, hdq, h, 1)?;
13078                (a, b, c, gt)
13079            }
13080            None => {
13081                if e.uses_q8_1_fast(&fa.wq)
13082                    && e.uses_q8_1_fast(&fa.wk)
13083                    && e.uses_q8_1_fast(&fa.wv)
13084                    && e.uses_q8_1_fast(gw)
13085                {
13086                    let (hq, hdq) = e.quantize_q8_1(h, 1, n_embd)?;
13087                    let (a, b, c) = match e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, &hq, &hdq)? {
13088                        Some(t3) => t3,
13089                        None => (
13090                            e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
13091                            e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
13092                            e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?,
13093                        ),
13094                    };
13095                    let gt = e.matmul_pre(gw, &hq, &hdq, h, 1)?;
13096                    (a, b, c, gt)
13097                } else {
13098                    (
13099                        e.matmul(&fa.wq, h, 1)?,
13100                        e.matmul(&fa.wk, h, 1)?,
13101                        e.matmul(&fa.wv, h, 1)?,
13102                        e.matmul(gw, h, 1)?,
13103                    )
13104                }
13105            }
13106        };
13107
13108        let mut q = e.uninit(nh * hd)?;
13109        e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, nh, eps)?;
13110        let mut k = e.uninit(nkv * hd)?;
13111        e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, nkv, eps)?;
13112        let ff = if swa {
13113            None
13114        } else {
13115            self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
13116        };
13117        #[cfg(debug_assertions)]
13118        if let Some(ff) = ff {
13119            crate::debug_assert_tensor_stream_device(
13120                ff,
13121                &e.stream(),
13122                "step35_decode_attn.rope_freqs",
13123            );
13124        }
13125        e.rope_neox2(&mut q, &mut k, pos_d, hd, n_rot, nh, nkv, 1, rbase, 1.0, ff)?;
13126
13127        if std::env::var("MEMRA_NOFA").is_ok() {
13128            return Err(
13129                "MEMRA_NOFA (naive f32 SDPA) is incompatible with the quantized KV \
13130                        cache; unset MEMRA_NOFA to use fa_decode"
13131                    .into(),
13132            );
13133        }
13134        let kvl = cache.kv[il].as_mut().unwrap();
13135        let next_len = kvl.len + 1;
13136        let (off, t_kv) = if swa && next_len > win {
13137            (next_len - win, win)
13138        } else {
13139            (0, next_len)
13140        };
13141        let write_row = e.prepare_kv_append(kvl, off & !31usize, 1)?;
13142        e.append_kv_quantized(
13143            &k,
13144            &v0,
13145            &mut kvl.k,
13146            &mut kvl.v,
13147            write_row,
13148            kvl.kv_dim_k,
13149            kvl.kv_dim_v,
13150            kvl.k_tok_bytes,
13151            kvl.v_tok_bytes,
13152            crate::Engine::kv_fp8_on(),
13153        )?;
13154        kvl.len = next_len;
13155        let physical = kvl.physical_rows(off, off + t_kv)?;
13156        let k_view = e.view_u8_range(
13157            &kvl.k,
13158            physical.start * kvl.k_tok_bytes,
13159            physical.end * kvl.k_tok_bytes,
13160        );
13161        let v_view = e.view_u8_range(
13162            &kvl.v,
13163            physical.start * kvl.v_tok_bytes,
13164            physical.end * kvl.v_tok_bytes,
13165        );
13166        let mut attn = e.uninit(nh * hd)?;
13167        e.fa_decode_kvmod(
13168            &q,
13169            &k_view,
13170            &v_view,
13171            &mut attn,
13172            hd,
13173            nh,
13174            nkv,
13175            t_kv,
13176            scale,
13177            kvl.k_tok_bytes,
13178            kvl.v_tok_bytes,
13179            crate::Engine::kv_fp8_on(),
13180        )?;
13181
13182        let mut ag = e.uninit(nh * hd)?;
13183        e.attn_head_gate(&attn, &gt, &mut ag, None, hd, nh, 1)?;
13184        Ok(e.matmul(&fa.wo, &ag, 1)?)
13185    }
13186}
13187
13188// ===================================================================================== //
13189//  gemma-4 E4B (per-layer embeddings + KV-sharing) — FIRST-LIGHT forward.               //
13190//  Dedicated simple path (Stage-B matmuls, per-row causal attention over the quantized  //
13191//  cache) so the tuned 26B/31B paths stay untouched. Arch: research/gemma4-bringup/     //
13192//  e4b-arch-map.md; llama reference: src/models/gemma4.cpp (E4B arms). Wired: forward   //
13193//  (prefill logits), prime (tokenwise-equivalent batched trunk), eager decode_step_h.   //
13194//  NOT wired: dc/graph serving arms, verify/spec, chunked prime (HANDOVER-E4B.md).      //
13195// ===================================================================================== //
13196impl HybridModel {
13197    pub fn is_gemma4_e4b(&self) -> bool {
13198        self.gemma4_aux.as_ref().is_some_and(|a| a.e4b.is_some())
13199    }
13200
13201    /// E4B per-layer geometry: nh/nkv derive from the layer's OWN tensor shapes (the E4B GGUF
13202    /// ships a scalar head_count_kv; swa q 8x256 / kv 2x256, global q 4x512 / kv 1x512).
13203    /// KV-shared layers report the SHARE TARGET's kv count (their wk IS the target's tensor).
13204    fn gemma4_e4b_geom(&self, il: usize) -> (usize, usize, usize, f32, f32, bool) {
13205        let g = self.cfg.gemma4.as_ref().unwrap();
13206        let swa = g.swa_pattern[il];
13207        let hd = if swa {
13208            g.key_length_swa
13209        } else {
13210            g.key_length_global
13211        } as usize;
13212        let Mixer::Full(fa) = &self.layers[il].mixer else {
13213            panic!("e4b layer {il} not full-attn")
13214        };
13215        let nh = fa.wq.out_features() / hd;
13216        let nkv = fa.wk.out_features() / hd;
13217        (
13218            hd,
13219            nkv,
13220            nh,
13221            if swa {
13222                g.rope_base_swa
13223            } else {
13224                g.rope_base_global
13225            },
13226            1.0,
13227            swa,
13228        )
13229    }
13230
13231    /// E4B KV-share target: Some(target) for the trailing shared layers, None = own cache.
13232    fn gemma4_e4b_kv_target(&self, il: usize) -> Option<usize> {
13233        self.layers[il]
13234            .gemma4
13235            .as_ref()
13236            .and_then(|b| b.e4b.as_ref())
13237            .and_then(|e4| e4.kv_share.map(|t| t as usize))
13238    }
13239
13240    /// E4B prologue: inp_pl [t][n_layer][n_epl] =
13241    ///   ( gather(per_layer_tok_embd, tok)*sqrt(n_epl)
13242    ///   + rms_norm(model_proj . x_scaled * 1/sqrt(n_embd), proj_norm) ) * 1/sqrt(2)
13243    /// (llama gemma4.cpp build_inp_per_layer + project_per_layer_inputs, exact order).
13244    fn gemma4_e4b_inp_pl(
13245        &self,
13246        e: &Engine,
13247        tokens: &[u32],
13248        x_scaled: &CudaSlice<f32>,
13249        t: usize,
13250    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13251        let tok_d = e.stream().clone_htod(&tokens.to_vec())?;
13252        self.gemma4_e4b_inp_pl_dev(e, &tok_d, x_scaled, t)
13253    }
13254
13255    /// Device-token prologue core (dc arm shares it: token ids never touch the host).
13256    fn gemma4_e4b_inp_pl_dev(
13257        &self,
13258        e: &Engine,
13259        tok_d: &CudaSlice<u32>,
13260        x_scaled: &CudaSlice<f32>,
13261        t: usize,
13262    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13263        let aux = self.gemma4_aux.as_ref().unwrap();
13264        let m = aux.e4b.as_ref().unwrap();
13265        let n_embd = self.cfg.n_embd as usize;
13266        let n_layer = self.layers.len();
13267        let width = m.n_epl * n_layer;
13268        let tbl = m.tok_tbl_gpu.get_or_init(|| {
13269            e.upload_u8(&m.tok_embd_bytes)
13270                .expect("e4b per-layer token table upload")
13271        });
13272        let mut a =
13273            e.embed_gather_device_td(tbl, tok_d, t, width, m.tok_embd_qt, m.tok_embd_row_bytes)?;
13274        e.scale_inplace(&mut a, (m.n_epl as f32).sqrt(), t * width)?;
13275        let mut p = e.matmul(&m.model_proj, x_scaled, t)?;
13276        e.scale_inplace(&mut p, 1.0 / (n_embd as f32).sqrt(), t * width)?;
13277        let mut pn = e.uninit(t * width)?;
13278        e.rms_norm(
13279            &p,
13280            m.proj_norm.float_data(),
13281            &mut pn,
13282            m.n_epl,
13283            t * n_layer,
13284            self.cfg.rms_eps,
13285        )?;
13286        let mut out = e.uninit(t * width)?;
13287        e.add_scale(&a, &pn, 1.0 / 2f32.sqrt(), &mut out, t * width)?;
13288        Ok(out)
13289    }
13290
13291    /// E4B attention (t-wide, causal, per-row fa over the QUANTIZED cache — first-light
13292    /// correctness path; the fa rows arms come later). Own-KV layers project+norm+rope k/v
13293    /// and append t rows; KV-shared layers are Q-only over the target layer's cache (which
13294    /// already holds this forward's rows — the target runs earlier in the stack).
13295    #[allow(clippy::too_many_arguments)]
13296    fn gemma4_e4b_attn(
13297        &self,
13298        e: &Engine,
13299        il: usize,
13300        hq: &CudaSlice<i8>,
13301        hdq: &CudaSlice<f32>,
13302        pos_d: &CudaSlice<i32>,
13303        t: usize,
13304        cache: &mut Cache,
13305        dc_bucket: Option<usize>,
13306    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13307        let (hd, nkv, nh, base, scale, swa) = self.gemma4_e4b_geom(il);
13308        let eps = self.cfg.rms_eps;
13309        let aux = self.gemma4_aux.as_ref().unwrap();
13310        let ones = aux.ones(e);
13311        #[cfg(debug_assertions)]
13312        crate::debug_assert_tensor_stream_device(ones, &e.stream(), "gemma4_e4b_attn.ones");
13313        let Mixer::Full(fa) = &self.layers[il].mixer else {
13314            unreachable!()
13315        };
13316        // pre-quantized layer input (E4B fusion port, 2026-07-12): ONE norm+quant feeds
13317        // wq/wk/wv via matmul_pre — the first-light path quantized the same h three times
13318        // (the profile's 342 quantize_q8_1/token; glue = 26% of the 6.1ms token).
13319        let h0 = e.zeros(0)?;
13320        let h = &h0;
13321
13322        let ff = if swa {
13323            None
13324        } else {
13325            Some(
13326                aux.rope_freqs(e)
13327                    .expect("e4b global rope needs rope_freqs.weight"),
13328            )
13329        };
13330        #[cfg(debug_assertions)]
13331        if let Some(ff) = ff {
13332            crate::debug_assert_tensor_stream_device(ff, &e.stream(), "gemma4_e4b_attn.rope_freqs");
13333        }
13334        let share = self.gemma4_e4b_kv_target(il);
13335        // Own-KV arms keep the f32 (k, v) alive for the prime fa arm below.
13336        let mut kv_f32: Option<(CudaSlice<f32>, CudaSlice<f32>)> = None;
13337        let mut q;
13338        if let Some(_tgt) = share {
13339            let q0 = e.matmul_pre(&fa.wq, hq, hdq, h, t)?;
13340            q = e.uninit(t * nh * hd)?;
13341            // Q-only through the same fused norm+rope kernel (rk = 0: the k/v segments are
13342            // empty; q0 stands in for the unused k/v pointers).
13343            let mut kdummy = e.uninit(1)?;
13344            let mut vdummy = e.uninit(1)?;
13345            e.rms_norm_qkv_rope(
13346                &q0,
13347                &q0,
13348                &q0,
13349                fa.q_norm.float_data(),
13350                fa.q_norm.float_data(),
13351                ones,
13352                &mut q,
13353                &mut kdummy,
13354                &mut vdummy,
13355                hd,
13356                nh * t,
13357                0,
13358                pos_d,
13359                nh,
13360                1,
13361                base,
13362                1.0,
13363                ff,
13364                eps,
13365            )?;
13366        } else {
13367            // wave-4b: ONE concat matvec (wq|wk|wv) at t == 1 when the cat tensor exists;
13368            // else the fused3 grid launch; else per-matvec. The cat output is contiguous
13369            // q|k|v rows — the cat norm+rope twin consumes it directly.
13370            let e4bits = self.layers[il].gemma4.as_ref().and_then(|g| g.e4b.as_ref());
13371            let cat = e4bits.and_then(|e4| e4.qkv_cat.as_ref());
13372            q = e.uninit(t * nh * hd)?;
13373            let mut k = e.uninit(t * nkv * hd)?;
13374            let mut v = e.uninit(t * nkv * hd)?;
13375            if t == 1 && cat.is_some() {
13376                let qkv0 = e.matmul_pre(cat.unwrap(), hq, hdq, h, 1)?;
13377                e.rms_norm_qkv_rope_cat(
13378                    &qkv0,
13379                    fa.q_norm.float_data(),
13380                    fa.k_norm.float_data(),
13381                    ones,
13382                    &mut q,
13383                    &mut k,
13384                    &mut v,
13385                    hd,
13386                    nh,
13387                    nkv,
13388                    pos_d,
13389                    nh,
13390                    nkv,
13391                    base,
13392                    1.0,
13393                    ff,
13394                    eps,
13395                )?;
13396            } else {
13397                let (q0, k0, v0) = match if t == 1 {
13398                    e.matmul_q4_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hdq)?
13399                } else {
13400                    // E4B verify f3 port (2026-07-14): the 31B segmented-grid batched qkv
13401                    // on E4B's real-V triple — same MEMRA_F2B seam, bit-identical per row.
13402                    static F2B_QKV: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13403                    if *F2B_QKV.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0")) {
13404                        e.matmul_q4_fused3_batched(&fa.wq, &fa.wk, &fa.wv, hq, hdq, t)?
13405                    } else {
13406                        None
13407                    }
13408                } {
13409                    Some(triple) => triple,
13410                    None => (
13411                        e.matmul_pre(&fa.wq, hq, hdq, h, t)?,
13412                        e.matmul_pre(&fa.wk, hq, hdq, h, t)?,
13413                        e.matmul_pre(&fa.wv, hq, hdq, h, t)?,
13414                    ), // E4B: real v (K != V)
13415                };
13416                // wave-3 fold: q/k/v norms + q/k rope in ONE launch (rope math verbatim on
13417                // the normed rows; V ones-rms, never roped).
13418                e.rms_norm_qkv_rope(
13419                    &q0,
13420                    &k0,
13421                    &v0,
13422                    fa.q_norm.float_data(),
13423                    fa.k_norm.float_data(),
13424                    ones,
13425                    &mut q,
13426                    &mut k,
13427                    &mut v,
13428                    hd,
13429                    nh * t,
13430                    nkv * t,
13431                    pos_d,
13432                    nh,
13433                    nkv,
13434                    base,
13435                    1.0,
13436                    ff,
13437                    eps,
13438                )?;
13439            }
13440            let kvl = cache.kv[il].as_mut().unwrap();
13441            // class flag must match the cache dims (the g-threading sweep hardcoded `false`
13442            // here and the wkv default corrupted E4B: q8_0 bytes into an e4m3 cache — the
13443            // degenerate tok-0 stream, 2026-07-12).
13444            let cls = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
13445            if dc_bucket.is_some() {
13446                // DC arm (graph serving): append at the len_d slot, advance the counter
13447                // in-stream — replay-correct, no host len in the launch args. Host mirrors
13448                // are NOT touched here (the replay loop owns them; a bump at capture-record
13449                // time would double-count the capture iteration).
13450                debug_assert!(t == 1);
13451                // wave 5c: append + len_d inc fused (one launch; single-block ordering).
13452                e.append_kv_quantized_row_dc_inc(
13453                    &k,
13454                    &v,
13455                    &mut kvl.k,
13456                    &mut kvl.v,
13457                    &mut kvl.len_d,
13458                    kvl.kv_dim_k,
13459                    kvl.kv_dim_v,
13460                    kvl.k_tok_bytes,
13461                    kvl.v_tok_bytes,
13462                    cls,
13463                )?;
13464            } else {
13465                e.append_kv_quantized_rows(
13466                    &k,
13467                    &v,
13468                    &mut kvl.k,
13469                    &mut kvl.v,
13470                    kvl.len,
13471                    t,
13472                    kvl.kv_dim_k,
13473                    kvl.kv_dim_v,
13474                    kvl.k_tok_bytes,
13475                    kvl.v_tok_bytes,
13476                    cls,
13477                )?;
13478                kvl.len += t;
13479            }
13480            kv_f32 = Some((k, v));
13481        }
13482        // attention: per-row causal fa over the (own or target) quantized cache. The cache
13483        // already contains this forward's rows in both arms; row i attends [.., base+i].
13484        let kvl_idx = share.unwrap_or(il);
13485        let kvl = cache.kv[kvl_idx].as_ref().unwrap();
13486        let base_len = kvl.len - t; // pre-append length (target appended this forward too)
13487        let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
13488        let mut attn = e.uninit(t * nh * hd)?;
13489        // PRIME FA ARMS (perf lane 3, completed 2026-07-31 — the H100 board pinned E4B
13490        // prefill at 482 tok/s: the per-row loop ran the DECODE kernel once per token per
13491        // layer). Fresh-prompt (base_len == 0) batched attention, one launch per layer:
13492        //   - own-KV hd256 under the window: f32 fa_prefill (full causal exact — 26B pattern)
13493        //   - own-KV hd256 above the window (swa): fa_prefill_w windowed twin (12B pattern)
13494        //   - own-KV hd512 globals: fa_prefill_hd512 (12B/31B globals pattern)
13495        //   - KV-shared layers (no f32 k/v): fa_prefill_view over the target's QUANTIZED
13496        //     rows (the T=K verify kernel; the target appended this forward's rows already).
13497        //     swa-shared above the window keeps the per-row loop (no windowed view twin).
13498        // Same numeric class split as the 26B prime (f32/batched prefill vs per-row
13499        // quantized decode); run-gen argmax + chat gates arbitrate. MEMRA_NOFA=1 reverts all.
13500        if t > 1 && base_len == 0 && std::env::var("MEMRA_NOFA").is_err() {
13501            if let Some((kf, vf)) = &kv_f32 {
13502                if hd == 256 && t <= win {
13503                    e.fa_prefill(&q, kf, vf, &mut attn, hd, nh, nkv, t, t, scale, true)?;
13504                    return Ok(e.matmul(&fa.wo, &attn, t)?);
13505                }
13506                if hd == 256 && swa && t > win {
13507                    e.fa_prefill_w(&q, kf, vf, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
13508                    return Ok(e.matmul(&fa.wo, &attn, t)?);
13509                }
13510                if hd == 512 && !swa {
13511                    e.fa_prefill_hd512(&q, kf, vf, &mut attn, hd, nh, nkv, t, t, scale, true)?;
13512                    return Ok(e.matmul(&fa.wo, &attn, t)?);
13513                }
13514            } else if share.is_some() {
13515                let g = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
13516                let k_view = e.view_u8(&kvl.k, kvl.k.len());
13517                let v_view = e.view_u8(&kvl.v, kvl.v.len());
13518                if hd == 256 && (!swa || t <= win) {
13519                    // inline-dequant quantized-view prefill (fa_prefill_q stamps 256/128)
13520                    e.fa_prefill_view(
13521                        &q,
13522                        &k_view,
13523                        &v_view,
13524                        &mut attn,
13525                        hd,
13526                        nh,
13527                        nkv,
13528                        t,
13529                        t,
13530                        scale,
13531                        true,
13532                        kvl.k_tok_bytes,
13533                        kvl.v_tok_bytes,
13534                        g,
13535                    )?;
13536                    return Ok(e.matmul(&fa.wo, &attn, t)?);
13537                }
13538                // remaining shared classes (swa above the window; hd512 globals): dequant
13539                // the target's t rows ONCE to f32, then the same f32 twins as own-KV.
13540                let kv_dim = nkv * hd;
13541                let mut kf = e.uninit(t * kv_dim)?;
13542                let mut vf = e.uninit(t * kv_dim)?;
13543                e.fa_dequant_kv_view_f32(
13544                    &k_view,
13545                    &v_view,
13546                    &mut kf,
13547                    &mut vf,
13548                    kv_dim,
13549                    kv_dim,
13550                    t,
13551                    kvl.k_tok_bytes,
13552                    kvl.v_tok_bytes,
13553                    g,
13554                )?;
13555                if hd == 512 {
13556                    e.fa_prefill_hd512(&q, &kf, &vf, &mut attn, hd, nh, nkv, t, t, scale, true)?;
13557                } else {
13558                    e.fa_prefill_w(&q, &kf, &vf, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
13559                }
13560                return Ok(e.matmul(&fa.wo, &attn, t)?);
13561            }
13562        }
13563        if let Some(bucket) = dc_bucket {
13564            // DC arm (t == 1, under-window regime — the generate gate enforces it): ONE
13565            // fa_decode_dc over the live counter. len_d already advanced past this token
13566            // (t_kv = len_d[0], the cache.rs contract). KV-shared layers read the target's
13567            // counter (advanced when the target ran earlier in the stack).
13568            assert!(t == 1);
13569            // hd512 globals: eager (kvmod) picks the SCALAR unified below the fa512 floor,
13570            // and under the window every live t_kv sits below it — cap the capture bucket
13571            // under the floor so fa_decode_dc bakes the same scalar symbol, or the graph
13572            // rides the dpl16 twin and its numeric class diverges from the dc-eager stream
13573            // (E4B-GRAPH-GATE 2/64, 2026-07-12).
13574            let bucket = if hd == 512 && win <= crate::fa512_min_tkv() {
13575                bucket.min(crate::fa512_min_tkv().saturating_sub(1))
13576            } else {
13577                bucket
13578            };
13579            let k_view = e.view_u8(&kvl.k, kvl.k.len());
13580            let v_view = e.view_u8(&kvl.v, kvl.v.len());
13581            let g = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
13582            // Weight prefetch (SOTA item 3, 2026-07-13): wo's decode plane prefetched into
13583            // L2 across the fa window — fa reads KV only, the weight DRAM lanes are idle
13584            // there (E4B valid window +0.65%: 196.8 vs 195.6). Value-free scheduling op,
13585            // captured into the dc graph like any other launch. Extending the cascade to
13586            // the ffn gate/up planes measured NEGATIVE (193.9 vs 195.8 — 29MB/layer floods
13587            // the fill path and evicts still-hot lines); wo-only is the shipped shape.
13588            // MEMRA_WPF=0 rollback seam.
13589            if crate::Engine::wpf_level() >= 1 {
13590                e.prefetch_weight_l2(&fa.wo)?;
13591            }
13592            // wave 5b: the combine emits the wo input q8 pair directly — the standalone
13593            // quantize launch + the f32 attn round-trip fold away (t=1 fast path only).
13594            if e.uses_q8_1_fast(&fa.wo) {
13595                let mut oq = e.alloc_i8_uninit(nh * hd)?;
13596                let mut od = e.zeros(nh * hd / 32)?;
13597                e.fa_decode_dc_q8(
13598                    &q,
13599                    &k_view,
13600                    &v_view,
13601                    &mut attn,
13602                    hd,
13603                    nh,
13604                    nkv,
13605                    &kvl.len_d,
13606                    bucket,
13607                    scale,
13608                    kvl.k_tok_bytes,
13609                    kvl.v_tok_bytes,
13610                    g,
13611                    Some((&mut oq, &mut od)),
13612                )?;
13613                return Ok(e.matmul_pre(&fa.wo, &oq, &od, &attn, t)?);
13614            }
13615            e.fa_decode_dc(
13616                &q,
13617                &k_view,
13618                &v_view,
13619                &mut attn,
13620                hd,
13621                nh,
13622                nkv,
13623                &kvl.len_d,
13624                bucket,
13625                scale,
13626                kvl.k_tok_bytes,
13627                kvl.v_tok_bytes,
13628                g,
13629            )?;
13630            return Ok(e.matmul(&fa.wo, &attn, t)?);
13631        }
13632        for i in 0..t {
13633            let avail = base_len + i + 1;
13634            let (off_tok, t_kv) = if swa && avail > win {
13635                (avail - win, win)
13636            } else {
13637                (0, avail)
13638            };
13639            let k_view = e.view_u8_range(
13640                &kvl.k,
13641                off_tok * kvl.k_tok_bytes,
13642                (off_tok + t_kv) * kvl.k_tok_bytes,
13643            );
13644            let v_view = e.view_u8_range(
13645                &kvl.v,
13646                off_tok * kvl.v_tok_bytes,
13647                (off_tok + t_kv) * kvl.v_tok_bytes,
13648            );
13649            let qv = e.view(&q, t * nh * hd);
13650            let q_row = qv.slice(i * nh * hd..(i + 1) * nh * hd);
13651            let mut q_one = e.uninit(nh * hd)?;
13652            e.copy_view_into(&mut q_one, 0, &q_row, nh * hd)?;
13653            let mut a_one = e.uninit(nh * hd)?;
13654            // read class MUST match the append class (globals are e4m3 under gkv): the
13655            // swa-only flag decoded global-layer e4m3 bytes as q8_0/q5_1 — attention over
13656            // garbage on EVERY E4B global layer, the cross-mode maxdiff-30 root (2026-07-12).
13657            e.fa_decode_kvmod(
13658                &q_one,
13659                &k_view,
13660                &v_view,
13661                &mut a_one,
13662                hd,
13663                nh,
13664                nkv,
13665                t_kv,
13666                scale,
13667                kvl.k_tok_bytes,
13668                kvl.v_tok_bytes,
13669                (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on()),
13670            )?;
13671            e.copy_into(&mut attn, i * nh * hd, &a_one, nh * hd)?;
13672        }
13673        Ok(e.matmul(&fa.wo, &attn, t)?)
13674    }
13675
13676    /// E4B trunk: embed -> prologue -> layers (attn + dense ffn via the 31B tail_core + the
13677    /// per-layer-embedding tail + layer scale) -> output_norm. Returns (softcapped logits
13678    /// device [t, n_vocab], pre-output_norm hidden [t, n_embd]). Appends t rows per own-KV
13679    /// layer; does NOT advance cache.pos (caller owns pos).
13680    fn gemma4_e4b_trunk(
13681        &self,
13682        e: &Engine,
13683        tokens: &[u32],
13684        pos0: usize,
13685        cache: &mut Cache,
13686        head_last: bool,
13687    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13688        let n_embd = self.cfg.n_embd as usize;
13689        let t = tokens.len();
13690        let pos: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
13691        let pos_d = e.htod_i32(&pos)?;
13692        let mut x = e.htod(&self.embd.gather(n_embd, tokens))?;
13693        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
13694        let inp_pl = self.gemma4_e4b_inp_pl(e, tokens, &x, t)?;
13695        self.gemma4_e4b_trunk_core(e, x, inp_pl, &pos_d, t, cache, None, true, head_last)
13696    }
13697
13698    /// Layer stack + head over prebuilt (x_scaled, inp_pl, device pos) — everything below
13699    /// here is device-driven, so the dc arm shares it verbatim (stream identity with the
13700    /// eager chain by construction: SAME functions, not twins).
13701    fn gemma4_e4b_trunk_core(
13702        &self,
13703        e: &Engine,
13704        x_in: CudaSlice<f32>,
13705        inp_pl: CudaSlice<f32>,
13706        pos_d: &CudaSlice<i32>,
13707        t: usize,
13708        cache: &mut Cache,
13709        dc_bucket: Option<usize>,
13710        cap_logits: bool,
13711        head_last: bool,
13712    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13713        let n_embd = self.cfg.n_embd as usize;
13714        let eps = self.cfg.rms_eps;
13715        let n_layer = self.layers.len();
13716        let mut x = x_in;
13717        let aux_e4b = self.gemma4_aux.as_ref().unwrap().e4b.as_ref().unwrap();
13718        let n_epl = aux_e4b.n_epl;
13719
13720        // cross-layer fusion (2026-07-12 port of the 26B/31B trunk structure): each layer's
13721        // closing add_scale also EMITS the next layer's attn-normed input pre-quantized
13722        // q8_1 (add_scale_rms_norm_q8_1); the LAST layer emits through output_norm so the
13723        // head rides matmul_pre too. First layer's pair comes from a standalone fused
13724        // norm+quant.
13725        let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
13726        for il in 0..n_layer {
13727            let layer = &self.layers[il];
13728            let (hq, hdq) = match h_carry.take() {
13729                Some(p) => p,
13730                None => e.rms_norm_q8_1(&x, layer.attn_norm.float_data(), n_embd, t, eps)?,
13731            };
13732            let o = self.gemma4_e4b_attn(e, il, &hq, &hdq, pos_d, t, cache, dc_bucket)?;
13733            // dense ffn tail with the post-attn norm FOLDED into its entry (one launch for
13734            // rms(o, post_attn_norm) + residual add + ffn_norm — glue-fusion lane).
13735            let bits = layer.gemma4.as_ref().unwrap();
13736            let e4b = bits.e4b.as_ref().expect("e4b layer bits");
13737            // glue wave 5: the tail DEFERS its post_ffw norm — the FFN exit fuses
13738            // rms(f0, post_ffw) + residual add + q8 emit into ONE launch (rms_pre_add_q8_1),
13739            // killing the rms_norm + add_q8_1 pair per layer. T-GENERIC since 2026-07-13:
13740            // the fused single-phase reduction is NOT FP-order-identical to the unfused
13741            // rms_norm+add pair (the original "bit-identical chain" claim was FALSE — the
13742            // t==1 gate left the batched VERIFY on the unfused chain and split E4B verify
13743            // from decode by logit maxdiff ~0.45, the greedy tie-flip at depth: gate 135/256.
13744            // NOFUSE bisect: disabling ONLY this fusion -> verify maxdiff 0.000e0). With the
13745            // gate dropped, decode AND verify ride the same fused chain — parity by
13746            // construction, VERIFY-GATE 0.000e0.
13747            let fuse_exit = e.uses_q8_1_fast(&e4b.inp_gate);
13748            let (sn, attn_out) = self.gemma4_layer_tail_core_pn(
13749                e,
13750                layer,
13751                &o,
13752                &x,
13753                t,
13754                Some(layer.post_attn_norm.float_data()),
13755                fuse_exit,
13756            )?;
13757            let mut resid = e.uninit(t * n_embd)?;
13758            // per-layer-embedding tail: resid += rms(proj . (gelu(inp_gate . resid) * inp_pl[il]))
13759            // wave-2: the residual add emits q8_1 alongside — inp_gate rides matmul_pre.
13760            // (PLE one-block mega-fusion PROBED NEGATIVE 2026-07-13: argmax-correct but
13761            // 126 vs 189 tok/s — one SM pulling 0.74MB of weights loses to the multi-block
13762            // launch chain it replaced; jsonl row. Kernel deleted per doctrine.)
13763            let g = if fuse_exit {
13764                // sn here = RAW f0 (post_ffw deferred).
13765                let (rq, rd) = e.rms_pre_add_q8_1(
13766                    &sn,
13767                    bits.post_ffw_norm.float_data(),
13768                    &attn_out,
13769                    &mut resid,
13770                    n_embd,
13771                    t,
13772                    self.cfg.rms_eps,
13773                )?;
13774                e.matmul_pre(&e4b.inp_gate, &rq, &rd, &resid, t)?
13775            } else {
13776                e.add(&sn, &attn_out, &mut resid, t * n_embd)?;
13777                e.matmul(&e4b.inp_gate, &resid, t)?
13778            };
13779            let mut act = e.uninit(t * n_epl)?;
13780            let y = if t == 1 && e.uses_q8_1_fast(&e4b.proj) {
13781                let ipv = e.view(&inp_pl, n_epl * n_layer);
13782                let row = ipv.slice(il * n_epl..(il + 1) * n_epl);
13783                let (aq, ad) = e.gelu_tanh_mul_q8_1(&g, &row, &mut act, n_epl, 1)?;
13784                e.matmul_pre(&e4b.proj, &aq, &ad, &act, t)?
13785            } else {
13786                let mut inp_this = e.uninit(t * n_epl)?;
13787                e.copy_rows_strided(
13788                    &inp_pl,
13789                    &mut inp_this,
13790                    n_epl,
13791                    t,
13792                    n_epl * n_layer,
13793                    il * n_epl,
13794                )?;
13795                e.gelu_tanh_mul(&g, &inp_this, &mut act, t * n_epl)?;
13796                e.matmul(&e4b.proj, &act, t)?
13797            };
13798            // rms(y, post_norm) + (yn + resid)*layer_scale + next-layer norm+quant emit,
13799            // ONE launch (glue-fusion lane; last layer emits through output_norm).
13800            let next_norm = if il + 1 < n_layer {
13801                self.layers[il + 1].attn_norm.float_data()
13802            } else {
13803                self.output_norm.float_data()
13804            };
13805            let mut xn = e.uninit(t * n_embd)?;
13806            let pair = e.rms_pre_add_scale_rms_norm_q8_1(
13807                &y,
13808                e4b.post_norm.float_data(),
13809                &resid,
13810                bits.layer_scale,
13811                next_norm,
13812                &mut xn,
13813                n_embd,
13814                t,
13815                eps,
13816            )?;
13817            h_carry = Some(pair);
13818            x = xn;
13819        }
13820        // the head consumes the last layer's fused (output_norm) emit. head_last callers
13821        // (prime, last_only forward) need only the final row's logits — the all-T head is
13822        // t*n_vocab of discarded work (E4B prime: ~134ms GEMM + a 2.26GB dtoh kept 1 row).
13823        let (oq, odq) = h_carry.take().unwrap();
13824        let h0 = e.zeros(0)?;
13825        let hm = if head_last { 1 } else { t };
13826        let (hq, hd) = if head_last && t > 1 {
13827            let mut q1 = e.uninit_i8(n_embd)?;
13828            e.dtod_copy_view_i8(&oq.slice((t - 1) * n_embd..t * n_embd), &mut q1)?;
13829            let nb = n_embd / 32;
13830            let mut d1 = e.uninit(nb)?;
13831            e.dtod_copy_view(&odq.slice((t - 1) * nb..t * nb), &mut d1)?;
13832            (q1, d1)
13833        } else {
13834            (oq, odq)
13835        };
13836        let mut ld = e.matmul_pre(&self.output, &hq, &hd, &h0, hm)?;
13837        // softcap is strictly monotonic — greedy (argmax-only) consumers skip it, matching
13838        // the 26B/31B dc precedent (their dc head goes matmul -> argmax with no cap).
13839        // Logit-returning callers (host logits / spec prime) keep the capped emit.
13840        if cap_logits {
13841            let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
13842            e.softcap(&mut ld, cap, hm * self.output.out_features())?;
13843        }
13844        self.gemma4_suppress(e, &mut ld, hm)?; // mask both capped and argmax-only consumers
13845        Ok((ld, x))
13846    }
13847
13848    /// E4B batched VERIFY (device tokens, the spec round's t=K+1 step): t rows through the
13849    /// e4b trunk (per-row causal attention; own-KV layers append t rows host-len, KV-shared
13850    /// layers ride their targets), per-row device argmax + the POST-output_norm hidden
13851    /// stack (the drafter's h convention). Advances cache.pos/kvl.len by t — the spec
13852    /// round rolls back rejected rows (shared layers have no KvLayer, so the plain rewind
13853    /// covers exactly the layers that appended).
13854    pub fn gemma4_e4b_decode_step_t_am_dev(
13855        &self,
13856        e: &Engine,
13857        tok_d: &CudaSlice<u32>,
13858        t: usize,
13859        pos0: usize,
13860        cache: &mut Cache,
13861    ) -> Result<(CudaSlice<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13862        let n_embd = self.cfg.n_embd as usize;
13863        let eps = self.cfg.rms_eps;
13864        let pos: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
13865        let pos_d = e.htod_i32(&pos)?;
13866        let embd_gpu = self
13867            .embd_gpu
13868            .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload"));
13869        let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
13870        let mut x = e.embed_gather_device_td(embd_gpu, tok_d, t, n_embd, qt, rb)?;
13871        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
13872        let inp_pl = self.gemma4_e4b_inp_pl_dev(e, tok_d, &x, t)?;
13873        let (ld, xp) =
13874            self.gemma4_e4b_trunk_core(e, x, inp_pl, &pos_d, t, cache, None, true, false)?;
13875        // softcap is monotonic — the per-row argmax is invariant to it (the trunk's head
13876        // emit is already capped, matching the eager chain bit-for-bit).
13877        let n_vocab = self.output.out_features();
13878        let mut vam = e.stream().alloc_zeros::<u32>(t)?;
13879        for i in 0..t {
13880            e.argmax_token_device_col(&ld, i, n_vocab, &mut vam, i)?;
13881        }
13882        let mut hn = e.uninit(t * n_embd)?;
13883        e.rms_norm(&xp, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
13884        cache.pos += t;
13885        Ok((vam, hn))
13886    }
13887
13888    /// E4B verify + host logits + POST-output_norm hidden stack (the short-prompt spec
13889    /// prime path — mirror of `gemma4_decode_step_t_h`).
13890    pub(crate) fn gemma4_e4b_decode_step_t_h(
13891        &self,
13892        e: &Engine,
13893        tokens: &[u32],
13894        pos0: usize,
13895        cache: &mut Cache,
13896    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13897        let n_embd = self.cfg.n_embd as usize;
13898        let eps = self.cfg.rms_eps;
13899        let t = tokens.len();
13900        let (ld, xp) = self.gemma4_e4b_trunk(e, tokens, pos0, cache, false)?;
13901        let mut hn = e.uninit(t * n_embd)?;
13902        e.rms_norm(&xp, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
13903        cache.pos += t;
13904        Ok((e.dtoh(&ld)?, hn))
13905    }
13906
13907    /// E4B GRAPH-CAPTURABLE dc step: same trunk as the dc step but token_d is updated IN
13908    /// PLACE (self-feeding replay) and every launch arg is a device counter — pos from
13909    /// pos_d (inc'd in-stream), KV slots from len_d (advanced in-stream), attention from
13910    /// fa_decode_dc at `bucket`. Host mirrors (cache.pos / kvl.len) advance in the caller's
13911    /// replay loop. UNDER-WINDOW regime only (the caller gates pos + budget < window).
13912    pub fn gemma4_e4b_decode_step_dcg(
13913        &self,
13914        e: &Engine,
13915        token_d: &mut CudaSlice<u32>,
13916        pos_d: &mut CudaSlice<i32>,
13917        embd_gpu: &CudaSlice<u8>,
13918        embd_qt: i32,
13919        embd_rb: usize,
13920        cache: &mut Cache,
13921        n_vocab: usize,
13922        bucket: usize,
13923    ) -> Result<(), Box<dyn std::error::Error>> {
13924        let n_embd = self.cfg.n_embd as usize;
13925        let mut x = e.embed_gather_device(embd_gpu, token_d, n_embd, embd_qt, embd_rb)?;
13926        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
13927        let inp_pl = self.gemma4_e4b_inp_pl_dev(e, token_d, &x, 1)?;
13928        let (ld, _x) =
13929            self.gemma4_e4b_trunk_core(e, x, inp_pl, pos_d, 1, cache, Some(bucket), false, false)?;
13930        e.argmax_token_device_into(&ld, token_d, n_vocab)?;
13931        e.inc_seqlen(pos_d)?;
13932        Ok(())
13933    }
13934
13935    /// E4B DEVICE-COUNTER decode step (the dc serving arm): token id rides `token_d`, the
13936    /// greedy argmax lands in the returned device buffer — 4B/token host traffic. The layer
13937    /// stack is gemma4_e4b_trunk_core, i.e. the SAME functions the eager chain runs (stream
13938    /// identity by construction, not by twin-kernel parity). Host KV mirrors advance like the
13939    /// 26B dc-eager arm (window views are host math); len_d stays synced by the caller's
13940    /// entry sync + the appends here don't read it. Graph capture is NOT wired (no
13941    /// cap_bucket_max) — the E4B graph arc comes after the perf gates.
13942    #[allow(clippy::too_many_arguments)]
13943    pub fn gemma4_e4b_decode_step_dc(
13944        &self,
13945        e: &Engine,
13946        token_d: &CudaSlice<u32>,
13947        pos_d: &mut CudaSlice<i32>,
13948        embd_gpu: &CudaSlice<u8>,
13949        embd_qt: i32,
13950        embd_rb: usize,
13951        cache: &mut Cache,
13952        n_vocab: usize,
13953    ) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
13954        let n_embd = self.cfg.n_embd as usize;
13955        let eps = self.cfg.rms_eps;
13956        let mut x = e.embed_gather_device(embd_gpu, token_d, n_embd, embd_qt, embd_rb)?;
13957        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
13958        let inp_pl = self.gemma4_e4b_inp_pl_dev(e, token_d, &x, 1)?;
13959        let (ld, _x) =
13960            self.gemma4_e4b_trunk_core(e, x, inp_pl, pos_d, 1, cache, None, false, false)?;
13961        let mut tok_out = e.stream().alloc_zeros::<u32>(1)?;
13962        e.argmax_token_device_into(&ld, &mut tok_out, n_vocab)?;
13963        e.inc_seqlen(pos_d)?;
13964        cache.pos += 1;
13965        let _ = eps;
13966        Ok(tok_out)
13967    }
13968
13969    /// E4B eager T=1 decode step (decode_step_h contract): returns (softcapped logits host,
13970    /// pre-output_norm hidden). Advances cache.pos.
13971    pub(crate) fn gemma4_e4b_decode_step_h(
13972        &self,
13973        e: &Engine,
13974        token: u32,
13975        cache: &mut Cache,
13976    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13977        let (ld, x) = self.gemma4_e4b_trunk(e, &[token], cache.pos, cache, false)?;
13978        let logits = e.dtoh(&ld)?;
13979        cache.pos += 1;
13980        Ok((logits, x))
13981    }
13982
13983    /// E4B batched prime (prime_cache contract: (last-row logits host, h_seed pre-norm last
13984    /// row, hidden stack)). First-light: the t-wide trunk (per-row attention) — correct, not
13985    /// fast; the prefill fa arms come later.
13986    pub(crate) fn gemma4_e4b_prime(
13987        &self,
13988        e: &Engine,
13989        tokens: &[u32],
13990        cache: &mut Cache,
13991    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13992        // Err, not assert (2026-08-07, lane/gemma4-serve-gaps): same served-chunked-prompt
13993        // process-kill as gemma4_prime — refuse per-request.
13994        if cache.pos != 0 {
13995            return Err(
13996                "e4b prime is fresh-prompt only (v0) — prime the full prompt in one \
13997                        call or decode tokenwise"
13998                    .into(),
13999            );
14000        }
14001        let n_embd = self.cfg.n_embd as usize;
14002        let t = tokens.len();
14003        let (ld, x) = self.gemma4_e4b_trunk(e, tokens, 0, cache, true)?;
14004        cache.pos += t;
14005        let last = e.dtoh(&ld)?; // head_last: ld is already the final row only
14006        let xv = e.view(&x, t * n_embd);
14007        let row = xv.slice((t - 1) * n_embd..t * n_embd);
14008        let mut h_seed = e.uninit(n_embd)?;
14009        e.copy_view_into(&mut h_seed, 0, &row, n_embd)?;
14010        Ok((last, h_seed, x))
14011    }
14012
14013    /// E4B prefill logits (forward contract — no persistent cache; scratch cache internally).
14014    pub(crate) fn gemma4_e4b_forward(
14015        &self,
14016        e: &Engine,
14017        tokens: &[u32],
14018        last_only: bool,
14019    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
14020        let mut cache = Cache::new(e, &self.cfg, tokens.len() + 8)?;
14021        let (ld, _x) = self.gemma4_e4b_trunk(e, tokens, 0, &mut cache, last_only)?;
14022        Ok(e.dtoh(&ld)?) // head_last already reduced to the final row when last_only
14023    }
14024}
14025
14026#[cfg(test)]
14027mod prime_chunk_schedule_tests {
14028    use super::{
14029        PRIME_MIN_T, PRIME_PIPE_MIN_CHUNK, dynamic_prime_chunk_ranges, fixed_prime_chunk_ranges,
14030        fixed_prime_chunk_ranges_for_ring,
14031    };
14032
14033    fn sizes(ranges: &[(usize, usize)]) -> Vec<usize> {
14034        ranges.iter().map(|(start, end)| end - start).collect()
14035    }
14036
14037    fn auto_chunk(t: usize) -> usize {
14038        t.div_ceil(8).max(PRIME_PIPE_MIN_CHUNK).min(4096)
14039    }
14040
14041    #[test]
14042    fn fixed_schedule_retains_measured_geometry() {
14043        assert_eq!(
14044            sizes(&fixed_prime_chunk_ranges(461, 128)),
14045            vec![128, 128, 128, 77]
14046        );
14047        assert_eq!(
14048            sizes(&fixed_prime_chunk_ranges(1833, 230)),
14049            vec![230, 230, 230, 230, 230, 230, 230, 223]
14050        );
14051        assert_eq!(sizes(&fixed_prime_chunk_ranges(4096, 512)), vec![512; 8]);
14052        let capped = sizes(&fixed_prime_chunk_ranges_for_ring(8200, 4096, true));
14053        assert_eq!(capped, vec![4096, 4088, 16]);
14054        assert!(capped.iter().all(|&rows| rows <= 4096));
14055        assert_eq!(
14056            sizes(&fixed_prime_chunk_ranges_for_ring(4100, 4096, false)),
14057            vec![4100],
14058            "flag-off schedule remains byte-for-byte the legacy monolithic tail",
14059        );
14060    }
14061
14062    #[test]
14063    fn dynamic_schedule_matches_registered_shapes() {
14064        let cases = [
14065            (461, vec![64, 141, 132, 124]),
14066            (1833, vec![115, 269, 260, 252, 244, 237, 231, 225]),
14067            (4096, vec![256, 602, 580, 563, 545, 531, 516, 503]),
14068        ];
14069        for (t, expected) in cases {
14070            let chunk = auto_chunk(t);
14071            let fixed = fixed_prime_chunk_ranges(t, chunk);
14072            assert_eq!(
14073                sizes(&dynamic_prime_chunk_ranges(t, chunk, &fixed)),
14074                expected
14075            );
14076        }
14077    }
14078
14079    #[test]
14080    fn dynamic_schedule_covers_exactly_and_shrinks_after_fill() {
14081        for t in 256..=8192 {
14082            let chunk = auto_chunk(t);
14083            let fixed = fixed_prime_chunk_ranges(t, chunk);
14084            let dynamic = dynamic_prime_chunk_ranges(t, chunk, &fixed);
14085            assert_eq!(dynamic.len(), fixed.len(), "T={t}");
14086            assert_eq!(dynamic.first().unwrap().0, 0, "T={t}");
14087            assert_eq!(dynamic.last().unwrap().1, t, "T={t}");
14088            for pair in dynamic.windows(2) {
14089                assert_eq!(pair[0].1, pair[1].0, "T={t}");
14090            }
14091            assert!(
14092                dynamic
14093                    .iter()
14094                    .all(|(start, end)| end - start >= PRIME_MIN_T),
14095                "T={t} sizes={:?}",
14096                sizes(&dynamic)
14097            );
14098            if dynamic.len() >= 3 {
14099                let chunk_sizes = sizes(&dynamic);
14100                assert!(
14101                    chunk_sizes[0] < chunk_sizes[1],
14102                    "T={t} sizes={chunk_sizes:?}"
14103                );
14104                assert!(
14105                    chunk_sizes[1..].windows(2).all(|pair| pair[0] >= pair[1]),
14106                    "T={t} sizes={chunk_sizes:?}"
14107                );
14108            }
14109        }
14110    }
14111}
14112
14113#[cfg(test)]
14114mod page_prefetch_tests {
14115    use super::{
14116        grouped_worker_prefetch_position, page_prefetch_positions,
14117        page_prefetch_window_from_values, worker_prefetch_positions,
14118    };
14119
14120    #[test]
14121    fn page_prefetch_window_keeps_existing_opt_in_default() {
14122        assert_eq!(page_prefetch_window_from_values(false, None), 0);
14123        assert_eq!(page_prefetch_window_from_values(false, Some("8")), 0);
14124        assert_eq!(page_prefetch_window_from_values(true, None), 1);
14125        assert_eq!(page_prefetch_window_from_values(true, Some("bad")), 1);
14126        assert_eq!(page_prefetch_window_from_values(true, Some("0")), 0);
14127        assert_eq!(page_prefetch_window_from_values(true, Some("8")), 8);
14128    }
14129
14130    #[test]
14131    fn rolling_page_prefetch_advises_each_future_expert_once() {
14132        let advised: Vec<_> = (0..7)
14133            .flat_map(|position| page_prefetch_positions(position, 7, 3))
14134            .collect();
14135        assert_eq!(advised, vec![1, 2, 3, 4, 5, 6]);
14136
14137        let one_ahead: Vec<_> = (0..4)
14138            .flat_map(|position| page_prefetch_positions(position, 4, 1))
14139            .collect();
14140        assert_eq!(one_ahead, vec![1, 2, 3]);
14141        assert!(page_prefetch_positions(0, 4, 0).is_empty());
14142    }
14143
14144    #[test]
14145    fn grouped_worker_prefetch_primes_first_then_each_known_next_once() {
14146        assert_eq!(grouped_worker_prefetch_position(0, None), None);
14147        let positions: Vec<_> = std::iter::once(grouped_worker_prefetch_position(4, None).unwrap())
14148            .chain(
14149                (0..4).filter_map(|position| grouped_worker_prefetch_position(4, Some(position))),
14150            )
14151            .collect();
14152        assert_eq!(positions, vec![0, 1, 2, 3]);
14153        assert_eq!(grouped_worker_prefetch_position(1, Some(0)), None);
14154    }
14155
14156    #[test]
14157    fn rolling_worker_prefetch_primes_current_and_each_future_expert_once() {
14158        let queued: Vec<_> = (0..8)
14159            .flat_map(|position| worker_prefetch_positions(position, 8, 5))
14160            .collect();
14161        assert_eq!(queued, (0..8).collect::<Vec<_>>());
14162
14163        let one_at_a_time: Vec<_> = (0..4)
14164            .flat_map(|position| worker_prefetch_positions(position, 4, 1))
14165            .collect();
14166        assert_eq!(one_at_a_time, vec![0, 1, 2, 3]);
14167        assert!(worker_prefetch_positions(0, 4, 0).is_empty());
14168    }
14169}
14170
14171pub struct G4DcSlots {
14172    x: CudaSlice<f32>,
14173    xn: CudaSlice<f32>,
14174    cur: CudaSlice<f32>,
14175    hq: CudaSlice<i8>,
14176    hd_: CudaSlice<f32>,
14177    q0: CudaSlice<f32>,
14178    k0: CudaSlice<f32>,
14179    v0: CudaSlice<f32>,
14180    q: CudaSlice<f32>,
14181    k: CudaSlice<f32>,
14182    v: CudaSlice<f32>,
14183    attn: CudaSlice<f32>,
14184    o: CudaSlice<f32>,
14185    attn_out: CudaSlice<f32>,
14186    zsh: CudaSlice<f32>,
14187    zq: CudaSlice<i8>,
14188    zd: CudaSlice<f32>,
14189    gate: CudaSlice<f32>,
14190    up: CudaSlice<f32>,
14191    act: CudaSlice<f32>,
14192    actq: CudaSlice<i8>,
14193    actd: CudaSlice<f32>,
14194    f0: CudaSlice<f32>,
14195    sn: CudaSlice<f32>,
14196    hn: CudaSlice<f32>,
14197    logits: CudaSlice<f32>,
14198}