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            // chunked prime writes tap rows at the chunk's absolute offset
894            if let Some(taps) = cache.dflash_taps.as_mut() {
895                taps.base = start;
896            }
897            let (l, hs, x) =
898                self.prime_chunk(e, &tokens[start..end], cache, seq_end, start, overlay)?;
899            e.copy_into(&mut hiddens, start * n_embd, &x, (end - start) * n_embd)?;
900            last = Some((l, hs));
901        }
902        let (logits, h_seed) = last.unwrap();
903        Ok((logits, h_seed, hiddens))
904    }
905
906    /// PP-2 chunk scheduler: stage 1 of chunk N and stage 0 of chunk N+1 are both queued
907    /// before N's epilogue D2H drains the last-stage stream. Arithmetic is unchanged:
908    /// every chunk still runs the same two `prime_layers` ranges, boundary copy, output
909    /// norm, lm head, and caller hidden-stack copy as the serial split.
910    fn prime_cache_pp2_pipelined(
911        &self,
912        e: &Engine,
913        tokens: &[u32],
914        cache: &mut Cache,
915        seq_end: usize,
916        ranges: &[(usize, usize)],
917        fence: &[usize],
918    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
919        debug_assert_eq!(fence.len(), 3);
920        debug_assert!(ranges.len() >= 2);
921        let rt = crate::pp::PpNRt::get(e)?;
922        assert_eq!(
923            rt.n_stages(),
924            2,
925            "prime pipeline requires exactly two PP stages"
926        );
927        let n_embd = self.cfg.n_embd as usize;
928        let t = tokens.len();
929        let initial_base = cache.pos;
930        let caller_stream = e.stream();
931
932        // #87 reverse publication before any new stage allocation, then prewarm both
933        // boundary slots while the stage streams are otherwise empty. Lazy-growing slot B
934        // after stage 1(N) is queued would synchronize that stream and erase the first
935        // overlap on a two-chunk prompt.
936        rt.fence_stages_behind(&caller_stream)?;
937        let max_payload = ranges.iter().map(|(s, e)| (e - s) * n_embd).max().unwrap();
938        rt.prepare_overlap_slots(0, max_payload)?;
939
940        let mut hiddens = e.uninit(t * n_embd)?;
941        let mut last: Option<(Vec<f32>, CudaSlice<f32>)> = None;
942        let mut stage_caches = PrimeCacheStages::new(cache, fence[1]);
943        let (cache0, cache1) = stage_caches.parts();
944        let (first_start, first_end) = ranges[0];
945        let mut slot = self.prime_pp2_stage0_enqueue(
946            e,
947            rt,
948            &tokens[first_start..first_end],
949            cache0,
950            seq_end,
951            fence,
952            initial_base + first_start,
953            true,
954        )?;
955        cache0.pos = initial_base + first_end;
956
957        for (i, &(start, end)) in ranges.iter().enumerate() {
958            let base = initial_base + start;
959            debug_assert_eq!(
960                cache1.pos, base,
961                "stage 1 must drain chunks in original position order"
962            );
963            let (out, next_slot) = if let Some(&(next_start, next_end)) = ranges.get(i + 1) {
964                let next_base = initial_base + next_start;
965                debug_assert_eq!(
966                    cache0.pos, next_base,
967                    "stage 0 must issue chunks in original position order"
968                );
969                let cache0_stage = &mut *cache0;
970                // Step's MoE router readback synchronizes once per layer. Two CUDA streams
971                // on one host thread therefore serialize even if the calls are ordered as
972                // a pipeline. Drive the disjoint stage caches from two scoped host threads:
973                // stage 1 consumes slot N while stage 0 produces slot N+1.
974                std::thread::scope(|scope| -> Result<_, Box<dyn std::error::Error>> {
975                    let stage0 = scope.spawn(move || -> Result<usize, String> {
976                        let next = self
977                            .prime_pp2_stage0_enqueue(
978                                e,
979                                rt,
980                                &tokens[next_start..next_end],
981                                cache0_stage,
982                                seq_end,
983                                fence,
984                                next_base,
985                                true,
986                            )
987                            .map_err(|err| err.to_string())?;
988                        cache0_stage.pos = initial_base + next_end;
989                        Ok(next)
990                    });
991                    let x = self.prime_pp2_stage1_enqueue(
992                        e,
993                        rt,
994                        slot,
995                        end - start,
996                        cache1,
997                        seq_end,
998                        fence,
999                        base,
1000                        true,
1001                    )?;
1002                    let out = {
1003                        rt.bind_stage(1)?;
1004                        let _st1 = rt.enter(1);
1005                        let e1 = rt.engine(1, e);
1006                        self.prime_chunk_epilogue(e1, x, end - start, cache1)?
1007                    };
1008                    let next = stage0
1009                        .join()
1010                        .map_err(|_| "pipeprime stage-0 host walker panicked")?
1011                        .map_err(|err| -> Box<dyn std::error::Error> { err.into() })?;
1012                    Ok((out, Some(next)))
1013                })?
1014            } else {
1015                let x = self.prime_pp2_stage1_enqueue(
1016                    e,
1017                    rt,
1018                    slot,
1019                    end - start,
1020                    cache1,
1021                    seq_end,
1022                    fence,
1023                    base,
1024                    true,
1025                )?;
1026                let out = {
1027                    rt.bind_stage(1)?;
1028                    let _st1 = rt.enter(1);
1029                    let e1 = rt.engine(1, e);
1030                    self.prime_chunk_epilogue(e1, x, end - start, cache1)?
1031                };
1032                (out, None)
1033            };
1034
1035            rt.publish_to(1, &caller_stream)?;
1036            e.copy_into(&mut hiddens, start * n_embd, &out.2, (end - start) * n_embd)?;
1037            last = Some((out.0, out.1));
1038            crate::pp::PRIME_SPLIT_CHUNKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1039
1040            if let Some(next) = next_slot {
1041                // The caller copy above reads a stage-1 allocation. Before stage 1 of the
1042                // next chunk can allocate/reuse blocks, mirror #87's body-entry fence.
1043                // Stage 0(N+1) is already queued before this wait is appended, so its
1044                // overlap with stage 1(N) is preserved.
1045                rt.fence_stages_behind(&caller_stream)?;
1046                slot = next;
1047            }
1048        }
1049
1050        debug_assert_eq!(cache0.pos, initial_base + t);
1051        debug_assert_eq!(cache1.pos, initial_base + t);
1052        let (logits, h_seed) = last.unwrap();
1053        Ok((logits, h_seed, hiddens))
1054    }
1055
1056    /// One prime chunk: the full layer stack over `tokens`, continuing from the cache's current
1057    /// state (`cache.pos` = tokens already primed; 0 = fresh). Positions/RoPE are absolute
1058    /// task #21 de-broadcast: the q/k head count PREP must emit so the scan's consumers
1059    /// agree. Compact (num_k) ONLY when the scan will take the chunked+mma route AND
1060    /// num_k == num_v/2 (the engine-side hint mirrors this exact formula); everything
1061    /// else (s128 verify tier, chunked-off, mma-off) keeps the broadcast layout.
1062    fn gdn_hk(e: &Engine, t: usize, num_v: usize, num_k: usize) -> usize {
1063        if Engine::gdn_db_on()
1064            && Engine::gdn_chunked_enabled()
1065            && t >= 16
1066            && e.gdn_mma_enabled(Engine::gdn_chunk_size())
1067            && num_k * 2 == num_v
1068        {
1069            num_k
1070        } else {
1071            num_v
1072        }
1073    }
1074
1075    /// task #17: gate for the fused fp16-operand epilogues (silu_mul/gated_rmsnorm/sig_mul
1076    /// `_f16out` twins). Bit-identical class (the twins emit the cvt kernel's exact halves),
1077    /// but seam-gated (MEMRA_F16OUT=0) so a bit-check can arbitrate, and OFF under verify-exact
1078    /// (matmul_group skips the f16 lane there; the twins must not resurrect it).
1079    fn f16out_on(e: &Engine, t: usize) -> bool {
1080        crate::f16_ffi::pp_f16_enabled()
1081            && t >= 16
1082            && !e.verify_exact_on()
1083            && std::env::var("MEMRA_F16OUT").as_deref() != Ok("0")
1084    }
1085
1086    /// (cache.pos + i). Returns (last-row logits, h_seed, this chunk's hidden stack [T, n_embd]).
1087    /// See HybridModel::prime_slabs — the eager prime's resident trunk transients.
1088    /// PER-DEVICE since lane/pp-leverb (2026-08-08): the map is keyed by the allocating
1089    /// engine's CUDA ordinal — under the prime stage split each stage's range walks through
1090    /// its OWN slabs on its own device (a dev0 slab dereferenced by a dev1 kernel would be
1091    /// a peer read per GEMM operand, the exact class Lever B removes). Single-device rigs
1092    /// see one entry, byte-identical behavior.
1093    pub fn prime_slabs_get(
1094        &self,
1095        e: &Engine,
1096        t: usize,
1097        n_embd: usize,
1098        n_ff_max: usize,
1099    ) -> Result<std::sync::Arc<std::sync::Mutex<PrimeSlabs>>, Box<dyn std::error::Error>> {
1100        let mut slabs = self.prime_slabs.lock().unwrap();
1101        let dev = e.ctx().ordinal();
1102        let need_new = match slabs.get(&dev) {
1103            None => true,
1104            Some(sl) => sl.lock().unwrap().t_cap < t,
1105        };
1106        if need_new {
1107            slabs.insert(
1108                dev,
1109                std::sync::Arc::new(std::sync::Mutex::new(PrimeSlabs {
1110                    t_cap: t,
1111                    h: e.uninit(t * n_embd)?,
1112                    x1: e.uninit(t * n_embd)?,
1113                    z: e.uninit(t * n_embd)?,
1114                    act: e.uninit(t * n_ff_max)?,
1115                    xa: e.uninit(t * n_embd)?,
1116                    xb: e.uninit(t * n_embd)?,
1117                    h16: e.alloc_u8_uninit(t * n_embd * 2)?,
1118                    z16: e.alloc_u8_uninit(t * n_embd * 2)?,
1119                    gate: e.uninit(t * n_ff_max)?,
1120                    up: e.uninit(t * n_ff_max)?,
1121                    ffn_out: e.uninit(t * n_embd)?,
1122                    seg_glue: Vec::new(),
1123                    mixed: e.uninit(t * n_embd)?,
1124                    seg_mid: Vec::new(),
1125                    seg_t: 0,
1126                })),
1127            );
1128        }
1129        Ok(slabs.get(&dev).expect("prime slab inserted").clone())
1130    }
1131
1132    /// `seq_end` = the whole REQUEST's absolute end position (`cache.pos + prompt_len` at
1133    /// `prime_cache` entry), NOT this chunk's end. Chunk-size-invariant by construction; step35's
1134    /// SWA arm selects on it (see `step35_attn_pre_wo`). Every other arch ignores it.
1135    fn prime_chunk(
1136        &self,
1137        e: &Engine,
1138        tokens: &[u32],
1139        cache: &mut Cache,
1140        seq_end: usize,
1141        chunk_off: usize,
1142        overlay: Option<&crate::vision::EmbedOverlay>,
1143    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
1144        if crate::pp::pp_host_bounce_active()
1145            && (self.cfg.gemma4.is_some() || !crate::pp::prime_pp_on())
1146        {
1147            return Err(
1148                "prime_chunk: refused with MEMRA_PP_HOST_BOUNCE=1 because this configuration \
1149                 has no active prime stage split and would peer-read remote weights; keep \
1150                 MEMRA_PRIME_PP enabled and use a PP-prime-supported model"
1151                    .into(),
1152            );
1153        }
1154        // LEVER B (lane/pp-leverb, 2026-08-08): the ppN door for the chunked prime. With the
1155        // door open + per-stage streams, each chunk walks its layer ranges on the OWNING
1156        // stage's engine/device (the anatomy receipt this kills: dev1 ran ZERO prefill
1157        // kernels; stage-1 trunk weights were peer-read = 22% of the pp4096 wall).
1158        // MEMRA_PRIME_PP=0 = the unsplit rollback (also the prime-split-gate's reference
1159        // arm — prime deliberately keeps NO refuse_unsplit_if_remote, see pp.rs). The
1160        // MEMRA_PP_STREAMS=0 seam keeps the unsplit walk too: in that regime the sharded
1161        // loader is off and there is nothing remote to split for.
1162        if self.cfg.gemma4.is_none() && !crate::pp::pp2_streams_off() && crate::pp::prime_pp_on() {
1163            if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
1164                if overlay.is_some() {
1165                    return Err("vision embedding overlay + PP prime unsupported (v1); \
1166                         run single-device or MEMRA_PRIME_PP=0"
1167                        .into());
1168                }
1169                return self.prime_chunk_ppn(e, tokens, cache, seq_end, &fence);
1170            }
1171        }
1172        if crate::pp::pp_host_bounce_active() {
1173            return Err(
1174                "prime_chunk: MEMRA_PP_HOST_BOUNCE=1 found no valid prime stage split; \
1175                 refusing an unsplit remote-weight walk"
1176                    .into(),
1177            );
1178        }
1179        let t = tokens.len();
1180        let base = cache.pos;
1181        debug_assert!(
1182            seq_end >= base + t,
1183            "prime_chunk: seq_end must cover this chunk"
1184        );
1185        let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
1186        let pos_d = e.htod_i32(&pos)?;
1187
1188        let mut x_embed = self.embed(e, tokens)?; // [T, n_embd]
1189        if let Some(ov) = overlay {
1190            // Mixed-embedding splice: image rows overwrite the pad-token embeddings that
1191            // fall inside this chunk's prompt-relative window [chunk_off, chunk_off+t).
1192            // Images larger than one prime chunk straddle boundaries, hence the clipping.
1193            let n_embd = self.cfg.n_embd as usize;
1194            for &(pos, row_off, n_rows) in &ov.spans {
1195                let lo = pos.max(chunk_off);
1196                let hi = (pos + n_rows).min(chunk_off + t);
1197                if lo < hi {
1198                    let src_row = row_off + (lo - pos);
1199                    let view = ov
1200                        .rows
1201                        .slice(src_row * n_embd..(src_row + (hi - lo)) * n_embd);
1202                    e.copy_view_into(
1203                        &mut x_embed,
1204                        (lo - chunk_off) * n_embd,
1205                        &view,
1206                        (hi - lo) * n_embd,
1207                    )?;
1208                }
1209            }
1210        }
1211        let x = self.prime_layers(
1212            e,
1213            x_embed,
1214            0,
1215            self.layers.len(),
1216            &pos_d,
1217            t,
1218            base,
1219            cache,
1220            seq_end,
1221        )?;
1222        self.prime_chunk_epilogue(e, x, t, cache)
1223    }
1224
1225    /// PRIME RANGE SUBGRAPH (lane/pp-leverb, 2026-08-08): layers `[lo, hi)` of the chunked
1226    /// prime walk — `prime_chunk`'s trunk loop extracted verbatim to the
1227    /// `decode_layers_eager(lo, hi)` / `verify_layers(lo, hi)` contract: enters with a
1228    /// MATERIALIZED `[T, n_embd]` residual, exits with the range's final residual
1229    /// materialized (cloned out of the slab). At `lo=0, hi=n_layers` — the unsplit call —
1230    /// the launch sequence is byte-identical to the pre-extraction body. Range semantics:
1231    ///   - the cross-layer [down-add + NEXT attn-norm] fusion is range-LOCAL (`il + 1 < hi`):
1232    ///     layer `hi`'s attn_norm belongs to the NEXT stage's device, so the range ends with
1233    ///     the plain add (materialize) and the next stage hoists its own first norm — the
1234    ///     kernel-check-pinned `add_rms_norm == add then rms_norm` identity, the same law
1235    ///     the decode split rests on (`prime-split-gate` arbitrates end-to-end);
1236    ///   - prime slabs are PER-DEVICE (`prime_slabs_get` keys on the engine ordinal), so
1237    ///     each stage walks through its own resident transients;
1238    ///   - the S-glue/S-mid capture path requires the FULL range (its lookahead fuses
1239    ///     `self.layers[il+1]` unconditionally) — `use_seg` gains `lo == 0 && hi == n_layers`.
1240    #[allow(clippy::too_many_arguments)]
1241    fn prime_layers(
1242        &self,
1243        e: &Engine,
1244        x_in: CudaSlice<f32>,
1245        lo: usize,
1246        hi: usize,
1247        pos_d: &CudaSlice<i32>,
1248        t: usize,
1249        base: usize,
1250        cache: &mut Cache,
1251        seq_end: usize,
1252    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1253        let cfg = &self.cfg;
1254        let n_embd = cfg.n_embd as usize;
1255        let eps = cfg.rms_eps;
1256        // task #14: fuse the fp16 GEMM-operand emission into the trunk norms (kills the
1257        // standalone convert launches). Only when the f16 lane serves and T reaches the
1258        // GEMM tier; bit-identical either way.
1259        let f16fuse = crate::f16_ffi::pp_f16_enabled() && t >= 16;
1260        // PRIME SLABS (piecewise foundation): trunk transients in resident buffers —
1261        // ~224 fewer alloc/free calls per prime and FROZEN Lt operand addresses
1262        // (nvjet's alignment-variant selection becomes run-to-run stable). Every slab is
1263        // fully overwritten before use; x ping-pongs xa<->xb; the hidden-stack return
1264        // clones the final x (the slab cannot leave). Non-slab fallback: MEMRA_PRIME_SLABS=0.
1265        let n_ff_max = self
1266            .layers
1267            .iter()
1268            .map(|l| match &l.ffn {
1269                crate::hybrid::Ffn::Dense { ffn_gate, .. } => ffn_gate.out_features(),
1270                _ => n_embd,
1271            })
1272            .max()
1273            .unwrap_or(n_embd)
1274            .max(n_embd);
1275        let use_slabs = std::env::var("MEMRA_PRIME_SLABS").as_deref() != Ok("0");
1276        let slab = if use_slabs {
1277            Some(self.prime_slabs_get(e, t, n_embd, n_ff_max)?)
1278        } else {
1279            None
1280        };
1281        let mut slab_guard = slab.as_ref().map(|sl| sl.lock().unwrap());
1282        let mut x_own; // fallback storage when slabs are off
1283        type SlabRefs<'a> = (
1284            &'a mut CudaSlice<f32>,
1285            &'a mut CudaSlice<f32>,
1286            &'a mut CudaSlice<f32>,
1287            &'a mut CudaSlice<f32>,
1288            &'a mut CudaSlice<u8>,
1289            &'a mut CudaSlice<u8>,
1290            &'a mut CudaSlice<f32>,
1291            &'a mut CudaSlice<f32>,
1292            &'a mut CudaSlice<f32>,
1293        );
1294        let (mut x_cur, mut x_nxt, sl): (
1295            &mut CudaSlice<f32>,
1296            &mut CudaSlice<f32>,
1297            Option<SlabRefs>,
1298        );
1299        let mut seg: Option<(
1300            &mut Vec<Option<cudarc::driver::CudaGraph>>,
1301            &mut Vec<Option<cudarc::driver::CudaGraph>>,
1302            &mut CudaSlice<f32>,
1303            &mut usize,
1304        )> = None;
1305        let mut x_own2;
1306        match slab_guard.as_mut() {
1307            Some(g) => {
1308                let slabs = &mut **g;
1309                e.copy_into(&mut slabs.xa, 0, &x_in, t * n_embd)?;
1310                let PrimeSlabs {
1311                    xa,
1312                    xb,
1313                    h,
1314                    x1,
1315                    z,
1316                    act,
1317                    h16,
1318                    z16,
1319                    gate,
1320                    up,
1321                    ffn_out,
1322                    seg_glue,
1323                    mixed,
1324                    seg_mid,
1325                    seg_t,
1326                    ..
1327                } = slabs;
1328                x_cur = xa;
1329                x_nxt = xb;
1330                seg = Some((seg_glue, seg_mid, mixed, seg_t));
1331                sl = Some((h, x1, z, act, h16, z16, gate, up, ffn_out));
1332            }
1333            None => {
1334                x_own = x_in;
1335                x_own2 = e.uninit(t * n_embd)?;
1336                x_cur = &mut x_own;
1337                x_nxt = &mut x_own2;
1338                sl = None;
1339            }
1340        }
1341        let mut alloc_h;
1342        let mut alloc_x1;
1343        let mut alloc_z;
1344        let mut alloc_act;
1345        let mut alloc_h16;
1346        let mut alloc_z16;
1347        let mut alloc_gate;
1348        let mut alloc_up;
1349        let mut alloc_fo;
1350        let (h, x1, z, act): (
1351            &mut CudaSlice<f32>,
1352            &mut CudaSlice<f32>,
1353            &mut CudaSlice<f32>,
1354            &mut CudaSlice<f32>,
1355        );
1356        let (h16, z16): (&mut CudaSlice<u8>, &mut CudaSlice<u8>);
1357        let (sl_gate, sl_up, sl_fo): (
1358            &mut CudaSlice<f32>,
1359            &mut CudaSlice<f32>,
1360            &mut CudaSlice<f32>,
1361        );
1362        match sl {
1363            Some((a, b, c, d, e16, f16b, g, u, fo)) => {
1364                h = a;
1365                x1 = b;
1366                z = c;
1367                act = d;
1368                h16 = e16;
1369                z16 = f16b;
1370                sl_gate = g;
1371                sl_up = u;
1372                sl_fo = fo;
1373            }
1374            None => {
1375                alloc_h = e.uninit(t * n_embd)?;
1376                alloc_x1 = e.uninit(t * n_embd)?;
1377                alloc_z = e.uninit(t * n_embd)?;
1378                alloc_act = e.uninit(t * n_ff_max)?;
1379                alloc_h16 = e.alloc_u8_uninit(t * n_embd * 2)?;
1380                alloc_z16 = e.alloc_u8_uninit(t * n_embd * 2)?;
1381                alloc_gate = e.uninit(t * n_ff_max)?;
1382                alloc_up = e.uninit(t * n_ff_max)?;
1383                alloc_fo = e.uninit(t * n_embd)?;
1384                h = &mut alloc_h;
1385                x1 = &mut alloc_x1;
1386                z = &mut alloc_z;
1387                act = &mut alloc_act;
1388                h16 = &mut alloc_h16;
1389                z16 = &mut alloc_z16;
1390                sl_gate = &mut alloc_gate;
1391                sl_up = &mut alloc_up;
1392                sl_fo = &mut alloc_fo;
1393            }
1394        }
1395        // piecewise increment 3: layer-0 norm hoisted; the per-layer tail fuses
1396        // [down-add + NEXT layer's attn-norm] into one captured S-glue segment
1397        // (all-slab IO, zero in-graph allocations). Capture happens lazily on the
1398        // first prime at this t (capture does not execute -> launch right after).
1399        let n_layers = self.layers.len();
1400        // OPT-IN (2026-07-26 interleaved verdict): 2-kernel segments measured NET
1401        // -0.7%-to-neutral under the interleaved A/B protocol — one cuGraphLaunch costs
1402        // about what two kernel submissions do. The earlier "+0.9%" was cross-run clock
1403        // drift (the repo's interleaved-A/B law exists for exactly this). Larger segments
1404        // (S-prep/S-attn, 7-9 kernels) remain the open hypothesis; the core-split
1405        // machinery stays (byte-identical) as their foundation.
1406        // step35 is excluded: the core-split path calls `full_attn_prime_core_inner`, which is
1407        // the GENERIC attn core (uniform n_head, rope_dim_count, no window, no head-wise gate).
1408        // step35 rides its own mixer through the normal per-layer arm below.
1409        let use_seg = f16fuse
1410            && seg.is_some()
1411            && self.cfg.step35.is_none()
1412            && lo == 0
1413            && hi == n_layers
1414            && std::env::var("MEMRA_PRIME_SEG").as_deref() == Ok("1");
1415        if let Some((sg, sm, _, st)) = seg.as_mut() {
1416            if **st != t {
1417                sg.clear();
1418                sg.extend((0..n_layers).map(|_| None));
1419                sm.clear();
1420                sm.extend((0..n_layers).map(|_| None));
1421                **st = t;
1422            }
1423        }
1424        {
1425            let layer_lo = &self.layers[lo];
1426            if f16fuse {
1427                e.rms_norm_f16out(
1428                    x_cur,
1429                    layer_lo.attn_norm.float_data(),
1430                    h,
1431                    h16,
1432                    n_embd,
1433                    t,
1434                    eps,
1435                )?;
1436            } else {
1437                e.rms_norm(x_cur, layer_lo.attn_norm.float_data(), h, n_embd, t, eps)?;
1438            }
1439        }
1440        for il in lo..hi {
1441            let layer = &self.layers[il];
1442            let hx16 = if f16fuse { Some(&*h16) } else { None };
1443            if use_seg {
1444                // core-split path: projections -> _inner core -> out-GEMM INTO the mixed
1445                // slab (no copies) -> S-mid segment [add + post-norm] as one graph launch.
1446                let (pre, pre16, w_out) = match &layer.mixer {
1447                    Mixer::Full(fa) => {
1448                        let g3 = match hx16 {
1449                            Some(xh) => e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], h, xh, t)?,
1450                            None => e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?,
1451                        };
1452                        let (pre, pre16) =
1453                            self.full_attn_prime_core_inner(e, fa, g3, &pos_d, t, cache, il)?;
1454                        (pre, pre16, &fa.wo)
1455                    }
1456                    Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
1457                    Mixer::Linear(la) => {
1458                        let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
1459                        let g4 = match hx16 {
1460                            Some(xh) => e.matmul_group_xh(&ws, h, xh, t)?,
1461                            None => e.matmul_group(&ws, h, t)?,
1462                        };
1463                        let (pre, pre16) =
1464                            self.linear_attn_prime_core_pad_inner(e, la, g4, t, cache, il, None)?;
1465                        (pre, pre16, &la.ssm_out)
1466                    }
1467                };
1468                {
1469                    let (_, sm, mslab, _) = seg.as_mut().unwrap();
1470                    let pre_n = pre.len() / t;
1471                    let xh_pre = match pre16 {
1472                        Some(x) => x,
1473                        None => e.f16_act(&pre, t * pre_n, pre_n)?,
1474                    };
1475                    if !e.try_f16_gemm_pre_into(w_out, &xh_pre, t, mslab)? {
1476                        let y = e.matmul(w_out, &pre, t)?;
1477                        e.copy_into(mslab, 0, &y, t * n_embd)?;
1478                    }
1479                    if sm[il].is_none() {
1480                        use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
1481                        let w_post = layer.post_attn_norm.float_data();
1482                        e.stream().synchronize()?;
1483                        e.stream()
1484                            .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
1485                        let r = (|| -> Result<(), Box<dyn std::error::Error>> {
1486                            e.add(x_cur, mslab, x1, t * n_embd)?;
1487                            e.rms_norm_f16out(x1, w_post, z, z16, n_embd, t, eps)?;
1488                            Ok(())
1489                        })();
1490                        let g = e.stream().end_capture(
1491                            CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH);
1492                        r?;
1493                        sm[il] = Some(g?.ok_or("S-mid capture produced no graph")?);
1494                    }
1495                    sm[il].as_ref().unwrap().launch()?;
1496                }
1497            } else {
1498                let mixed = match &layer.mixer {
1499                    Mixer::Full(fa) => {
1500                        self.full_attn_prime(e, fa, h, hx16, &pos_d, t, cache, il, seq_end)?
1501                    }
1502                    Mixer::Linear(la) => self.linear_attn_prime(e, la, h, hx16, t, cache, il)?,
1503                    Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
1504                };
1505                if f16fuse {
1506                    // round 28: residual+norm in ONE kernel (add_rms_norm precedent,
1507                    // bit-identical) — the standalone add pass disappears.
1508                    e.add_rms_norm_f16out(
1509                        x_cur,
1510                        &mixed,
1511                        layer.post_attn_norm.float_data(),
1512                        x1,
1513                        z,
1514                        z16,
1515                        n_embd,
1516                        t,
1517                        eps,
1518                    )?;
1519                } else {
1520                    e.add(x_cur, &mixed, x1, t * n_embd)?;
1521                    e.rms_norm(x1, layer.post_attn_norm.float_data(), z, n_embd, t, eps)?;
1522                }
1523            }
1524            let zx16 = if f16fuse { Some(&*z16) } else { None };
1525            match &layer.ffn {
1526                crate::hybrid::Ffn::Dense {
1527                    ffn_gate,
1528                    ffn_up,
1529                    ffn_down,
1530                } => {
1531                    let n_ff = ffn_gate.out_features();
1532                    // gate/up INTO boundary slabs (piecewise increment 2); fall back to
1533                    // the allocating group + copy when a mirror is missing.
1534                    let mut into_ok = false;
1535                    if let Some(xh) = zx16 {
1536                        into_ok = e.try_f16_gemm_pre_into(ffn_gate, xh, t, sl_gate)?
1537                            && e.try_f16_gemm_pre_into(ffn_up, xh, t, sl_up)?;
1538                    }
1539                    if !into_ok {
1540                        let mut g2 = match zx16 {
1541                            Some(xh) => e.matmul_group_xh(&[ffn_gate, ffn_up], z, xh, t)?,
1542                            None => e.matmul_group(&[ffn_gate, ffn_up], z, t)?,
1543                        };
1544                        let up_y = g2.pop().unwrap();
1545                        let gate_y = g2.pop().unwrap();
1546                        e.copy_into(sl_gate, 0, &gate_y, t * n_ff)?;
1547                        e.copy_into(sl_up, 0, &up_y, t * n_ff)?;
1548                    }
1549                    // task #17: the silu arm's f16out twin emits the down GEMM's fp16
1550                    // operand in-epilogue; non-silu activations keep the standalone convert.
1551                    // silu_mul_f16out is PLAIN silu(gate)*up — a clamped layer (step35 dense
1552                    // blocks under a live swiglu_clamp_shexp) must take the ffn_act_lim arm.
1553                    let d_lim = self.cfg.clamp_shexp_at(il as u32);
1554                    let act16 = if Self::f16out_on(e, t) && self.cfg.m3.is_none() && d_lim.is_none()
1555                    {
1556                        let mut a16 = e.alloc_u8_uninit(t * n_ff * 2)?;
1557                        e.silu_mul_f16out(sl_gate, sl_up, act, &mut a16, t * n_ff)?;
1558                        Some(a16)
1559                    } else {
1560                        Self::ffn_act_lim(
1561                            e,
1562                            &self.cfg,
1563                            sl_gate,
1564                            sl_up,
1565                            1.0,
1566                            1.0,
1567                            d_lim,
1568                            act,
1569                            t * n_ff,
1570                        )?;
1571                        None
1572                    };
1573                    // down GEMM into the ffn_out slab (f16 arm; fallback copies)
1574                    let xh_act = match act16 {
1575                        Some(x) => x,
1576                        None => e.f16_act(act, t * n_ff, n_ff)?,
1577                    };
1578                    if !e.try_f16_gemm_pre_into(ffn_down, &xh_act, t, sl_fo)? {
1579                        let y = e.matmul(ffn_down, &*act, t)?;
1580                        e.copy_into(sl_fo, 0, &y, t * n_embd)?;
1581                    }
1582                }
1583                crate::hybrid::Ffn::Moe(m) => {
1584                    let y = self.moe_ffn_il_prefill(e, m, z, t, il as u16)?;
1585                    e.copy_into(sl_fo, 0, &y, t * n_embd)?;
1586                }
1587            }
1588            if use_seg && il + 1 < hi {
1589                // S-glue segment: [add + next attn-norm(+f16out)] — one cuGraphLaunch
1590                let w_next = self.layers[il + 1].attn_norm.float_data();
1591                let (sg, _, _, _) = seg.as_mut().unwrap();
1592                if sg[il].is_none() {
1593                    use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
1594                    e.stream().synchronize()?;
1595                    e.stream()
1596                        .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
1597                    let r = (|| -> Result<(), Box<dyn std::error::Error>> {
1598                        e.add(x1, sl_fo, x_nxt, t * n_embd)?;
1599                        e.rms_norm_f16out(x_nxt, w_next, h, h16, n_embd, t, eps)?;
1600                        Ok(())
1601                    })();
1602                    let g = e.stream().end_capture(
1603                        CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
1604                    );
1605                    r?;
1606                    sg[il] = Some(g?.ok_or("S-glue capture produced no graph")?);
1607                }
1608                sg[il].as_ref().unwrap().launch()?;
1609            } else {
1610                if il + 1 < hi {
1611                    let w_next = self.layers[il + 1].attn_norm.float_data();
1612                    if f16fuse {
1613                        e.add_rms_norm_f16out(x1, sl_fo, w_next, x_nxt, h, h16, n_embd, t, eps)?;
1614                    } else {
1615                        e.add(x1, sl_fo, x_nxt, t * n_embd)?;
1616                        e.rms_norm(x_nxt, w_next, h, n_embd, t, eps)?;
1617                    }
1618                } else {
1619                    e.add(x1, sl_fo, x_nxt, t * n_embd)?;
1620                }
1621            }
1622            // CHUNK-INVARIANCE BISECT SEAM (lane/chunk-invariance, 2026-08-05): dump the
1623            // per-layer post-residual hidden for the LAST row of this chunk, keyed by
1624            // absolute position, so two runs at different MEMRA_PRIME_CHUNK can be diffed
1625            // layer-by-layer to find the FIRST diverging layer. Diagnostic only —
1626            // unset (the default) costs one OnceLock read per layer.
1627            if let Some(path) = Self::prime_trace_path() {
1628                let row = (base + t - 1) as usize;
1629                let host = e.dtoh(x_nxt)?;
1630                let last = &host[(t - 1) * n_embd..t * n_embd];
1631                use std::io::Write as _;
1632                let mut f = std::fs::OpenOptions::new()
1633                    .create(true)
1634                    .append(true)
1635                    .open(path)?;
1636                let mut h64: u64 = 0xcbf29ce484222325;
1637                for v in last {
1638                    h64 ^= v.to_bits() as u64;
1639                    h64 = h64.wrapping_mul(0x100000001b3);
1640                }
1641                writeln!(
1642                    f,
1643                    "{{\"pos\":{row},\"layer\":{il},\"t\":{t},\"base\":{base},\
1644                             \"hash\":\"{h64:016x}\",\"v0\":{:.9e},\"v1\":{:.9e},\"v2\":{:.9e}}}",
1645                    last[0], last[1], last[2]
1646                )?;
1647            }
1648            // dflash/dspark tap (no-op when no sink armed): post-layer residual rows for
1649            // drafter conditioning — the qwen twin of the gemma4 tap sites.
1650            self.dflash_tap(e, cache, il, x_nxt, t)?;
1651            std::mem::swap(&mut x_cur, &mut x_nxt);
1652        }
1653        // hidden-stack return: clone the final x out of the slab
1654        let mut x = e.uninit(t * n_embd)?;
1655        e.copy_into(&mut x, 0, x_cur, t * n_embd)?;
1656        drop(slab_guard);
1657        Ok(x)
1658    }
1659
1660    /// The prime chunk's tail — h_seed + output_norm + one-row lm head + cache.pos advance —
1661    /// shared verbatim by the unsplit walk and the last stage of the ppN walk (`e` = the
1662    /// engine that produced `x`, i.e. the last stage's under the split; output_norm/output
1663    /// were loaded through that engine by the sharded loader, hybrid.rs `e_head`).
1664    fn prime_chunk_epilogue(
1665        &self,
1666        e: &Engine,
1667        x: CudaSlice<f32>,
1668        t: usize,
1669        cache: &mut Cache,
1670    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
1671        let n_embd = self.cfg.n_embd as usize;
1672        let eps = self.cfg.rms_eps;
1673        // h_seed = LAST row of x BEFORE output_norm (MTP-PLAN §A default) or AFTER it
1674        // (MEMRA_SPEC_HPOST — reference convention; hn is computed just below either way, so
1675        // the post-norm copy happens after hn exists).
1676        let mut h_seed = e.uninit(n_embd)?;
1677        if !crate::spec::spec_hpost() {
1678            e.copy_view_into(
1679                &mut h_seed,
1680                0,
1681                &x.slice((t - 1) * n_embd..t * n_embd),
1682                n_embd,
1683            )?;
1684        }
1685        // last-row logits, exactly like forward_last (norm all T — per-row op — then lm_head on 1 row).
1686        let mut hn = e.uninit(t * n_embd)?;
1687        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
1688        if crate::spec::spec_hpost() {
1689            e.copy_view_into(
1690                &mut h_seed,
1691                0,
1692                &hn.slice((t - 1) * n_embd..t * n_embd),
1693                n_embd,
1694            )?;
1695        }
1696        let last = e.view(&hn, t * n_embd);
1697        let last_row = last.slice((t - 1) * n_embd..t * n_embd);
1698        let mut hlast = e.uninit(n_embd)?;
1699        e.copy_view_into(&mut hlast, 0, &last_row, n_embd)?;
1700        let logits = e.matmul(&self.output, &hlast, 1)?;
1701        cache.pos += t;
1702        // Hidden stack handed to generate_spec as prompt_h: pre-norm x (default) or the full
1703        // post-norm stack hn (MEMRA_SPEC_HPOST).
1704        Ok((
1705            e.dtoh(&logits)?,
1706            h_seed,
1707            if crate::spec::spec_hpost() { hn } else { x },
1708        ))
1709    }
1710
1711    /// THE PRIME STAGE SPLIT (lane/pp-leverb, 2026-08-08): one prime chunk as N stage
1712    /// subgraphs, each range on ITS OWN engine/stream/device, with the `[T, n_embd]`
1713    /// boundary handoff at every fence cut — the prime twin of `decode_step_h_ppn` /
1714    /// `decode_step_t_core_ppn`, and the kill for the anatomy's two receipts: stage-1 trunk
1715    /// weights stop being peer-read (22% of the pp4096 wall) and dev1 stops running zero
1716    /// prefill kernels. Structure mirrors the verify split exactly:
1717    ///   stage 0        `rt.enter(0)` → per-stage pos_d + embed (the table lives with
1718    ///                  stage 0) → `prime_layers(fence[0], fence[1])` → `rt.tx`
1719    ///   middle stages  `rt.rx` → per-stage pos_d → range → `rt.tx`
1720    ///   last stage     `rt.rx` → range → the shared epilogue (output_norm + head live
1721    ///                  there via the sharded loader) → `publish_to`
1722    /// Laws inherited (not relearned): `fence_stages_behind` at entry (#87 — a previous
1723    /// round's stage-freed buffers must not be reused under the caller's queued reads);
1724    /// per-stage `pos_d` (allocated/consumed/freed on one stream); per-stage Engines from
1725    /// `PpNRt` (shared-scratch race); EXIT PUBLICATION for the device-resident returns
1726    /// (h_seed + hidden stack live on the last stage — the caller's stream must wait).
1727    /// KV/MoE locality falls out: each range's KV appends run on the owning stage
1728    /// (`pp::new_cache` placed the buffers there), each stage engine owns its own SLRU
1729    /// pool (per-Engine `moe_cache`, sized on ITS device — the SGLang #33666 law), and the
1730    /// slab-local MoE arm's `DevExps.dev` gate now matches on stage-1 layers too.
1731    /// EXACTNESS: the split adds zero deviation by construction (same kernels, same bytes,
1732    /// boundary = straight f32 copy); `prime-split-gate` (ppsplit) arbitrates bit-for-bit
1733    /// and its liveness counter is bumped here — the gate goes green with this function.
1734    fn prime_chunk_ppn(
1735        &self,
1736        e: &Engine,
1737        tokens: &[u32],
1738        cache: &mut Cache,
1739        seq_end: usize,
1740        fence: &[usize],
1741    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
1742        let rt = crate::pp::PpNRt::get(e)?;
1743        let n_st = fence.len() - 1;
1744        assert_eq!(
1745            rt.n_stages(),
1746            n_st,
1747            "PpNRt stage count {} != fence stages {n_st}",
1748            rt.n_stages()
1749        );
1750        let n_embd = self.cfg.n_embd as usize;
1751        let t = tokens.len();
1752        let base = cache.pos;
1753        debug_assert!(
1754            seq_end >= base + t,
1755            "prime_chunk_ppn: seq_end must cover this chunk"
1756        );
1757        let payload = t * n_embd;
1758        // The caller's ambient stream, captured BEFORE any enter() pushes a stage stream
1759        // (inside a stage scope e.stream() IS the stage stream and the exit wait would
1760        // self-order into a no-op) — the decode_step_t_core_ppn pattern.
1761        let caller_stream = e.stream();
1762        rt.fence_stages_behind(&caller_stream)?;
1763
1764        if n_st == 2 {
1765            let slot =
1766                self.prime_pp2_stage0_enqueue(e, rt, tokens, cache, seq_end, fence, base, false)?;
1767            let x =
1768                self.prime_pp2_stage1_enqueue(e, rt, slot, t, cache, seq_end, fence, base, false)?;
1769            let out = {
1770                rt.bind_stage(1)?;
1771                let _st1 = rt.enter(1);
1772                let e1 = rt.engine(1, e);
1773                self.prime_chunk_epilogue(e1, x, t, cache)?
1774            };
1775            rt.publish_to(1, &caller_stream)?;
1776            crate::pp::PRIME_SPLIT_CHUNKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1777            return Ok(out);
1778        }
1779
1780        let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
1781
1782        // ---- STAGE 0: embed + layers [0, fence[1]) + boundary-0 TX ----
1783        let mut slot = {
1784            let _st0 = rt.enter(0);
1785            let e0 = rt.engine(0, e);
1786            let pos_d = e0.htod_i32(&pos)?;
1787            let x = self.embed(e0, tokens)?;
1788            let x =
1789                self.prime_layers(e0, x, fence[0], fence[1], &pos_d, t, base, cache, seq_end)?;
1790            rt.tx(0, &x, payload)?
1791            // x + pos_d drop here: freed stream-ordered on stage-0's stream after use.
1792        };
1793
1794        // ---- MIDDLE STAGES: RX boundary s-1 -> range -> TX boundary s ----
1795        for s in 1..n_st - 1 {
1796            let _st = rt.enter(s);
1797            let es = rt.engine(s, e);
1798            let pos_d = es.htod_i32(&pos)?;
1799            let x = rt.rx(s - 1, slot, payload)?;
1800            let x = self.prime_layers(
1801                es,
1802                x,
1803                fence[s],
1804                fence[s + 1],
1805                &pos_d,
1806                t,
1807                base,
1808                cache,
1809                seq_end,
1810            )?;
1811            slot = rt.tx(s, &x, payload)?;
1812        }
1813
1814        // ---- LAST STAGE: RX + final range + the shared epilogue ----
1815        let _stl = rt.enter(n_st - 1);
1816        let el = rt.engine(n_st - 1, e);
1817        let pos_d = el.htod_i32(&pos)?;
1818        let x = rt.rx(n_st - 2, slot, payload)?;
1819        let x = self.prime_layers(
1820            el,
1821            x,
1822            fence[n_st - 1],
1823            fence[n_st],
1824            &pos_d,
1825            t,
1826            base,
1827            cache,
1828            seq_end,
1829        )?;
1830        let out = self.prime_chunk_epilogue(el, x, t, cache)?;
1831        // EXIT PUBLICATION: h_seed + the hidden stack are device-resident on the last
1832        // stage's stream; the caller resumes on its own stream (chunk-loop copy_into /
1833        // generate_spec's prompt_h consumer). The logits dtoh above already drained the
1834        // stage stream host-side, but the law is stated in events, not in a dtoh side
1835        // effect a later deferred form would remove.
1836        rt.publish_to(n_st - 1, &caller_stream)?;
1837        crate::pp::PRIME_SPLIT_CHUNKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1838        Ok(out)
1839    }
1840
1841    fn prime_pp2_stage0_enqueue(
1842        &self,
1843        e: &Engine,
1844        rt: &crate::pp::PpNRt,
1845        tokens: &[u32],
1846        cache: &mut Cache,
1847        seq_end: usize,
1848        fence: &[usize],
1849        base: usize,
1850        pipelined: bool,
1851    ) -> Result<usize, Box<dyn std::error::Error>> {
1852        let t = tokens.len();
1853        let n_embd = self.cfg.n_embd as usize;
1854        let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
1855        rt.bind_stage(0)?;
1856        let _st0 = rt.enter(0);
1857        let e0 = rt.engine(0, e);
1858        let pos_d = e0.htod_i32(&pos)?;
1859        let x = self.embed(e0, tokens)?;
1860        let _overlap = pipelined.then(crate::pp::enter_prime_pipe_stage);
1861        let x = self.prime_layers(e0, x, fence[0], fence[1], &pos_d, t, base, cache, seq_end)?;
1862        if pipelined {
1863            rt.tx_pipelined(0, &x, t * n_embd)
1864        } else {
1865            rt.tx(0, &x, t * n_embd)
1866        }
1867    }
1868
1869    fn prime_pp2_stage1_enqueue(
1870        &self,
1871        e: &Engine,
1872        rt: &crate::pp::PpNRt,
1873        slot: usize,
1874        t: usize,
1875        cache: &mut Cache,
1876        seq_end: usize,
1877        fence: &[usize],
1878        base: usize,
1879        pipelined: bool,
1880    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1881        let n_embd = self.cfg.n_embd as usize;
1882        let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
1883        rt.bind_stage(1)?;
1884        let _st1 = rt.enter(1);
1885        let e1 = rt.engine(1, e);
1886        let pos_d = e1.htod_i32(&pos)?;
1887        let x = rt.rx(0, slot, t * n_embd)?;
1888        let _overlap = pipelined.then(crate::pp::enter_prime_pipe_stage);
1889        self.prime_layers(e1, x, fence[1], fence[2], &pos_d, t, base, cache, seq_end)
1890    }
1891
1892    /// CAPTURE-SAFE prime trunk (task #14 increment 1, ARCHITECTURE-H100.md design v2):
1893    /// prime_chunk's layer stack with every capture hazard hoisted — `x` is the
1894    /// PRE-EMBEDDED input in a STABLE graph-input buffer (embed + its 8MB htod stay
1895    /// eager, one launch), `pos_d` is a baked device param (fresh prime = 0..T, constant
1896    /// per bucket), logits/h_seed/hidden stay DEVICE-resident (no dtoh), and cache.pos is
1897    /// NOT advanced (host state — the replay wrapper owns it). Body mirrors prime_chunk
1898    /// (the prime-graph-gate pins them together). KNOWN smoke-scope gap: append host-len
1899    /// bookkeeping still runs on the host per call — the real replay path moves the write
1900    /// slot to the len_d device counter (increment 3).
1901    /// GRAPH-OUTPUT CONTRACT: results are COPIED into caller-provided stable buffers
1902    /// (`logits_out` [n_vocab], `h_seed_out` [n_embd]) — every internal allocation drops
1903    /// INSIDE the capture region (alloc+free node pairs). Retaining an in-capture
1904    /// allocation across end_capture makes instantiate throw INVALID_VALUE (smoke finding
1905    /// 2026-07-26), and under AUTO_FREE_ON_LAUNCH its address wouldn't survive a launch
1906    /// anyway — the decode GraphSession's pre-allocated-output pattern is the law here.
1907    pub fn prime_chunk_captured(
1908        &self,
1909        e: &Engine,
1910        x_in: &CudaSlice<f32>,
1911        pos_d: &CudaSlice<i32>,
1912        t: usize,
1913        cache: &mut Cache,
1914        len_d: &CudaSlice<i32>,
1915        logits_out: &mut CudaSlice<f32>,
1916        h_seed_out: &mut CudaSlice<f32>,
1917    ) -> Result<(), Box<dyn std::error::Error>> {
1918        let cfg = &self.cfg;
1919        let n_embd = cfg.n_embd as usize;
1920        let eps = cfg.rms_eps;
1921        let f16fuse = crate::f16_ffi::pp_f16_enabled() && t >= 16;
1922        let mut x = e.uninit(t * n_embd)?;
1923        e.copy_into(&mut x, 0, x_in, t * n_embd)?;
1924        for (il, layer) in self.layers.iter().enumerate() {
1925            let mut h = e.uninit(t * n_embd)?;
1926            let mut hx16: Option<CudaSlice<u8>> = None;
1927            if f16fuse {
1928                let mut b16 = e.alloc_u8_uninit(t * n_embd * 2)?;
1929                e.rms_norm_f16out(
1930                    &x,
1931                    layer.attn_norm.float_data(),
1932                    &mut h,
1933                    &mut b16,
1934                    n_embd,
1935                    t,
1936                    eps,
1937                )?;
1938                hx16 = Some(b16);
1939            } else {
1940                e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
1941            }
1942            let mixed = match &layer.mixer {
1943                // The captured prime is ONE unchunked bucket over a FRESH cache (pos == 0), so the
1944                // request ends at t. If bucketed capture ever composes with chunking, seq_end must
1945                // come from the caller (see step35_attn_pre_wo's doc note).
1946                Mixer::Full(fa) => {
1947                    self.full_attn_prime(e, fa, &h, hx16.as_ref(), pos_d, t, cache, il, t)?
1948                }
1949                Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
1950                Mixer::Linear(la) => {
1951                    let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
1952                    let g4 = match hx16.as_ref() {
1953                        Some(xh) => e.matmul_group_xh(&ws, &h, xh, t)?,
1954                        None => e.matmul_group(&ws, &h, t)?,
1955                    };
1956                    self.linear_attn_prime_core_pad(e, la, g4, t, cache, il, Some(len_d))?
1957                }
1958            };
1959            let mut x1 = e.uninit(t * n_embd)?;
1960            e.add(&x, &mixed, &mut x1, t * n_embd)?;
1961            let mut z = e.uninit(t * n_embd)?;
1962            let mut zx16: Option<CudaSlice<u8>> = None;
1963            if f16fuse {
1964                let mut b16 = e.alloc_u8_uninit(t * n_embd * 2)?;
1965                e.rms_norm_f16out(
1966                    &x1,
1967                    layer.post_attn_norm.float_data(),
1968                    &mut z,
1969                    &mut b16,
1970                    n_embd,
1971                    t,
1972                    eps,
1973                )?;
1974                zx16 = Some(b16);
1975            } else {
1976                e.rms_norm(
1977                    &x1,
1978                    layer.post_attn_norm.float_data(),
1979                    &mut z,
1980                    n_embd,
1981                    t,
1982                    eps,
1983                )?;
1984            }
1985            let ffn_out = match &layer.ffn {
1986                crate::hybrid::Ffn::Dense {
1987                    ffn_gate,
1988                    ffn_up,
1989                    ffn_down,
1990                } => {
1991                    let n_ff = ffn_gate.out_features();
1992                    let mut g2 = match &zx16 {
1993                        Some(xh) => e.matmul_group_xh(&[ffn_gate, ffn_up], &z, xh, t)?,
1994                        None => e.matmul_group(&[ffn_gate, ffn_up], &z, t)?,
1995                    };
1996                    let up = g2.pop().unwrap();
1997                    let gate = g2.pop().unwrap();
1998                    let mut act = e.uninit(t * n_ff)?;
1999                    // dense FFN keys off the SHEXP clamp array — see forward()'s note.
2000                    Self::ffn_act_lim(
2001                        e,
2002                        &self.cfg,
2003                        &gate,
2004                        &up,
2005                        1.0,
2006                        1.0,
2007                        self.cfg.clamp_shexp_at(il as u32),
2008                        &mut act,
2009                        t * n_ff,
2010                    )?;
2011                    e.matmul(ffn_down, &act, t)?
2012                }
2013                crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_prefill(e, m, &z, t, il as u16)?,
2014            };
2015            let mut x2 = e.uninit(t * n_embd)?;
2016            e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
2017            x = x2;
2018        }
2019        // device-indexed TRUE last row (pads sit past it in a bucketed graph)
2020        if !crate::spec::spec_hpost() {
2021            e.row_gather_dev(&x, h_seed_out, len_d, n_embd)?;
2022        }
2023        let mut hn = e.uninit(t * n_embd)?;
2024        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
2025        if crate::spec::spec_hpost() {
2026            e.row_gather_dev(&hn, h_seed_out, len_d, n_embd)?;
2027        }
2028        let mut hlast = e.uninit(n_embd)?;
2029        e.row_gather_dev(&hn, &mut hlast, len_d, n_embd)?;
2030        let logits = e.matmul(&self.output, &hlast, 1)?;
2031        let nv = logits.len();
2032        e.copy_into(logits_out, 0, &logits, nv)?;
2033        Ok(())
2034    }
2035
2036    fn step35_prime_batch_on() -> bool {
2037        std::env::var("MEMRA_STEP35_PRIME_BATCH").as_deref() != Ok("0")
2038    }
2039
2040    /// Step35 cross-request prime range: weight-streaming work runs once at `m=sum(T)`;
2041    /// sequence-scoped attention/KV work stays on each request's own cache and positions.
2042    #[allow(clippy::too_many_arguments)]
2043    fn step35_prime_batch_layers(
2044        &self,
2045        e: &Engine,
2046        mut x: CudaSlice<f32>,
2047        lo: usize,
2048        hi: usize,
2049        ts: &[usize],
2050        offs: &[usize],
2051        pos_ds: &[CudaSlice<i32>],
2052        caches: &mut [&mut Cache],
2053    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2054        let cfg = &self.cfg;
2055        let n_embd = cfg.n_embd as usize;
2056        let eps = cfg.rms_eps;
2057        let b = ts.len();
2058        let total: usize = ts.iter().sum();
2059        let f16fuse = crate::f16_ffi::pp_f16_enabled() && total >= 16;
2060
2061        let split = |e: &Engine,
2062                     y: &CudaSlice<f32>,
2063                     dim: usize|
2064         -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
2065            let mut out = Vec::with_capacity(b);
2066            for s in 0..b {
2067                let mut ys = e.uninit(ts[s] * dim)?;
2068                e.copy_view_into(
2069                    &mut ys,
2070                    0,
2071                    &y.slice(offs[s] * dim..(offs[s] + ts[s]) * dim),
2072                    ts[s] * dim,
2073                )?;
2074                out.push(ys);
2075            }
2076            Ok(out)
2077        };
2078
2079        for il in lo..hi {
2080            let layer = &self.layers[il];
2081            let Mixer::Full(fa) = &layer.mixer else {
2082                return Err(format!("step35 layer {il} is not full-attn — corrupt config").into());
2083            };
2084
2085            let mut h = e.uninit(total * n_embd)?;
2086            let mut hx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
2087            if f16fuse {
2088                e.rms_norm_f16out(
2089                    &x,
2090                    layer.attn_norm.float_data(),
2091                    &mut h,
2092                    &mut hx16,
2093                    n_embd,
2094                    total,
2095                    eps,
2096                )?;
2097            } else {
2098                e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, total, eps)?;
2099            }
2100
2101            // Q/K/V + gate projections and wo are batched weight streams. The step35 core
2102            // remains per sequence, so its partial RoPE, SWA view, KV append, and gate
2103            // application stay verbatim.
2104            let gate_w = fa
2105                .attn_gate
2106                .as_ref()
2107                .ok_or("step35 layer is missing attn_gate.weight (head-wise attention gate)")?;
2108            let mut g4 = if f16fuse {
2109                e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv, gate_w], &h, &hx16, total)?
2110            } else {
2111                e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv, gate_w], &h, total)?
2112            };
2113            let gate = g4.pop().unwrap();
2114            let mut parts: Vec<Vec<CudaSlice<f32>>> =
2115                (0..b).map(|_| Vec::with_capacity(3)).collect();
2116            for (w, y) in [&fa.wq, &fa.wk, &fa.wv].iter().zip(g4) {
2117                for (s, ys) in split(e, &y, w.out_features())?.into_iter().enumerate() {
2118                    parts[s].push(ys);
2119                }
2120            }
2121            let gates = split(e, &gate, gate_w.out_features())?;
2122            let geometry = self.step35_geom(il);
2123            let hd = geometry.head_dim_k as usize;
2124            let nh = geometry.n_head as usize;
2125            let mut ag_cat = e.uninit(total * nh * hd)?;
2126            for (s, (g3s, gate)) in parts.into_iter().zip(gates).enumerate() {
2127                let ag = self.step35_attn_pre_wo(
2128                    e,
2129                    fa,
2130                    g3s,
2131                    None,
2132                    Some(&gate),
2133                    &pos_ds[s],
2134                    ts[s],
2135                    Some(&mut *caches[s]),
2136                    il,
2137                    ts[s],
2138                )?;
2139                e.copy_into(&mut ag_cat, offs[s] * nh * hd, &ag, ts[s] * nh * hd)?;
2140            }
2141            let mixed = e.matmul(&fa.wo, &ag_cat, total)?;
2142
2143            let mut x1 = e.uninit(total * n_embd)?;
2144            let mut z = e.uninit(total * n_embd)?;
2145            let mut zx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
2146            if f16fuse {
2147                e.add_rms_norm_f16out(
2148                    &x,
2149                    &mixed,
2150                    layer.post_attn_norm.float_data(),
2151                    &mut x1,
2152                    &mut z,
2153                    &mut zx16,
2154                    n_embd,
2155                    total,
2156                    eps,
2157                )?;
2158            } else {
2159                e.add(&x, &mixed, &mut x1, total * n_embd)?;
2160                e.rms_norm(
2161                    &x1,
2162                    layer.post_attn_norm.float_data(),
2163                    &mut z,
2164                    n_embd,
2165                    total,
2166                    eps,
2167                )?;
2168            }
2169
2170            let ffn_out = match &layer.ffn {
2171                crate::hybrid::Ffn::Dense {
2172                    ffn_gate,
2173                    ffn_up,
2174                    ffn_down,
2175                } => {
2176                    let n_ff = ffn_gate.out_features();
2177                    let mut g2 = if f16fuse {
2178                        e.matmul_group_xh(&[ffn_gate, ffn_up], &z, &zx16, total)?
2179                    } else {
2180                        e.matmul_group(&[ffn_gate, ffn_up], &z, total)?
2181                    };
2182                    let up = g2.pop().unwrap();
2183                    let gate = g2.pop().unwrap();
2184                    let mut act = e.uninit(total * n_ff)?;
2185                    let d_lim = cfg.clamp_shexp_at(il as u32);
2186                    if Self::f16out_on(e, total) && cfg.m3.is_none() && d_lim.is_none() {
2187                        let mut a16 = e.alloc_u8_uninit(total * n_ff * 2)?;
2188                        e.silu_mul_f16out(&gate, &up, &mut act, &mut a16, total * n_ff)?;
2189                        match e.try_f16_gemm_pre(ffn_down, &a16, total)? {
2190                            Some(y) => y,
2191                            None => e.matmul(ffn_down, &act, total)?,
2192                        }
2193                    } else {
2194                        Self::ffn_act_lim(
2195                            e,
2196                            cfg,
2197                            &gate,
2198                            &up,
2199                            1.0,
2200                            1.0,
2201                            d_lim,
2202                            &mut act,
2203                            total * n_ff,
2204                        )?;
2205                        e.matmul(ffn_down, &act, total)?
2206                    }
2207                }
2208                crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, total, il as u16)?,
2209            };
2210            let mut x2 = e.uninit(total * n_embd)?;
2211            e.add(&x1, &ffn_out, &mut x2, total * n_embd)?;
2212            x = x2;
2213        }
2214        Ok(x)
2215    }
2216
2217    fn step35_prime_batch_epilogue(
2218        &self,
2219        e: &Engine,
2220        x: CudaSlice<f32>,
2221        ts: &[usize],
2222        offs: &[usize],
2223        caches: &mut [&mut Cache],
2224    ) -> Result<Vec<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
2225        let n_embd = self.cfg.n_embd as usize;
2226        let total: usize = ts.iter().sum();
2227        let mut hn = e.uninit(total * n_embd)?;
2228        e.rms_norm(
2229            &x,
2230            self.output_norm.float_data(),
2231            &mut hn,
2232            n_embd,
2233            total,
2234            self.cfg.rms_eps,
2235        )?;
2236
2237        let hidden_src = if crate::spec::spec_hpost() { &hn } else { &x };
2238        let mut out = Vec::with_capacity(ts.len());
2239        for s in 0..ts.len() {
2240            let mut hidden = e.uninit(ts[s] * n_embd)?;
2241            e.copy_view_into(
2242                &mut hidden,
2243                0,
2244                &hidden_src.slice(offs[s] * n_embd..(offs[s] + ts[s]) * n_embd),
2245                ts[s] * n_embd,
2246            )?;
2247            let last0 = (offs[s] + ts[s] - 1) * n_embd;
2248            let mut h_seed = e.uninit(n_embd)?;
2249            e.copy_view_into(
2250                &mut h_seed,
2251                0,
2252                &hidden_src.slice(last0..last0 + n_embd),
2253                n_embd,
2254            )?;
2255            // Exactness-first: the serial reference runs the output head at m=1.
2256            let mut hlast = e.uninit(n_embd)?;
2257            e.copy_view_into(&mut hlast, 0, &hn.slice(last0..last0 + n_embd), n_embd)?;
2258            let logits = e.dtoh(&e.matmul(&self.output, &hlast, 1)?)?;
2259            caches[s].pos += ts[s];
2260            out.push((logits, h_seed, hidden));
2261        }
2262        Ok(out)
2263    }
2264
2265    fn step35_prime_cache_batch(
2266        &self,
2267        e: &Engine,
2268        prompts: &[&[u32]],
2269        caches: &mut [&mut Cache],
2270    ) -> Result<Vec<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
2271        if crate::pp::pp_host_bounce_active()
2272            && (!crate::pp::prime_pp_on() || crate::pp::pp_cuts(self.layers.len()).is_none())
2273        {
2274            return Err(
2275                "step35_prime_cache_batch: MEMRA_PP_HOST_BOUNCE=1 requires a valid prime \
2276                 stage split; refusing an unsplit remote-weight walk"
2277                    .into(),
2278            );
2279        }
2280        if !Self::step35_prime_batch_on() {
2281            return Err("step35 batched prime is disabled (MEMRA_STEP35_PRIME_BATCH=0)".into());
2282        }
2283        if caches.iter().any(|c| c.pos != 0) {
2284            return Err(
2285                "step35 batched prime currently supports complete fresh prompts only; \
2286                 continuation/tick chunks require per-request queued_after"
2287                    .into(),
2288            );
2289        }
2290
2291        let ts: Vec<usize> = prompts.iter().map(|p| p.len()).collect();
2292        for &t in &ts {
2293            assert!(
2294                t >= PRIME_MIN_T,
2295                "step35 batched prime needs T >= {PRIME_MIN_T}"
2296            );
2297        }
2298        for (s, c) in caches.iter().enumerate() {
2299            assert!(
2300                ts[s] <= c.max_ctx,
2301                "step35 batched prime exceeds cache max_ctx"
2302            );
2303        }
2304        let offs: Vec<usize> = ts
2305            .iter()
2306            .scan(0usize, |a, &t| {
2307                let o = *a;
2308                *a += t;
2309                Some(o)
2310            })
2311            .collect();
2312        let total: usize = ts.iter().sum();
2313        let payload = total * self.cfg.n_embd as usize;
2314        let cat_tokens: Vec<u32> = prompts.iter().flat_map(|p| p.iter().copied()).collect();
2315        let positions: Vec<Vec<i32>> = ts.iter().map(|&t| (0..t as i32).collect()).collect();
2316        let upload_positions =
2317            |e: &Engine| -> Result<Vec<CudaSlice<i32>>, Box<dyn std::error::Error>> {
2318                positions
2319                    .iter()
2320                    .map(|p| e.htod_i32(p))
2321                    .collect::<Result<_, _>>()
2322            };
2323
2324        static ONCE: std::sync::Once = std::sync::Once::new();
2325        ONCE.call_once(|| {
2326            eprintln!(
2327                "[step35-prime-batch] first concat prime: B={} tokens={total}",
2328                prompts.len()
2329            );
2330        });
2331
2332        let out = if !crate::pp::pp2_streams_off() && crate::pp::prime_pp_on() {
2333            if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
2334                let rt = crate::pp::PpNRt::get(e)?;
2335                let n_st = fence.len() - 1;
2336                assert_eq!(
2337                    rt.n_stages(),
2338                    n_st,
2339                    "step35 prime batch stage count mismatch"
2340                );
2341                let caller_stream = e.stream();
2342                rt.fence_stages_behind(&caller_stream)?;
2343
2344                let mut slot = {
2345                    let _st0 = rt.enter(0);
2346                    let e0 = rt.engine(0, e);
2347                    let pos_ds = upload_positions(e0)?;
2348                    let x = self.embed(e0, &cat_tokens)?;
2349                    let x = self.step35_prime_batch_layers(
2350                        e0, x, fence[0], fence[1], &ts, &offs, &pos_ds, caches,
2351                    )?;
2352                    rt.tx(0, &x, payload)?
2353                };
2354                for s in 1..n_st - 1 {
2355                    let _st = rt.enter(s);
2356                    let es = rt.engine(s, e);
2357                    let pos_ds = upload_positions(es)?;
2358                    let x = rt.rx(s - 1, slot, payload)?;
2359                    let x = self.step35_prime_batch_layers(
2360                        es,
2361                        x,
2362                        fence[s],
2363                        fence[s + 1],
2364                        &ts,
2365                        &offs,
2366                        &pos_ds,
2367                        caches,
2368                    )?;
2369                    slot = rt.tx(s, &x, payload)?;
2370                }
2371
2372                let _stl = rt.enter(n_st - 1);
2373                let el = rt.engine(n_st - 1, e);
2374                let pos_ds = upload_positions(el)?;
2375                let x = rt.rx(n_st - 2, slot, payload)?;
2376                let x = self.step35_prime_batch_layers(
2377                    el,
2378                    x,
2379                    fence[n_st - 1],
2380                    fence[n_st],
2381                    &ts,
2382                    &offs,
2383                    &pos_ds,
2384                    caches,
2385                )?;
2386                let out = self.step35_prime_batch_epilogue(el, x, &ts, &offs, caches)?;
2387                rt.publish_to(n_st - 1, &caller_stream)?;
2388                crate::pp::STEP35_PRIME_BATCH_SPLITS
2389                    .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2390                out
2391            } else {
2392                let pos_ds = upload_positions(e)?;
2393                let x = self.embed(e, &cat_tokens)?;
2394                let x = self.step35_prime_batch_layers(
2395                    e,
2396                    x,
2397                    0,
2398                    self.layers.len(),
2399                    &ts,
2400                    &offs,
2401                    &pos_ds,
2402                    caches,
2403                )?;
2404                self.step35_prime_batch_epilogue(e, x, &ts, &offs, caches)?
2405            }
2406        } else {
2407            let pos_ds = upload_positions(e)?;
2408            let x = self.embed(e, &cat_tokens)?;
2409            let x = self.step35_prime_batch_layers(
2410                e,
2411                x,
2412                0,
2413                self.layers.len(),
2414                &ts,
2415                &offs,
2416                &pos_ds,
2417                caches,
2418            )?;
2419            self.step35_prime_batch_epilogue(e, x, &ts, &offs, caches)?
2420        };
2421        crate::pp::STEP35_PRIME_BATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2422        Ok(out)
2423    }
2424
2425    /// Cross-request BATCHED fresh prime (task #13, design in ARCHITECTURE-H100.md): the
2426    /// trunk's token-parallel ops (embed, norms, adds, ffn, projection GROUPS) run once on
2427    /// the CONCATENATION of B sequences — GEMMs at m = sum_T, the continuous-batching win
2428    /// the serving-lane bench attributed the remaining vLLM gap to. The stateful mixer
2429    /// CORES (QK-norm/RoPE/FA/append, conv/GDN scans) run per sequence on split projection
2430    /// buffers (D2D row copies; the mixers' own out-projections stay per-seq this
2431    /// increment). CONTINUATION primes (increment (b), 2026-07-30): cache.pos > 0 seqs
2432    /// batch the projections/FFN/lm_head exactly like fresh; the mixer cores take the
2433    /// per-seq CONTINUATION arms (Full: core_inner with carried pos_d + fa_prefill_view
2434    /// over the quantized past; Linear: the stateful pad_view twin — the same state
2435    /// carry the chunked single-seq prime rides). The fresh-only favl/gdn-vl fast paths
2436    /// stay byte-identical (gated on !carried). gemma4 models have no continuation
2437    /// prime (v0 monolithic fresh) — carried gemma4 batches return Err (caller falls
2438    /// back to single-chunk serving).
2439    /// NUMERIC CONFIG: a concat GEMM tiles K differently than per-seq GEMMs — same class
2440    /// as every prefill GEMM change; prime_batch_gate arbitrates (argmax + stream battery).
2441    pub fn prime_cache_batch(
2442        &self,
2443        e: &Engine,
2444        prompts: &[&[u32]],
2445        caches: &mut [&mut Cache],
2446    ) -> Result<Vec<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
2447        let cfg = &self.cfg;
2448        let n_embd = cfg.n_embd as usize;
2449        let eps = cfg.rms_eps;
2450        let b = prompts.len();
2451        assert!(b >= 1 && b == caches.len());
2452        let pos0s: Vec<usize> = caches.iter().map(|c| c.pos).collect();
2453        let carried = pos0s.iter().any(|&p| p > 0);
2454        // gemma4: refuse UNCONDITIONALLY (2026-08-07, lane/gemma4-serve-gaps). The old guard
2455        // covered only `carried` — two concurrent FRESH gemma4 prompts batched into the
2456        // generic concat attn core below (uniform geometry, no per-layer swa window, no
2457        // softcapped head): compiles, runs, wrong logits. Same silent-wrong class as the
2458        // step35 refusal beneath. The per-sequence `gemma4_prime` is the supported prefill.
2459        if cfg.gemma4.is_some() {
2460            return Err(
2461                "prime_cache_batch: gemma4 has no batched prime core (per-layer \
2462                        swa/global geometry, softcapped head) — use gemma4_prime per sequence"
2463                    .into(),
2464            );
2465        }
2466        // Step35 has a dedicated concat walk: the generic core below cannot express its
2467        // per-layer geometry/SWA/head gate, and under PP the dedicated path is stage-scoped.
2468        if cfg.step35.is_some() {
2469            return self.step35_prime_cache_batch(e, prompts, caches);
2470        }
2471        let ts: Vec<usize> = prompts.iter().map(|p| p.len()).collect();
2472        for &t in &ts {
2473            assert!(
2474                t >= PRIME_MIN_T,
2475                "prime_cache_batch needs T >= {PRIME_MIN_T}"
2476            );
2477        }
2478        for (s, c) in caches.iter().enumerate() {
2479            assert!(
2480                c.pos + ts[s] <= c.max_ctx,
2481                "prime_cache_batch: prompt exceeds cache max_ctx"
2482            );
2483        }
2484        let total: usize = ts.iter().sum();
2485        let offs: Vec<usize> = ts
2486            .iter()
2487            .scan(0usize, |a, &t| {
2488                let o = *a;
2489                *a += t;
2490                Some(o)
2491            })
2492            .collect();
2493        // per-seq positions (fresh: 0..T_s; continuation: pos0..pos0+T_s)
2494        let pos_ds: Vec<CudaSlice<i32>> = ts
2495            .iter()
2496            .zip(&pos0s)
2497            .map(|(&t, &p0)| e.htod_i32(&(p0 as i32..(p0 + t) as i32).collect::<Vec<_>>()))
2498            .collect::<Result<_, _>>()?;
2499        // split a concat [total, dim] buffer into per-seq copies
2500        let split = |e: &Engine,
2501                     y: &CudaSlice<f32>,
2502                     dim: usize|
2503         -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
2504            let mut out = Vec::with_capacity(b);
2505            for s in 0..b {
2506                let mut ys = e.uninit(ts[s] * dim)?;
2507                e.copy_view_into(
2508                    &mut ys,
2509                    0,
2510                    &y.slice(offs[s] * dim..(offs[s] + ts[s]) * dim),
2511                    ts[s] * dim,
2512                )?;
2513                out.push(ys);
2514            }
2515            Ok(out)
2516        };
2517
2518        let cat_tokens: Vec<u32> = prompts.iter().flat_map(|p| p.iter().copied()).collect();
2519        let mut x = self.embed(e, &cat_tokens)?; // [total, n_embd]
2520        for (il, layer) in self.layers.iter().enumerate() {
2521            let mut h = e.uninit(total * n_embd)?;
2522            let mut hx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
2523            e.rms_norm_f16out(
2524                &x,
2525                layer.attn_norm.float_data(),
2526                &mut h,
2527                &mut hx16,
2528                n_embd,
2529                total,
2530                eps,
2531            )?;
2532            // mixer: projection GROUP on the concat (m = total), stateful core per seq
2533            let mut mixed = e.uninit(total * n_embd)?;
2534            match &layer.mixer {
2535                Mixer::Full(fa) => {
2536                    let g3 = e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], &h, &hx16, total)?;
2537                    // task #18 (attn side): the WHOLE attn core is varlen for fresh gated
2538                    // batches — split/QK-norm/RoPE/append (attn_pre_vl8, view inputs: the
2539                    // q/k/v split copies vanish) + ONE varlen FA. Per-block math identical
2540                    // everywhere (bit-gateable); MEMRA_FA_VL=0 or a non-bf16kv config falls
2541                    // back to the per-seq dispatch.
2542                    let geometry = self.cfg.full_attention_geometry_at(il as u32);
2543                    let (n_head, n_head_kv, head_dim) = (
2544                        geometry.n_head as usize,
2545                        geometry.n_head_kv as usize,
2546                        geometry.head_dim_k as usize,
2547                    );
2548                    let fa_scale = geometry.attention_scale();
2549                    let use_favl = !carried
2550                        && (2..=8).contains(&b)
2551                        && (head_dim == 256 || head_dim == 128)
2552                        && geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ
2553                        && std::env::var("MEMRA_NOFA").is_err()
2554                        && std::env::var("MEMRA_FA_FLOOR").is_err()
2555                        && std::env::var("MEMRA_FA_PP_W2").as_deref() != Ok("1")
2556                        && std::env::var("MEMRA_FA_BF16KV").as_deref() != Ok("0")
2557                        && std::env::var("MEMRA_FA_VL").as_deref() != Ok("0");
2558                    if use_favl {
2559                        let (qf_w, kf_w, vf_w) = (
2560                            fa.wq.out_features(),
2561                            fa.wk.out_features(),
2562                            fa.wv.out_features(),
2563                        );
2564                        // Same bounds contract as `Engine::q_gate_split`, applied to the varlen
2565                        // twin's PER-TOKEN stride. `attn_pre_vl8` takes raw device pointers so it
2566                        // cannot check its own extents; `qf_w` is the wq out-features that set
2567                        // them, and `q_gate_split_vl` reads 2*head_dim per head out of it.
2568                        memra_gguf::config::check_fused_q_gate_extent(qf_w, head_dim, n_head, 1)?;
2569                        struct APre {
2570                            q: CudaSlice<f32>,
2571                            gate: Option<CudaSlice<f32>>,
2572                            qn: CudaSlice<f32>,
2573                            kn: CudaSlice<f32>,
2574                        }
2575                        let mut aps = Vec::with_capacity(b);
2576                        for &t in ts.iter().take(b) {
2577                            aps.push(APre {
2578                                q: e.uninit(t * n_head * head_dim)?,
2579                                gate: Some(e.uninit(t * n_head * head_dim)?),
2580                                qn: e.uninit(t * n_head * head_dim)?,
2581                                kn: e.uninit(t * n_head_kv * head_dim)?,
2582                            });
2583                        }
2584                        let (kv_dim_k, kv_dim_v, ktb, vtb) = {
2585                            let kvl = caches[0].kv[il].as_ref().unwrap();
2586                            (kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes)
2587                        };
2588                        let pargs: Vec<crate::AttnPreVl> = (0..b)
2589                            .map(|s| {
2590                                let (o, t) = (offs[s], ts[s]);
2591                                let kvl = caches[s].kv[il].as_ref().unwrap();
2592                                assert!(
2593                                    kvl.len == 0 && kvl.len + t <= caches[s].max_ctx,
2594                                    "prime_cache_batch attn vl: fresh + capacity"
2595                                );
2596                                crate::AttnPreVl {
2597                                    qf: e.addr_f32v(&g3[0].slice(o * qf_w..(o + t) * qf_w)),
2598                                    kf: e.addr_f32v(&g3[1].slice(o * kf_w..(o + t) * kf_w)),
2599                                    vf: e.addr_f32v(&g3[2].slice(o * vf_w..(o + t) * vf_w)),
2600                                    q: e.addr_f32(&aps[s].q),
2601                                    gate: e.addr_f32(aps[s].gate.as_ref().unwrap()),
2602                                    qn: e.addr_f32(&aps[s].qn),
2603                                    kn: e.addr_f32(&aps[s].kn),
2604                                    kc: e.addr_u8(&kvl.k),
2605                                    vc: e.addr_u8(&kvl.v),
2606                                    t: t as i32,
2607                                    pad: 0,
2608                                }
2609                            })
2610                            .collect();
2611                        e.attn_pre_vl8(
2612                            &pargs,
2613                            fa.q_norm.float_data(),
2614                            fa.k_norm.float_data(),
2615                            head_dim,
2616                            geometry.n_rot as usize,
2617                            n_head,
2618                            n_head_kv,
2619                            self.cfg.rms_eps,
2620                            geometry.rope_base,
2621                            1.0,
2622                            kv_dim_k,
2623                            kv_dim_v,
2624                            ktb,
2625                            vtb,
2626                        )?;
2627                        for s in 0..b {
2628                            let kvl = caches[s].kv[il].as_mut().unwrap();
2629                            kvl.len += ts[s];
2630                            let new_len = kvl.len as i32;
2631                            e.set_i32_one(&mut kvl.len_d, new_len)?;
2632                        }
2633                        let mut attns = Vec::with_capacity(b);
2634                        let mut mirrors = Vec::with_capacity(b);
2635                        for &t in ts.iter().take(b) {
2636                            attns.push(e.uninit(t * n_head * head_dim)?);
2637                            let n = t * n_head_kv * head_dim;
2638                            mirrors.push((e.alloc_u8_uninit(n * 2)?, e.alloc_u8_uninit(n * 2)?));
2639                        }
2640                        // FA3 batched twin (round 31): TMA-swizzled wgmma vl when the
2641                        // promoted single-seq config is on; else the mma favl.
2642                        let fa3_on = match std::env::var("MEMRA_FA3").as_deref() {
2643                            Ok("0") => false,
2644                            Ok("1") => true,
2645                            _ => cfg!(memra_hopper_mma),
2646                        };
2647                        if fa3_on {
2648                            let mut q16s = Vec::with_capacity(b);
2649                            let mut v16s = Vec::with_capacity(b);
2650                            for s in 0..b {
2651                                let t = ts[s];
2652                                let mut q16 = e.alloc_u8_uninit(t * n_head * head_dim * 2)?;
2653                                e.f32_to_bf16_into(&aps[s].qn, &mut q16, t * n_head * head_dim)?;
2654                                let mut k16 = e.alloc_u8_uninit(t * n_head_kv * head_dim * 2)?;
2655                                e.f32_to_bf16_into(&aps[s].kn, &mut k16, t * n_head_kv * head_dim)?;
2656                                let mut v16 = e.alloc_u8_uninit(t * n_head_kv * head_dim * 2)?;
2657                                e.f32_to_bf16_v(
2658                                    &g3[2].slice(offs[s] * vf_w..(offs[s] + t) * vf_w),
2659                                    &mut v16,
2660                                    t * n_head_kv * head_dim,
2661                                )?;
2662                                q16s.push(q16);
2663                                v16s.push((k16, v16));
2664                            }
2665                            let mut qp = [core::ptr::null::<core::ffi::c_void>(); 8];
2666                            let mut kp = qp;
2667                            let mut vp = qp;
2668                            let mut op = [core::ptr::null_mut::<f32>(); 8];
2669                            let mut tsv = [0i32; 8];
2670                            for s in 0..b {
2671                                qp[s] = e.addr_u8(&q16s[s]) as *const core::ffi::c_void;
2672                                kp[s] = e.addr_u8(&v16s[s].0) as *const core::ffi::c_void;
2673                                vp[s] = e.addr_u8(&v16s[s].1) as *const core::ffi::c_void;
2674                                op[s] = e.addr_f32(&attns[s]) as *mut f32;
2675                                tsv[s] = ts[s] as i32;
2676                            }
2677                            let rc = unsafe {
2678                                crate::fa3_vl_raw(
2679                                    qp.as_ptr(),
2680                                    kp.as_ptr(),
2681                                    vp.as_ptr(),
2682                                    op.as_ptr(),
2683                                    tsv.as_ptr(),
2684                                    b as i32,
2685                                    n_head as i32,
2686                                    n_head_kv as i32,
2687                                    head_dim as i32,
2688                                    fa_scale,
2689                                    e.stream().cu_stream() as *mut core::ffi::c_void,
2690                                )
2691                            };
2692                            if rc != 0 {
2693                                return Err(format!("memra_fa3_vl rc={rc}").into());
2694                            }
2695                        } else {
2696                            let fargs: Vec<crate::FaSeqVl> = (0..b)
2697                                .map(|s| crate::FaSeqVl {
2698                                    q: e.addr_f32(&aps[s].qn),
2699                                    k16: e.addr_u8(&mirrors[s].0),
2700                                    v16: e.addr_u8(&mirrors[s].1),
2701                                    o: e.addr_f32(&attns[s]),
2702                                    kf: e.addr_f32(&aps[s].kn),
2703                                    vf: e.addr_f32v(
2704                                        &g3[2].slice(offs[s] * vf_w..(offs[s] + ts[s]) * vf_w),
2705                                    ),
2706                                    t: ts[s] as i32,
2707                                    pad: 0,
2708                                })
2709                                .collect();
2710                            e.fa_prefill_vl8(&fargs, head_dim, n_head, n_head_kv, fa_scale)?;
2711                        }
2712                        for (s, attn) in attns.into_iter().enumerate() {
2713                            let (attn_g, ag16) = self.full_attn_prime_post_fa(
2714                                e,
2715                                attn,
2716                                &aps[s].gate,
2717                                ts[s],
2718                                n_head,
2719                                head_dim,
2720                            )?;
2721                            let mut done = false;
2722                            if let Some(xh) = &ag16 {
2723                                done = e.try_f16_gemm_pre_into_off(
2724                                    &fa.wo,
2725                                    xh,
2726                                    ts[s],
2727                                    &mut mixed,
2728                                    offs[s] * n_embd,
2729                                )?;
2730                            }
2731                            if !done {
2732                                let m = e.matmul(&fa.wo, &attn_g, ts[s])?;
2733                                e.copy_into(&mut mixed, offs[s] * n_embd, &m, ts[s] * n_embd)?;
2734                            }
2735                        }
2736                    } else {
2737                        let mut parts: Vec<Vec<CudaSlice<f32>>> =
2738                            (0..b).map(|_| Vec::new()).collect();
2739                        for (w, y) in [&fa.wq, &fa.wk, &fa.wv].iter().zip(g3) {
2740                            for (s, ys) in split(e, &y, w.out_features())?.into_iter().enumerate() {
2741                                parts[s].push(ys);
2742                            }
2743                        }
2744                        for (s, g3s) in parts.into_iter().enumerate() {
2745                            // task #16 gather removal: wo writes into `mixed` at offs[s] directly.
2746                            let (attn_g, ag16) = self.full_attn_prime_core_inner(
2747                                e, fa, g3s, &pos_ds[s], ts[s], caches[s], il,
2748                            )?;
2749                            let mut done = false;
2750                            if let Some(xh) = &ag16 {
2751                                done = e.try_f16_gemm_pre_into_off(
2752                                    &fa.wo,
2753                                    xh,
2754                                    ts[s],
2755                                    &mut mixed,
2756                                    offs[s] * n_embd,
2757                                )?;
2758                            }
2759                            if !done {
2760                                let m = e.matmul(&fa.wo, &attn_g, ts[s])?;
2761                                e.copy_into(&mut mixed, offs[s] * n_embd, &m, ts[s] * n_embd)?;
2762                            }
2763                        }
2764                    }
2765                }
2766                Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
2767                Mixer::Linear(la) => {
2768                    // task #16: NO split copies (cores read row-offset views of the concat
2769                    // outputs; out-GEMMs write into `mixed` at offs[s]). task #18: the core
2770                    // itself is BATCHED — per-seq prep/K1-K3, then ONE varlen K4 + ONE
2771                    // varlen K5 launch for all sequences.
2772                    let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
2773                    let g4 = e.matmul_group_xh(&ws, &h, &hx16, total)?;
2774                    let outs =
2775                        self.linear_attn_prime_core_batch(e, la, &g4, &offs, &ts, caches, il)?;
2776                    for (s, (gn, gn16)) in outs.into_iter().enumerate() {
2777                        let (o, t) = (offs[s], ts[s]);
2778                        let mut done = false;
2779                        if let Some(xh) = &gn16 {
2780                            done = e.try_f16_gemm_pre_into_off(
2781                                &la.ssm_out,
2782                                xh,
2783                                t,
2784                                &mut mixed,
2785                                o * n_embd,
2786                            )?;
2787                        }
2788                        if !done {
2789                            let m = e.matmul(&la.ssm_out, &gn, t)?;
2790                            e.copy_into(&mut mixed, o * n_embd, &m, t * n_embd)?;
2791                        }
2792                    }
2793                }
2794            }
2795            let mut x1 = e.uninit(total * n_embd)?;
2796            let mut z = e.uninit(total * n_embd)?;
2797            let mut zx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
2798            e.add_rms_norm_f16out(
2799                &x,
2800                &mixed,
2801                layer.post_attn_norm.float_data(),
2802                &mut x1,
2803                &mut z,
2804                &mut zx16,
2805                n_embd,
2806                total,
2807                eps,
2808            )?;
2809            let ffn_out = match &layer.ffn {
2810                crate::hybrid::Ffn::Dense {
2811                    ffn_gate,
2812                    ffn_up,
2813                    ffn_down,
2814                } => {
2815                    let n_ff = ffn_gate.out_features();
2816                    let mut g2 = e.matmul_group_xh(&[ffn_gate, ffn_up], &z, &zx16, total)?;
2817                    let up = g2.pop().unwrap();
2818                    let gate = g2.pop().unwrap();
2819                    let mut act = e.uninit(total * n_ff)?;
2820                    // task #17 (batch trunk): silu twin emits the down GEMM's fp16 operand
2821                    // in-epilogue (nsys round-26: this trunk still paid 32 cvt passes).
2822                    // A clamped layer must skip the plain-SiLU twin (see prime_chunk's note).
2823                    let d_lim = self.cfg.clamp_shexp_at(il as u32);
2824                    if Self::f16out_on(e, total) && self.cfg.m3.is_none() && d_lim.is_none() {
2825                        let mut a16 = e.alloc_u8_uninit(total * n_ff * 2)?;
2826                        e.silu_mul_f16out(&gate, &up, &mut act, &mut a16, total * n_ff)?;
2827                        match e.try_f16_gemm_pre(ffn_down, &a16, total)? {
2828                            Some(y) => y,
2829                            None => e.matmul(ffn_down, &act, total)?,
2830                        }
2831                    } else {
2832                        Self::ffn_act_lim(
2833                            e,
2834                            &self.cfg,
2835                            &gate,
2836                            &up,
2837                            1.0,
2838                            1.0,
2839                            d_lim,
2840                            &mut act,
2841                            total * n_ff,
2842                        )?;
2843                        e.matmul(ffn_down, &act, total)?
2844                    }
2845                }
2846                crate::hybrid::Ffn::Moe(m) => {
2847                    self.moe_ffn_il_prefill(e, m, &z, total, il as u16)?
2848                }
2849            };
2850            let mut x2 = e.uninit(total * n_embd)?;
2851            e.add(&x1, &ffn_out, &mut x2, total * n_embd)?;
2852            x = x2;
2853        }
2854        // epilogue per seq (identical math to prime_chunk: norm all rows, lm_head on last row)
2855        let mut hn = e.uninit(total * n_embd)?;
2856        e.rms_norm(
2857            &x,
2858            self.output_norm.float_data(),
2859            &mut hn,
2860            n_embd,
2861            total,
2862            eps,
2863        )?;
2864        // batched lm_head (nsys round-26: B sequential m=1 matvecs re-read the 600MB
2865        // lm_head weight B times — 2.2ms at B=6): gather the B last rows and run ONE
2866        // m=B GEMM (f16 lane; falls back to the per-seq matvec when no mirror). The
2867        // f16-vs-mmvq logits delta is the prefill GEMM numeric class — prime_batch_gate's
2868        // argmax battery arbitrates, same as every other prefill GEMM change.
2869        let mut hcat = e.uninit(b * n_embd)?;
2870        for s in 0..b {
2871            let last0 = (offs[s] + ts[s] - 1) * n_embd;
2872            e.copy_view_into(
2873                &mut hcat,
2874                s * n_embd,
2875                &hn.slice(last0..last0 + n_embd),
2876                n_embd,
2877            )?;
2878        }
2879        let logits_cat = if b >= 2 {
2880            e.try_f16_gemm(&self.output, &hcat, b)?
2881        } else {
2882            None
2883        };
2884        let logits_host: Option<Vec<f32>> = match &logits_cat {
2885            Some(lc) => Some(e.dtoh(lc)?),
2886            None => None,
2887        };
2888        let n_vocab = self.output.out_features();
2889        let mut hidden_all = if crate::spec::spec_hpost() {
2890            split(e, &hn, n_embd)?
2891        } else {
2892            split(e, &x, n_embd)?
2893        };
2894        let mut out = Vec::with_capacity(b);
2895        for s in 0..b {
2896            let last0 = (offs[s] + ts[s] - 1) * n_embd;
2897            let mut h_seed = e.uninit(n_embd)?;
2898            if !crate::spec::spec_hpost() {
2899                e.copy_view_into(&mut h_seed, 0, &x.slice(last0..last0 + n_embd), n_embd)?;
2900            } else {
2901                e.copy_view_into(&mut h_seed, 0, &hn.slice(last0..last0 + n_embd), n_embd)?;
2902            }
2903            let logits = match &logits_host {
2904                Some(lh) => lh[s * n_vocab..(s + 1) * n_vocab].to_vec(),
2905                None => {
2906                    let mut hlast = e.uninit(n_embd)?;
2907                    e.copy_view_into(&mut hlast, 0, &hn.slice(last0..last0 + n_embd), n_embd)?;
2908                    e.dtoh(&e.matmul(&self.output, &hlast, 1)?)?
2909                }
2910            };
2911            caches[s].pos += ts[s];
2912            out.push((logits, h_seed, hidden_all.remove(0)));
2913        }
2914        Ok(out)
2915    }
2916
2917    /// `full_attn` (batched prefill mixer) + the cache side-effect: append the T post-RoPE K/V
2918    /// rows into the resident quantized KV cache (q8_0 K / q5_1 V) and advance len/len_d. Row
2919    /// bytes are BIT-IDENTICAL to the decode append (same per-warp quant kernel per row; the
2920    /// batched `append_kv_quantized_rows` runs that exact warp math on a (block, token) grid).
2921    /// The attention itself is unchanged prefill math (fa_prefill over the f32 K/V).
2922    ///
2923    /// `seq_end` = the ABSOLUTE end position of the WHOLE prime request (`cache.pos + prompt_len`
2924    /// measured BEFORE the chunk loop starts), i.e. a chunk-size-invariant property of the request.
2925    /// Only step35's SWA arm reads it (`step35_attn_pre_wo`'s doc note explains why keying on the
2926    /// chunk's own `t_kv` made the output depend on MEMRA_PRIME_CHUNK); every other arch ignores it.
2927    #[allow(clippy::too_many_arguments)]
2928    fn full_attn_prime(
2929        &self,
2930        e: &Engine,
2931        fa: &FullAttnLayer,
2932        h: &CudaSlice<f32>,
2933        hx: Option<&CudaSlice<u8>>,
2934        pos_d: &CudaSlice<i32>,
2935        t: usize,
2936        cache: &mut Cache,
2937        il: usize,
2938        seq_end: usize,
2939    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2940        if self.cfg.step35.is_some() {
2941            return self.step35_attn_prime(e, fa, h, hx, pos_d, t, cache, il, seq_end);
2942        }
2943        // PROJ/CORE SPLIT (task #13, 2026-07-26): the q/k/v group projection is hoisted so
2944        // the cross-request batch driver can run it at m = sum_T over concatenated tokens;
2945        // this single-seq path composes proj+core identically (byte-for-byte the old body).
2946        // task #14: `hx` = the norm-fused fp16 twin of `h` (skips the convert launch).
2947        let g3 = match hx {
2948            Some(xh) => e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], h, xh, t)?,
2949            None => e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?,
2950        };
2951        self.full_attn_prime_core(e, fa, g3, pos_d, t, cache, il)
2952    }
2953
2954    /// Everything after the q/k/v projections (split/QK-norm/RoPE/FA/gate/append + wo).
2955    /// `g3` = matmul_group([wq, wk, wv]) output rows for THIS sequence.
2956    /// Composes inner + wo (== the pre-split core; every existing caller unchanged).
2957    fn full_attn_prime_core(
2958        &self,
2959        e: &Engine,
2960        fa: &FullAttnLayer,
2961        g3: Vec<CudaSlice<f32>>,
2962        pos_d: &CudaSlice<i32>,
2963        t: usize,
2964        cache: &mut Cache,
2965        il: usize,
2966    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2967        let (attn_g, ag16) = self.full_attn_prime_core_inner(e, fa, g3, pos_d, t, cache, il)?;
2968        if let Some(xh) = &ag16 {
2969            if let Some(y) = e.try_f16_gemm_pre(&fa.wo, xh, t)? {
2970                return Ok(y);
2971            }
2972        }
2973        Ok(e.matmul(&fa.wo, &attn_g, t)?)
2974    }
2975
2976    fn full_attn_prime_core_inner(
2977        &self,
2978        e: &Engine,
2979        fa: &FullAttnLayer,
2980        g3: Vec<CudaSlice<f32>>,
2981        pos_d: &CudaSlice<i32>,
2982        t: usize,
2983        cache: &mut Cache,
2984        il: usize,
2985    ) -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
2986        let cfg = &self.cfg;
2987        let geometry = cfg.full_attention_geometry_at(il as u32);
2988        let n_head = geometry.n_head as usize;
2989        let n_head_kv = geometry.n_head_kv as usize;
2990        let head_dim = geometry.head_dim_k as usize;
2991        let scale = geometry.attention_scale();
2992        let (pre, base_len) = self.full_attn_prime_pre_fa(e, fa, g3, pos_d, t, cache, il)?;
2993        let AttnPre { q, k, v, gate } = pre;
2994        let mut attn = e.uninit(t * n_head * head_dim)?;
2995        self.full_attn_prime_fa_dispatch(
2996            e, &q, &k, &v, &mut attn, base_len, t, cache, il, head_dim, n_head, n_head_kv, scale,
2997        )?;
2998        self.full_attn_prime_post_fa(e, attn, &gate, t, n_head, head_dim)
2999    }
3000
3001    /// task #18 (attn side): projections tail through KV append — everything before the
3002    /// attention kernel. Returns the post-rope q/k, v, optional out-gate, and the KV rows
3003    /// present BEFORE this chunk's append (base_len; 0 == fresh).
3004    #[allow(clippy::type_complexity)]
3005    fn full_attn_prime_pre_fa(
3006        &self,
3007        e: &Engine,
3008        fa: &FullAttnLayer,
3009        mut g3: Vec<CudaSlice<f32>>,
3010        pos_d: &CudaSlice<i32>,
3011        t: usize,
3012        cache: &mut Cache,
3013        il: usize,
3014    ) -> Result<(AttnPre, usize), Box<dyn std::error::Error>> {
3015        let cfg = &self.cfg;
3016        let geometry = cfg.full_attention_geometry_at(il as u32);
3017        let n_head = geometry.n_head as usize;
3018        let n_head_kv = geometry.n_head_kv as usize;
3019        let head_dim = geometry.head_dim_k as usize;
3020        let eps = cfg.rms_eps;
3021
3022        // qwen35 fuses [q|gate] per head in wq (2*head_dim stride); M3/Hy3 have NO output gate
3023        // (attention_output_gate=false) — wq out = n_head*head_dim exactly, and q_gate_split
3024        // would read 2x out of bounds. `gated` keys both the split and the sigmoid epilogue.
3025        let gated = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
3026        let v = g3.pop().unwrap();
3027        let mut k = g3.pop().unwrap();
3028        let qf = g3.pop().unwrap();
3029        let (mut q, gate) = if gated {
3030            let mut q = e.uninit(t * n_head * head_dim)?;
3031            let mut gate = e.uninit(t * n_head * head_dim)?;
3032            e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, t)?;
3033            (q, Some(gate))
3034        } else {
3035            (qf, None)
3036        };
3037
3038        let mut qn = e.uninit(t * n_head * head_dim)?;
3039        e.rms_norm(
3040            &q,
3041            fa.q_norm.float_data(),
3042            &mut qn,
3043            head_dim,
3044            n_head * t,
3045            eps,
3046        )?;
3047        q = qn;
3048        let mut kn = e.uninit(t * n_head_kv * head_dim)?;
3049        e.rms_norm(
3050            &k,
3051            fa.k_norm.float_data(),
3052            &mut kn,
3053            head_dim,
3054            n_head_kv * t,
3055            eps,
3056        )?;
3057        k = kn;
3058        let rope_dims = geometry.n_rot as usize;
3059        e.rope_neox(
3060            &mut q,
3061            pos_d,
3062            head_dim,
3063            rope_dims,
3064            n_head,
3065            t,
3066            geometry.rope_base,
3067            1.0,
3068        )?;
3069        e.rope_neox(
3070            &mut k,
3071            pos_d,
3072            head_dim,
3073            rope_dims,
3074            n_head_kv,
3075            t,
3076            geometry.rope_base,
3077            1.0,
3078        )?;
3079
3080        // CACHE SIDE-EFFECT: append the T post-rope K/V token rows (token-major [T, kv_dim] ==
3081        // the cache row layout) quantized into cache.kv[il], then advance len + device len_d.
3082        {
3083            let kvl = cache.kv[il].as_mut().unwrap();
3084            assert!(kvl.len + t <= cache.max_ctx, "prime_cache: KV overflow");
3085            e.append_kv_quantized_rows(
3086                &k,
3087                &v,
3088                &mut kvl.k,
3089                &mut kvl.v,
3090                kvl.len,
3091                t,
3092                kvl.kv_dim_k,
3093                kvl.kv_dim_v,
3094                kvl.k_tok_bytes,
3095                kvl.v_tok_bytes,
3096                crate::Engine::kv_fp8_on(),
3097            )?;
3098            kvl.len += t;
3099            let new_len = kvl.len as i32;
3100            e.set_i32_one(&mut kvl.len_d, new_len)?;
3101        }
3102
3103        let base_len = {
3104            let kvl = cache.kv[il].as_ref().unwrap();
3105            kvl.len - t // KV rows present BEFORE this chunk's append above
3106        };
3107        Ok((AttnPre { q, k, v, gate }, base_len))
3108    }
3109
3110    /// batched prefill attention. FRESH prime (no past KV): unchanged forward_last math over
3111    /// the f32 K/V of this batch. CONTINUATION chunk (past KV present): the chunk's queries
3112    /// must attend to [0 .. base+t) — run fa_prefill_view over the resident QUANTIZED cache
3113    /// (the spec-verify pattern; kernel's causal mask offsets by T_kv-T). Numerically this
3114    /// reads q8_0/q5_1-dequantized K/V for the past AND the current chunk — the same class as
3115    /// decode reading the cache; the run-gen/first-16 battery is the accuracy authority.
3116    #[allow(clippy::too_many_arguments)]
3117    fn full_attn_prime_fa_dispatch(
3118        &self,
3119        e: &Engine,
3120        q: &CudaSlice<f32>,
3121        k: &CudaSlice<f32>,
3122        v: &CudaSlice<f32>,
3123        attn: &mut CudaSlice<f32>,
3124        base_len: usize,
3125        t: usize,
3126        cache: &mut Cache,
3127        il: usize,
3128        head_dim: usize,
3129        n_head: usize,
3130        n_head_kv: usize,
3131        scale: f32,
3132    ) -> Result<(), Box<dyn std::error::Error>> {
3133        // GRAIN-FREE CHUNK-INVARIANCE FIX (lane/chunkinv-flip, 2026-08-05): the base_len == 0
3134        // f32 special case (fa_prefill over this batch's f32 K/V) is DROPPED. pre_fa appends
3135        // the chunk's quantized rows into cache.kv[il] BEFORE this dispatch, so chunk 0 can
3136        // attend through the quantized cache exactly like every later chunk (quantize-then-
3137        // attend). One numeric class for every row => the chunk size cannot decide where a
3138        // precision edge falls, so chunked prefill is reduction-order-stable with NO grain
3139        // knob (the stronger fix VERDICT.md filed; supersedes the MEMRA_PRIME_INVARIANT door's
3140        // pin-the-boundary approach).
3141        // MEMRA_PRIME_F32CHUNK0=1 is the ROLLBACK SEAM to the old arithmetic (chunk 0 attends
3142        // f32 K/V) — flags-doctrine rollback door AND the chunkinv gate's canary injection:
3143        // with the fix unconditional, only re-introducing the class edge can prove the gate
3144        // still detects the mechanism. Never on in a measured default run.
3145        if base_len == 0 && std::env::var("MEMRA_PRIME_F32CHUNK0").as_deref() == Ok("1") {
3146            if std::env::var("MEMRA_NOFA").is_ok() || !(head_dim == 256 || head_dim == 128) {
3147                e.sdpa_naive(
3148                    q, k, v, attn, head_dim, n_head, n_head_kv, t, t, scale, true,
3149                )?;
3150            } else {
3151                e.fa_prefill(
3152                    q, k, v, attn, head_dim, n_head, n_head_kv, t, t, scale, true,
3153                )?;
3154            }
3155            return Ok(());
3156        }
3157        let kvl = cache.kv[il].as_ref().unwrap();
3158        let t_kv = base_len + t;
3159        let k_view = e.view_u8(&kvl.k, t_kv * kvl.k_tok_bytes);
3160        let v_view = e.view_u8(&kvl.v, t_kv * kvl.v_tok_bytes);
3161        // fa_prefill_q/_qw twins are stamped for head_dim 256 (qwen35) and 128 (M3) only;
3162        // other dims (and MEMRA_NOFA) take the naive quantized-view SDPA — SAME cache bytes,
3163        // same numeric class, so the uniform contract holds on the fallback too.
3164        if std::env::var("MEMRA_NOFA").is_ok() || !(head_dim == 256 || head_dim == 128) {
3165            e.sdpa_naive_quantized_view(
3166                q,
3167                &k_view,
3168                &v_view,
3169                attn,
3170                head_dim,
3171                n_head,
3172                n_head_kv,
3173                t,
3174                t_kv,
3175                scale,
3176                true,
3177                kvl.k_tok_bytes,
3178                kvl.v_tok_bytes,
3179            )?;
3180            return Ok(());
3181        }
3182        // ARC B (2026-07-05): dequant-once workspace, DEFAULT ON. fa_prefill_q's inline
3183        // dequant re-reads+re-dequants the whole quantized KV stream from every one of the
3184        // T/64 x n_head CTAs (64x+ redundant at chunk=4096; 30.5% of the 32k prime wall).
3185        // fa_prefill_view_ws dequants K/V ONCE into a resident bf16 workspace then runs the
3186        // bit-identical bf16 twin (fa_prefill_qw) — same staged values, same FP order, token-
3187        // identical output (gate: MEMRA_PRIME_CHUNK=4096 ws-on vs ws-off vs monolithic).
3188        // MEMRA_PRIME_DEQW=0 reverts to the inline-dequant kernel.
3189        let deqw = std::env::var("MEMRA_PRIME_DEQW")
3190            .map(|v| v != "0")
3191            .unwrap_or(true);
3192        if deqw {
3193            e.fa_prefill_view_ws(
3194                q,
3195                &k_view,
3196                &v_view,
3197                attn,
3198                head_dim,
3199                n_head,
3200                n_head_kv,
3201                t,
3202                t_kv,
3203                scale,
3204                true,
3205                kvl.k_tok_bytes,
3206                kvl.v_tok_bytes,
3207                crate::Engine::kv_fp8_on(),
3208            )?;
3209        } else {
3210            e.fa_prefill_view(
3211                q,
3212                &k_view,
3213                &v_view,
3214                attn,
3215                head_dim,
3216                n_head,
3217                n_head_kv,
3218                t,
3219                t_kv,
3220                scale,
3221                true,
3222                kvl.k_tok_bytes,
3223                kvl.v_tok_bytes,
3224                crate::Engine::kv_fp8_on(),
3225            )?;
3226        }
3227        Ok(())
3228    }
3229
3230    /// task #17: sig_mul_f16out fuses [sigmoid + mul + f16 convert] into one launch
3231    /// (bit-identical composition) and hands wo its fp16 operand directly.
3232    fn full_attn_prime_post_fa(
3233        &self,
3234        e: &Engine,
3235        attn: CudaSlice<f32>,
3236        gate: &Option<CudaSlice<f32>>,
3237        t: usize,
3238        n_head: usize,
3239        head_dim: usize,
3240    ) -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
3241        let (attn_g, ag16) = match gate {
3242            Some(gate) => {
3243                let n = t * n_head * head_dim;
3244                let mut ag = e.uninit(n)?;
3245                if Self::f16out_on(e, t) {
3246                    let mut a16 = e.alloc_u8_uninit(n * 2)?;
3247                    e.sig_mul_f16out(&attn, gate, &mut ag, &mut a16, n)?;
3248                    (ag, Some(a16))
3249                } else {
3250                    let mut gsig = e.uninit(n)?;
3251                    e.sigmoid(gate, &mut gsig, n)?;
3252                    e.mul(&attn, &gsig, &mut ag, n)?;
3253                    (ag, None)
3254                }
3255            }
3256            None => (attn, None),
3257        };
3258        Ok((attn_g, ag16))
3259    }
3260
3261    /// STATEFUL batched linear-attention prime: `linear_attn`'s prefill-dispatch pass (normal
3262    /// `e.matmul` — GEMM at m>=16 — plus the prefill repack/L2/glog kernels) but with the state
3263    /// carried THROUGH the cache like the spec verify does: carried-ring conv
3264    /// (ssm_conv1d_tm_state writes the final ring back) + ONE gdn_scan from cache.recur[il]'s
3265    /// current state (zero at a fresh prime) whose final state ping-pongs back into the cache.
3266    /// Wiring mirrors `linear_attn_verify_t` (spec.rs); dispatch mirrors `linear_attn` (prefill).
3267    fn linear_attn_prime(
3268        &self,
3269        e: &Engine,
3270        la: &LinearAttnLayer,
3271        h: &CudaSlice<f32>,
3272        hx: Option<&CudaSlice<u8>>,
3273        t: usize,
3274        cache: &mut Cache,
3275        il: usize,
3276    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3277        // PROJ/CORE SPLIT (task #13): see full_attn_prime — same hoist for the GDN 4-tuple.
3278        let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
3279        let g4 = match hx {
3280            Some(xh) => e.matmul_group_xh(&ws, h, xh, t)?,
3281            None => e.matmul_group(&ws, h, t)?,
3282        };
3283        self.linear_attn_prime_core(e, la, g4, t, cache, il)
3284    }
3285
3286    /// Everything after the GDN 4-tuple projections (conv/chunk stack/gated norm + ssm_out).
3287    fn linear_attn_prime_core(
3288        &self,
3289        e: &Engine,
3290        la: &LinearAttnLayer,
3291        mut g4: Vec<CudaSlice<f32>>,
3292        t: usize,
3293        cache: &mut Cache,
3294        il: usize,
3295    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3296        self.linear_attn_prime_core_pad(e, la, g4.drain(..).collect(), t, cache, il, None)
3297    }
3298
3299    /// task #14: `pad_len` = device true length for PADDED prime graphs. Pads become
3300    /// identity GDN steps (gdn_pad_mask zeroes beta/g_log past the true length) and the
3301    /// conv ring writes back from the true tail. None = classic path, byte-identical.
3302    #[allow(clippy::too_many_arguments)]
3303    fn linear_attn_prime_core_pad_inner(
3304        &self,
3305        e: &Engine,
3306        la: &LinearAttnLayer,
3307        mut g4: Vec<CudaSlice<f32>>,
3308        t: usize,
3309        cache: &mut Cache,
3310        il: usize,
3311        pad_len: Option<&CudaSlice<i32>>,
3312    ) -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
3313        // shim over the view twin (task #16): full-range views of the owned buffers.
3314        let ssm = self.cfg.ssm.as_ref().unwrap();
3315        let d_state = ssm.state_size as usize;
3316        let num_k = ssm.group_count as usize;
3317        let num_v = ssm.time_step_rank as usize;
3318        let key_dim = d_state * num_k;
3319        let value_dim = d_state * num_v;
3320        let conv_dim = key_dim * 2 + value_dim;
3321        let alpha = g4.pop().unwrap(); // [T, num_v]
3322        let beta_raw = g4.pop().unwrap(); // [T, num_v]
3323        let z = g4.pop().unwrap(); // [T, value_dim]
3324        let qkv_mixed = g4.pop().unwrap(); // [T, conv_dim] token-major
3325        self.linear_attn_prime_core_pad_view(
3326            e,
3327            la,
3328            &qkv_mixed.slice(0..t * conv_dim),
3329            &z.slice(0..t * value_dim),
3330            &beta_raw.slice(0..t * num_v),
3331            &alpha.slice(0..t * num_v),
3332            t,
3333            cache,
3334            il,
3335            pad_len,
3336        )
3337    }
3338
3339    /// task #18: the GDN prep stage (conv + repack + l2 x2 + sigmoid + glog + pad-mask) —
3340    /// shared verbatim by the per-seq scan path and the varlen batched path.
3341    #[allow(clippy::too_many_arguments)]
3342    fn linear_attn_gdn_prep(
3343        &self,
3344        e: &Engine,
3345        la: &LinearAttnLayer,
3346        qkv_mixed: &cudarc::driver::CudaView<f32>,
3347        beta_raw: &cudarc::driver::CudaView<f32>,
3348        alpha: &cudarc::driver::CudaView<f32>,
3349        t: usize,
3350        cache: &mut Cache,
3351        il: usize,
3352        pad_len: Option<&CudaSlice<i32>>,
3353    ) -> Result<GdnPrep, Box<dyn std::error::Error>> {
3354        let cfg = &self.cfg;
3355        let ssm = cfg.ssm.as_ref().unwrap();
3356        let d_state = ssm.state_size as usize; // 128
3357        let num_k = ssm.group_count as usize; // 16
3358        let num_v = ssm.time_step_rank as usize; // 32
3359        let d_conv = ssm.conv_kernel as usize; // 4
3360        let key_dim = d_state * num_k; // 2048
3361        let value_dim = d_state * num_v; // 4096
3362        let conv_dim = key_dim * 2 + value_dim; // 8192
3363        let eps = cfg.rms_eps;
3364        debug_assert!(
3365            t >= d_conv - 1,
3366            "stateful conv needs T >= pad (PRIME_MIN_T gates)"
3367        );
3368
3369        // conv with CARRIED ring state + ring roll (state read + final-window write-back).
3370        // task #18 conv-fuse (default ON): conv + SiLU + repack in one pass — the 67MB
3371        // conv_out intermediate and its transposed re-read disappear (values bit-identical).
3372        // MEMRA_CONV_FUSE=0 reverts to the two-kernel chain.
3373        let rl = cache.recur[il].as_mut().unwrap();
3374        let hk = Self::gdn_hk(e, t, num_v, num_k);
3375        let conv_fuse = std::env::var("MEMRA_CONV_FUSE").as_deref() != Ok("0");
3376        let hk = if conv_fuse { hk } else { num_v }; // de-broadcast rides the fused conv
3377        let mut q_g = e.uninit(d_state * hk * t)?;
3378        let mut k_g = e.uninit(d_state * hk * t)?;
3379        let mut v_g = e.uninit(d_state * num_v * t)?;
3380        if conv_fuse {
3381            e.ssm_conv1d_gdn_state_pad(
3382                qkv_mixed,
3383                &mut rl.conv_state,
3384                la.ssm_conv1d.float_data(),
3385                &mut q_g,
3386                &mut k_g,
3387                &mut v_g,
3388                conv_dim,
3389                t,
3390                d_conv,
3391                d_state,
3392                num_v,
3393                num_k,
3394                key_dim,
3395                hk,
3396                pad_len,
3397            )?;
3398        } else {
3399            let mut conv_out = e.uninit(conv_dim * t)?; // [conv_dim, T] channel-major, SiLU
3400            e.ssm_conv1d_tm_state_pad_v(
3401                qkv_mixed,
3402                &mut rl.conv_state,
3403                la.ssm_conv1d.float_data(),
3404                &mut conv_out,
3405                conv_dim,
3406                t,
3407                d_conv,
3408                pad_len,
3409            )?;
3410            e.qkv_to_gdn_repack(
3411                &conv_out, &mut q_g, &mut k_g, &mut v_g, d_state, num_v, num_k, key_dim, t,
3412            )?;
3413        }
3414        let mut q_l2 = e.uninit(d_state * hk * t)?;
3415        // mirror-fold (round 35): q's bf16 twin (wgmma K45/K2 A-operand) in-epilogue too.
3416        // Emitted only where a consumer exists (the wgmma config) — on other arches the
3417        // alloc + epilogue stores would be pure waste.
3418        let qb16 = if Engine::l2_v2_on(d_state) && e.gdn_wgmma_on(32) {
3419            let mut qb = e.alloc_u8_uninit(d_state * hk * t * 2)?;
3420            e.l2_norm_pp(&q_g, &mut q_l2, Some(&mut qb), d_state, hk * t, eps)?;
3421            Some(qb)
3422        } else {
3423            e.l2_norm_pp(&q_g, &mut q_l2, None, d_state, hk * t, eps)?;
3424            None
3425        };
3426        let mut k_l2 = e.uninit(d_state * hk * t)?;
3427        // mirror-fold: emit k's bf16 twin (the chunked-scan kb16 mirror) in-epilogue
3428        let kb16 = if Engine::l2_v2_on(d_state) {
3429            let mut kb = e.alloc_u8_uninit(d_state * hk * t * 2)?;
3430            e.l2_norm_pp(&k_g, &mut k_l2, Some(&mut kb), d_state, hk * t, eps)?;
3431            Some(kb)
3432        } else {
3433            e.l2_norm_pp(&k_g, &mut k_l2, None, d_state, hk * t, eps)?;
3434            None
3435        };
3436        let mut beta = e.uninit(t * num_v)?;
3437        e.sigmoid_v(beta_raw, &mut beta, t * num_v)?;
3438        let mut g_log = e.uninit(t * num_v)?;
3439        e.gdn_glog_v(
3440            alpha,
3441            la.ssm_dt.float_data(),
3442            la.ssm_a.float_data(),
3443            &mut g_log,
3444            num_v,
3445            t,
3446        )?;
3447        if let Some(len_d) = pad_len {
3448            e.gdn_pad_mask(&mut beta, &mut g_log, len_d, num_v, t)?;
3449        }
3450        Ok(GdnPrep {
3451            hk,
3452            q_l2,
3453            k_l2,
3454            v_g,
3455            beta,
3456            g_log,
3457            kb16,
3458            qb16,
3459        })
3460    }
3461
3462    /// task #18: batched GDN mixer core — per-seq prep + K1-K3, then ONE varlen K4 and
3463    /// ONE varlen K5 launch for all B sequences (full-machine occupancy vs B underfilled
3464    /// per-seq trains). Per-block math identical to the per-seq path (bit-gateable);
3465    /// MEMRA_GDN_VL=0 or a non-mma/non-C32 config falls back to the per-seq loop.
3466    #[allow(clippy::too_many_arguments)]
3467    fn linear_attn_prime_core_batch(
3468        &self,
3469        e: &Engine,
3470        la: &LinearAttnLayer,
3471        g4: &[CudaSlice<f32>],
3472        offs: &[usize],
3473        ts: &[usize],
3474        caches: &mut [&mut Cache],
3475        il: usize,
3476    ) -> Result<Vec<(CudaSlice<f32>, Option<CudaSlice<u8>>)>, Box<dyn std::error::Error>> {
3477        let ssm = self.cfg.ssm.as_ref().unwrap();
3478        let d_state = ssm.state_size as usize;
3479        let num_k = ssm.group_count as usize;
3480        let num_v = ssm.time_step_rank as usize;
3481        let key_dim = d_state * num_k;
3482        let value_dim = d_state * num_v;
3483        let conv_dim = key_dim * 2 + value_dim;
3484        let eps = self.cfg.rms_eps;
3485        let scale = 1.0 / (d_state as f32).sqrt();
3486        let b = ts.len();
3487        let c = Engine::gdn_chunk_size();
3488        // CONTINUATION batches (increment (b)) take the per-seq stateful pad_view twin
3489        // below — the varlen K4/K5 chain is fresh-only (zero initial state assumed).
3490        let carried = caches.iter().any(|c| c.pos > 0);
3491        let use_vl = !carried
3492            && (2..=8).contains(&b)
3493            && Engine::gdn_chunked_enabled()
3494            && ts.iter().all(|&t| t >= 16)
3495            && e.gdn_mma_enabled(c)
3496            && std::env::var("MEMRA_GDN_VL").as_deref() != Ok("0");
3497        if !use_vl {
3498            return (0..b)
3499                .map(|s| {
3500                    let (o, t) = (offs[s], ts[s]);
3501                    self.linear_attn_prime_core_pad_view(
3502                        e,
3503                        la,
3504                        &g4[0].slice(o * conv_dim..(o + t) * conv_dim),
3505                        &g4[1].slice(o * value_dim..(o + t) * value_dim),
3506                        &g4[2].slice(o * num_v..(o + t) * num_v),
3507                        &g4[3].slice(o * num_v..(o + t) * num_v),
3508                        t,
3509                        caches[s],
3510                        il,
3511                        None,
3512                    )
3513                })
3514                .collect();
3515        }
3516        // increment 3: the ENTIRE core is varlen — allocs only per seq, then 13 launches
3517        // total for the whole batch: [conv, ring, repack, l2(q+k), gate-prep] +
3518        // [k-mirror, K1, K2, K3, w-mirror, K4, K5] + [gated-norm tail].
3519        struct SeqBufs {
3520            conv_out: CudaSlice<f32>,
3521            q_g: CudaSlice<f32>,
3522            k_g: CudaSlice<f32>,
3523            v_g: CudaSlice<f32>,
3524            q_l2: CudaSlice<f32>,
3525            k_l2: CudaSlice<f32>,
3526            beta: CudaSlice<f32>,
3527            g_log: CudaSlice<f32>,
3528            gn: CudaSlice<f32>,
3529            gn16: CudaSlice<u8>,
3530        }
3531        let d_conv = ssm.conv_kernel as usize;
3532        let f16o = Self::f16out_on(e, 16);
3533        let hk = Self::gdn_hk(e, 16, num_v, num_k); // vl path is always chunked+mma
3534        let mut sb = Vec::with_capacity(b);
3535        let mut pres = Vec::with_capacity(b);
3536        for &t in ts.iter().take(b) {
3537            sb.push(SeqBufs {
3538                conv_out: e.uninit(conv_dim * t)?,
3539                q_g: e.uninit(d_state * hk * t)?,
3540                k_g: e.uninit(d_state * hk * t)?,
3541                v_g: e.uninit(d_state * num_v * t)?,
3542                q_l2: e.uninit(d_state * hk * t)?,
3543                k_l2: e.uninit(d_state * hk * t)?,
3544                beta: e.uninit(t * num_v)?,
3545                g_log: e.uninit(t * num_v)?,
3546                gn: e.uninit(d_state * num_v * t)?,
3547                gn16: e.alloc_u8_uninit(d_state * num_v * t * 2)?,
3548            });
3549            pres.push(e.gdn_chunk_alloc(num_v, t, c, hk)?);
3550        }
3551        let prep_args: Vec<crate::GdnPrepVl> = (0..b)
3552            .map(|s| {
3553                let (o, t) = (offs[s], ts[s]);
3554                let rl = caches[s].recur[il].as_ref().unwrap();
3555                crate::GdnPrepVl {
3556                    qkv: e.addr_f32v(&g4[0].slice(o * conv_dim..(o + t) * conv_dim)),
3557                    conv_state: e.addr_f32(&rl.conv_state),
3558                    conv_out: e.addr_f32(&sb[s].conv_out),
3559                    q_g: e.addr_f32(&sb[s].q_g),
3560                    k_g: e.addr_f32(&sb[s].k_g),
3561                    v_g: e.addr_f32(&sb[s].v_g),
3562                    q_l2: e.addr_f32(&sb[s].q_l2),
3563                    k_l2: e.addr_f32(&sb[s].k_l2),
3564                    beta_raw: e.addr_f32v(&g4[2].slice(o * num_v..(o + t) * num_v)),
3565                    alpha: e.addr_f32v(&g4[3].slice(o * num_v..(o + t) * num_v)),
3566                    beta: e.addr_f32(&sb[s].beta),
3567                    g_log: e.addr_f32(&sb[s].g_log),
3568                    o: e.addr_f32(&pres[s].o),
3569                    z: e.addr_f32v(&g4[1].slice(o * value_dim..(o + t) * value_dim)),
3570                    gn: e.addr_f32(&sb[s].gn),
3571                    gn16: e.addr_u8(&sb[s].gn16),
3572                    kb16: if Engine::l2_v2_on(d_state) {
3573                        e.addr_u8(&pres[s].kb16)
3574                    } else {
3575                        0
3576                    },
3577                    qb16: if Engine::l2_v2_on(d_state) && e.gdn_wgmma_on(c) {
3578                        e.addr_u8(&pres[s].qb16)
3579                    } else {
3580                        0
3581                    },
3582                    t: t as i32,
3583                    pad: 0,
3584                }
3585            })
3586            .collect();
3587        let args: Vec<crate::GdnSeqVl> = (0..b)
3588            .map(|s| {
3589                let rl = caches[s].recur[il].as_ref().unwrap();
3590                crate::GdnSeqVl {
3591                    kb16: e.addr_u8(&pres[s].kb16),
3592                    gcum: e.addr_f32(&pres[s].gcum),
3593                    beta: e.addr_f32(&sb[s].beta),
3594                    u: e.addr_f32(&pres[s].u),
3595                    wb16: e.addr_u8(&pres[s].wb16),
3596                    y: e.addr_u8(&pres[s].y16),
3597                    ssnap: e.addr_u8(&pres[s].ssnap16),
3598                    state_in: e.addr_f32(&rl.ssm_state),
3599                    state_out: e.addr_f32(&rl.ssm_state_alt),
3600                    q: e.addr_f32(&sb[s].q_l2),
3601                    p: e.addr_f32(&pres[s].p),
3602                    o: e.addr_f32(&pres[s].o),
3603                    k: e.addr_f32(&sb[s].k_l2),
3604                    v: e.addr_f32(&sb[s].v_g),
3605                    g: e.addr_f32(&sb[s].g_log),
3606                    a: e.addr_f32(&pres[s].a),
3607                    w: e.addr_f32(&pres[s].w),
3608                    t: ts[s] as i32,
3609                    nc: pres[s].nc as i32,
3610                }
3611            })
3612            .collect();
3613        e.gdn_prep_vl8(
3614            &prep_args,
3615            la.ssm_conv1d.float_data(),
3616            la.ssm_dt.float_data(),
3617            la.ssm_a.float_data(),
3618            conv_dim,
3619            d_conv,
3620            d_state,
3621            num_v,
3622            num_k,
3623            key_dim,
3624            hk,
3625            eps,
3626        )?;
3627        // mirror-fold (round 27): l2 v2 emits kb16 in-epilogue; K3's store emits wb16 —
3628        // both standalone mirror launches vanish on the default config.
3629        if !Engine::l2_v2_on(d_state) {
3630            e.gdn_mirror_vl8(&args, num_v, 0, hk)?;
3631        }
3632        // task #22: wgmma-fused vl twins — qb16 mirrors + GdnWVl8 side-struct.
3633        let wq8: Option<crate::GdnWVl8> = if e.gdn_wgmma_on(c) {
3634            // qb16 emitted by the vl l2 mirror-fold when l2-v2 serves; bulk cvt otherwise
3635            if !Engine::l2_v2_on(d_state) {
3636                for s in 0..b {
3637                    e.f32_to_bf16_into(&sb[s].q_l2, &mut pres[s].qb16, d_state * hk * ts[s])?;
3638                }
3639            }
3640            let mut wa = [crate::GdnWVl::default(); 8];
3641            for s in 0..b {
3642                wa[s] = crate::GdnWVl {
3643                    qb16: e.addr_u8(&pres[s].qb16),
3644                    pb16: e.addr_u8(&pres[s].pb16),
3645                };
3646            }
3647            Some(crate::GdnWVl8(wa))
3648        } else {
3649            None
3650        };
3651        e.gdn_chunk_k123_vl8(&args, num_v, hk, wq8.as_ref())?;
3652        e.gdn_chunk_vl8(&args, num_v, scale, hk, wq8.as_ref())?;
3653        if f16o {
3654            e.gdn_tail_vl8(&prep_args, la.ssm_norm.float_data(), d_state, num_v, eps)?;
3655        }
3656        // per-seq state swap (+ non-f16out tail fallback)
3657        let mut out = Vec::with_capacity(b);
3658        for (s, bufs) in sb.into_iter().enumerate() {
3659            let rl = caches[s].recur[il].as_mut().unwrap();
3660            std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
3661            let (o, t) = (offs[s], ts[s]);
3662            let SeqBufs { mut gn, gn16, .. } = bufs;
3663            if f16o {
3664                out.push((gn, Some(gn16)));
3665            } else {
3666                let z_v = g4[1].slice(o * value_dim..(o + t) * value_dim);
3667                e.gated_rmsnorm_zv(
3668                    &pres[s].o,
3669                    la.ssm_norm.float_data(),
3670                    &z_v,
3671                    &mut gn,
3672                    d_state,
3673                    num_v * t,
3674                    eps,
3675                )?;
3676                out.push((gn, None));
3677            }
3678        }
3679        Ok(out)
3680    }
3681
3682    /// task #16: view-consuming GDN prime core — the batched prime hands row-offset
3683    /// views of the CONCAT projection outputs directly (no per-seq split copies).
3684    /// Same kernels, same values, byte-identical to the Vec shim above.
3685    #[allow(clippy::too_many_arguments)]
3686    fn linear_attn_prime_core_pad_view(
3687        &self,
3688        e: &Engine,
3689        la: &LinearAttnLayer,
3690        qkv_mixed: &cudarc::driver::CudaView<f32>,
3691        z: &cudarc::driver::CudaView<f32>,
3692        beta_raw: &cudarc::driver::CudaView<f32>,
3693        alpha: &cudarc::driver::CudaView<f32>,
3694        t: usize,
3695        cache: &mut Cache,
3696        il: usize,
3697        pad_len: Option<&CudaSlice<i32>>,
3698    ) -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
3699        let cfg = &self.cfg;
3700        let ssm = cfg.ssm.as_ref().unwrap();
3701        let d_state = ssm.state_size as usize; // 128
3702        let num_v = ssm.time_step_rank as usize; // 32
3703        let eps = cfg.rms_eps;
3704        let scale = 1.0 / (d_state as f32).sqrt();
3705
3706        let prep =
3707            self.linear_attn_gdn_prep(e, la, qkv_mixed, beta_raw, alpha, t, cache, il, pad_len)?;
3708
3709        // ONE gdn_scan over T from the cache's CURRENT state (zero at fresh prime); the final
3710        // state lands in the spare buffer and ping-pongs back (stable resident pointers, the
3711        // decode-determinism discipline from linear_attn_decode_inner). A4: `gdn_scan_prefill`
3712        // dispatches the chunked WY form under MEMRA_GDN_CHUNKED (prefill-only seam; decode +
3713        // verify keep the sequential kernel).
3714        let mut o = e.uninit(d_state * num_v * t)?;
3715        let rl = cache.recur[il].as_mut().unwrap();
3716        {
3717            let crate::cache::RecurLayer {
3718                ssm_state,
3719                ssm_state_alt,
3720                ..
3721            } = rl;
3722            e.gdn_scan_prefill(
3723                &prep.q_l2,
3724                &prep.k_l2,
3725                &prep.v_g,
3726                &prep.g_log,
3727                &prep.beta,
3728                prep.kb16.as_ref(),
3729                prep.qb16.as_ref(),
3730                ssm_state,
3731                ssm_state_alt,
3732                &mut o,
3733                num_v,
3734                t,
3735                scale,
3736                prep.hk,
3737            )?;
3738        }
3739        std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
3740
3741        // gated RMSNorm + out projection (prefill dispatch). task #17: the f16out twin also
3742        // emits the ssm_out GEMM's fp16 operand in-epilogue (kills the standalone convert).
3743        let mut gn = e.uninit(d_state * num_v * t)?;
3744        let gn16 = if Self::f16out_on(e, t) {
3745            let mut g16 = e.alloc_u8_uninit(d_state * num_v * t * 2)?;
3746            e.gated_rmsnorm_f16out_zv(
3747                &o,
3748                la.ssm_norm.float_data(),
3749                z,
3750                &mut gn,
3751                &mut g16,
3752                d_state,
3753                num_v * t,
3754                eps,
3755            )?;
3756            Some(g16)
3757        } else {
3758            e.gated_rmsnorm_zv(
3759                &o,
3760                la.ssm_norm.float_data(),
3761                z,
3762                &mut gn,
3763                d_state,
3764                num_v * t,
3765                eps,
3766            )?;
3767            None
3768        };
3769        Ok((gn, gn16))
3770    }
3771
3772    /// Task #15 core-split wrapper: composes inner + ssm_out (== the old core).
3773    #[allow(clippy::too_many_arguments)]
3774    fn linear_attn_prime_core_pad(
3775        &self,
3776        e: &Engine,
3777        la: &LinearAttnLayer,
3778        g4: Vec<CudaSlice<f32>>,
3779        t: usize,
3780        cache: &mut Cache,
3781        il: usize,
3782        pad_len: Option<&CudaSlice<i32>>,
3783    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3784        let (gn, gn16) = self.linear_attn_prime_core_pad_inner(e, la, g4, t, cache, il, pad_len)?;
3785        if let Some(xh) = &gn16 {
3786            if let Some(y) = e.try_f16_gemm_pre(&la.ssm_out, xh, t)? {
3787                return Ok(y);
3788            }
3789        }
3790        Ok(e.matmul(&la.ssm_out, &gn, t)?)
3791    }
3792
3793    /// Full-attention mixer with QK-norm, partial RoPE, sigmoid output gate (qwen35 :257-336).
3794    ///
3795    /// `il` = layer index: step35 needs it (per-layer n_head / rope width / window / gate) and
3796    /// routes to its own mixer. Every other arch ignores it (uniform geometry).
3797    pub fn full_attn(
3798        &self,
3799        e: &Engine,
3800        fa: &FullAttnLayer,
3801        h: &CudaSlice<f32>,
3802        pos_d: &CudaSlice<i32>,
3803        t: usize,
3804        il: usize,
3805    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3806        if self.cfg.step35.is_some() {
3807            return self.step35_attn(e, fa, h, pos_d, t, il);
3808        }
3809        let cfg = &self.cfg;
3810        let _n_embd = cfg.n_embd as usize;
3811        let geometry = cfg.full_attention_geometry_at(il as u32);
3812        let n_head = geometry.n_head as usize;
3813        let n_head_kv = geometry.n_head_kv as usize;
3814        let head_dim = geometry.head_dim_k as usize;
3815        let eps = cfg.rms_eps;
3816        let scale = geometry.attention_scale();
3817
3818        // qwen35: wq output = head_dim*2*n_head (fused [q|gate] per head). M3/Hy3: NO output
3819        // gate — wq out = n_head*head_dim, no split (see prime-path note).
3820        let gated = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
3821        // grouped: one f16 activation convert feeds q/k/v (matmul_group)
3822        let mut g3 = e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?;
3823        let v = g3.pop().unwrap();
3824        let mut k = g3.pop().unwrap();
3825        let qf = g3.pop().unwrap();
3826        let (mut q, gate) = if gated {
3827            let mut q = e.uninit(t * n_head * head_dim)?;
3828            let mut gate = e.uninit(t * n_head * head_dim)?;
3829            e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, t)?;
3830            (q, Some(gate))
3831        } else {
3832            (qf, None)
3833        };
3834
3835        // QK-norm (per head_dim row), then partial RoPE.
3836        let mut qn = e.uninit(t * n_head * head_dim)?;
3837        e.rms_norm(
3838            &q,
3839            fa.q_norm.float_data(),
3840            &mut qn,
3841            head_dim,
3842            n_head * t,
3843            eps,
3844        )?;
3845        q = qn;
3846        let mut kn = e.uninit(t * n_head_kv * head_dim)?;
3847        e.rms_norm(
3848            &k,
3849            fa.k_norm.float_data(),
3850            &mut kn,
3851            head_dim,
3852            n_head_kv * t,
3853            eps,
3854        )?;
3855        k = kn;
3856        let rope_dims = geometry.n_rot as usize;
3857        e.rope_neox(
3858            &mut q,
3859            pos_d,
3860            head_dim,
3861            rope_dims,
3862            n_head,
3863            t,
3864            geometry.rope_base,
3865            1.0,
3866        )?;
3867        e.rope_neox(
3868            &mut k,
3869            pos_d,
3870            head_dim,
3871            rope_dims,
3872            n_head_kv,
3873            t,
3874            geometry.rope_base,
3875            1.0,
3876        )?;
3877
3878        // SDPA
3879        let mut attn = e.uninit(t * n_head * head_dim)?;
3880        // hand-written FlashAttention prefill (head_dim 256/128 stamped twins). MEMRA_NOFA
3881        // falls back to naive sdpa.
3882        if std::env::var("MEMRA_NOFA").is_ok() || !(head_dim == 256 || head_dim == 128) {
3883            // head_dim gate: see prime-path note (fa_prefill is stamped at 256 and 128 only).
3884            e.sdpa_naive(
3885                &q, &k, &v, &mut attn, head_dim, n_head, n_head_kv, t, t, scale, true,
3886            )?;
3887        } else {
3888            e.fa_prefill(
3889                &q, &k, &v, &mut attn, head_dim, n_head, n_head_kv, t, t, scale, true,
3890            )?;
3891        }
3892
3893        // output gate: attn * sigmoid(gate) — qwen35 only (M3 has no gate).
3894        let attn_g = match &gate {
3895            Some(gate) => {
3896                let mut gsig = e.uninit(t * n_head * head_dim)?;
3897                e.sigmoid(gate, &mut gsig, t * n_head * head_dim)?;
3898                let mut ag = e.uninit(t * n_head * head_dim)?;
3899                e.mul(&attn, &gsig, &mut ag, t * n_head * head_dim)?;
3900                ag
3901            }
3902            None => attn,
3903        };
3904
3905        // o projection
3906        let o = e.matmul(&fa.wo, &attn_g, t)?;
3907        Ok(o)
3908    }
3909
3910    /// Linear-attention (Gated DeltaNet) mixer (qwen35 :338-470).
3911    pub fn linear_attn(
3912        &self,
3913        e: &Engine,
3914        la: &LinearAttnLayer,
3915        h: &CudaSlice<f32>,
3916        t: usize,
3917    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3918        let cfg = &self.cfg;
3919        let _n_embd = cfg.n_embd as usize;
3920        let ssm = cfg.ssm.as_ref().unwrap();
3921        let d_state = ssm.state_size as usize; // 128
3922        let num_k = ssm.group_count as usize; // 16
3923        let num_v = ssm.time_step_rank as usize; // 32
3924        let d_conv = ssm.conv_kernel as usize; // 4
3925        let head_k = d_state;
3926        let head_v = d_state;
3927        let key_dim = head_k * num_k; // 2048
3928        let value_dim = head_v * num_v; // 4096
3929        let conv_dim = key_dim * 2 + value_dim; // 8192
3930        let eps = cfg.rms_eps;
3931        let scale = 1.0 / (d_state as f32).sqrt();
3932
3933        // projections
3934        // grouped: one f16 activation convert feeds all four projections (matmul_group)
3935        let mut g4 = e.matmul_group(
3936            &[&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha],
3937            h,
3938            t,
3939        )?;
3940        let alpha = g4.pop().unwrap(); // [T, num_v]
3941        let beta_raw = g4.pop().unwrap(); // [T, num_v]
3942        let z = g4.pop().unwrap(); // [T, value_dim]
3943        let qkv_mixed = g4.pop().unwrap(); // [T, conv_dim] token-major
3944
3945        // conv + GDN repack, FUSED (2026-07-03): ssm_conv1d_gdn reads qkv_mixed [T, conv_dim]
3946        // token-major DIRECTLY (causal window rows t-pad..t, rows<0 = zero prefill state), applies
3947        // the 8-tap conv + SiLU, and scatters straight into the GDN [d_state, num_v, T] q/k/v
3948        // layout with the modulo head-repeat. Replaces transpose + zeros + conv_left_pad +
3949        // ssm_conv1d + qkv_to_gdn_repack (5 launches, conv_in/conv_out scratch + a 16MB@T=512
3950        // round-trip). BIT-IDENTICAL accumulation and scatter mapping.
3951        let _ = (head_k, head_v);
3952        let mut q_g = e.uninit(d_state * num_v * t)?;
3953        let mut k_g = e.uninit(d_state * num_v * t)?;
3954        let mut v_g = e.uninit(d_state * num_v * t)?;
3955        e.ssm_conv1d_gdn(
3956            &qkv_mixed,
3957            la.ssm_conv1d.float_data(),
3958            &mut q_g,
3959            &mut k_g,
3960            &mut v_g,
3961            conv_dim,
3962            t,
3963            d_conv,
3964            d_state,
3965            num_v,
3966            num_k,
3967            key_dim,
3968        )?;
3969        // L2-norm q,k per (head_dim) row — rows are contiguous d_state in q_g.
3970        let mut q_l2 = e.uninit(d_state * num_v * t)?;
3971        e.l2_norm(&q_g, &mut q_l2, d_state, num_v * t, eps)?;
3972        let mut k_l2 = e.uninit(d_state * num_v * t)?;
3973        e.l2_norm(&k_g, &mut k_l2, d_state, num_v * t, eps)?;
3974        let v_gd = v_g;
3975
3976        // beta = sigmoid(beta_raw) ; g_log = a * softplus(alpha + dt). Both need [num_v, T] layout
3977        // (g[t*num_v + h]). beta_raw/alpha are [T, num_v] token-major == that layout already.
3978        let mut beta = e.uninit(t * num_v)?;
3979        e.sigmoid(&beta_raw, &mut beta, t * num_v)?;
3980        // gdn_glog expects alpha [H,T] with alpha[t*H+h] and dt_bias/a [H] — matches token-major [T,num_v].
3981        let mut g_log = e.uninit(t * num_v)?;
3982        e.gdn_glog(
3983            &alpha,
3984            la.ssm_dt.float_data(),
3985            la.ssm_a.float_data(),
3986            &mut g_log,
3987            num_v,
3988            t,
3989        )?;
3990
3991        // GDN scan (A4: gdn_scan_prefill dispatches chunked WY under MEMRA_GDN_CHUNKED)
3992        let state_in = e.zeros(d_state * d_state * num_v)?; // zero state (prefill)
3993        let mut state_out = e.zeros(d_state * d_state * num_v)?;
3994        let mut o = e.uninit(d_state * num_v * t)?;
3995        e.gdn_scan_prefill(
3996            &q_l2,
3997            &k_l2,
3998            &v_gd,
3999            &g_log,
4000            &beta,
4001            None,
4002            None,
4003            &state_in,
4004            &mut state_out,
4005            &mut o,
4006            num_v,
4007            t,
4008            scale,
4009            num_v,
4010        )?;
4011
4012        // gated RMSNorm: dst = RMSNorm(o, ssm_norm[head_v]) * silu(z). o is [d_state, num_v, T];
4013        // rows of head_v=d_state, nrows = num_v*T. z must match row layout: z is [T, value_dim] token-major
4014        // = [T, num_v*head_v]; per (t, vh) the head_v slice is contiguous -> rows align as (t*num_v+vh).
4015        // o rows are (t*num_v+vh) too. Good.
4016        let mut gn = e.uninit(d_state * num_v * t)?;
4017        e.gated_rmsnorm(
4018            &o,
4019            la.ssm_norm.float_data(),
4020            &z,
4021            &mut gn,
4022            d_state,
4023            num_v * t,
4024            eps,
4025        )?;
4026
4027        // ssm_out projection: gn is [d_state, num_v, T] = [value_dim, T] viewed token-major as [T, value_dim]?
4028        // gn layout: (t*num_v+vh)*d_state + i  == token t, then (vh,i) = channel vh*d_state+i. That's
4029        // token-major [T, value_dim]. linear wants [T, in=value_dim]. Good.
4030        let out = e.matmul(&la.ssm_out, &gn, t)?;
4031        Ok(out)
4032    }
4033}
4034
4035impl HybridModel {
4036    /// MoE FFN (EDGE-1). z: [T, n_embd] (already post-attention-normed). Returns moe_out [T, n_embd].
4037    /// Node-for-node vs llama.cpp build_moe_ffn + qwen35moe::build_layer_ffn.
4038    ///
4039    /// `il` is the trunk layer index — the residency-cache key prefix (a gate-expert of layer 3 is a
4040    /// different 860160-byte block than the same expert of layer 7).
4041    ///
4042    /// Routing: host softmax+sort (default) OR the fused router kernel (MEMRA_FUSED_ROUTER).
4043    /// Dispatch: stage-every-token into 3 scratch slots (default) OR the SLRU residency cache
4044    /// (MEMRA_MOE_CACHE). The cache-HIT weight path is bit-identical to stage-every-token (§B.3).
4045    /// Convenience wrapper used by the hybrid trunk/MTP loops: pulls dims + max-block from `self`.
4046    pub fn moe_ffn_il(
4047        &self,
4048        e: &Engine,
4049        m: &MoeWeights,
4050        z: &CudaSlice<f32>,
4051        t: usize,
4052        il: u16,
4053    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4054        Self::moe_ffn_inner(e, m, z, None, t, &self.cfg, il, self.max_moe_block(), false)
4055    }
4056
4057    /// Prefill twin: Step35 promotes expert-grouped dispatch by default while decode/spec callers
4058    /// keep `moe_ffn_il` and therefore retain their existing dispatch class.
4059    pub fn moe_ffn_il_prefill(
4060        &self,
4061        e: &Engine,
4062        m: &MoeWeights,
4063        z: &CudaSlice<f32>,
4064        t: usize,
4065        il: u16,
4066    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4067        Self::moe_ffn_inner(e, m, z, None, t, &self.cfg, il, self.max_moe_block(), true)
4068    }
4069
4070    /// Decode-path twin with a PRE-QUANTIZED z (from add_rms_norm_zq8): threads (zq, zd) into the
4071    /// t=1 dev arm so the per-layer standalone quantize_q8_1 launch folds away. Identical bytes
4072    /// (the fused kernel reproduces quantize_q8_1 exactly); every other path ignores the pair.
4073    pub fn moe_ffn_il_zq8(
4074        &self,
4075        e: &Engine,
4076        m: &MoeWeights,
4077        z: &CudaSlice<f32>,
4078        zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
4079        t: usize,
4080        il: u16,
4081    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4082        Self::moe_ffn_inner(e, m, z, zq8, t, &self.cfg, il, self.max_moe_block(), false)
4083    }
4084
4085    /// MoE FFN (EDGE-1), source-/model-agnostic. z: [T, n_embd] (already post-attention-normed).
4086    /// Returns moe_out [T, n_embd]. Node-for-node vs llama.cpp build_moe_ffn. Shared by the hybrid
4087    /// (qwen35moe, shared expert present) and the dense-attention MoE (OLMoE, no shared expert) paths;
4088    /// `cfg.moe` supplies the dims and the optional shexp fields decide whether step 3 runs.
4089    ///
4090    /// `il` is the layer index — the residency-cache key prefix. `max_block` is the global max expert
4091    /// stride (fixed cache-slot size); pass `self.max_moe_block()`.
4092    pub(crate) fn moe_ffn(
4093        e: &Engine,
4094        m: &MoeWeights,
4095        z: &CudaSlice<f32>,
4096        t: usize,
4097        cfg: &ModelConfig,
4098        il: u16,
4099        max_block: usize,
4100    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4101        Self::moe_ffn_inner(e, m, z, None, t, cfg, il, max_block, false)
4102    }
4103
4104    #[allow(clippy::too_many_arguments)]
4105    pub(crate) fn moe_ffn_inner(
4106        e: &Engine,
4107        m: &MoeWeights,
4108        z: &CudaSlice<f32>,
4109        zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
4110        t: usize,
4111        cfg: &ModelConfig,
4112        il: u16,
4113        max_block: usize,
4114        prefill: bool,
4115    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4116        let worker_io = crate::spill_pread::worker_enabled();
4117        let epoch_lfu = std::env::var_os("MEMRA_MOE_LFU_DECAY").is_some();
4118        if Engine::moe_cache_enabled() && (worker_io || epoch_lfu) {
4119            e.with_moe_cache(max_block, |cache, _| {
4120                cache.begin_forward_epoch(il, t);
4121                if worker_io {
4122                    cache.begin_worker_scope();
4123                }
4124                Ok(())
4125            })?;
4126        }
4127        if Self::sigmoid_resident_dev_eligible(e, m, cfg) {
4128            let moe = cfg.moe.as_ref().unwrap();
4129            let n_expert = moe.expert_count as usize;
4130            let n_used = moe.expert_used_count as usize;
4131            let sigmoid = cfg.sigmoid_router().unwrap();
4132            let logits = Self::moe_router_logits(e, m, z, t, cfg)?;
4133            Self::trace_sigmoid_router_logits(e, il, t, n_expert, n_used, &logits, m, sigmoid)?;
4134            return Self::moe_ffn_sigmoid_dev(e, m, z, zq8, &logits, t, cfg, il, sigmoid);
4135        }
4136        // Expert-grouped dispatch for prefill (T>1). MEMRA_MOE_GROUPED explicitly opts the
4137        // current caller into this research arm; the naked default stays on the established path.
4138        if t > 1 && moe_grouped_enabled(cfg, prefill) {
4139            let grouped_out = Self::moe_ffn_grouped(e, m, z, t, cfg, il, max_block)?;
4140            // MEMRA_MOE_GATE: byte-identity comparison vs sequential path.
4141            // The grouped q8 path uses the same row-wise quantize/dot/activation/FMA programs as
4142            // sequential dispatch; mixed or q8-disabled layouts fall back to the matching f32
4143            // path. A mismatch is therefore a correctness failure, not an accepted numeric class.
4144            if std::env::var("MEMRA_MOE_GATE").is_ok() {
4145                let seq_out = Self::moe_ffn_sequential(e, m, z, t, cfg, il, max_block)?;
4146                let g_host = e.dtoh(&grouped_out)?;
4147                let s_host = e.dtoh(&seq_out)?;
4148                let g_bytes: &[u8] = unsafe {
4149                    std::slice::from_raw_parts(g_host.as_ptr() as *const u8, g_host.len() * 4)
4150                };
4151                let s_bytes: &[u8] = unsafe {
4152                    std::slice::from_raw_parts(s_host.as_ptr() as *const u8, s_host.len() * 4)
4153                };
4154                if g_bytes == s_bytes {
4155                    println!("moe-gate il={il} t={t} BYTE-IDENTICAL");
4156                } else {
4157                    let diffs = g_host
4158                        .iter()
4159                        .zip(s_host.iter())
4160                        .enumerate()
4161                        .filter(|(_, (a, b))| a != b)
4162                        .count();
4163                    let maxdiff = g_host
4164                        .iter()
4165                        .zip(s_host.iter())
4166                        .map(|(a, b)| (a - b).abs())
4167                        .fold(0.0f32, f32::max);
4168                    panic!(
4169                        "moe-gate il={il} t={t} MISMATCH: {diffs}/{} elems differ, maxdiff={maxdiff:.6e}",
4170                        g_host.len()
4171                    );
4172                }
4173            }
4174            return Ok(grouped_out);
4175        }
4176        Self::moe_ffn_sequential_zq8(e, m, z, zq8, t, cfg, il, max_block)
4177    }
4178
4179    fn sigmoid_resident_dev_eligible(e: &Engine, m: &MoeWeights, cfg: &ModelConfig) -> bool {
4180        let Some(moe) = cfg.moe.as_ref() else {
4181            return false;
4182        };
4183        // Cached once per process: this predicate runs per MoE layer per decode step, and five
4184        // uncached env reads were ~300+ environ scans/step on the exact path this arm de-hosts.
4185        static OBSERVATION_MODE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4186        let observation_mode = *OBSERVATION_MODE.get_or_init(|| {
4187            std::env::var("MEMRA_MOE_STATS").is_ok()
4188                || std::env::var("MEMRA_MOE_TRACE").is_ok()
4189                || std::env::var("MEMRA_MOE_WEIGHT_TRACE").is_ok()
4190                || std::env::var("MEMRA_MOE_INPUT_TRACE_DIR").is_ok()
4191                || std::env::var("MEMRA_MOE_GATE").is_ok()
4192        });
4193        cfg.step35.is_some()
4194            && sigmoid_router_enabled()
4195            && moe_dev_enabled()
4196            && moe_slab_enabled()
4197            && !observation_mode
4198            && moe.expert_used_count <= 8
4199            && m.has_uniform_expert_layout()
4200            && m.gate_exps.macros.is_none()
4201            && m.up_exps.macros.is_none()
4202            && m.down_exps.macros.is_none()
4203            && !m.has_macros
4204            && moe_q8_enabled()
4205            && q8_expert_supported(m.gate_exps.qtype)
4206            && q8_expert_supported(m.up_exps.qtype)
4207            && q8_expert_supported(m.down_exps.qtype)
4208            && m.dev_exps
4209                .as_ref()
4210                .is_some_and(|dev| dev.dev == e.ctx().ordinal())
4211    }
4212
4213    /// Sequential (per-token) MoE FFN -- the original path. Factored out for the gate comparison.
4214    pub(crate) fn moe_ffn_sequential(
4215        e: &Engine,
4216        m: &MoeWeights,
4217        z: &CudaSlice<f32>,
4218        t: usize,
4219        cfg: &ModelConfig,
4220        il: u16,
4221        max_block: usize,
4222    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4223        Self::moe_ffn_sequential_zq8(e, m, z, None, t, cfg, il, max_block)
4224    }
4225
4226    /// One router-logit selector shared by sequential and grouped dispatch. Real prefill must use
4227    /// the row-wise GEMV when exact routing is enabled: cuBLASLt's reduction changes with `m`,
4228    /// which makes expert selection depend on the caller's chunk or concat-batch shape.
4229    fn moe_router_logits(
4230        e: &Engine,
4231        m: &MoeWeights,
4232        z: &CudaSlice<f32>,
4233        t: usize,
4234        cfg: &ModelConfig,
4235    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4236        if t < PRIME_MIN_T {
4237            // Decode and speculative verify use one fixed per-row reduction program.
4238            if crate::router_kernel_on() {
4239                e.router_gemv(
4240                    m.gate_inp.float_data(),
4241                    z,
4242                    cfg.n_embd as usize,
4243                    m.gate_exps.n_expert,
4244                    t,
4245                )
4246            } else {
4247                e.matmul_decode_exact(&m.gate_inp, z, t)
4248            }
4249        } else if crate::router_prefill_exact_on() && crate::router_kernel_on() {
4250            e.router_gemv(
4251                m.gate_inp.float_data(),
4252                z,
4253                cfg.n_embd as usize,
4254                m.gate_exps.n_expert,
4255                t,
4256            )
4257        } else {
4258            e.matmul(&m.gate_inp, z, t)
4259        }
4260    }
4261
4262    /// Append the host-visible router selection for one layer/forward when calibration tracing is
4263    /// enabled. Both sequential and expert-grouped prefill must call this after routing so the
4264    /// trace is independent of the dispatch optimization selected for the forward.
4265    fn trace_moe_routes(
4266        il: u16,
4267        t: usize,
4268        sel_all: &[u32],
4269        weights: &[f32],
4270    ) -> Result<(), Box<dyn std::error::Error>> {
4271        use std::io::Write as _;
4272        if let Ok(path) = std::env::var("MEMRA_MOE_TRACE") {
4273            let mut f = std::fs::OpenOptions::new()
4274                .create(true)
4275                .append(true)
4276                .open(path)?;
4277            let ids: Vec<String> = sel_all.iter().map(|s| s.to_string()).collect();
4278            writeln!(f, "{} {} {}", il, t, ids.join(","))?;
4279        }
4280        if let Ok(path) = std::env::var("MEMRA_MOE_WEIGHT_TRACE") {
4281            let mut f = std::fs::OpenOptions::new()
4282                .create(true)
4283                .append(true)
4284                .open(path)?;
4285            let pairs: Vec<String> = sel_all
4286                .iter()
4287                .zip(weights)
4288                .map(|(expert, weight)| format!("{expert}:{weight:.9}"))
4289                .collect();
4290            writeln!(f, "{} {} {}", il, t, pairs.join(","))?;
4291        }
4292        Ok(())
4293    }
4294
4295    #[allow(clippy::too_many_arguments)]
4296    fn trace_sigmoid_router_logits(
4297        e: &Engine,
4298        il: u16,
4299        t: usize,
4300        n_expert: usize,
4301        n_used: usize,
4302        logits: &CudaSlice<f32>,
4303        m: &MoeWeights,
4304        (scaling_factor, route_norm): (f32, bool),
4305    ) -> Result<(), Box<dyn std::error::Error>> {
4306        if !crate::sigrouter_contract::served_logit_trace_enabled() || t != 1 {
4307            return Ok(());
4308        }
4309        let logits = e.dtoh(logits)?;
4310        let active: Vec<u8> = m
4311            .active_experts
4312            .as_ref()
4313            .map(|mask| mask.iter().map(|&enabled| u8::from(enabled)).collect())
4314            .unwrap_or_else(|| vec![1; n_expert]);
4315        let bias = m.exp_probs_b.clone().unwrap_or_else(|| vec![0.0; n_expert]);
4316        crate::sigrouter_contract::capture_served_logits(
4317            il as u32,
4318            t,
4319            n_expert,
4320            n_used,
4321            scaling_factor,
4322            route_norm,
4323            &active,
4324            &bias,
4325            &logits,
4326        )?;
4327        Ok(())
4328    }
4329
4330    /// Append the f32 input to one MoE layer for offline, layerwise calibration. This diagnostic
4331    /// intentionally performs a DtoH copy and is therefore disabled unless an explicit fresh trace
4332    /// directory is supplied. Each layer owns one payload file; index.jsonl records byte offsets so
4333    /// a validator can prove request/layer coverage before the trace is used for pruning or healing.
4334    fn trace_moe_input(
4335        e: &Engine,
4336        il: u16,
4337        t: usize,
4338        n_embd: usize,
4339        z: &CudaSlice<f32>,
4340    ) -> Result<(), Box<dyn std::error::Error>> {
4341        use std::io::Write as _;
4342        let Ok(dir) = std::env::var("MEMRA_MOE_INPUT_TRACE_DIR") else {
4343            return Ok(());
4344        };
4345        let host = e.dtoh(z)?;
4346        if host.len() != t * n_embd {
4347            return Err(format!(
4348                "MoE input trace shape mismatch at layer {il}: got {} values, expected {}x{}",
4349                host.len(),
4350                t,
4351                n_embd
4352            )
4353            .into());
4354        }
4355        let bytes = unsafe {
4356            std::slice::from_raw_parts(
4357                host.as_ptr().cast::<u8>(),
4358                host.len() * std::mem::size_of::<f32>(),
4359            )
4360        };
4361        let state = MOE_INPUT_TRACE_WRITER.get_or_init(|| std::sync::Mutex::new(None));
4362        let mut state = state
4363            .lock()
4364            .map_err(|_| "MoE input trace writer lock is poisoned")?;
4365        if state.is_none() {
4366            let dir = std::path::PathBuf::from(&dir);
4367            std::fs::create_dir_all(&dir)?;
4368            let index = std::fs::OpenOptions::new()
4369                .create(true)
4370                .append(true)
4371                .open(dir.join("index.jsonl"))?;
4372            *state = Some(MoeInputTraceWriter {
4373                dir,
4374                index,
4375                payloads: std::collections::HashMap::new(),
4376            });
4377        }
4378        let writer = state.as_mut().unwrap();
4379        if writer.dir != std::path::Path::new(&dir) {
4380            return Err("MEMRA_MOE_INPUT_TRACE_DIR changed after capture started".into());
4381        }
4382        let file_name = format!("layer-{il:03}.f32");
4383        if !writer.payloads.contains_key(&il) {
4384            let payload = std::fs::OpenOptions::new()
4385                .create(true)
4386                .append(true)
4387                .open(writer.dir.join(&file_name))?;
4388            let offset = payload.metadata()?.len();
4389            writer.payloads.insert(il, (payload, offset));
4390        }
4391        let (payload, offset) = writer.payloads.get_mut(&il).unwrap();
4392        let row_offset = *offset;
4393        payload.write_all(bytes)?;
4394        *offset += bytes.len() as u64;
4395        writeln!(
4396            writer.index,
4397            "{{\"format\":\"memra-moe-input-trace-v1\",\"layer\":{il},\"tokens\":{t},\
4398             \"hidden_size\":{n_embd},\"file\":\"{file_name}\",\"offset\":{row_offset},\
4399             \"payload_bytes\":{}}}",
4400            bytes.len()
4401        )?;
4402        Ok(())
4403    }
4404
4405    #[allow(clippy::too_many_arguments)]
4406    pub(crate) fn moe_ffn_sequential_zq8(
4407        e: &Engine,
4408        m: &MoeWeights,
4409        z: &CudaSlice<f32>,
4410        zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
4411        t: usize,
4412        cfg: &ModelConfig,
4413        il: u16,
4414        max_block: usize,
4415    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4416        use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
4417        let moe = cfg.moe.as_ref().unwrap();
4418        let n_embd = cfg.n_embd as usize; // 2048 (gate/up in_f, down out_f)
4419        let n_expert = moe.expert_count as usize; // 256
4420        let n_used = moe.expert_used_count as usize; // 8
4421        let n_ff_exp = moe.expert_ff_length as usize; // 512 (gate/up out_f, down in_f)
4422
4423        // verify the HostExps dims match cfg (catches a wrong-file / transpose mixup)
4424        debug_assert_eq!(m.gate_exps.in_f, n_embd);
4425        debug_assert_eq!(m.gate_exps.out_f, n_ff_exp);
4426        debug_assert_eq!(m.down_exps.in_f, n_ff_exp); // down is TRANSPOSED: in=512
4427        debug_assert_eq!(m.down_exps.out_f, n_embd); //                     out=2048
4428        debug_assert_eq!(m.gate_exps.n_expert, n_expert);
4429
4430        // step35 PER-LAYER SwiGLU clamp (None on every other arch and on every unclamped layer).
4431        // Routed experts and the shared expert read SEPARATE arrays — never share one value.
4432        let lim_exp = cfg.clamp_exp_at(il as u32);
4433        let lim_shexp = cfg.clamp_shexp_at(il as u32);
4434        let use_cache = Engine::moe_cache_enabled();
4435        let uniform_experts = m.has_uniform_expert_layout();
4436        let moe_q8 = uniform_experts
4437            && moe_q8_enabled()
4438            && q8_expert_supported(m.gate_exps.qtype)
4439            && q8_expert_supported(m.up_exps.qtype)
4440            && q8_expert_supported(m.down_exps.qtype);
4441        // Experimental secondary backend: complete experts already resident in the SLRU stay on
4442        // CUDA; any expert missing one or more projections runs from the original host GGUF bytes
4443        // through llama.cpp CPU quant dots. Small-t only keeps decode and speculative verification
4444        // in the same numeric/dispatch class; real prefill remains on the established GPU path and
4445        // seeds the residency cache. An explicit library path is the build/runtime gate, so naked
4446        // commands and CI have no llama.cpp or OpenMP dependency.
4447        let cpu_expert_requested = crate::cpu_experts::configured();
4448        if cpu_expert_requested && (cfg.hy3.is_none() || cfg.m3.is_some()) {
4449            return Err(std::io::Error::other(
4450                "MEMRA_CPU_EXPERT_LIB is experimental and currently gated to Hy3",
4451            )
4452            .into());
4453        }
4454        let cpu_hybrid = cpu_expert_requested && t < PRIME_MIN_T && m.dev_exps.is_none();
4455        // A dynamic SLRU changes which complete experts run on CUDA versus llama.cpp CPU dots.
4456        // Those backends are each deterministic but are different numeric configurations, so a
4457        // later prefill eviction can change greedy output. Freeze after the first real prefill;
4458        // decode/spec reads the fixed resident set, while later prefill misses use transient GPU
4459        // staging below and cannot change backend assignment.
4460        let freeze_cpu_residency = cpu_expert_requested
4461            && std::env::var("MEMRA_CPU_EXPERT_FREEZE_CACHE").as_deref() == Ok("1");
4462        let caller_warms_before_freeze = std::env::var("MEMRA_CPU_EXPERT_FREEZE_WARMUP_TOKENS")
4463            .ok()
4464            .and_then(|value| value.parse::<usize>().ok())
4465            .is_some_and(|tokens| tokens > 0);
4466        if cpu_hybrid && freeze_cpu_residency && !caller_warms_before_freeze {
4467            e.freeze_moe_cache();
4468        }
4469        let cache_frozen = use_cache && e.moe_cache_frozen();
4470        let cache_dispatch = use_cache && (!cache_frozen || cpu_hybrid);
4471
4472        // 1. ROUTER: one selector is shared with expert-grouped prefill so changing dispatch
4473        // cannot change logits, selected expert ids, or routing weights.
4474        let logits = Self::moe_router_logits(e, m, z, t, cfg)?;
4475        if let Some(sig) = cfg.sigmoid_router() {
4476            Self::trace_sigmoid_router_logits(e, il, t, n_expert, n_used, &logits, m, sig)?;
4477        }
4478
4479        // LAUNCH-STRUCTURE STAGE 3 (2026-07-05, MEMRA_MOE_DEV default ON, =0 rollback): ZERO-DtoH
4480        // device-dispatch when this layer's expert blocks are ALL cache-resident. The fused
4481        // router's sel/w stay ON DEVICE; the expert weight pointers come from the per-layer
4482        // device table of fixed slot addresses; gate/up/silu + down/fma run as the same TWO
4483        // launches per token as gdec. Removes the per-layer router DtoH + stream sync — the
4484        // per-token host stall that dominated the 35B decode wall after stages 1+2.
4485        // BIT-IDENTITY: the router kernel is selection-exact vs the host oracle (kernel-check
4486        // tie gate) and the _dev matvec twins reproduce the gdec kernels' exact FP chains; the
4487        // only difference is where sel/w/pointers are READ from (device instead of params).
4488        // Residency: one-shot PREWARM force-admits the layer while free slots cover it
4489        // (MEMRA_MOE_PREWARM=0 -> organic residency, dev path fires when the SLRU fills).
4490        // Any non-resident layer falls through to host routing + the gdec/sequential path.
4491        // FITS-VRAM RESIDENT EXPERTS (2026-07-06): the layer's expert slabs are device-resident
4492        // (load-time decision) — fire the zero-DtoH dev path unconditionally with the prebuilt
4493        // pointer row. No cache, no dispatch, no residency check: the llama full-offload regime
4494        // (it measured 169.55 vs the cache path's 28.5 on the local 35B — the residency-gate
4495        // all-or-nothing fallback was the 6x). BIT-IDENTITY: same _dev kernels, same math; only
4496        // the pointer table's provenance differs (slab base+stride vs SLRU slot addresses).
4497        // MoE PREFILL PAIR-BATCH (2026-07-06, the 16x pp hole): t>1 on resident experts — ONE
4498        // launch per proj covers ALL (token,expert) pairs (grid.y=pair, warp-per-row), replacing
4499        // the per-expert loop (256 experts x 3-4 launches x tiny m_e). Scatter is slot-ordered
4500        // per token (the sequential-axpy bit-identity class). Requires q8-supported qtypes +
4501        // resident slabs. MEMRA_MOE_PAIRS=0 rollback.
4502        // t >= PRIME_MIN_T only (2026-07-06 exactness fix, verify-probe proof): spec VERIFY
4503        // batches (t = 2..K+2) previously rode these pairs kernels while T=1 decode rode the
4504        // dev_q8 loop — different FP chains, verify-T2 logit maxdiff 2.6e-1 vs eager -> greedy
4505        // flips at tight margins -> 35B real-prompt spec self-consistency FAIL (the 27B
4506        // "verify must be kernel-DISPATCH-identical to decode" lesson, MoE edition). Small-t
4507        // now rides the dev loop below (same kernels per token as decode); pairs serves real
4508        // prefill (t >= 16, where spec never verifies).
4509        // sigmoid-router archs must NOT enter the pairs/dev arms: those dispatch arms still route
4510        // via the fused SOFTMAX device router (moe_router_topk) — silently wrong experts (the
4511        // M3 gate-MISMATCH 74602-vs-92 lesson, 2026-07-07). Sigmoid routing has its own device
4512        // top-k kernel; MEMRA_SIG_ROUTER=0 selects the host sigmoid oracle instead.
4513        // Per-expert macro-scales (compressed-tensors NVFP4 ST class, e.g. unsloth 35B-A3B):
4514        // the device-dispatch kernels (pairs/dev) do NOT fold them — those checkpoints must
4515        // ride the macro-aware sequential/staged paths below or every expert output is off by
4516        // its global scale (~3e4x, measured garbage 2026-07-16). GGUF experts: macros None.
4517        let no_exp_macros = m.gate_exps.macros.is_none()
4518            && m.up_exps.macros.is_none()
4519            && m.down_exps.macros.is_none();
4520        // A clamped-SwiGLU layer (step35 43/44) also cannot ride pairs: moe_pairs_silu_mul's
4521        // epilogue is plain silu(gate)*up with no clamp form, and moe_ffn_pairs takes no `il`
4522        // so it cannot even see the per-layer limit.
4523        if cfg.sigmoid_router().is_none()
4524            && cfg.m3.is_none()
4525            && cfg.hy3.is_none()
4526            && !cfg.swiglu_clamped_at(il as u32)
4527            && no_exp_macros
4528            && t >= PRIME_MIN_T
4529            && m.dev_exps.is_some()
4530            && moe_q8_enabled()
4531            && q8_expert_supported(m.gate_exps.qtype)
4532            && q8_expert_supported(m.up_exps.qtype)
4533            && q8_expert_supported(m.down_exps.qtype)
4534            && std::env::var("MEMRA_MOE_PAIRS")
4535                .map(|v| v != "0")
4536                .unwrap_or(true)
4537            && std::env::var("MEMRA_MOE_STATS").is_err()
4538        {
4539            return Self::moe_ffn_pairs(e, m, z, &logits, t, cfg);
4540        }
4541
4542        // t < PRIME_MIN_T: moe_ffn_dev loops tokens serially (1 launch-pair per token) — the
4543        // decode path AND the spec-verify path (dispatch parity = exactness; see pairs gate
4544        // above). Serial launches are fine at t<=10 (K+2); real prefill never lands here.
4545        // moe_ffn_dev routes via the FUSED SOFTMAX device router (moe_router_topk). Sigmoid
4546        // routing has a separate device top-k kernel, but this dev arm cannot consume its sel/w
4547        // contract, so sigmoid-router arches must NOT enter it: with MOE_CACHE=1 M3 silently
4548        // routed softmax = wrong experts (gate MISMATCH 74602 vs 92, caught 2026-07-07).
4549        // MEMRA_SIG_ROUTER=0 selects the host sigmoid oracle for the supported sigmoid path.
4550        // macro-carrying experts are handled inside the dev path now (epilogue fold + w-scale);
4551        // pairs/gdec/csr keep their macro gates until their kernels grow the fold.
4552        // step35 (2026-08-06) is the third sigmoid-router arch and the deny above was written as
4553        // two arch names, not as the mechanism — so `dev_ok` let it into moe_ffn_dev's
4554        // moe_router_topk (softmax, no exp_probs_b bias, no expert_weights_scale) = silently
4555        // wrong experts, the same failure the pairs gate at :2254 already blocks by predicate.
4556        // Keyed off sigmoid_router() so arch #4 is denied by construction.
4557        // The clamped-SwiGLU layers must also fall through: every moe_ffn_dev kernel's fused
4558        // epilogue is plain silu(gate)*up (moe_gate_up_silu8_dev_*), which has no clamp form.
4559        let dev_ok = uniform_experts
4560            && cfg.sigmoid_router().is_none()
4561            && cfg.m3.is_none()
4562            && cfg.hy3.is_none()
4563            && !cfg.swiglu_clamped_at(il as u32);
4564        // Observation modes must route through the host-visible selection below. Otherwise a fully
4565        // resident layer returns through device dispatch before its trace/stats row is recorded,
4566        // silently biasing calibration toward only non-resident layers on large-VRAM machines.
4567        let observe_routes = std::env::var("MEMRA_MOE_STATS").is_ok()
4568            || std::env::var("MEMRA_MOE_TRACE").is_ok()
4569            || std::env::var("MEMRA_MOE_WEIGHT_TRACE").is_ok()
4570            || std::env::var("MEMRA_MOE_INPUT_TRACE_DIR").is_ok();
4571        if dev_ok
4572            && t < PRIME_MIN_T
4573            && m.dev_exps.is_some()
4574            && n_used <= 8
4575            && moe_dev_enabled()
4576            && !observe_routes
4577        {
4578            return Self::moe_ffn_dev(e, m, z, zq8, &logits, t, cfg, il, max_block);
4579        }
4580        if dev_ok && use_cache && n_used <= 8 && moe_dev_enabled() && !observe_routes {
4581            let row_ok = e.with_moe_cache(max_block, |c, eng| {
4582                if moe_prewarm_enabled() {
4583                    c.prewarm_layer(il, m, eng)?;
4584                }
4585                Ok(c.layer_dev_row(il, n_expert, eng)?.is_some())
4586            })?;
4587            if row_ok {
4588                return Self::moe_ffn_dev(e, m, z, zq8, &logits, t, cfg, il, max_block);
4589            }
4590        }
4591
4592        // Per-token (sel[8], w[8]) — sigmoid host oracle (M3/Hy3), else fused-router/softmax.
4593        let (sel_all, w_all, routed_cpu_input) = if let Some(sig) = cfg.sigmoid_router() {
4594            if cpu_hybrid {
4595                let (sel, w, input) = Self::moe_route_sigmoid_with_input(
4596                    e,
4597                    &logits,
4598                    z,
4599                    t,
4600                    n_expert,
4601                    n_used,
4602                    m.exp_probs_b.as_deref(),
4603                    sig,
4604                    m.active_experts.as_deref(),
4605                )?;
4606                (sel, w, Some(input))
4607            } else {
4608                let (sel, w) =
4609                    Self::moe_route_sigmoid_cfg(e, &logits, t, n_expert, n_used, m, sig)?;
4610                (sel, w, None)
4611            }
4612        } else {
4613            let (sel, w) =
4614                Self::moe_route_cfg(e, &logits, t, n_expert, n_used, m.active_experts.as_deref())?;
4615            (sel, w, None)
4616        };
4617        crate::moesd::record_host_routes(il, n_expert, n_used, &sel_all)?;
4618
4619        // MEMRA_MOE_TRACE=<path>: append one line per (layer, step) with the selected expert ids —
4620        // offline analysis derives the decode working set + step-to-step reuse (the go/no-go
4621        // measurement for resident-expert tiering; see rig5090.jsonl 2026-07-07 pinned-tier row).
4622        Self::trace_moe_routes(il, t, &sel_all, &w_all)?;
4623        Self::trace_moe_input(e, il, t, n_embd, z)?;
4624
4625        // Worker reads ordinarily overlap NVMe with the current expert but leave their H2D copies
4626        // demand-serialized on the compute stream. Hy3 decode already has a safe owner-thread
4627        // boundary here: sigmoid routing read `logits` back to the host, which synchronized all
4628        // earlier-layer compute. Submit the complete selected set first, then reserve only victims
4629        // outside that set and queue its H2D copies on the copy stream. Dispatch inserts an event
4630        // wait for each pending block, so later copies can overlap the earlier expert kernels while
4631        // the selected weights and numerical kernels remain unchanged. Restrict this experiment to
4632        // T=1; batched forwards can have token-local consumers still in flight between selections.
4633        // The CPU backend reads mmap-backed model bytes directly. Positioned-read worker buffers
4634        // are CUDA write-combined staging allocations and are intentionally not CPU-compute inputs;
4635        // do not spend NVMe bandwidth filling them for a layer whose misses go to CPU.
4636        let worker_disk_prefetch =
4637            cache_dispatch && crate::spill_pread::worker_enabled() && !cpu_hybrid;
4638        let promote_worker_h2d =
4639            t == 1 && worker_disk_prefetch && crate::spill_pread::copy_h2d_enabled();
4640        if promote_worker_h2d {
4641            let mut selected_blocks = Vec::with_capacity(n_used * 3);
4642            for &ex in sel_all.iter().take(n_used) {
4643                let ex = ex as u16;
4644                selected_blocks.extend([
4645                    BlockId::new(il, PROJ_GATE, ex),
4646                    BlockId::new(il, PROJ_UP, ex),
4647                    BlockId::new(il, PROJ_DOWN, ex),
4648                ]);
4649            }
4650            for &ex in sel_all.iter().take(n_used) {
4651                Self::moe_prefetch_disk_expert(e, il, ex as usize, m, max_block, &selected_blocks)?;
4652            }
4653            e.with_moe_cache(max_block, |cache, eng| {
4654                cache.promote_worker_reads_at_safe_boundary(
4655                    &selected_blocks,
4656                    &selected_blocks,
4657                    eng,
4658                )?;
4659                Ok(())
4660            })?;
4661        }
4662
4663        // MEMRA_MOE_STATS: per-layer routing stats for the A2 (expert-grouped prefill) baseline —
4664        // per-token expert-id entropy, active-expert coverage, tokens-per-expert group sizes.
4665        if t > 1 && std::env::var("MEMRA_MOE_STATS").is_ok() {
4666            let mut cnt = vec![0u32; n_expert];
4667            for &s in sel_all.iter() {
4668                cnt[s as usize] += 1;
4669            }
4670            let total = sel_all.len() as f64;
4671            let mut h = 0.0f64;
4672            let mut active = 0usize;
4673            for &c in &cnt {
4674                if c > 0 {
4675                    active += 1;
4676                    let p = c as f64 / total;
4677                    h -= p * p.log2();
4678                }
4679            }
4680            let maxc = cnt.iter().copied().max().unwrap_or(0);
4681            println!(
4682                "moe-stats il={} t={} assignments={} active={}/{} entropy={:.3}b (max {:.3}b) mean_tok_per_active={:.2} max_tok_per_expert={}",
4683                il,
4684                t,
4685                sel_all.len(),
4686                active,
4687                n_expert,
4688                h,
4689                (n_expert as f64).log2(),
4690                total / active.max(1) as f64,
4691                maxc
4692            );
4693        }
4694
4695        // LAUNCH-STRUCTURE STAGE 2 (2026-07-05): moe_out memset elision on the gdec path.
4696        // moe_down8_fma_f32 FULLY overwrites its token row (dst[o] = the in-kernel FMA chain that
4697        // starts at 0.0f — numerically the axpy-into-zeroed-row chain), so when the grouped-decode
4698        // path fires the upfront `e.zeros(t*n_embd)` memset is pure launch churn. Allocate uninit
4699        // when gdec CAN fire (any t — decode t=1 AND the spec verify t=K+1 route here per token)
4700        // and lazily zero ONLY the row of a token that falls through to the sequential axpy loop.
4701        // BIT-IDENTITY: unchanged — every row is either fully overwritten (gdec) or
4702        // zeroed-then-accumulated exactly as before (fallback).
4703        // `!swiglu_clamped_at`: the grouped-decode kernels' fused epilogue is plain
4704        // silu(gate)*up (see the m3 note at the call sites below) — a clamped layer must fall
4705        // through to the sequential loop's ffn_act_lim. Hoisted into the fire predicate so the
4706        // uninit/memset invariant at :2418 stays consistent with what actually dispatches.
4707        let gdec_may_fire = uniform_experts
4708            && use_cache
4709            && n_used <= 8
4710            && gdec_enabled()
4711            && !cfg.swiglu_clamped_at(il as u32);
4712        // SLAB-LOCAL RESIDENT ARM (lane/pp-leverb 2026-08-08): the fits-VRAM resident slabs
4713        // (`dev_exps`, sized per PP device by cx-503b) were consumed ONLY by the pairs/dev
4714        // arms — every one of which DENIES sigmoid-router archs (step35/M3/Hy3). So on those
4715        // archs the slabs were uploaded but never read, and every expert went through the
4716        // SLRU (37 GB H2D staging per pp4096 prime on the Step SKU; and post-cx-503b the
4717        // dead slabs additionally STARVE the SLRU, which sizes itself on free-VRAM-after-
4718        // residents). This arm reads the slabs through the SAME kernels the SLRU arm runs —
4719        // gdec's fused pair when eligible, the per-expert qmatvec twins otherwise — with the
4720        // slab base + ex*stride as the pointer provenance (the exact bit-identity class
4721        // `moe_ffn_dev` documents for its resident-vs-SLRU arms; bytes are the same HostExps
4722        // bytes either way). LOCALITY GATE: `d.dev == e.ctx().ordinal()` — a slab on another
4723        // device must NOT be dereferenced (m=1 peer reads are the measured 34-150x class,
4724        // strictly worse than staging); under PP-2 without the prime walker this admits
4725        // stage-0 layers on dev0 and leaves stage-1 layers on the SLRU, which is the honest
4726        // interim shape. MEMRA_MOE_SLAB=0 = rollback/A-B seam (read per call).
4727        let slab_local = m
4728            .dev_exps
4729            .as_ref()
4730            .filter(|d| !d.gu_il && moe_slab_enabled() && d.dev == e.ctx().ordinal());
4731        let slab_bases = slab_local.map(|d| {
4732            use cudarc::driver::DevicePtr;
4733            let s = e.stream();
4734            let (pg, _g0) = d.gate.device_ptr(&s);
4735            let (pu, _g1) = d.up.device_ptr(&s);
4736            let (pd, _g2) = d.down.device_ptr(&s);
4737            (pg as u64, pu as u64, pd as u64)
4738        });
4739        // Fused-pair eligibility mirrors the gdec call sites exactly (plain-SiLU epilogue,
4740        // no macros, m3 clamp excluded, q8-supported qtypes, <=8 experts) — INCLUDING
4741        // `gdec_enabled()`: the pair IS the gdec kernel pair, so MEMRA_MOE_GDEC=0 must
4742        // disable it here too. That seam is also the exactness localizer: with GDEC=0 both
4743        // provenances run the SAME per-expert qmatvec kernels (slab base+stride vs SLRU
4744        // slot) and must be BIT-IDENTICAL — the true provenance-only pair — while GDEC=1
4745        // compares the fused-pair class against the SLRU's hit/miss MIX (gdec for
4746        // all-resident tokens, staged loop for misses), which is a dispatch-class
4747        // comparison, not a provenance one.
4748        let slab_fused_may_fire = slab_bases.is_some()
4749            && n_used <= 8
4750            && gdec_enabled()
4751            && !cfg.swiglu_clamped_at(il as u32)
4752            && cfg.m3.is_none()
4753            && no_exp_macros
4754            && moe_q8;
4755        // moe_out memset elision: BOTH full-row-overwrite arms (gdec + slab fused) allocate
4756        // uninit; a token that falls through to any accumulating loop zeroes its own row.
4757        let mut moe_out = if gdec_may_fire || slab_fused_may_fire {
4758            e.uninit(t * n_embd)?
4759        } else {
4760            e.zeros(t * n_embd)?
4761        };
4762        // The router readback above already established a host boundary. Copy each small-t hidden
4763        // row once so CPU miss experts can start while the owner thread queues resident GPU work.
4764        let cpu_input = if cpu_hybrid {
4765            Some(routed_cpu_input.ok_or("CPU expert routing did not return the MoE input")?)
4766        } else {
4767            None
4768        };
4769
4770        // GPU scratch: one slot per proj, big enough for ONE expert (default stage-every-token path).
4771        // STAGE 2: LAZY — allocated only if the no-cache staging path actually runs (under
4772        // MEMRA_MOE_CACHE they were 3 dead ~1MB alloc_zeros + memset + free per layer per token,
4773        // measured ~123 memsets/token of the decode wall).
4774        let g_len = m.gate_exps.max_expert_bytes(); // 860160 for the uniform 35B gate
4775        let u_len = m.up_exps.max_expert_bytes(); // 860160 for the uniform 35B up
4776        let d_len = m.down_exps.max_expert_bytes(); // 1114112 for the uniform 35B down
4777        let mut scratch_g: Option<CudaSlice<u8>> = None;
4778        let mut scratch_u: Option<CudaSlice<u8>> = None;
4779        let mut scratch_d: Option<CudaSlice<u8>> = None;
4780        // `max_block` (the GLOBAL max expert stride across all layers) is passed in — the cache slots
4781        // are FIXED-ADDRESS and must fit any layer's block (UD/dynamic GGUFs vary quant per layer).
4782
4783        // EDGE-1 §C.2/C.3: the optional pipeline queues the next selected expert's MISS blocks on
4784        // the copy stream before launching the current expert's compute. Pending slots stay invisible
4785        // to cache hits until dispatch inserts the completion-event wait; current gate/up/down ids are
4786        // protected from eviction. This changes scheduling only, never the GGUF bytes or GEMM path.
4787        let page_window = moe_page_prefetch_window();
4788
4789        // 2. PER TOKEN: routed-expert loop. The ONE dispatch change vs Stage-1: a resident slot
4790        //    (cache HIT, no H2D) OR a staged slot (MISS) feeds the SAME unchanged qmatvec_view.
4791        for tok in 0..t {
4792            let sel = &sel_all[tok * n_used..(tok + 1) * n_used];
4793            let w = &w_all[tok * n_used..(tok + 1) * n_used];
4794            let zt = z.slice(tok * n_embd..(tok + 1) * n_embd); // CudaView<f32>
4795            let mut tok_q8: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
4796
4797            // STAGE-2 GROUPED DECODE (2026-07-04, MEMRA_MOE_GDEC default ON, =0 rollback): fold
4798            // this token's whole routed-expert FFN (8x gate/up/silu + 8x down/axpy = 40 launches)
4799            // into TWO launches via expert-pointer indirection over the fixed-address cache slots.
4800            // Fires only when ALL 3*n_used blocks are ALREADY cache-resident (pure-HIT: zero
4801            // memcpy, zero admission, so no slot can move under the collected pointers) — any
4802            // miss falls through to the sequential loop below, which admits as before. In steady
4803            // state on a fully-resident rig every token-layer takes the grouped path.
4804            // BIT-IDENTITY: each in-kernel dot reproduces qmatvec_f32's exact reduction; SiLU is
4805            // silu_mul_f32's exact expression; the down accumulation is a slot-ordered
4806            // __fmaf_rn chain == the sequential axpy_f32 chain (MEMRA_MOE_GDEC_GATE compares).
4807            // cfg.m3: the grouped kernels' fused epilogues are plain SiLU — M3's swigluoai must
4808            // NOT take them until the kernels grow the clamped variant. NVFP4 experts carry
4809            // per-expert macro-scales the fused kernels don't fold — those fall through too.
4810            let no_macros = m.gate_exps.macros.is_none()
4811                && m.up_exps.macros.is_none()
4812                && m.down_exps.macros.is_none();
4813            // SLAB-LOCAL FUSED PAIR (lane/pp-leverb 2026-08-08): the gdec launch pair
4814            // (moe_gate_up_silu8_q8 + moe_down8_fma_q8 — the EXACT kernels, same FP chains)
4815            // with pointers computed from the resident slab base + ex*stride instead of
4816            // collected SLRU slot addresses. No cache lock, no residency predicate — the
4817            // slab holds every expert by construction, so this arm never falls through
4818            // (gdec's P(all 24 resident) ≈ 0.37 coin-flip and the miss path's ~49-launch
4819            // staging both die). Bit-identity class: pointer provenance only, the same
4820            // slab-vs-SLRU equivalence moe_ffn_dev documents. Ordered BEFORE gdec: when a
4821            // slab exists it is strictly better (no lock, no miss).
4822            if slab_fused_may_fire {
4823                let (pg, pu, pd) = slab_bases.unwrap();
4824                let mut gp = [0u64; 8];
4825                let mut up = [0u64; 8];
4826                let mut dp = [0u64; 8];
4827                for (j, &ex) in sel.iter().enumerate() {
4828                    let ex = ex as usize;
4829                    gp[j] = pg + (ex * m.gate_exps.expert_stride) as u64;
4830                    up[j] = pu + (ex * m.up_exps.expert_stride) as u64;
4831                    dp[j] = pd + (ex * m.down_exps.expert_stride) as u64;
4832                }
4833                let mut wv = [0f32; 8];
4834                wv[..n_used].copy_from_slice(w);
4835                if tok_q8.is_none() {
4836                    tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
4837                }
4838                let (zq, zd) = tok_q8.as_ref().unwrap();
4839                let act = e.moe_gate_up_silu8_q8(
4840                    crate::WPtr8(gp),
4841                    crate::WPtr8(up),
4842                    zq,
4843                    zd,
4844                    n_embd,
4845                    n_ff_exp,
4846                    n_used,
4847                    m.gate_exps.qtype,
4848                    m.up_exps.qtype,
4849                    m.gate_exps.row_bytes,
4850                    m.up_exps.row_bytes,
4851                )?;
4852                let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
4853                let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
4854                e.moe_down8_fma_q8(
4855                    crate::WPtr8(dp),
4856                    crate::F32x8(wv),
4857                    &aq2,
4858                    &ad2,
4859                    &mut dst,
4860                    n_ff_exp,
4861                    n_embd,
4862                    n_used,
4863                    m.down_exps.qtype,
4864                    m.down_exps.row_bytes,
4865                )?;
4866                continue;
4867            }
4868            if gdec_may_fire && moe_q8 && cfg.m3.is_none() && no_macros {
4869                if tok_q8.is_none() {
4870                    tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
4871                }
4872                let (zq, zd) = tok_q8.as_ref().unwrap();
4873                if Self::moe_gdec_token_q8(
4874                    e,
4875                    m,
4876                    il,
4877                    max_block,
4878                    zq,
4879                    zd,
4880                    sel,
4881                    w,
4882                    &mut moe_out,
4883                    tok,
4884                    n_embd,
4885                    n_ff_exp,
4886                    n_used,
4887                )? {
4888                    continue;
4889                }
4890            } else if gdec_may_fire
4891                && cfg.m3.is_none()
4892                && no_macros
4893                && Self::moe_gdec_token(
4894                    e,
4895                    m,
4896                    il,
4897                    max_block,
4898                    &zt,
4899                    sel,
4900                    w,
4901                    &mut moe_out,
4902                    tok,
4903                    n_embd,
4904                    n_ff_exp,
4905                    n_used,
4906                )?
4907            {
4908                continue;
4909            }
4910
4911            // STAGE 2 memset-elision invariant: moe_out was allocated UNINIT when gdec or the
4912            // slab pair could fire. This token fell through to a sequential axpy loop, which
4913            // ACCUMULATES — zero its row first (row-sized memset; other rows are owned by the
4914            // full-overwrite arms). slab_fused_may_fire never actually reaches here (its arm
4915            // has no fallible predicate), included for the allocation invariant's symmetry.
4916            if gdec_may_fire || slab_fused_may_fire {
4917                let mut row = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
4918                e.memset_zeros_view(&mut row)?;
4919            }
4920
4921            // Fiddler-style split at whole-expert granularity. Partial residency is deliberately a
4922            // CPU assignment: transferring the missing projection would reintroduce the PCIe/NVMe
4923            // stall this path exists to remove, while mixing projections would require another
4924            // activation round-trip. Weight addresses remain valid until this worker is joined at
4925            // the bottom of the token scope.
4926            let mut cpu_mask = vec![false; sel.len()];
4927            let cpu_worker = if let Some(host_input) = cpu_input.as_ref() {
4928                let gpu_resident = if use_cache {
4929                    e.with_moe_cache(max_block, |cache, _| {
4930                        Ok(sel
4931                            .iter()
4932                            .map(|&expert| {
4933                                let expert = expert as u16;
4934                                [PROJ_GATE, PROJ_UP, PROJ_DOWN]
4935                                    .into_iter()
4936                                    .filter(|&projection| {
4937                                        cache
4938                                            .resident(BlockId::new(il, projection, expert))
4939                                            .is_some()
4940                                    })
4941                                    .count()
4942                            })
4943                            .collect::<Vec<_>>())
4944                    })?
4945                } else {
4946                    vec![0; sel.len()]
4947                };
4948                let mut cpu_selected = Vec::new();
4949                for (index, (&expert, &route_weight)) in sel.iter().zip(w).enumerate() {
4950                    if gpu_resident[index] != 3 {
4951                        cpu_mask[index] = true;
4952                        crate::cpu_experts::record_incomplete_gpu_residency(gpu_resident[index]);
4953                        let expert = expert as usize;
4954                        cpu_selected.push((expert, route_weight));
4955                    }
4956                }
4957                if crate::cpu_experts::predictor_enabled() {
4958                    // Fire-and-forget lookahead: the predictor worker scores layers il+1..
4959                    // from this layer's MoE input and prefetches predicted-and-missing
4960                    // experts into the companion RAM cache. Never blocks this thread.
4961                    let row = &host_input[tok * n_embd..(tok + 1) * n_embd];
4962                    crate::cpu_experts::predictor_submit(il, row);
4963                }
4964                if cpu_selected.is_empty() {
4965                    None
4966                } else {
4967                    let row = &host_input[tok * n_embd..(tok + 1) * n_embd];
4968                    let job = crate::cpu_experts::prepare_job(m, il, &cpu_selected, row)
4969                        .map_err(std::io::Error::other)?;
4970                    Some(crate::cpu_experts::submit(job).map_err(std::io::Error::other)?)
4971                }
4972            } else {
4973                None
4974            };
4975
4976            let worker_window = worker_disk_prefetch
4977                .then(worker_prefetch_window)
4978                .unwrap_or(0);
4979            for (j, &ex) in sel.iter().enumerate() {
4980                if cpu_mask[j] {
4981                    continue;
4982                }
4983                let ex = ex as usize;
4984                // PER-EXPERT SLAB READ (lane/pp-leverb 2026-08-08): the layers the fused
4985                // pair above excludes — step35's CLAMPED layers 43/44 (ffn_act_lim has no
4986                // fused form) and macro-carrying artifacts — still have their bytes in the
4987                // local resident slab. Same kernels as the SLRU arms (`qmatvec_expert_q8` /
4988                // `qmatvec_view`), same ffn_act_lim/macro folds; provenance = slab base +
4989                // ex*stride. No dispatch lock, no admission, no prefetch — nothing to miss.
4990                if let Some(d) = slab_local {
4991                    let gl = m.gate_exps.expert_layout(ex);
4992                    let ul = m.up_exps.expert_layout(ex);
4993                    let dl = m.down_exps.expert_layout(ex);
4994                    let (g0, u0, d0) = (
4995                        ex * m.gate_exps.expert_stride,
4996                        ex * m.up_exps.expert_stride,
4997                        ex * m.down_exps.expert_stride,
4998                    );
4999                    let (gate, up) = if moe_q8 {
5000                        if tok_q8.is_none() {
5001                            tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
5002                        }
5003                        let (zq, zd) = tok_q8.as_ref().unwrap();
5004                        (
5005                            e.qmatvec_expert_q8(
5006                                &d.gate,
5007                                g0..g0 + gl.len,
5008                                zq,
5009                                zd,
5010                                1,
5011                                m.gate_exps.in_f,
5012                                m.gate_exps.out_f,
5013                                gl.qtype,
5014                                gl.row_bytes,
5015                            )?,
5016                            e.qmatvec_expert_q8(
5017                                &d.up,
5018                                u0..u0 + ul.len,
5019                                zq,
5020                                zd,
5021                                1,
5022                                m.up_exps.in_f,
5023                                m.up_exps.out_f,
5024                                ul.qtype,
5025                                ul.row_bytes,
5026                            )?,
5027                        )
5028                    } else {
5029                        (
5030                            e.qmatvec_view(
5031                                &d.gate,
5032                                g0..g0 + gl.len,
5033                                &zt,
5034                                1,
5035                                m.gate_exps.in_f,
5036                                m.gate_exps.out_f,
5037                                gl.qtype,
5038                                gl.row_bytes,
5039                            )?,
5040                            e.qmatvec_view(
5041                                &d.up,
5042                                u0..u0 + ul.len,
5043                                &zt,
5044                                1,
5045                                m.up_exps.in_f,
5046                                m.up_exps.out_f,
5047                                ul.qtype,
5048                                ul.row_bytes,
5049                            )?,
5050                        )
5051                    };
5052                    let mut act = e.uninit(n_ff_exp)?;
5053                    Self::ffn_act_lim(
5054                        e,
5055                        cfg,
5056                        &gate,
5057                        &up,
5058                        m.gate_exps.macro_scale(ex),
5059                        m.up_exps.macro_scale(ex),
5060                        lim_exp,
5061                        &mut act,
5062                        n_ff_exp,
5063                    )?;
5064                    let y = if moe_q8 {
5065                        let (aq2, ad2) = e.quantize_q8_1(&act, 1, n_ff_exp)?;
5066                        e.qmatvec_expert_q8(
5067                            &d.down,
5068                            d0..d0 + dl.len,
5069                            &aq2,
5070                            &ad2,
5071                            1,
5072                            m.down_exps.in_f,
5073                            m.down_exps.out_f,
5074                            dl.qtype,
5075                            dl.row_bytes,
5076                        )?
5077                    } else {
5078                        let actv = act.slice(0..n_ff_exp);
5079                        e.qmatvec_view(
5080                            &d.down,
5081                            d0..d0 + dl.len,
5082                            &actv,
5083                            1,
5084                            m.down_exps.in_f,
5085                            m.down_exps.out_f,
5086                            dl.qtype,
5087                            dl.row_bytes,
5088                        )?
5089                    };
5090                    let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
5091                    e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
5092                    continue;
5093                }
5094                for next in page_prefetch_positions(j, sel.len(), page_window) {
5095                    Self::moe_prefetch_host_expert(sel[next] as usize, m);
5096                }
5097                let keep = [
5098                    crate::moe_cache::BlockId::new(il, crate::moe_cache::PROJ_GATE, ex as u16),
5099                    crate::moe_cache::BlockId::new(il, crate::moe_cache::PROJ_UP, ex as u16),
5100                    crate::moe_cache::BlockId::new(il, crate::moe_cache::PROJ_DOWN, ex as u16),
5101                ];
5102                if worker_disk_prefetch && worker_window > 0 {
5103                    for next in worker_prefetch_positions(j, sel.len(), worker_window) {
5104                        Self::moe_prefetch_disk_expert(
5105                            e,
5106                            il,
5107                            sel[next] as usize,
5108                            m,
5109                            max_block,
5110                            &keep,
5111                        )?;
5112                    }
5113                } else if cache_dispatch
5114                    && !cpu_hybrid
5115                    && moe_prefetch_enabled()
5116                    && j + 1 < sel.len()
5117                {
5118                    let next = sel[j + 1] as usize;
5119                    Self::moe_prefetch_expert(e, il, next, m, max_block, &keep)?;
5120                }
5121                let [gate_q8, up_q8, down_q8] = [moe_q8; 3];
5122                if cache_dispatch && (gate_q8 || up_q8 || down_q8) {
5123                    // dp4a EXPERT PATH (MEMRA_MOE_Q8): quantize z-row once per token. Mixed expert
5124                    // layouts stay on the metadata-aware f32 path.
5125                    if (gate_q8 || up_q8) && tok_q8.is_none() {
5126                        tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
5127                    }
5128                    let gate = if gate_q8 {
5129                        let (zq, zd) = tok_q8.as_ref().unwrap();
5130                        Self::moe_cached_gemm_q8(e, il, PROJ_GATE, ex, m, max_block, zq, zd)?
5131                    } else {
5132                        Self::moe_cached_gemm(e, il, PROJ_GATE, ex, m, max_block, &zt)?
5133                    };
5134                    let up = if up_q8 {
5135                        let (zq, zd) = tok_q8.as_ref().unwrap();
5136                        Self::moe_cached_gemm_q8(e, il, PROJ_UP, ex, m, max_block, zq, zd)?
5137                    } else {
5138                        Self::moe_cached_gemm(e, il, PROJ_UP, ex, m, max_block, &zt)?
5139                    };
5140                    let mut act = e.uninit(n_ff_exp)?;
5141                    Self::ffn_act_lim(
5142                        e,
5143                        cfg,
5144                        &gate,
5145                        &up,
5146                        m.gate_exps.macro_scale(ex),
5147                        m.up_exps.macro_scale(ex),
5148                        lim_exp,
5149                        &mut act,
5150                        n_ff_exp,
5151                    )?;
5152                    let y = if down_q8 {
5153                        let (aq2, ad2) = e.quantize_q8_1(&act, 1, n_ff_exp)?;
5154                        Self::moe_cached_gemm_q8(e, il, PROJ_DOWN, ex, m, max_block, &aq2, &ad2)?
5155                    } else {
5156                        let actv = act.slice(0..n_ff_exp);
5157                        Self::moe_cached_gemm(e, il, PROJ_DOWN, ex, m, max_block, &actv)?
5158                    };
5159                    let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
5160                    // down-proj macro folds into the accumulate weight (1.0 for non-macro archs).
5161                    e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
5162                } else if cache_dispatch {
5163                    // SLRU residency cache: per-projection, dispatch the block (HIT => resident slot,
5164                    // MISS => staged slot) then run the SAME unchanged qmatvec_view from that slot.
5165                    // The bytes the kernel reads are byte-for-byte the same GGUF block (§B.3); the
5166                    // only difference between HIT and MISS is whether the memcpy_htod ran.
5167                    let gate = Self::moe_cached_gemm(e, il, PROJ_GATE, ex, m, max_block, &zt)?;
5168                    let up = Self::moe_cached_gemm(e, il, PROJ_UP, ex, m, max_block, &zt)?;
5169                    let mut act = e.uninit(n_ff_exp)?; // activation fully overwrites
5170                    Self::ffn_act_lim(
5171                        e,
5172                        cfg,
5173                        &gate,
5174                        &up,
5175                        m.gate_exps.macro_scale(ex),
5176                        m.up_exps.macro_scale(ex),
5177                        lim_exp,
5178                        &mut act,
5179                        n_ff_exp,
5180                    )?;
5181                    let actv = act.slice(0..n_ff_exp);
5182                    let y = Self::moe_cached_gemm(e, il, PROJ_DOWN, ex, m, max_block, &actv)?;
5183                    let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
5184                    // down-proj macro folds into the accumulate weight (post-matmul linear scale).
5185                    e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
5186                } else if cache_frozen {
5187                    // A later prompt prime must not change the CPU/GPU assignment frozen after the
5188                    // first prime. Reuse every fixed resident projection directly and stage only a
5189                    // true miss through the ordinary scratch slot. This preserves the established
5190                    // f32-dequant numeric path while avoiding a full-bank reread on every prime.
5191                    let gate = Self::moe_frozen_gemm(
5192                        e,
5193                        il,
5194                        PROJ_GATE,
5195                        ex,
5196                        m,
5197                        max_block,
5198                        &zt,
5199                        &mut scratch_g,
5200                        g_len,
5201                    )?;
5202                    let up = Self::moe_frozen_gemm(
5203                        e,
5204                        il,
5205                        PROJ_UP,
5206                        ex,
5207                        m,
5208                        max_block,
5209                        &zt,
5210                        &mut scratch_u,
5211                        u_len,
5212                    )?;
5213                    let mut act = e.uninit(n_ff_exp)?;
5214                    Self::ffn_act_lim(
5215                        e,
5216                        cfg,
5217                        &gate,
5218                        &up,
5219                        m.gate_exps.macro_scale(ex),
5220                        m.up_exps.macro_scale(ex),
5221                        lim_exp,
5222                        &mut act,
5223                        n_ff_exp,
5224                    )?;
5225                    let actv = act.slice(0..n_ff_exp);
5226                    let y = Self::moe_frozen_gemm(
5227                        e,
5228                        il,
5229                        PROJ_DOWN,
5230                        ex,
5231                        m,
5232                        max_block,
5233                        &actv,
5234                        &mut scratch_d,
5235                        d_len,
5236                    )?;
5237                    let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
5238                    e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
5239                } else {
5240                    // Stage-1: stage gate/up/down for expert `ex` into the scratch slots, then GEMM.
5241                    // Lazy scratch: first no-cache expert allocates the 3 slots (uninit — stage_expert
5242                    // fully overwrites the byte range the GEMM reads).
5243                    if scratch_g.is_none() {
5244                        scratch_g = Some(e.alloc_u8_uninit(g_len)?);
5245                        scratch_u = Some(e.alloc_u8_uninit(u_len)?);
5246                        scratch_d = Some(e.alloc_u8_uninit(d_len)?);
5247                    }
5248                    let (sg, su, sd) = (
5249                        scratch_g.as_mut().unwrap(),
5250                        scratch_u.as_mut().unwrap(),
5251                        scratch_d.as_mut().unwrap(),
5252                    );
5253                    let gl = m.gate_exps.expert_layout(ex);
5254                    let ul = m.up_exps.expert_layout(ex);
5255                    let dl = m.down_exps.expert_layout(ex);
5256                    e.stage_expert(m.gate_exps.expert_bytes(ex), sg, 0)?;
5257                    let gate = e.qmatvec_view(
5258                        sg,
5259                        0..gl.len,
5260                        &zt,
5261                        1,
5262                        m.gate_exps.in_f,
5263                        m.gate_exps.out_f,
5264                        gl.qtype,
5265                        gl.row_bytes,
5266                    )?;
5267
5268                    e.stage_expert(m.up_exps.expert_bytes(ex), su, 0)?;
5269                    let up = e.qmatvec_view(
5270                        su,
5271                        0..ul.len,
5272                        &zt,
5273                        1,
5274                        m.up_exps.in_f,
5275                        m.up_exps.out_f,
5276                        ul.qtype,
5277                        ul.row_bytes,
5278                    )?;
5279
5280                    let mut act = e.uninit(n_ff_exp)?; // activation fully overwrites
5281                    Self::ffn_act_lim(
5282                        e,
5283                        cfg,
5284                        &gate,
5285                        &up,
5286                        m.gate_exps.macro_scale(ex),
5287                        m.up_exps.macro_scale(ex),
5288                        lim_exp,
5289                        &mut act,
5290                        n_ff_exp,
5291                    )?;
5292
5293                    e.stage_expert(m.down_exps.expert_bytes(ex), sd, 0)?;
5294                    let actv = act.slice(0..n_ff_exp);
5295                    let y = e.qmatvec_view(
5296                        sd,
5297                        0..dl.len,
5298                        &actv,
5299                        1,
5300                        m.down_exps.in_f,
5301                        m.down_exps.out_f,
5302                        dl.qtype,
5303                        dl.row_bytes,
5304                    )?;
5305
5306                    let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
5307                    e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
5308                }
5309            }
5310            if let Some(worker) = cpu_worker {
5311                let cpu_output = worker.wait().map_err(std::io::Error::other)?;
5312                let cpu_output = e.htod(&cpu_output)?;
5313                let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
5314                e.axpy_into(&cpu_output, 1.0, &mut dst, n_embd)?;
5315            }
5316            if cpu_hybrid && !cache_frozen && cpu_expert_profile_admit_enabled() {
5317                for (j, &ex) in sel.iter().enumerate() {
5318                    if cpu_mask[j] {
5319                        Self::moe_profile_admit_expert(e, il, ex as usize, m, max_block)?;
5320                    }
5321                }
5322            }
5323        }
5324
5325        // 3. SHARED EXPERT (ALWAYS-ON, no routing) on the SAME z — qwen35moe only. OLMoE and most
5326        //    vanilla MoE have NO shared expert (the shexp tensors are absent / `None`); skip it then.
5327        // gate_inp_shexp is OPTIONAL: qwen35moe gates the shared expert (sigmoid(gate_inp) x sh);
5328        // MiniMax-M3 (DeepSeek-V3 class) has NO shexp gate — the shared expert adds directly.
5329        if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
5330            (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
5331        {
5332            let n_ff_sh = gate_shexp.out_features(); // 512
5333            // Q8 TRUNK-FUSION (decode t=1): gate_shexp+up_shexp are Q8_0 same-shape on the 35B —
5334            // ONE fused2 launch (also folds the two per-matmul re-quantizes of z into one).
5335            // Bit-identical per (tensor,row); falls back to the two matmul calls when ineligible.
5336            // Small-t (spec verify 2..15) rides matmul_decode_exact so shexp FP chains match the
5337            // t==1 decode chain per column (cuBLASLt n-dependence + dp4a-vs-mmvq class); real
5338            // prefill keeps the batched matmul. Activation routes through ffn_act (SiLU for
5339            // softmax archs, clamped swigluoai for M3 — identical to silu_mul when cfg.m3 is None).
5340            let verify_t = t > 1 && t < PRIME_MIN_T;
5341            let (sg_gate, sg_up) = if t == 1 {
5342                match e.matmul_q8_fused2_x(gate_shexp, up_shexp, z)? {
5343                    Some(pair) => pair,
5344                    None => (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?),
5345                }
5346            } else if verify_t {
5347                (
5348                    e.matmul_decode_exact(gate_shexp, z, t)?,
5349                    e.matmul_decode_exact(up_shexp, z, t)?,
5350                )
5351            } else {
5352                (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?) // [T, 512] each
5353            };
5354            let mut sa = e.uninit(t * n_ff_sh)?; // activation fully overwrites
5355            Self::ffn_act_lim(
5356                e,
5357                cfg,
5358                &sg_gate,
5359                &sg_up,
5360                1.0,
5361                1.0,
5362                lim_shexp,
5363                &mut sa,
5364                t * n_ff_sh,
5365            )?;
5366            let sh = if verify_t {
5367                e.matmul_decode_exact(down_shexp, &sa, t)?
5368            } else {
5369                e.matmul(down_shexp, &sa, t)?
5370            }; // [T, n_embd]
5371
5372            // shexp gate: qwen35moe sigmoid-gates via ffn_gate_inp_shexp (1-D ne=[n_embd] ->
5373            // out_f=1); M3 has no gate tensor -> weight 1.0. Decode + verify ride the fused
5374            // sigmoid-dot kernel (one fold order for both chains; kills the per-layer
5375            // cuBLASLt m=1 splitK GEMM — 40x/step, ~10% of the H100 q35 decode step).
5376            // SERVE ISOLATION (lane/concat-prime-exact, 2026-08-02): PREFILL rides it too.
5377            // This out_f=1 cuBLASLt GEMV is the SECOND m-dependent op in the trunk (probed:
5378            // rows [0,19) move by 1.07e-4 between m=74 and m=75 while sigmoid_dot_rows is
5379            // BIT-IDENTICAL — allw-shexpgate-o35b.log). The gate multiplies the shared
5380            // expert's contribution into every token's residual, so under cross-request
5381            // concat prefill a session's hidden state depended on its co-arrivals' token
5382            // count. MEMRA_ROUTER_PREFILL_EXACT=0 reverts to the batched cuBLASLt linear.
5383            let g = match &m.gate_inp_shexp {
5384                Some(gate_inp_shexp) => {
5385                    if t < PRIME_MIN_T || crate::router_prefill_exact_on() {
5386                        e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
5387                    } else {
5388                        let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
5389                        let mut g = e.uninit(t)?; // sigmoid fully overwrites
5390                        e.sigmoid(&gs, &mut g, t)?;
5391                        g
5392                    }
5393                }
5394                None => e.htod(&vec![1.0f32; t])?,
5395            };
5396            // moe_out[r, :] += sh[r, :] * g[r]   (per-token scalar gate; g=1 ungated)
5397            e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
5398        }
5399
5400        Ok(moe_out)
5401    }
5402
5403    /// Stage-1 (no-cache) per-DECODE-TOKEN H2D bytes: every routed block re-staged every layer every
5404    /// token = sum over MoE layers of n_used * (gate+up+down expert_stride). The §D.4 PCIe baseline.
5405    pub fn stage1_h2d_per_token(&self) -> u64 {
5406        use crate::hybrid::Ffn;
5407        let n_used = self
5408            .cfg
5409            .moe
5410            .as_ref()
5411            .map(|m| m.expert_used_count as u64)
5412            .unwrap_or(0);
5413        let mut bytes = 0u64;
5414        for l in self.layers.iter() {
5415            if let Ffn::Moe(m) = &l.ffn {
5416                bytes += n_used
5417                    * (m.gate_exps.max_expert_bytes()
5418                        + m.up_exps.max_expert_bytes()
5419                        + m.down_exps.max_expert_bytes()) as u64;
5420            }
5421        }
5422        bytes
5423    }
5424
5425    /// Largest expert block (bytes) across ALL MoE layers + the MTP head — the fixed cache slot size.
5426    /// UD/dynamic GGUFs quant different layers differently, so `expert_stride` varies per layer; the
5427    /// residency cache slots are fixed-address and must fit any block, so size to this global max.
5428    pub(crate) fn max_moe_block(&self) -> usize {
5429        use crate::hybrid::Ffn;
5430        let mut mx = 0usize;
5431        let mut scan = |ffn: &Ffn| {
5432            if let Ffn::Moe(m) = ffn {
5433                mx = mx
5434                    .max(m.gate_exps.max_expert_bytes())
5435                    .max(m.up_exps.max_expert_bytes())
5436                    .max(m.down_exps.max_expert_bytes());
5437            }
5438        };
5439        for l in self.layers.iter() {
5440            scan(&l.ffn);
5441        }
5442        if let Some(mtp) = self.mtp.as_ref() {
5443            scan(&mtp.ffn);
5444        }
5445        mx
5446    }
5447
5448    /// Exact retained projection lengths for cache class sizing. Pruned ids keep router positions
5449    /// but have no bytes and therefore consume no residency slot.
5450    pub(crate) fn moe_cache_block_sizes(&self) -> Vec<usize> {
5451        use crate::hybrid::Ffn;
5452        let mut sizes = Vec::new();
5453        let mut scan = |ffn: &Ffn| {
5454            let Ffn::Moe(m) = ffn else { return };
5455            for ex in 0..m.gate_exps.n_expert {
5456                if m.active_experts.as_ref().is_some_and(|active| !active[ex]) {
5457                    continue;
5458                }
5459                for exps in [&m.gate_exps, &m.up_exps, &m.down_exps] {
5460                    let len = exps.expert_layout(ex).len;
5461                    if len > 0 {
5462                        sizes.push(len);
5463                    }
5464                }
5465            }
5466        };
5467        for layer in &self.layers {
5468            scan(&layer.ffn);
5469        }
5470        if let Some(mtp) = &self.mtp {
5471            scan(&mtp.ffn);
5472        }
5473        sizes
5474    }
5475
5476    /// Persist the frozen residency set so a later process can restage it directly and skip
5477    /// the profiling warmup. Plain text: a versioned header binding slot geometry, then one
5478    /// `layer proj ex` triple per line. A mismatched or stale profile is rejected at load
5479    /// (header check) or degrades to fewer restaged blocks (per-id checks); either way the
5480    /// post-freeze argmax gate still validates the serving assignment.
5481    pub fn save_cpu_expert_residency_profile(
5482        &self,
5483        e: &Engine,
5484        path: &std::path::Path,
5485    ) -> Result<(), Box<dyn std::error::Error>> {
5486        let Some(ids) = e.export_moe_residency() else {
5487            return Err("no MoE residency cache to persist".into());
5488        };
5489        let mut body = format!(
5490            "memra-freeze-profile v1 max_block={} blocks={}\n",
5491            self.max_moe_block(),
5492            ids.len()
5493        );
5494        for (layer, proj, ex) in &ids {
5495            body.push_str(&format!("{layer} {proj} {ex}\n"));
5496        }
5497        let tmp = path.with_extension("tmp");
5498        std::fs::write(&tmp, body)?;
5499        std::fs::rename(&tmp, path)?;
5500        println!(
5501            "[moe-cache] freeze profile saved: {} blocks -> {}",
5502            ids.len(),
5503            path.display()
5504        );
5505        Ok(())
5506    }
5507
5508    /// Restage a saved freeze profile and freeze immediately, skipping the profiling warmup.
5509    /// Returns false (leaving the cache untouched for a normal warmup) when the profile is
5510    /// missing or its header does not match this model's slot geometry.
5511    pub fn restore_cpu_expert_residency_profile(
5512        &self,
5513        e: &Engine,
5514        path: &std::path::Path,
5515    ) -> Result<bool, Box<dyn std::error::Error>> {
5516        use crate::hybrid::Ffn;
5517        use crate::moe_cache::BlockId;
5518        let Ok(content) = std::fs::read_to_string(path) else {
5519            return Ok(false);
5520        };
5521        let mut lines = content.lines();
5522        let Some(header) = lines.next() else {
5523            return Ok(false);
5524        };
5525        let expected = format!("memra-freeze-profile v1 max_block={}", self.max_moe_block());
5526        if !header.starts_with(&expected) {
5527            println!(
5528                "[moe-cache] freeze profile ignored (geometry mismatch): {}",
5529                path.display()
5530            );
5531            return Ok(false);
5532        }
5533        let mut by_layer: std::collections::HashMap<u16, Vec<BlockId>> =
5534            std::collections::HashMap::new();
5535        for line in lines {
5536            let mut fields = line.split_whitespace();
5537            let (Some(layer), Some(proj), Some(ex)) = (fields.next(), fields.next(), fields.next())
5538            else {
5539                continue;
5540            };
5541            let (Ok(layer), Ok(proj), Ok(ex)) =
5542                (layer.parse::<u16>(), proj.parse::<u8>(), ex.parse::<u16>())
5543            else {
5544                continue;
5545            };
5546            by_layer
5547                .entry(layer)
5548                .or_default()
5549                .push(BlockId::new(layer, proj, ex));
5550        }
5551        let requested: usize = by_layer.values().map(Vec::len).sum();
5552        if requested == 0 {
5553            return Ok(false);
5554        }
5555        let max_block = self.max_moe_block();
5556        let mut restaged = 0usize;
5557        let mut stage_layer =
5558            |layer_index: u16, ffn: &Ffn| -> Result<(), Box<dyn std::error::Error>> {
5559                let Ffn::Moe(m) = ffn else { return Ok(()) };
5560                let Some(ids) = by_layer.get(&layer_index) else {
5561                    return Ok(());
5562                };
5563                e.with_moe_cache(max_block, |cache, eng| {
5564                    for id in ids {
5565                        if cache.restage_block(*id, m, eng)? {
5566                            restaged += 1;
5567                        }
5568                    }
5569                    Ok(())
5570                })
5571            };
5572        for (index, layer) in self.layers.iter().enumerate() {
5573            stage_layer(index as u16, &layer.ffn)?;
5574        }
5575        if let Some(mtp) = self.mtp.as_ref() {
5576            stage_layer(u16::MAX, &mtp.ffn)?;
5577        }
5578        e.freeze_moe_cache();
5579        println!(
5580            "[moe-cache] freeze profile restored: {restaged}/{requested} blocks restaged from {}",
5581            path.display()
5582        );
5583        Ok(true)
5584    }
5585
5586    /// Freeze the heterogeneous CPU/GPU split after the caller's discarded profile warmup.
5587    pub fn freeze_cpu_expert_residency(
5588        &self,
5589        e: &Engine,
5590    ) -> Result<(), Box<dyn std::error::Error>> {
5591        e.freeze_moe_cache();
5592        Ok(())
5593    }
5594
5595    /// FFN activation dispatch: swigluoai (clamped, alpha/limit) when cfg.m3 says so, else the
5596    /// standard SiLU*up. One seam so every FFN site (dense, routed expert, shared expert) follows
5597    /// the model's activation exactly.
5598    ///
5599    /// NO-`il` FORM: cannot apply step35's PER-LAYER SwiGLU clamp. Only call it from a site whose
5600    /// layer provably has no live limit (dense-FFN layers, MTP blocks) — `ffn_act_lim` is the
5601    /// form for anything that can land on a clamped layer.
5602    pub fn ffn_act(
5603        e: &Engine,
5604        cfg: &ModelConfig,
5605        gate: &CudaSlice<f32>,
5606        up: &CudaSlice<f32>,
5607        act: &mut CudaSlice<f32>,
5608        n: usize,
5609    ) -> Result<(), Box<dyn std::error::Error>> {
5610        Self::ffn_act_scaled(e, cfg, gate, up, 1.0, 1.0, act, n)
5611    }
5612
5613    /// ffn_act with per-tensor post-matmul macro-scales folded in (gs/us == 1.0 -> identical
5614    /// float ops to ffn_act; used by the ModelOpt NVFP4 expert path where each expert tensor
5615    /// carries a `weight_scale_2`). Same no-`il` contract as `ffn_act`.
5616    #[allow(clippy::too_many_arguments)]
5617    pub(crate) fn ffn_act_scaled(
5618        e: &Engine,
5619        cfg: &ModelConfig,
5620        gate: &CudaSlice<f32>,
5621        up: &CudaSlice<f32>,
5622        gs: f32,
5623        us: f32,
5624        act: &mut CudaSlice<f32>,
5625        n: usize,
5626    ) -> Result<(), Box<dyn std::error::Error>> {
5627        Self::ffn_act_lim(e, cfg, gate, up, gs, us, None, act, n)
5628    }
5629
5630    /// ffn_act_scaled + step35's PER-LAYER clamped SwiGLU. `limit`:
5631    ///   * `None`   -> the unclamped dispatch (every arch except step35's layers 43-44).
5632    ///   * `Some(l)`-> `min(silu(gate*gs), l) * clamp(up*us, +-l)` (llama-graph.cpp:2146/1751,
5633    ///                 non-DEEPSEEK4 branch). Callers source it from `cfg.clamp_exp_at(il)`
5634    ///                 (routed experts) or `cfg.clamp_shexp_at(il)` (shared expert) — the two
5635    ///                 arrays are SEPARATE and a layer can have one without the other.
5636    /// The `> 1e-6` eps gate lives in `clamp_exp_at`/`clamp_shexp_at`, so a `Some` here is
5637    /// already known live.
5638    #[allow(clippy::too_many_arguments)]
5639    pub(crate) fn ffn_act_lim(
5640        e: &Engine,
5641        cfg: &ModelConfig,
5642        gate: &CudaSlice<f32>,
5643        up: &CudaSlice<f32>,
5644        gs: f32,
5645        us: f32,
5646        limit: Option<f32>,
5647        act: &mut CudaSlice<f32>,
5648        n: usize,
5649    ) -> Result<(), Box<dyn std::error::Error>> {
5650        if let Some(m3) = cfg.m3.as_ref() {
5651            debug_assert!(
5652                limit.is_none(),
5653                "m3 swigluoai and step35 clamp are different archs"
5654            );
5655            return e.swigluoai_mul_scaled(
5656                gate,
5657                up,
5658                gs,
5659                us,
5660                m3.swiglu_alpha,
5661                m3.swiglu_limit,
5662                act,
5663                n,
5664            );
5665        }
5666        if let Some(l) = limit {
5667            return e.swiglu_clamped_mul_scaled(gate, up, gs, us, l, act, n);
5668        }
5669        if gs == 1.0 && us == 1.0 {
5670            return e.silu_mul(gate, up, act, n);
5671        }
5672        e.silu_mul_scaled(gate, up, gs, us, act, n)
5673    }
5674
5675    /// Routing for the whole batch: returns (sel [T*n_used] expert ids, w [T*n_used] renorm weights),
5676    /// token-major. Default = the Stage-1 host path (dtoh logits, softmax-256, stable DESC top-k,
5677    /// renorm). MEMRA_FUSED_ROUTER = the device kernel (§A) which reproduces the same numerics; we
5678    /// still dtoh the tiny [T,n_used] sel/w buffers (64 B/token vs 1 KB/token) — the host loop
5679    /// indexes HostExps.bytes on the CPU to choose the DMA source (§A.2 output staging).
5680    fn moe_route(
5681        e: &Engine,
5682        logits: &CudaSlice<f32>,
5683        t: usize,
5684        n_expert: usize,
5685        n_used: usize,
5686    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
5687        Self::moe_route_cfg(e, logits, t, n_expert, n_used, None)
5688    }
5689
5690    /// DeepSeek-V3-class sigmoid routing (Step-3.7, MiniMax-M3, Hy3, GLM-DSA). Reference:
5691    /// M3/Hy3 modeling code — scores = sigmoid(logits); selection over scores + expert bias
5692    /// (M3 `e_score_correction_bias` / Hy3 `expert_bias`, both surfaced as `exp_probs_b`);
5693    /// weights = un-biased scores of the selected experts, sum-normalized when `route_norm`,
5694    /// x scaling factor (M3 routed_scaling_factor 2.0 / Hy3 router_scaling_factor 2.826).
5695    /// Default uses one device kernel plus the same pinned sel/w readback contract as the fused
5696    /// softmax router. `MEMRA_SIG_ROUTER=0` restores the full-logit DtoH host oracle.
5697    #[allow(clippy::too_many_arguments)]
5698    fn moe_route_sigmoid_cfg(
5699        e: &Engine,
5700        logits: &CudaSlice<f32>,
5701        t: usize,
5702        n_expert: usize,
5703        n_used: usize,
5704        m: &MoeWeights,
5705        (sf, route_norm): (f32, bool),
5706    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
5707        if sigmoid_router_enabled() {
5708            return e.moe_router_sigmoid_topk_host(
5709                logits,
5710                t,
5711                n_expert,
5712                n_used,
5713                m.active_count(),
5714                &m.exp_probs_b_dev,
5715                &m.active_experts_dev,
5716                sf,
5717                route_norm,
5718            );
5719        }
5720        let lg = e.dtoh(logits)?;
5721        Self::moe_route_sigmoid_host(
5722            &lg,
5723            t,
5724            n_expert,
5725            n_used,
5726            m.exp_probs_b.as_deref(),
5727            sf,
5728            route_norm,
5729            m.active_experts.as_deref(),
5730        )
5731    }
5732
5733    /// Softmax router for qwen35moe/OLMoE. An active-mask overlay retains the host oracle because
5734    /// the existing softmax device kernel has no mask input.
5735    fn moe_route_cfg(
5736        e: &Engine,
5737        logits: &CudaSlice<f32>,
5738        t: usize,
5739        n_expert: usize,
5740        n_used: usize,
5741        active: Option<&[bool]>,
5742    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
5743        // LAUNCH-STRUCTURE STAGE 1 (2026-07-05): fused router DEFAULT ON (MEMRA_FUSED_ROUTER=0
5744        // rollback) via the single-sync pinned readback — softmax arch only.
5745        if active.is_none() && !matches!(std::env::var("MEMRA_FUSED_ROUTER").as_deref(), Ok("0")) {
5746            return e.moe_router_topk_host(logits, t, n_expert, n_used);
5747        }
5748        // Host oracle (the §D bit-identity reference).
5749        let lg = e.dtoh(logits)?; // [T*n_expert] host
5750        let mut sel = vec![0u32; t * n_used];
5751        let mut w_out = vec![0f32; t * n_used];
5752        for tok in 0..t {
5753            let row = &lg[tok * n_expert..(tok + 1) * n_expert];
5754            // softmax over ALL n_expert (stable: subtract max)
5755            let maxl = row
5756                .iter()
5757                .enumerate()
5758                .filter(|(i, _)| active.is_none_or(|mask| mask[*i]))
5759                .map(|(_, &x)| x)
5760                .fold(f32::NEG_INFINITY, f32::max);
5761            let mut probs = vec![0f32; n_expert];
5762            let mut den = 0f32;
5763            for i in 0..n_expert {
5764                if active.is_some_and(|mask| !mask[i]) {
5765                    continue;
5766                }
5767                let x = (row[i] - maxl).exp();
5768                probs[i] = x;
5769                den += x;
5770            }
5771            for p in probs.iter_mut() {
5772                *p /= den;
5773            }
5774            // stable DESC sort: prob DESC, ascending-index tiebreak.
5775            let mut idx: Vec<usize> = (0..n_expert)
5776                .filter(|&i| active.is_none_or(|mask| mask[i]))
5777                .collect();
5778            idx.sort_by(|&a, &b| probs[b].total_cmp(&probs[a]).then(a.cmp(&b)));
5779            let sl = &idx[..n_used];
5780            let mut wv: Vec<f32> = sl.iter().map(|&i| probs[i]).collect();
5781            let mut ws: f32 = wv.iter().sum();
5782            ws = ws.max(6.103515625e-5_f32); // F16 smallest normal, clamp BEFORE divide
5783            for x in wv.iter_mut() {
5784                *x /= ws;
5785            }
5786            for j in 0..n_used {
5787                sel[tok * n_used + j] = sl[j] as u32;
5788                w_out[tok * n_used + j] = wv[j];
5789            }
5790        }
5791        Ok((sel, w_out))
5792    }
5793
5794    #[allow(clippy::too_many_arguments)]
5795    fn moe_route_sigmoid_with_input(
5796        e: &Engine,
5797        logits: &CudaSlice<f32>,
5798        input: &CudaSlice<f32>,
5799        t: usize,
5800        n_expert: usize,
5801        n_used: usize,
5802        bias: Option<&[f32]>,
5803        (sf, route_norm): (f32, bool),
5804        active: Option<&[bool]>,
5805    ) -> Result<(Vec<u32>, Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
5806        let (lg, input) = e.dtoh_pair(logits, input)?;
5807        let (sel, w) =
5808            Self::moe_route_sigmoid_host(&lg, t, n_expert, n_used, bias, sf, route_norm, active)?;
5809        Ok((sel, w, input))
5810    }
5811
5812    /// Start the prediction-guided prefetch worker (MEMRA_MOE_PREFETCH=depth). Call after
5813    /// residency freeze: the worker filters against a static snapshot of the frozen HBM set.
5814    /// Builds a fully-owned per-layer table (host router copies via one-time DtoH, bias,
5815    /// active mask, prebuilt projection descriptors) so no model reference escapes.
5816    pub fn start_moe_prefetch_predictor(
5817        &self,
5818        e: &Engine,
5819        cfg: &ModelConfig,
5820    ) -> Result<(), Box<dyn std::error::Error>> {
5821        use crate::hybrid::Ffn;
5822        let Some(sig) = cfg.sigmoid_router() else {
5823            return Err("prefetch predictor requires a sigmoid-router arch".into());
5824        };
5825        let resident: std::collections::HashSet<(u16, u8, u16)> = e
5826            .export_moe_residency()
5827            .ok_or("prefetch predictor needs the frozen MoE residency cache")?
5828            .into_iter()
5829            .collect();
5830        let mut layers = Vec::new();
5831        for (index, layer) in self.layers.iter().enumerate() {
5832            let Ffn::Moe(m) = &layer.ffn else { continue };
5833            let crate::model::GpuTensor::Float { data, .. } = &m.gate_inp else {
5834                continue;
5835            };
5836            let router = e.dtoh(data)?;
5837            let n_expert = m.gate_exps.n_expert;
5838            let n_embd = m.gate_exps.in_f;
5839            if router.len() != n_embd * n_expert {
5840                continue;
5841            }
5842            let build = |exps: &crate::model::HostExps| {
5843                (0..n_expert)
5844                    .map(|expert| crate::cpu_experts::predictor_projection(exps, expert))
5845                    .collect::<Vec<_>>()
5846            };
5847            layers.push((
5848                index as u16,
5849                crate::cpu_experts::PredictLayerInit {
5850                    router,
5851                    bias: m.exp_probs_b.clone(),
5852                    active: m.active_experts.clone(),
5853                    n_embd,
5854                    n_used: cfg
5855                        .moe
5856                        .as_ref()
5857                        .map(|moe| moe.expert_used_count as usize)
5858                        .ok_or("prefetch predictor requires MoE config")?,
5859                    sig,
5860                    weights_n_expert: n_expert,
5861                    gate: build(&m.gate_exps),
5862                    up: build(&m.up_exps),
5863                    down: build(&m.down_exps),
5864                },
5865            ));
5866        }
5867        crate::cpu_experts::start_prefetch_predictor(layers, resident).map_err(|error| error.into())
5868    }
5869
5870    /// Sigmoid-routing oracle shared by the prefetch predictor and `kernel-check`: identical
5871    /// selection math to the rollback runtime, applied to host-computed logits.
5872    #[allow(clippy::too_many_arguments)]
5873    pub fn moe_route_sigmoid_host_public(
5874        logits: &[f32],
5875        t: usize,
5876        n_expert: usize,
5877        n_used: usize,
5878        bias: Option<&[f32]>,
5879        sf: f32,
5880        route_norm: bool,
5881        active: Option<&[bool]>,
5882    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
5883        Self::moe_route_sigmoid_host(logits, t, n_expert, n_used, bias, sf, route_norm, active)
5884    }
5885
5886    #[allow(clippy::too_many_arguments)]
5887    fn moe_route_sigmoid_host(
5888        lg: &[f32],
5889        t: usize,
5890        n_expert: usize,
5891        n_used: usize,
5892        bias: Option<&[f32]>,
5893        sf: f32,
5894        route_norm: bool,
5895        active: Option<&[bool]>,
5896    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
5897        let active_count = active
5898            .map(|mask| mask.iter().filter(|&&enabled| enabled).count())
5899            .unwrap_or(n_expert);
5900        crate::sigrouter_contract::validate_active_count(n_used, active_count)?;
5901        if lg.len() != t * n_expert {
5902            return Err(format!(
5903                "sigmoid router logits length mismatch: got {}, expected {}",
5904                lg.len(),
5905                t * n_expert,
5906            )
5907            .into());
5908        }
5909        let mut sel = vec![0u32; t * n_used];
5910        let mut w_out = vec![0f32; t * n_used];
5911        for tok in 0..t {
5912            let row = &lg[tok * n_expert..(tok + 1) * n_expert];
5913            let scores: Vec<f32> = row.iter().map(|&x| 1.0 / (1.0 + (-x).exp())).collect();
5914            // selection score = sigmoid + bias; weight = plain sigmoid.
5915            let selsc: Vec<f32> = match bias {
5916                Some(b) => scores.iter().zip(b).map(|(s, bb)| s + bb).collect(),
5917                None => scores.clone(),
5918            };
5919            let mut idx: Vec<usize> = (0..n_expert)
5920                .filter(|&i| active.is_none_or(|mask| mask[i]))
5921                .collect();
5922            idx.sort_by(|&a, &b| selsc[b].total_cmp(&selsc[a]).then(a.cmp(&b)));
5923            let sl = &idx[..n_used];
5924            let mut wv: Vec<f32> = sl.iter().map(|&i| scores[i]).collect();
5925            if route_norm {
5926                let ws: f32 = wv.iter().sum::<f32>().max(1e-20);
5927                for x in wv.iter_mut() {
5928                    *x = *x / ws * sf;
5929                }
5930            } else {
5931                for x in wv.iter_mut() {
5932                    *x *= sf;
5933                }
5934            }
5935            for j in 0..n_used {
5936                sel[tok * n_used + j] = sl[j] as u32;
5937                w_out[tok * n_used + j] = wv[j];
5938            }
5939        }
5940        Ok((sel, w_out))
5941    }
5942
5943    /// Step-3.7 resident-slab dispatch: sigmoid top-k outputs remain on device and feed the
5944    /// existing slot-ordered q8 expert kernels directly. Mixed layouts, spill, remote slabs,
5945    /// macro-scaled experts, and observation modes are denied by the caller.
5946    #[allow(clippy::too_many_arguments)]
5947    fn moe_ffn_sigmoid_dev(
5948        e: &Engine,
5949        m: &MoeWeights,
5950        z: &CudaSlice<f32>,
5951        zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
5952        logits: &CudaSlice<f32>,
5953        t: usize,
5954        cfg: &ModelConfig,
5955        il: u16,
5956        (scaling_factor, route_norm): (f32, bool),
5957    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5958        let moe = cfg.moe.as_ref().unwrap();
5959        let n_embd = cfg.n_embd as usize;
5960        let n_expert = moe.expert_count as usize;
5961        let n_used = moe.expert_used_count as usize;
5962        let n_ff_exp = moe.expert_ff_length as usize;
5963        let dev = m.dev_exps.as_ref().unwrap();
5964        debug_assert!(cfg.step35.is_some());
5965        debug_assert_eq!(dev.dev, e.ctx().ordinal());
5966        debug_assert!(m.has_uniform_expert_layout());
5967        debug_assert!(!m.has_macros);
5968
5969        let (sel_d, w_d) = e.moe_router_sigmoid_topk(
5970            logits,
5971            t,
5972            n_expert,
5973            n_used,
5974            m.active_count(),
5975            &m.exp_probs_b_dev,
5976            &m.active_experts_dev,
5977            scaling_factor,
5978            route_norm,
5979        )?;
5980        crate::moesd::record_device_routes(e, il, n_expert, n_used, &sel_d)?;
5981        let (gate_row_bytes, up_row_bytes) = if dev.gu_il {
5982            let combined = m.gate_exps.row_bytes + m.up_exps.row_bytes;
5983            (combined, combined)
5984        } else {
5985            (m.gate_exps.row_bytes, m.up_exps.row_bytes)
5986        };
5987        let (zq, zd) = match (t, zq8) {
5988            (1, Some((q, d))) => (q.clone(), d.clone()),
5989            _ => e.quantize_q8_1(z, t, n_embd)?,
5990        };
5991        let n_pairs = t * n_used;
5992        let mut moe_out = if cfg.clamp_exp_at(il as u32).is_some() {
5993            // The final Step layers retain the established separate gate/up -> clamp -> down
5994            // arithmetic. Pair rows are derived from token position; selected expert ids and
5995            // routing weights remain the device router's buffers throughout.
5996            let pair_tok: Vec<i32> = (0..n_pairs).map(|pair| (pair / n_used) as i32).collect();
5997            let pair_tok_d = e.htod_i32(&pair_tok)?;
5998            let gate = e.moe_pairs_matvec_q8(
5999                &dev.ptr_row,
6000                0,
6001                &pair_tok_d,
6002                &sel_d,
6003                &zq,
6004                &zd,
6005                n_embd,
6006                n_ff_exp,
6007                n_expert,
6008                n_pairs,
6009                m.gate_exps.qtype,
6010                gate_row_bytes,
6011            )?;
6012            let up = e.moe_pairs_matvec_q8(
6013                &dev.ptr_row,
6014                1,
6015                &pair_tok_d,
6016                &sel_d,
6017                &zq,
6018                &zd,
6019                n_embd,
6020                n_ff_exp,
6021                n_expert,
6022                n_pairs,
6023                m.up_exps.qtype,
6024                up_row_bytes,
6025            )?;
6026            let mut act = e.uninit(n_pairs * n_ff_exp)?;
6027            Self::ffn_act_lim(
6028                e,
6029                cfg,
6030                &gate,
6031                &up,
6032                1.0,
6033                1.0,
6034                cfg.clamp_exp_at(il as u32),
6035                &mut act,
6036                n_pairs * n_ff_exp,
6037            )?;
6038            let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
6039            let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
6040            let pair_self_d = e.htod_i32(&pair_self)?;
6041            let down = e.moe_pairs_matvec_q8(
6042                &dev.ptr_row,
6043                2,
6044                &pair_self_d,
6045                &sel_d,
6046                &aq2,
6047                &ad2,
6048                n_ff_exp,
6049                n_embd,
6050                n_expert,
6051                n_pairs,
6052                m.down_exps.qtype,
6053                m.down_exps.row_bytes,
6054            )?;
6055            let tok_off: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
6056            let tok_ids: Vec<i32> = (0..n_pairs as i32).collect();
6057            let tok_off_d = e.htod_i32(&tok_off)?;
6058            let tok_ids_d = e.htod_i32(&tok_ids)?;
6059            let mut output = e.uninit(t * n_embd)?;
6060            e.moe_pairs_scatter(&down, &w_d, &tok_off_d, &tok_ids_d, &mut output, t, n_embd)?;
6061            output
6062        } else {
6063            let act = e.moe_gate_up_silu8_dev_q8_rows(
6064                &dev.ptr_row,
6065                &sel_d,
6066                &zq,
6067                &zd,
6068                t,
6069                n_embd,
6070                n_ff_exp,
6071                n_used,
6072                n_expert,
6073                m.gate_exps.qtype,
6074                m.up_exps.qtype,
6075                gate_row_bytes,
6076                up_row_bytes,
6077                &m.dev_macros,
6078            )?;
6079            let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
6080            let mut output = e.uninit(t * n_embd)?;
6081            e.moe_down8_fma_dev_q8_rows_g(
6082                &dev.ptr_row,
6083                &sel_d,
6084                &w_d,
6085                &aq2,
6086                &ad2,
6087                &mut output,
6088                t,
6089                n_ff_exp,
6090                n_embd,
6091                n_used,
6092                n_expert,
6093                m.down_exps.qtype,
6094                m.down_exps.row_bytes,
6095            )?;
6096            output
6097        };
6098
6099        if std::env::var("MEMRA_SIG_ROUTER_DISPATCH_TRACE").as_deref() == Ok("1") {
6100            eprintln!(
6101                "[sigrouter-dev] layer={il} tokens={t} experts={n_expert} used={n_used} clamp={} gu_il={}",
6102                cfg.clamp_exp_at(il as u32).is_some(),
6103                dev.gu_il,
6104            );
6105        }
6106        Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut moe_out)?;
6107        Ok(moe_out)
6108    }
6109
6110    /// LAUNCH-STRUCTURE STAGE 3: the ZERO-DtoH fully-resident MoE FFN. Caller guarantees the
6111    /// layer's device pointer row exists (checked under the cache lock). Router top-k runs on
6112    /// device; sel/w are consumed by the `_dev` matvec twins directly; NOTHING crosses PCIe.
6113    /// Same numerics as the fused-router + gdec chain (kernel-level bit-identity, see the
6114    /// MoE PREFILL PAIR-BATCH: host routing (sel/w like the sequential path), then 5 launches
6115    /// TOTAL per layer (quantize z, gate-pairs, up-pairs, silu, act-quantize, down-pairs,
6116    /// scatter) regardless of T or expert count. Bit-identity class: per (pair,row) dot =
6117    /// qmatvec_expert_q8 order; per-token accumulation slot-ordered (scatter kernel).
6118    fn moe_ffn_pairs(
6119        e: &Engine,
6120        m: &MoeWeights,
6121        z: &CudaSlice<f32>,
6122        logits: &CudaSlice<f32>,
6123        t: usize,
6124        cfg: &ModelConfig,
6125    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6126        let moe = cfg.moe.as_ref().unwrap();
6127        let n_embd = cfg.n_embd as usize;
6128        let n_expert = moe.expert_count as usize;
6129        let n_used = moe.expert_used_count as usize;
6130        let n_ff_exp = moe.expert_ff_length as usize;
6131        // This arm has no `il` and every kernel in it fuses PLAIN silu(gate)*up, so a clamped
6132        // model must never reach it. The caller's gate at `moe_ffn_sequential_zq8` denies per
6133        // layer via `swiglu_clamped_at(il)`; assert the whole-model form here so a future caller
6134        // that forgets the gate fails loudly in debug instead of returning wrong logits.
6135        debug_assert!(
6136            !cfg.swiglu_clamped_anywhere(),
6137            "moe_ffn_pairs has no per-layer clamp: fused epilogues are plain SiLU"
6138        );
6139        let dev = m.dev_exps.as_ref().unwrap();
6140        // WALL-GAP ARC: interleaved gate/up slab strides (see moe_ffn_dev).
6141        let (rbg_d, rbu_d) = if dev.gu_il {
6142            let sxx = m.gate_exps.row_bytes + m.up_exps.row_bytes;
6143            (sxx, sxx)
6144        } else {
6145            (m.gate_exps.row_bytes, m.up_exps.row_bytes)
6146        };
6147
6148        let (sel_all, w_all) = Self::moe_route(e, logits, t, n_expert, n_used)?;
6149        let n_pairs = t * n_used;
6150        // pair arrays: pair p = (token p/n_used, slot p%n_used) — ALREADY slot-ordered per token,
6151        // so the CSR is trivial: tok_pair_off[tok] = tok*n_used, ids identity.
6152        let pair_tok: Vec<i32> = (0..n_pairs).map(|p| (p / n_used) as i32).collect();
6153        let pair_ex: Vec<i32> = sel_all.iter().map(|&x| x as i32).collect();
6154        let pair_w: Vec<f32> = w_all.clone();
6155        let tok_off: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
6156        let tok_ids: Vec<i32> = (0..n_pairs as i32).collect();
6157        let pt = e.htod_i32(&pair_tok)?;
6158        let px = e.htod_i32(&pair_ex)?;
6159        let pw = e.htod(&pair_w)?;
6160        let toff = e.htod_i32(&tok_off)?;
6161        let tids = e.htod_i32(&tok_ids)?;
6162
6163        // z quantized ONCE for all tokens; gate/up pair matvecs; silu; act quantize; down; scatter.
6164        // EXPERT-MAJOR CSR (rung 2): pairs grouped by expert -> the kernel reuses each weight
6165        // row across the expert's token group (llama-MMQ's core win). Host grouping is O(pairs).
6166        let mut by_ex: Vec<Vec<i32>> = vec![Vec::new(); n_expert];
6167        for p in 0..n_pairs {
6168            by_ex[pair_ex[p] as usize].push(p as i32);
6169        }
6170        let mut ex_ids: Vec<i32> = Vec::new();
6171        let mut ex_off: Vec<i32> = vec![0];
6172        let mut ex_pairs: Vec<i32> = Vec::with_capacity(n_pairs);
6173        for (ex, list) in by_ex.iter().enumerate() {
6174            if list.is_empty() {
6175                continue;
6176            }
6177            ex_ids.push(ex as i32);
6178            ex_pairs.extend_from_slice(list);
6179            ex_off.push(ex_pairs.len() as i32);
6180        }
6181        let n_active = ex_ids.len();
6182        let exi = e.htod_i32(&ex_ids)?;
6183        let exo = e.htod_i32(&ex_off)?;
6184        let exp_d = e.htod_i32(&ex_pairs)?;
6185        let _ = &px; // pair-major twin keeps it; em path uses CSR
6186
6187        // INT8-MMA EXPERT MMQ (MEMRA_MOE_MMA=1, opt-in): the m16n8k16.s8 tensor-core analog of the
6188        // _dec dp4a kernel (cu/mmq_iq_experts.cu). Same CSR grouping; per-expert matvec runs as a
6189        // 128x128-tile int8 MMA GEMM over the expert's token group. Weight IQ nibbles decode to int8
6190        // at tile-load + per-32 float scale; activation is q8_1_mmq (D4, same quant class as dp4a).
6191        // FP-ORDER differs from dp4a (MMA reduction) — logits SHIFT, gated on argmax/spec/closeness,
6192        // NOT byte-identity (like the W4A8 path). Requires IQ3_S/IQ4_XS + in_f % 256 == 0.
6193        // t >= 16 (GEMM_M-class rule): the MMA tile needs token volume (crossover ~200 tok/expert;
6194        // microbench: dp4a wins at tiny groups). ALSO an exactness requirement — spec verify
6195        // batches (t=2..K+2) must ride the dp4a path whose FP order matches the T=1 decode chain,
6196        // else K=1 self-consistency FAILs (caught 2026-07-06: MMA at T=2 flipped a verify argmax).
6197        // DEFAULT ON (2026-07-06, third flip — this time with the real culprit fixed): the
6198        // "MMA prime breaks spec" failure was the ROUTER's cuBLASLt n-dependence (d994271),
6199        // not MMA's own FP order — both this and the k-quant arms were innocent suspects whose
6200        // margin shifts surfaced the router bug. With the router decode-exact at verify t, the
6201        // full battery is green with MMA on (spec p1/p2/p3 PASS, raw K=1..8 PASS, argmax MATCH,
6202        // pp6257 2862 = 2.1x dec). t>=16 floor still required: verify batches must ride dp4a
6203        // (dispatch parity with the T=1 decode chain). MEMRA_MOE_MMA=0 rollback;
6204        // MEMRA_MOE_MMA_T overrides the floor (bisect seam).
6205        static MMA_T: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
6206        let mma_t = *MMA_T.get_or_init(|| {
6207            std::env::var("MEMRA_MOE_MMA_T")
6208                .ok()
6209                .and_then(|v| v.parse().ok())
6210                .unwrap_or(16)
6211        });
6212        let use_mma = std::env::var("MEMRA_MOE_MMA")
6213            .map(|v| v != "0")
6214            .unwrap_or(true)
6215            && t >= mma_t
6216            && q8_expert_dec_supported(m.gate_exps.qtype)
6217            && q8_expert_dec_supported(m.up_exps.qtype)
6218            && q8_expert_dec_supported(m.down_exps.qtype)
6219            && n_embd % 256 == 0
6220            && n_ff_exp % 256 == 0;
6221        // GROUPED f16 LANE admission (MEMRA_MOE_F16G, rounds 46-49): its own door, no longer a
6222        // subset of the MMQ arm — q35's k-quant stragglers (Q6_K/Q4_K down, one Q3_K gate/up
6223        // layer) fail q8_expert_dec_supported but dequant fine to f16, so f16g must be able to
6224        // take a layer the MMQ arm would reject. Same t >= mma_t floor as MMA: decode and
6225        // spec-verify batches must ride the dp4a path whose FP order matches the T=1 chain.
6226        // MODE 2 (sm_120a naked default, lane/f16g-default-rearb 2026-08-02): every layer
6227        // whose three projections pass f16g_proj_ok rides the sk visitor with direct tile
6228        // loaders — with IQ4_XS/IQ3_S direct coverage the sk arm beats the int8-MMA MMQ
6229        // tiles on the IQ-bank models too (q35 board-2048 +33.9%, KAT pp512 +46.7%).
6230        // AUTO-KQUANT (mode 3, MEMRA_MOE_F16G=3, lane/q4k-expert-prefill): admit f16g ONLY
6231        // where the MMA arm can't take the layer's QTYPES — the k-quant expert class whose
6232        // baseline is the per-pair _em fallback (Ornith-35B Q4_K: board-2048 3.14x). Its
6233        // MMA-capable carve-out (IQ3_S/IQ4_XS/Q4_0 banks on MMQ) was priced pre-IQ-direct;
6234        // it survives as the rollback seam. Keyed on qtype capability, NOT use_mma, so
6235        // MEMRA_MOE_MMA=0 stays a pure dp4a rollback seam.
6236        let mma_capable = q8_expert_dec_supported(m.gate_exps.qtype)
6237            && q8_expert_dec_supported(m.up_exps.qtype)
6238            && q8_expert_dec_supported(m.down_exps.qtype)
6239            && n_embd % 256 == 0
6240            && n_ff_exp % 256 == 0;
6241        let f16g_mode = crate::moe_f16g_mode();
6242        let f16g = f16g_mode != 0
6243            && t >= mma_t
6244            && (f16g_mode != 3 || !mma_capable)
6245            && f16g_proj_ok(m.gate_exps.qtype, n_embd)
6246            && f16g_proj_ok(m.up_exps.qtype, n_embd)
6247            && f16g_proj_ok(m.down_exps.qtype, n_ff_exp);
6248        if use_mma || f16g {
6249            // GROUPED f16 LANE (MEMRA_MOE_F16G, rounds 46-49, experimental door): dequant
6250            // the active experts ONCE per projection to f16 and run one grouped f16 GEMM over
6251            // the CSR groups. f16-mirror numeric class — argmax/spec gated before promotion.
6252            // NOTE: the pair-gather kernel needs pair p ordered EXPERT-MAJOR (ex_pairs order);
6253            // the y rows come back in that same CSR order, so gate/up/down all stay pair-major
6254            // in ex_pairs order — but moe_pairs_silu_mul and the scatter consume PAIR-ID order.
6255            // We therefore gather activations per ex_pairs and scatter y back through ex_pairs.
6256            let y_down = if f16g {
6257                // CSR order end-to-end: gather z rows by the pair's TOKEN (pair p's token is
6258                // p / n_used — the trivial CSR above), silu in CSR order (elementwise), one
6259                // permute at the very end back to pair-id order for the scatter.
6260                let csr_tok: Vec<i32> = ex_pairs.iter().map(|&p| p / n_used as i32).collect();
6261                let csr_tok_d = e.htod_i32(&csr_tok)?;
6262                let (z_f16, z_s) = e.moe_f16g_act(z, Some(&csr_tok_d), n_embd, n_pairs)?;
6263                let g_csr = e.moe_f16_grouped(
6264                    &dev.ptr_row,
6265                    0,
6266                    n_expert,
6267                    &exi,
6268                    &ex_off,
6269                    &exo,
6270                    &z_f16,
6271                    &z_s,
6272                    n_embd,
6273                    n_ff_exp,
6274                    n_active,
6275                    n_pairs,
6276                    m.gate_exps.qtype,
6277                    rbg_d,
6278                )?;
6279                let u_csr = e.moe_f16_grouped(
6280                    &dev.ptr_row,
6281                    1,
6282                    n_expert,
6283                    &exi,
6284                    &ex_off,
6285                    &exo,
6286                    &z_f16,
6287                    &z_s,
6288                    n_embd,
6289                    n_ff_exp,
6290                    n_active,
6291                    n_pairs,
6292                    m.up_exps.qtype,
6293                    rbu_d,
6294                )?;
6295                let act_csr = e.moe_pairs_silu_mul(&g_csr, &u_csr, n_pairs * n_ff_exp)?;
6296                let (a_f16, a_s) = e.moe_f16g_act(&act_csr, None, n_ff_exp, n_pairs)?;
6297                let d_csr = e.moe_f16_grouped(
6298                    &dev.ptr_row,
6299                    2,
6300                    n_expert,
6301                    &exi,
6302                    &ex_off,
6303                    &exo,
6304                    &a_f16,
6305                    &a_s,
6306                    n_ff_exp,
6307                    n_embd,
6308                    n_active,
6309                    n_pairs,
6310                    m.down_exps.qtype,
6311                    m.down_exps.row_bytes,
6312                )?;
6313                e.rows_permute(&d_csr, &exp_d, n_pairs, n_embd)?
6314            } else {
6315                // gate/up: activation = z, token-major over t tokens; pair_tok gathers the routed row.
6316                let z_scr = e.mmq_iq_quantize_act(z, n_embd, t)?;
6317                let gate = e.mmq_iq_experts(
6318                    &dev.ptr_row,
6319                    0,
6320                    n_expert,
6321                    &exi,
6322                    &exo,
6323                    &exp_d,
6324                    &pt,
6325                    &z_scr,
6326                    n_embd,
6327                    n_ff_exp,
6328                    n_active,
6329                    n_pairs,
6330                    t,
6331                    m.gate_exps.qtype,
6332                    rbg_d,
6333                )?;
6334                let up = e.mmq_iq_experts(
6335                    &dev.ptr_row,
6336                    1,
6337                    n_expert,
6338                    &exi,
6339                    &exo,
6340                    &exp_d,
6341                    &pt,
6342                    &z_scr,
6343                    n_embd,
6344                    n_ff_exp,
6345                    n_active,
6346                    n_pairs,
6347                    t,
6348                    m.up_exps.qtype,
6349                    rbu_d,
6350                )?;
6351                // down: activation = silu(gate)*up, pair-major [n_pairs, n_ff_exp]; pair_tok =
6352                // identity. FUSED ACT-EPILOGUE (default on): one launch computes the activation in
6353                // registers and writes ONLY the quantized scratch — the two-pass chain
6354                // (moe_pairs_silu_mul writes act f32, mmq_iq_quantize_act re-reads it) is the
6355                // MEMRA_MOE_FUSE_ACTQ=0 rollback. Scratch bytes are BYTE-IDENTICAL (kernel-check).
6356                let a_scr = if crate::moe_fuse_actq_on() {
6357                    e.mmq_iq_fused_act_quant(&gate, &up, n_ff_exp, n_pairs, 0)?
6358                } else {
6359                    let act = e.moe_pairs_silu_mul(&gate, &up, n_pairs * n_ff_exp)?;
6360                    e.mmq_iq_quantize_act(&act, n_ff_exp, n_pairs)?
6361                };
6362                let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
6363                let pself = e.htod_i32(&pair_self)?;
6364                e.mmq_iq_experts(
6365                    &dev.ptr_row,
6366                    2,
6367                    n_expert,
6368                    &exi,
6369                    &exo,
6370                    &exp_d,
6371                    &pself,
6372                    &a_scr,
6373                    n_ff_exp,
6374                    n_embd,
6375                    n_active,
6376                    n_pairs,
6377                    n_pairs,
6378                    m.down_exps.qtype,
6379                    m.down_exps.row_bytes,
6380                )?
6381            };
6382            let mut moe_out = e.uninit(t * n_embd)?;
6383            e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
6384            if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
6385                (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
6386            {
6387                let n_ff_sh = gate_shexp.out_features();
6388                let sg_gate = e.matmul(gate_shexp, z, t)?;
6389                let sg_up = e.matmul(up_shexp, z, t)?;
6390                let mut sa = e.uninit(t * n_ff_sh)?;
6391                Self::ffn_act(e, cfg, &sg_gate, &sg_up, &mut sa, t * n_ff_sh)?;
6392                let sh = e.matmul(down_shexp, &sa, t)?;
6393                // shexp gate: qwen35moe sigmoid-gates; M3 has no gate tensor -> weight 1.0.
6394                // SERVE ISOLATION (lane/concat-prime-exact): the fused sigmoid-dot is the
6395                // m-INVARIANT form — see the sequential arm's note. This is the PAIRS arm,
6396                // i.e. the one real prefill actually takes on a resident-expert MoE model,
6397                // so the concat-prime isolation fix has to land here as well.
6398                let g = match &m.gate_inp_shexp {
6399                    Some(gate_inp_shexp) if crate::router_prefill_exact_on() => {
6400                        e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
6401                    }
6402                    Some(gate_inp_shexp) => {
6403                        let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
6404                        let mut g = e.uninit(t)?;
6405                        e.sigmoid(&gs, &mut g, t)?;
6406                        g
6407                    }
6408                    None => e.htod(&vec![1.0f32; t])?,
6409                };
6410                e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
6411            }
6412            return Ok(moe_out);
6413        }
6414
6415        // DECODE-ONCE MMQ (rung 3, MEMRA_MOE_DEC=1 default-on): dequant each weight group once per
6416        // (row,group) then dp4a across the expert's tokens. _em re-decoded per token (NEUTRAL).
6417        let dec = std::env::var("MEMRA_MOE_DEC")
6418            .map(|v| v != "0")
6419            .unwrap_or(true);
6420        let matvec = |proj,
6421                      exi: &_,
6422                      exo: &_,
6423                      exp_d: &_,
6424                      pt: &_,
6425                      aq: &_,
6426                      ad: &_,
6427                      inf,
6428                      outf,
6429                      qtype,
6430                      rb|
6431         -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6432            // _dec's decode-once extractors are IQ-only; k-quant expert layers take the _em dot path.
6433            let dec = dec && q8_expert_dec_supported(qtype);
6434            if dec {
6435                e.moe_pairs_matvec_q8_dec(
6436                    &dev.ptr_row,
6437                    proj,
6438                    exi,
6439                    exo,
6440                    exp_d,
6441                    pt,
6442                    aq,
6443                    ad,
6444                    inf,
6445                    outf,
6446                    n_expert,
6447                    n_active,
6448                    n_pairs,
6449                    qtype,
6450                    rb,
6451                )
6452            } else {
6453                e.moe_pairs_matvec_q8_em(
6454                    &dev.ptr_row,
6455                    proj,
6456                    exi,
6457                    exo,
6458                    exp_d,
6459                    pt,
6460                    aq,
6461                    ad,
6462                    inf,
6463                    outf,
6464                    n_expert,
6465                    n_active,
6466                    n_pairs,
6467                    qtype,
6468                    rb,
6469                )
6470            }
6471        };
6472        let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
6473        let gate = matvec(
6474            0,
6475            &exi,
6476            &exo,
6477            &exp_d,
6478            &pt,
6479            &zq,
6480            &zd,
6481            n_embd,
6482            n_ff_exp,
6483            m.gate_exps.qtype,
6484            rbg_d,
6485        )?;
6486        let up = matvec(
6487            1,
6488            &exi,
6489            &exo,
6490            &exp_d,
6491            &pt,
6492            &zq,
6493            &zd,
6494            n_embd,
6495            n_ff_exp,
6496            m.up_exps.qtype,
6497            rbu_d,
6498        )?;
6499        let act = e.moe_pairs_silu_mul(&gate, &up, n_pairs * n_ff_exp)?;
6500        let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
6501        // down consumes PAIR-major activation rows: pair_tok = identity.
6502        let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
6503        let pself = e.htod_i32(&pair_self)?;
6504        let y_down = matvec(
6505            2,
6506            &exi,
6507            &exo,
6508            &exp_d,
6509            &pself,
6510            &aq2,
6511            &ad2,
6512            n_ff_exp,
6513            n_embd,
6514            m.down_exps.qtype,
6515            m.down_exps.row_bytes,
6516        )?;
6517        let mut moe_out = e.uninit(t * n_embd)?; // scatter fully overwrites per (token,col)
6518        e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
6519
6520        // SHARED EXPERT epilogue — same as the other paths.
6521        // gate_inp_shexp is OPTIONAL: qwen35moe gates the shared expert (sigmoid(gate_inp) x sh);
6522        // MiniMax-M3 (DeepSeek-V3 class) has NO shexp gate — the shared expert adds directly.
6523        if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
6524            (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
6525        {
6526            let n_ff_sh = gate_shexp.out_features();
6527            // These decode-exact forms are required by the new Step resident arm. Keep the
6528            // established grouped shared-expert program for every other architecture: widening
6529            // this to Gemma changed its speculative acceptance despite green argmax gates.
6530            let step_exact = cfg.step35.is_some();
6531            let verify_t = step_exact && t > 1 && t < PRIME_MIN_T;
6532            let (sg_gate, sg_up) = if step_exact && t == 1 {
6533                match e.matmul_q8_fused2_x(gate_shexp, up_shexp, z)? {
6534                    Some(pair) => pair,
6535                    None => (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?),
6536                }
6537            } else if verify_t {
6538                let mut fused = None;
6539                if crate::spec::spec_fused_t()
6540                    && (2..=4).contains(&t)
6541                    && e.uses_q8_1_fast(gate_shexp)
6542                    && e.uses_q8_1_fast(up_shexp)
6543                {
6544                    let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
6545                    fused = e.matmul_q8_fused2_t(gate_shexp, up_shexp, &zq, &zd, t)?;
6546                }
6547                match fused {
6548                    Some(pair) => pair,
6549                    None => (
6550                        e.matmul_decode_exact(gate_shexp, z, t)?,
6551                        e.matmul_decode_exact(up_shexp, z, t)?,
6552                    ),
6553                }
6554            } else {
6555                (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?)
6556            };
6557            let mut sa = e.uninit(t * n_ff_sh)?;
6558            e.silu_mul(&sg_gate, &sg_up, &mut sa, t * n_ff_sh)?;
6559            let sh = if verify_t {
6560                e.matmul_decode_exact(down_shexp, &sa, t)?
6561            } else {
6562                e.matmul(down_shexp, &sa, t)?
6563            };
6564            // shexp gate: qwen35moe sigmoid-gates; M3 has no gate tensor -> weight 1.0.
6565            // m-INVARIANT fused sigmoid-dot under router_prefill_exact (serve isolation,
6566            // lane/concat-prime-exact) — every shexp-gate arm shares the same form so a
6567            // dispatch choice cannot change bits.
6568            let g = match &m.gate_inp_shexp {
6569                Some(gate_inp_shexp) if crate::router_prefill_exact_on() => {
6570                    e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
6571                }
6572                Some(gate_inp_shexp) => {
6573                    let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
6574                    let mut g = e.uninit(t)?;
6575                    e.sigmoid(&gs, &mut g, t)?;
6576                    g
6577                }
6578                None => e.htod(&vec![1.0f32; t])?,
6579            };
6580            e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
6581        }
6582        Ok(moe_out)
6583    }
6584
6585    /// kernel headers); the shared-expert epilogue is byte-identical to moe_ffn_sequential's.
6586    #[allow(clippy::too_many_arguments)]
6587    #[allow(clippy::too_many_arguments)]
6588    fn moe_ffn_dev(
6589        e: &Engine,
6590        m: &MoeWeights,
6591        z: &CudaSlice<f32>,
6592        zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
6593        logits: &CudaSlice<f32>,
6594        t: usize,
6595        cfg: &ModelConfig,
6596        il: u16,
6597        max_block: usize,
6598    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6599        let moe = cfg.moe.as_ref().unwrap();
6600        let n_embd = cfg.n_embd as usize;
6601        let n_expert = moe.expert_count as usize;
6602        let n_used = moe.expert_used_count as usize;
6603        let n_ff_exp = moe.expert_ff_length as usize;
6604        // moe_router_topk is SOFTMAX-only (no exp_probs_b bias, no expert_weights_scale) and every
6605        // gate_up kernel here fuses PLAIN silu(gate)*up. `dev_ok` denies sigmoid-router archs and
6606        // clamped layers; assert both so a future caller that skips the gate fails loudly.
6607        debug_assert!(
6608            cfg.sigmoid_router().is_none(),
6609            "moe_ffn_dev routes SOFTMAX: a sigmoid-router arch would pick wrong experts"
6610        );
6611        debug_assert!(
6612            !cfg.swiglu_clamped_at(il as u32),
6613            "moe_ffn_dev's fused epilogue is plain SiLU: no clamped form"
6614        );
6615
6616        // device top-k: sel [t, n_used] i32, w [t, n_used] f32 — stays on device.
6617        let (sel_d, mut w_d) = e.moe_router_topk(logits, t, n_expert, n_used)?;
6618        // Down-projection macro fold (compressed-tensors NVFP4 artifacts): one tiny launch,
6619        // skipped entirely for macro-free experts (every k-quant GGUF).
6620        if m.has_macros {
6621            e.moe_w_scale_by_expert(&mut w_d, &sel_d, &m.dev_macros, n_expert, t * n_used)?;
6622        }
6623
6624        // moe_out rows are FULLY overwritten by moe_down8_fma_dev — uninit (stage-2 rule).
6625        let mut moe_out = e.uninit(t * n_embd)?;
6626
6627        // RESIDENT-EXPERTS arm: the pointer row comes from the load-time slab (no cache, no
6628        // lock). Same kernels/loop as the SLRU arm below — only the row's provenance differs.
6629        if let Some(dev) = m.dev_exps.as_ref() {
6630            // WALL-GAP ARC: interleaved gate/up slab (MEMRA_MOE_GU_IL) -> both projections use
6631            // the combined stride; up's base is offset in the ptr table. Down unchanged.
6632            let (rbg_d, rbu_d) = if dev.gu_il {
6633                let sxx = m.gate_exps.row_bytes + m.up_exps.row_bytes;
6634                (sxx, sxx)
6635            } else {
6636                (m.gate_exps.row_bytes, m.up_exps.row_bytes)
6637            };
6638            let q8 = moe_q8_enabled()
6639                && q8_expert_supported(m.gate_exps.qtype)
6640                && q8_expert_supported(m.up_exps.qtype)
6641                && q8_expert_supported(m.down_exps.qtype);
6642            // SMALL-M ROWS ARM (MEMRA_SPEC_M2, lane/spec-m2): batch the verify token loop —
6643            // ONE batched z-quantize + ONE gate_up rows launch + ONE act quantize + ONE down
6644            // rows launch (4 launches/layer, was 4t). BIT-IDENTICAL per token to the serial
6645            // loop below (rows twins = the _v/w8h2v per-token programs on a grid.z token axis;
6646            // quantize_q8_1 is per-32-block row-independent). Gated to the AUTO kernel modes —
6647            // a custom MEMRA_MOE_DEVQ8_GU/DOWN diagnostic run keeps the serial loop so the
6648            // dispatched kernel stays exactly the env-selected one — and to the w8h2v shape
6649            // (n_ff_exp==512, n_used<=8), the same contract the AUTO down dispatch keys on.
6650            let rows_arm = q8
6651                && t > 1
6652                && crate::spec::spec_m2()
6653                && n_ff_exp == 512
6654                && n_used <= 8
6655                && std::env::var("MEMRA_MOE_DEVQ8_GU")
6656                    .map(|v| v.is_empty() || v == "v")
6657                    .unwrap_or(true)
6658                && std::env::var("MEMRA_MOE_DEVQ8_DOWN")
6659                    .map(|v| v.is_empty() || v == "w8h2v")
6660                    .unwrap_or(true);
6661            // CSR EXPERT-DEDUP gate_up (verify-cost target #1, DEFAULT ON 2026-07-10): the
6662            // owner-scan kernel serves every (token, slot) pair of an expert from ONE block,
6663            // deduping the 38-40% duplicated weight-stream+decode the overlap probe measured
6664            // (gate_up 55.0 -> 39.7us/launch; +1.3-2.1% spec e2e p2, +0.6-1.7% p3, all K).
6665            // Bit-identical to the _rows twins (explicit-intrinsic accumulate — the ULP/fmad
6666            // lesson; =2 byte-compare verified zero diffs). down stays on _rows (CSR down
6667            // measured 23.5 -> 37.5us: 16-group rows can't amortize the serial pair loop).
6668            // MEMRA_MOE_CSR=0 rollback; =2 runs BOTH paths and byte-compares (debug).
6669            let csr_mode = std::env::var("MEMRA_MOE_CSR")
6670                .ok()
6671                .and_then(|v| v.parse::<i32>().ok())
6672                .unwrap_or(1);
6673            let csr_qt = |qt: i32| qt == crate::QT_IQ4_XS || qt == crate::QT_IQ3_S;
6674            let csr_arm = rows_arm
6675                && csr_mode > 0
6676                && t <= 10
6677                && csr_qt(m.gate_exps.qtype)
6678                && csr_qt(m.up_exps.qtype)
6679                && csr_qt(m.down_exps.qtype);
6680            if csr_arm {
6681                if csr_mode == 2 {
6682                    static ENGAGED: std::sync::Once = std::sync::Once::new();
6683                    ENGAGED.call_once(|| eprintln!("[memra] moe CSR byte-compare mode ON (t={t})"));
6684                }
6685                let n_pairs = t * n_used;
6686                let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
6687                let act = e.moe_gate_up_silu8_dev_q8_csr(
6688                    &dev.ptr_row,
6689                    &sel_d,
6690                    &zq,
6691                    &zd,
6692                    n_pairs,
6693                    n_embd,
6694                    n_ff_exp,
6695                    n_used,
6696                    n_expert,
6697                    m.gate_exps.qtype,
6698                    m.up_exps.qtype,
6699                    rbg_d,
6700                    rbu_d,
6701                )?;
6702                let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
6703                // down stays on the _rows twin — BOTH CSR down variants measured negative
6704                // (v1 serial pairs 23.5->37.5us; v2 warp-parallel+SMEM -14% e2e, K=8 -37%):
6705                // 16-group rows have too little decode to amortize any dedup structure.
6706                e.moe_down8_fma_dev_q8_rows(
6707                    &dev.ptr_row,
6708                    &sel_d,
6709                    &w_d,
6710                    &aq2,
6711                    &ad2,
6712                    &mut moe_out,
6713                    t,
6714                    n_ff_exp,
6715                    n_embd,
6716                    n_used,
6717                    n_expert,
6718                    m.down_exps.qtype,
6719                    m.down_exps.row_bytes,
6720                )?;
6721                if csr_mode == 2 {
6722                    // DEBUG BYTE-COMPARE: run the _rows twins on the same inputs, diff bits.
6723                    let act_r = e.moe_gate_up_silu8_dev_q8_rows(
6724                        &dev.ptr_row,
6725                        &sel_d,
6726                        &zq,
6727                        &zd,
6728                        t,
6729                        n_embd,
6730                        n_ff_exp,
6731                        n_used,
6732                        n_expert,
6733                        m.gate_exps.qtype,
6734                        m.up_exps.qtype,
6735                        rbg_d,
6736                        rbu_d,
6737                        &m.dev_macros,
6738                    )?;
6739                    let mut out_r = e.uninit(t * n_embd)?;
6740                    let (aq2r, ad2r) = e.quantize_q8_1(&act_r, n_pairs, n_ff_exp)?;
6741                    e.moe_down8_fma_dev_q8_rows(
6742                        &dev.ptr_row,
6743                        &sel_d,
6744                        &w_d,
6745                        &aq2r,
6746                        &ad2r,
6747                        &mut out_r,
6748                        t,
6749                        n_ff_exp,
6750                        n_embd,
6751                        n_used,
6752                        n_expert,
6753                        m.down_exps.qtype,
6754                        m.down_exps.row_bytes,
6755                    )?;
6756                    let (a1, a2) = (e.dtoh(&act)?, e.dtoh(&act_r)?);
6757                    let (o1, o2) = (e.dtoh(&moe_out)?, e.dtoh(&out_r)?);
6758                    let ba = a1
6759                        .iter()
6760                        .zip(&a2)
6761                        .filter(|(x, y)| x.to_bits() != y.to_bits())
6762                        .count();
6763                    let bo = o1
6764                        .iter()
6765                        .zip(&o2)
6766                        .filter(|(x, y)| x.to_bits() != y.to_bits())
6767                        .count();
6768                    if ba + bo > 0 {
6769                        eprintln!(
6770                            "[csr-check] il={il} t={t} ACT diffs={ba}/{} OUT diffs={bo}/{}",
6771                            a1.len(),
6772                            o1.len()
6773                        );
6774                        // First 4 differing ACT elements: (pair, o, csr, rows) + that pair's expert
6775                        let sel_h = e.dtoh_i32(&sel_d)?;
6776                        let mut shown = 0;
6777                        for (i, (x, y)) in a1.iter().zip(&a2).enumerate() {
6778                            if x.to_bits() != y.to_bits() && shown < 4 {
6779                                let (p, o) = (i / n_ff_exp, i % n_ff_exp);
6780                                let ex = sel_h[p];
6781                                let npx = sel_h.iter().filter(|&&v| v == ex).count();
6782                                eprintln!(
6783                                    "  ACT p={p} ex={ex} np={npx} o={o} csr={x:e} rows={y:e}"
6784                                );
6785                                shown += 1;
6786                            }
6787                        }
6788                        std::process::exit(3);
6789                    }
6790                }
6791            } else if rows_arm {
6792                // TEMP PROBE (MEMRA_MOE_OVERLAP=1): cross-token expert-activation overlap at
6793                // verify — sizes the CSR dedup win (unique experts vs t*n_used pairs).
6794                if std::env::var("MEMRA_MOE_OVERLAP").as_deref() == Ok("1") {
6795                    use std::sync::atomic::{AtomicU64, Ordering};
6796                    static PAIRS: AtomicU64 = AtomicU64::new(0);
6797                    static UNIQ: AtomicU64 = AtomicU64::new(0);
6798                    static CALLS: AtomicU64 = AtomicU64::new(0);
6799                    let sel_h = e.dtoh_i32(&sel_d)?;
6800                    let mut u: Vec<i32> = sel_h.clone();
6801                    u.sort_unstable();
6802                    u.dedup();
6803                    PAIRS.fetch_add(sel_h.len() as u64, Ordering::Relaxed);
6804                    UNIQ.fetch_add(u.len() as u64, Ordering::Relaxed);
6805                    let c = CALLS.fetch_add(1, Ordering::Relaxed) + 1;
6806                    if c % 480 == 0 {
6807                        let p = PAIRS.load(Ordering::Relaxed);
6808                        let q = UNIQ.load(Ordering::Relaxed);
6809                        eprintln!(
6810                            "[overlap] calls={c} pairs={p} unique={q} ratio={:.3} (t={t})",
6811                            q as f64 / p as f64
6812                        );
6813                    }
6814                }
6815                let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
6816                let act = e.moe_gate_up_silu8_dev_q8_rows(
6817                    &dev.ptr_row,
6818                    &sel_d,
6819                    &zq,
6820                    &zd,
6821                    t,
6822                    n_embd,
6823                    n_ff_exp,
6824                    n_used,
6825                    n_expert,
6826                    m.gate_exps.qtype,
6827                    m.up_exps.qtype,
6828                    rbg_d,
6829                    rbu_d,
6830                    &m.dev_macros,
6831                )?;
6832                let (aq2, ad2) = e.quantize_q8_1(&act, t * n_used, n_ff_exp)?;
6833                e.moe_down8_fma_dev_q8_rows(
6834                    &dev.ptr_row,
6835                    &sel_d,
6836                    &w_d,
6837                    &aq2,
6838                    &ad2,
6839                    &mut moe_out,
6840                    t,
6841                    n_ff_exp,
6842                    n_embd,
6843                    n_used,
6844                    n_expert,
6845                    m.down_exps.qtype,
6846                    m.down_exps.row_bytes,
6847                )?;
6848            } else {
6849                for tok in 0..t {
6850                    let zt = z.slice(tok * n_embd..(tok + 1) * n_embd);
6851                    let selt = sel_d.slice(tok * n_used..(tok + 1) * n_used);
6852                    let wt = w_d.slice(tok * n_used..(tok + 1) * n_used);
6853                    let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
6854                    if q8 {
6855                        let (zq, zd) = match (t, zq8) {
6856                            (1, Some((q, d))) => (q.clone(), d.clone()),
6857                            _ => e.quantize_q8_1_view(&zt, 1, n_embd)?,
6858                        };
6859                        let act = e.moe_gate_up_silu8_dev_q8(
6860                            &dev.ptr_row,
6861                            &selt,
6862                            &zq,
6863                            &zd,
6864                            n_embd,
6865                            n_ff_exp,
6866                            n_used,
6867                            n_expert,
6868                            m.gate_exps.qtype,
6869                            m.up_exps.qtype,
6870                            rbg_d,
6871                            rbu_d,
6872                            &m.dev_macros,
6873                        )?;
6874                        let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
6875                        e.moe_down8_fma_dev_q8(
6876                            &dev.ptr_row,
6877                            &selt,
6878                            &wt,
6879                            &aq2,
6880                            &ad2,
6881                            &mut dst,
6882                            n_ff_exp,
6883                            n_embd,
6884                            n_used,
6885                            n_expert,
6886                            m.down_exps.qtype,
6887                            m.down_exps.row_bytes,
6888                        )?;
6889                    } else {
6890                        let act = e.moe_gate_up_silu8_dev(
6891                            &dev.ptr_row,
6892                            &selt,
6893                            &zt,
6894                            n_embd,
6895                            n_ff_exp,
6896                            n_used,
6897                            n_expert,
6898                            m.gate_exps.qtype,
6899                            m.up_exps.qtype,
6900                            rbg_d,
6901                            rbu_d,
6902                            &m.dev_macros,
6903                        )?;
6904                        e.moe_down8_fma_dev(
6905                            &dev.ptr_row,
6906                            &selt,
6907                            &wt,
6908                            &act,
6909                            &mut dst,
6910                            n_ff_exp,
6911                            n_embd,
6912                            n_used,
6913                            n_expert,
6914                            m.down_exps.qtype,
6915                            m.down_exps.row_bytes,
6916                        )?;
6917                    }
6918                }
6919            }
6920        } else {
6921            // Launch under the cache lock: the row borrow lives as long as the closure, and the
6922            // lock covers only launch ISSUE (µs), same policy as moe_cached_gemm.
6923            // Q8 ARM PARITY (2026-07-06): the SLRU arm ran the f32-dequant _dev kernels only —
6924            // 80us/launch vs the q8 twins' 15us on the SAME shapes (fixed-build profile: 228
6925            // f32 launches = 36ms of the 64-tok window). Same q8 gate + kernels as the resident
6926            // arm above; MEMRA_MOE_Q8=0 restores the byte-identical f32 path.
6927            let q8 = moe_q8_enabled()
6928                && q8_expert_supported(m.gate_exps.qtype)
6929                && q8_expert_supported(m.up_exps.qtype)
6930                && q8_expert_supported(m.down_exps.qtype);
6931            e.with_moe_cache(max_block, |c, eng| {
6932                let row = c
6933                    .layer_dev_row(il, n_expert, eng)?
6934                    .ok_or("moe_ffn_dev: layer row vanished under the lock")?;
6935                for tok in 0..t {
6936                    let zt = z.slice(tok * n_embd..(tok + 1) * n_embd);
6937                    let selt = sel_d.slice(tok * n_used..(tok + 1) * n_used);
6938                    let wt = w_d.slice(tok * n_used..(tok + 1) * n_used);
6939                    let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
6940                    if q8 {
6941                        let (zq, zd) = match (t, zq8) {
6942                            (1, Some((q, d))) => (q.clone(), d.clone()),
6943                            _ => eng.quantize_q8_1_view(&zt, 1, n_embd)?,
6944                        };
6945                        let act = eng.moe_gate_up_silu8_dev_q8(
6946                            row,
6947                            &selt,
6948                            &zq,
6949                            &zd,
6950                            n_embd,
6951                            n_ff_exp,
6952                            n_used,
6953                            n_expert,
6954                            m.gate_exps.qtype,
6955                            m.up_exps.qtype,
6956                            m.gate_exps.row_bytes,
6957                            m.up_exps.row_bytes,
6958                            &m.dev_macros,
6959                        )?;
6960                        let (aq2, ad2) = eng.quantize_q8_1(&act, n_used, n_ff_exp)?;
6961                        eng.moe_down8_fma_dev_q8(
6962                            row,
6963                            &selt,
6964                            &wt,
6965                            &aq2,
6966                            &ad2,
6967                            &mut dst,
6968                            n_ff_exp,
6969                            n_embd,
6970                            n_used,
6971                            n_expert,
6972                            m.down_exps.qtype,
6973                            m.down_exps.row_bytes,
6974                        )?;
6975                    } else {
6976                        let act = eng.moe_gate_up_silu8_dev(
6977                            row,
6978                            &selt,
6979                            &zt,
6980                            n_embd,
6981                            n_ff_exp,
6982                            n_used,
6983                            n_expert,
6984                            m.gate_exps.qtype,
6985                            m.up_exps.qtype,
6986                            m.gate_exps.row_bytes,
6987                            m.up_exps.row_bytes,
6988                            &m.dev_macros,
6989                        )?;
6990                        eng.moe_down8_fma_dev(
6991                            row,
6992                            &selt,
6993                            &wt,
6994                            &act,
6995                            &mut dst,
6996                            n_ff_exp,
6997                            n_embd,
6998                            n_used,
6999                            n_expert,
7000                            m.down_exps.qtype,
7001                            m.down_exps.row_bytes,
7002                        )?;
7003                    }
7004                }
7005                // instrumentation parity with the host paths (3 blocks/expert-slot, all hits).
7006                c.hits += (t * 3 * n_used) as u64;
7007                Ok(())
7008            })?;
7009        }
7010
7011        // SHARED EXPERT epilogue — byte-identical to moe_ffn_sequential step 3 (incl. its Q8
7012        // TRUNK-FUSION arm: fused2 is bit-identical to the two matmul calls per (tensor,row)).
7013        // gate_inp_shexp is OPTIONAL: qwen35moe gates the shared expert (sigmoid(gate_inp) x sh);
7014        // MiniMax-M3 (DeepSeek-V3 class) has NO shexp gate — the shared expert adds directly.
7015        if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
7016            (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
7017        {
7018            let n_ff_sh = gate_shexp.out_features();
7019            // verify-t (2..15) decode-exact arm: this fn now serves the spec verify batches
7020            // (pairs gate moved to t>=PRIME_MIN_T), so the shexp chain must match t==1 per col.
7021            let verify_t = t > 1 && t < PRIME_MIN_T;
7022            let (sg_gate, sg_up) = if t == 1 {
7023                match e.matmul_q8_fused2_x(gate_shexp, up_shexp, z)? {
7024                    Some(pair) => pair,
7025                    None => (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?),
7026                }
7027            } else if verify_t {
7028                // VERIFY-TIER TRUNK FUSION (MEMRA_SPEC_FUSED_T, t=2-4): the shexp gate+up pair
7029                // rides one shared quantize + one fused2 batched launch instead of two
7030                // decode-exact calls. Bit-identical per (tensor,token,row) — see spec_fused_t.
7031                let mut fused = None;
7032                if crate::spec::spec_fused_t()
7033                    && (2..=4).contains(&t)
7034                    && e.uses_q8_1_fast(gate_shexp)
7035                    && e.uses_q8_1_fast(up_shexp)
7036                {
7037                    let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
7038                    fused = e.matmul_q8_fused2_t(gate_shexp, up_shexp, &zq, &zd, t)?;
7039                }
7040                match fused {
7041                    Some(pair) => pair,
7042                    None => (
7043                        e.matmul_decode_exact(gate_shexp, z, t)?,
7044                        e.matmul_decode_exact(up_shexp, z, t)?,
7045                    ),
7046                }
7047            } else {
7048                (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?)
7049            };
7050            let mut sa = e.uninit(t * n_ff_sh)?; // silu_mul fully overwrites
7051            e.silu_mul(&sg_gate, &sg_up, &mut sa, t * n_ff_sh)?;
7052            let sh = if verify_t {
7053                e.matmul_decode_exact(down_shexp, &sa, t)?
7054            } else {
7055                e.matmul(down_shexp, &sa, t)?
7056            };
7057            // shexp gate: qwen35moe sigmoid-gates; M3 has no gate tensor -> weight 1.0.
7058            // Same fused sigmoid-dot as moe_ffn_sequential step 3 (byte-identity contract
7059            // between the two arms; prefill keeps the batched cuBLASLt linear).
7060            let g = match &m.gate_inp_shexp {
7061                Some(gate_inp_shexp) => {
7062                    // router_prefill_exact_on(): the fused sigmoid-dot is m-INVARIANT and
7063                    // serves EVERY t (serve isolation, lane/concat-prime-exact).
7064                    if t < PRIME_MIN_T || crate::router_prefill_exact_on() {
7065                        e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
7066                    } else {
7067                        let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
7068                        let mut g = e.uninit(t)?;
7069                        e.sigmoid(&gs, &mut g, t)?;
7070                        g
7071                    }
7072                }
7073                None => e.htod(&vec![1.0f32; t])?,
7074            };
7075            e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
7076        }
7077
7078        Ok(moe_out)
7079    }
7080
7081    /// STAGE-2 GROUPED DECODE (2026-07-04): run ONE token's whole routed-expert FFN in TWO
7082    /// launches when every one of its 3*n_used blocks is ALREADY cache-resident. Returns
7083    /// Ok(true) if the grouped path ran (caller skips the sequential loop for this token);
7084    /// Ok(false) on ANY miss (caller falls through — the sequential loop stages/admits as
7085    /// before, so the NEXT occurrence takes the grouped path). Pointer safety: cache slots are
7086    /// fixed-address for the engine's lifetime and the pure-HIT path performs no admission, so
7087    /// the collected raw pointers cannot move between collection and launch (single-threaded
7088    /// decode; the lock is held only for collection, launches are stream-ordered after any
7089    /// prior same-stream staging writes).
7090    #[allow(clippy::too_many_arguments)]
7091    /// q8 twin of moe_gdec_token (dp4a arc): same residency check + 2-launch shape; the fused
7092    /// kernels consume the pre-quantized z-row and re-quantize act per slot batch.
7093    #[allow(clippy::too_many_arguments)]
7094    fn moe_gdec_token_q8(
7095        e: &Engine,
7096        m: &MoeWeights,
7097        il: u16,
7098        max_block: usize,
7099        zq: &CudaSlice<i8>,
7100        zd: &CudaSlice<f32>,
7101        sel: &[u32],
7102        w: &[f32],
7103        moe_out: &mut CudaSlice<f32>,
7104        tok: usize,
7105        n_embd: usize,
7106        n_ff_exp: usize,
7107        n_used: usize,
7108    ) -> Result<bool, Box<dyn std::error::Error>> {
7109        use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
7110        use cudarc::driver::DevicePtr;
7111        let ptrs = e.with_moe_cache(max_block, |c, eng| {
7112            let mut g = [0u64; 8];
7113            let mut u = [0u64; 8];
7114            let mut d = [0u64; 8];
7115            for (j, &ex) in sel.iter().enumerate() {
7116                let ex = ex as u16;
7117                let (Some(sg), Some(su), Some(sd)) = (
7118                    c.resident(BlockId::new(il, PROJ_GATE, ex)),
7119                    c.resident(BlockId::new(il, PROJ_UP, ex)),
7120                    c.resident(BlockId::new(il, PROJ_DOWN, ex)),
7121                ) else {
7122                    return Ok(None);
7123                };
7124                let __s = eng.stream();
7125                let (pg, _e0) = c.slot(sg).device_ptr(&__s);
7126                let (pu, _e1) = c.slot(su).device_ptr(&__s);
7127                let (pd, _e2) = c.slot(sd).device_ptr(&__s);
7128                g[j] = pg as u64;
7129                u[j] = pu as u64;
7130                d[j] = pd as u64;
7131            }
7132            if cpu_expert_profile_admit_enabled() && !c.is_frozen() {
7133                for &ex in sel {
7134                    let ex = ex as u16;
7135                    for proj in [PROJ_GATE, PROJ_UP, PROJ_DOWN] {
7136                        c.note_profile_hit(BlockId::new(il, proj, ex));
7137                    }
7138                }
7139            }
7140            c.hits += (3 * n_used) as u64;
7141            Ok(Some((g, u, d)))
7142        })?;
7143        let Some((g, u, d)) = ptrs else {
7144            return Ok(false);
7145        };
7146        let mut wv = [0f32; 8];
7147        wv[..n_used].copy_from_slice(w);
7148        let act = e.moe_gate_up_silu8_q8(
7149            crate::WPtr8(g),
7150            crate::WPtr8(u),
7151            zq,
7152            zd,
7153            n_embd,
7154            n_ff_exp,
7155            n_used,
7156            m.gate_exps.qtype,
7157            m.up_exps.qtype,
7158            m.gate_exps.row_bytes,
7159            m.up_exps.row_bytes,
7160        )?;
7161        // per-slot act quantize: [n_used, n_ff] rows in one quantize launch.
7162        let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
7163        let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
7164        e.moe_down8_fma_q8(
7165            crate::WPtr8(d),
7166            crate::F32x8(wv),
7167            &aq2,
7168            &ad2,
7169            &mut dst,
7170            n_ff_exp,
7171            n_embd,
7172            n_used,
7173            m.down_exps.qtype,
7174            m.down_exps.row_bytes,
7175        )?;
7176        Ok(true)
7177    }
7178
7179    fn moe_gdec_token(
7180        e: &Engine,
7181        m: &MoeWeights,
7182        il: u16,
7183        max_block: usize,
7184        zt: &cudarc::driver::CudaView<f32>,
7185        sel: &[u32],
7186        w: &[f32],
7187        moe_out: &mut CudaSlice<f32>,
7188        tok: usize,
7189        n_embd: usize,
7190        n_ff_exp: usize,
7191        n_used: usize,
7192    ) -> Result<bool, Box<dyn std::error::Error>> {
7193        use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
7194        use cudarc::driver::DevicePtr;
7195        // One lock hold: residency-check all 3*n_used blocks, collect raw slot pointers.
7196        let ptrs = e.with_moe_cache(max_block, |c, eng| {
7197            let mut g = [0u64; 8];
7198            let mut u = [0u64; 8];
7199            let mut d = [0u64; 8];
7200            for (j, &ex) in sel.iter().enumerate() {
7201                let ex = ex as u16;
7202                let (Some(sg), Some(su), Some(sd)) = (
7203                    c.resident(BlockId::new(il, PROJ_GATE, ex)),
7204                    c.resident(BlockId::new(il, PROJ_UP, ex)),
7205                    c.resident(BlockId::new(il, PROJ_DOWN, ex)),
7206                ) else {
7207                    return Ok(None);
7208                };
7209                let __s = eng.stream();
7210                let (pg, _e0) = c.slot(sg).device_ptr(&__s);
7211                let (pu, _e1) = c.slot(su).device_ptr(&__s);
7212                let (pd, _e2) = c.slot(sd).device_ptr(&__s);
7213                g[j] = pg as u64;
7214                u[j] = pu as u64;
7215                d[j] = pd as u64;
7216            }
7217            if cpu_expert_profile_admit_enabled() && !c.is_frozen() {
7218                for &ex in sel {
7219                    let ex = ex as u16;
7220                    for proj in [PROJ_GATE, PROJ_UP, PROJ_DOWN] {
7221                        c.note_profile_hit(BlockId::new(il, proj, ex));
7222                    }
7223                }
7224            }
7225            c.hits += (3 * n_used) as u64; // instrumentation parity with dispatch()
7226            Ok(Some((g, u, d)))
7227        })?;
7228        let Some((g, u, d)) = ptrs else {
7229            return Ok(false);
7230        };
7231        let mut wv = [0f32; 8];
7232        wv[..n_used].copy_from_slice(w);
7233        // 2 launches: (gate+up+silu) x8, then (down + slot-ordered FMA accumulate) x8.
7234        let act = e.moe_gate_up_silu8(
7235            crate::WPtr8(g),
7236            crate::WPtr8(u),
7237            zt,
7238            n_embd,
7239            n_ff_exp,
7240            n_used,
7241            m.gate_exps.qtype,
7242            m.up_exps.qtype,
7243            m.gate_exps.row_bytes,
7244            m.up_exps.row_bytes,
7245        )?;
7246        let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
7247        e.moe_down8_fma_into(
7248            crate::WPtr8(d),
7249            crate::F32x8(wv),
7250            &act,
7251            &mut dst,
7252            n_ff_exp,
7253            n_embd,
7254            n_used,
7255            m.down_exps.qtype,
7256            m.down_exps.row_bytes,
7257        )?;
7258        Ok(true)
7259    }
7260
7261    /// EDGE-1 §B.3: dispatch one expert projection through the SLRU cache, then run the SAME
7262    /// `qmatvec_view` from whichever slot it landed in (resident HIT or staged MISS). `x` is the
7263    /// sliced activation row. `proj` selects the gate/up/down HostExps tensor. Returns y = W_expert @ x.
7264    /// q8 twin of moe_cached_gemm: same dispatch/slot mechanics, dp4a expert kernel.
7265    fn moe_cached_gemm_q8(
7266        e: &Engine,
7267        il: u16,
7268        proj: u8,
7269        ex: usize,
7270        m: &MoeWeights,
7271        max_block: usize,
7272        aq: &CudaSlice<i8>,
7273        ad: &CudaSlice<f32>,
7274    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7275        use crate::moe_cache::{BlockId, DispatchSlot, PROJ_GATE, PROJ_UP};
7276        let exps = match proj {
7277            PROJ_GATE => &m.gate_exps,
7278            PROJ_UP => &m.up_exps,
7279            _ => &m.down_exps,
7280        };
7281        let layout = exps.expert_layout(ex);
7282        let id = BlockId::new(il, proj, ex as u16);
7283        let source = exps.expert_source(ex);
7284        e.with_moe_cache(max_block, |c, eng| {
7285            let slot = c.dispatch_source(id, source, eng)?;
7286            let DispatchSlot::Resident(sl) = slot;
7287            let buf = c.slot(sl);
7288            eng.qmatvec_expert_q8(
7289                buf,
7290                0..layout.len,
7291                aq,
7292                ad,
7293                1,
7294                exps.in_f,
7295                exps.out_f,
7296                layout.qtype,
7297                layout.row_bytes,
7298            )
7299        })
7300    }
7301
7302    fn moe_cached_gemm(
7303        e: &Engine,
7304        il: u16,
7305        proj: u8,
7306        ex: usize,
7307        m: &MoeWeights,
7308        max_block: usize,
7309        x: &cudarc::driver::CudaView<f32>,
7310    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7311        use crate::moe_cache::{BlockId, DispatchSlot, PROJ_GATE, PROJ_UP};
7312        let exps = match proj {
7313            PROJ_GATE => &m.gate_exps,
7314            PROJ_UP => &m.up_exps,
7315            _ => &m.down_exps,
7316        };
7317        let layout = exps.expert_layout(ex);
7318        let id = BlockId::new(il, proj, ex as u16);
7319        let source = exps.expert_source(ex);
7320        // dispatch under the lock (lookup/admit/memcpy-issue), then resolve the slot and GEMM.
7321        e.with_moe_cache(max_block, |c, eng| {
7322            let slot = c.dispatch_source(id, source, eng)?;
7323            // resolve the device buffer for this slot; the GEMM is enqueued on the compute stream
7324            // (the same stream the memcpy was issued on, so ordering holds without extra sync).
7325            let DispatchSlot::Resident(sl) = slot;
7326            let buf = c.slot(sl);
7327            eng.qmatvec_view(
7328                buf,
7329                0..layout.len,
7330                x,
7331                1,
7332                exps.in_f,
7333                exps.out_f,
7334                layout.qtype,
7335                layout.row_bytes,
7336            )
7337        })
7338    }
7339
7340    /// Populate the warmup cache for an expert whose current-token output intentionally ran on
7341    /// CPU. No GEMM is launched and callers invoke this only after the CPU result has completed,
7342    /// so the current forward's backend assignment and output remain unchanged.
7343    fn moe_profile_admit_expert(
7344        e: &Engine,
7345        il: u16,
7346        ex: usize,
7347        m: &MoeWeights,
7348        max_block: usize,
7349    ) -> Result<(), Box<dyn std::error::Error>> {
7350        use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
7351        e.with_moe_cache(max_block, |cache, eng| {
7352            for (proj, exps) in [
7353                (PROJ_GATE, &m.gate_exps),
7354                (PROJ_UP, &m.up_exps),
7355                (PROJ_DOWN, &m.down_exps),
7356            ] {
7357                let id = BlockId::new(il, proj, ex as u16);
7358                let _ = cache.dispatch_source(id, exps.expert_source(ex), eng)?;
7359            }
7360            Ok(())
7361        })
7362    }
7363
7364    /// Read a projection from the immutable residency set when present; otherwise use one
7365    /// transient slot without admitting or evicting anything. Used only by post-freeze prefill.
7366    #[allow(clippy::too_many_arguments)]
7367    fn moe_frozen_gemm(
7368        e: &Engine,
7369        il: u16,
7370        proj: u8,
7371        ex: usize,
7372        m: &MoeWeights,
7373        max_block: usize,
7374        x: &cudarc::driver::CudaView<f32>,
7375        scratch: &mut Option<CudaSlice<u8>>,
7376        scratch_len: usize,
7377    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7378        use crate::moe_cache::{BlockId, PROJ_GATE, PROJ_UP};
7379        let exps = match proj {
7380            PROJ_GATE => &m.gate_exps,
7381            PROJ_UP => &m.up_exps,
7382            _ => &m.down_exps,
7383        };
7384        let layout = exps.expert_layout(ex);
7385        let id = BlockId::new(il, proj, ex as u16);
7386        if let Some(output) = e.with_moe_cache(max_block, |cache, eng| {
7387            let Some(slot) = cache.resident(id) else {
7388                return Ok(None);
7389            };
7390            let buf = cache.slot(slot);
7391            Ok(Some(eng.qmatvec_view(
7392                buf,
7393                0..layout.len,
7394                x,
7395                1,
7396                exps.in_f,
7397                exps.out_f,
7398                layout.qtype,
7399                layout.row_bytes,
7400            )?))
7401        })? {
7402            return Ok(output);
7403        }
7404        if scratch.is_none() {
7405            *scratch = Some(e.alloc_u8_uninit(scratch_len)?);
7406        }
7407        let scratch = scratch.as_mut().unwrap();
7408        e.stage_expert(exps.expert_bytes(ex), scratch, 0)?;
7409        e.qmatvec_view(
7410            scratch,
7411            0..layout.len,
7412            x,
7413            1,
7414            exps.in_f,
7415            exps.out_f,
7416            layout.qtype,
7417            layout.row_bytes,
7418        )
7419    }
7420
7421    fn moe_prefetch_expert(
7422        e: &Engine,
7423        il: u16,
7424        ex: usize,
7425        m: &MoeWeights,
7426        max_block: usize,
7427        keep: &[crate::moe_cache::BlockId],
7428    ) -> Result<(), Box<dyn std::error::Error>> {
7429        use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
7430        e.with_moe_cache(max_block, |c, eng| {
7431            for (proj, exps) in [
7432                (PROJ_GATE, &m.gate_exps),
7433                (PROJ_UP, &m.up_exps),
7434                (PROJ_DOWN, &m.down_exps),
7435            ] {
7436                let id = BlockId::new(il, proj, ex as u16);
7437                let _ = c.prefetch_source(id, exps.expert_source(ex), keep, eng)?;
7438            }
7439            Ok(())
7440        })
7441    }
7442
7443    /// Worker-mode disk lookahead for grouped prefill. Memory sources are deliberately skipped so
7444    /// selecting `worker` changes only storage scheduling here; all CUDA work stays in dispatch.
7445    fn moe_prefetch_disk_expert(
7446        e: &Engine,
7447        il: u16,
7448        ex: usize,
7449        m: &MoeWeights,
7450        max_block: usize,
7451        keep: &[crate::moe_cache::BlockId],
7452    ) -> Result<(), Box<dyn std::error::Error>> {
7453        use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
7454        e.with_moe_cache(max_block, |c, eng| {
7455            for (proj, exps) in [
7456                (PROJ_GATE, &m.gate_exps),
7457                (PROJ_UP, &m.up_exps),
7458                (PROJ_DOWN, &m.down_exps),
7459            ] {
7460                let source = exps.expert_source(ex);
7461                if let crate::model::ExpertSource::Disk { .. } = &source {
7462                    let id = BlockId::new(il, proj, ex as u16);
7463                    let _ = c.prefetch_source(id, source, keep, eng)?;
7464                }
7465            }
7466            Ok(())
7467        })
7468    }
7469
7470    #[inline]
7471    fn moe_prefetch_host_expert(ex: usize, m: &MoeWeights) {
7472        let _ = m.gate_exps.prefetch_expert_pages(ex);
7473        let _ = m.up_exps.prefetch_expert_pages(ex);
7474        let _ = m.down_exps.prefetch_expert_pages(ex);
7475    }
7476}
7477
7478// ================================================================================================
7479// A2: EXPERT-GROUPED MoE PREFILL (MEMRA_MOE_GROUPED=1). Resident-case prototype.
7480//
7481// Instead of the per-token loop (T * 8 experts * 3 projections = 12024 individual m=1 matvecs),
7482// this groups tokens by expert and runs ONE matmul per active expert per projection at m=m_e.
7483// On a 501-token prefill with ~170 active experts, that's ~510 matmuls (vs 12024).
7484//
7485// EXACTNESS: per-token accumulation across its 8 experts is reordered (grouped processes experts
7486// in expert-id order, not the router's top-k order). To preserve bit-identity with the sequential
7487// loop, we use an 8-SLOT scheme: expert outputs are scattered into slots keyed by the token's
7488// top-k position (0..7), then reduced in that fixed order. This makes the f32 addition order
7489// identical to the per-token loop regardless of expert processing order.
7490//
7491// Memory: T * 8 * n_embd * 4 = 501 * 8 * 2048 * 4 = ~32 MB (slot buffer). Fine on 96GB.
7492// ================================================================================================
7493
7494impl HybridModel {
7495    /// Resident-slab fast path for host-routed grouped prefill. Unclamped layers batch the
7496    /// sequential fused q8 program over the token axis; clamped layers use the separate
7497    /// expert-major q8 chain so `ffn_act_lim` remains authoritative.
7498    #[allow(clippy::too_many_arguments)]
7499    fn moe_ffn_grouped_resident_q8(
7500        e: &Engine,
7501        m: &MoeWeights,
7502        z: &CudaSlice<f32>,
7503        t: usize,
7504        cfg: &ModelConfig,
7505        il: u16,
7506        sel_all: &[u32],
7507        w_all: &[f32],
7508        table: &CudaSlice<u64>,
7509        gu_il: bool,
7510    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7511        let moe = cfg.moe.as_ref().unwrap();
7512        let n_embd = cfg.n_embd as usize;
7513        let n_expert = moe.expert_count as usize;
7514        let n_used = moe.expert_used_count as usize;
7515        let n_ff_exp = moe.expert_ff_length as usize;
7516        let n_pairs = t * n_used;
7517        debug_assert_eq!(sel_all.len(), n_pairs);
7518        debug_assert_eq!(w_all.len(), n_pairs);
7519        debug_assert!(
7520            m.gate_exps.macros.is_none()
7521                && m.up_exps.macros.is_none()
7522                && m.down_exps.macros.is_none(),
7523            "resident grouped q8 does not fold per-expert macro scales",
7524        );
7525
7526        // The rows twins run the resident sequential program verbatim on grid.z = token:
7527        // fused gate/up/SiLU per slot, batched activation quantization, then the original
7528        // slot-ordered down/FMA chain. Routing remains the sigmoid selector above (device by
7529        // default; MEMRA_SIG_ROUTER=0 is the host oracle); these kernels consume sel/w only and
7530        // never enter the softmax router.
7531        if !cfg.swiglu_clamped_at(il as u32) {
7532            let sel: Vec<i32> = sel_all.iter().map(|&expert| expert as i32).collect();
7533            let sel_d = e.htod_i32(&sel)?;
7534            let w_d = e.htod(w_all)?;
7535            let (gate_row_bytes, up_row_bytes) = if gu_il {
7536                let combined = m.gate_exps.row_bytes + m.up_exps.row_bytes;
7537                (combined, combined)
7538            } else {
7539                (m.gate_exps.row_bytes, m.up_exps.row_bytes)
7540            };
7541            let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
7542            let act = e.moe_gate_up_silu8_dev_q8_rows(
7543                table,
7544                &sel_d,
7545                &zq,
7546                &zd,
7547                t,
7548                n_embd,
7549                n_ff_exp,
7550                n_used,
7551                n_expert,
7552                m.gate_exps.qtype,
7553                m.up_exps.qtype,
7554                gate_row_bytes,
7555                up_row_bytes,
7556                &m.dev_macros,
7557            )?;
7558            let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
7559            let mut moe_out = e.uninit(t * n_embd)?;
7560            e.moe_down8_fma_dev_q8_rows_g(
7561                table,
7562                &sel_d,
7563                &w_d,
7564                &aq2,
7565                &ad2,
7566                &mut moe_out,
7567                t,
7568                n_ff_exp,
7569                n_embd,
7570                n_used,
7571                n_expert,
7572                m.down_exps.qtype,
7573                m.down_exps.row_bytes,
7574            )?;
7575
7576            if std::env::var("MEMRA_MOE_STATS").is_ok() {
7577                let mut counts = vec![0usize; n_expert];
7578                for &expert in sel_all {
7579                    counts[expert as usize] += 1;
7580                }
7581                let mut sizes: Vec<usize> =
7582                    counts.into_iter().filter(|&count| count != 0).collect();
7583                sizes.sort_unstable();
7584                let mean = sizes.iter().sum::<usize>() as f64 / sizes.len().max(1) as f64;
7585                println!(
7586                    "moe-grouped il={il} t={t} dispatch=resident-q8-rows active={}/{} \
7587                     m_e: min={} median={} mean={mean:.1} max={}",
7588                    sizes.len(),
7589                    n_expert,
7590                    sizes.first().copied().unwrap_or(0),
7591                    sizes.get(sizes.len() / 2).copied().unwrap_or(0),
7592                    sizes.last().copied().unwrap_or(0),
7593                );
7594            }
7595            return Ok(moe_out);
7596        }
7597
7598        // Clamped layers keep gate/up, activation, and down as separate stages. The pair-major
7599        // matvec body is qmatvec_expert_q8 verbatim, while one launch covers every routed pair.
7600        // Pair ids stay in router slot order so scatter preserves the sequential FMA chain.
7601        let pair_tok: Vec<i32> = (0..n_pairs).map(|pair| (pair / n_used) as i32).collect();
7602        let pair_ex: Vec<i32> = sel_all.iter().map(|&expert| expert as i32).collect();
7603        let tok_off: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
7604        let tok_ids: Vec<i32> = (0..n_pairs as i32).collect();
7605
7606        let mut by_expert: Vec<Vec<i32>> = vec![Vec::new(); n_expert];
7607        for (pair, &expert) in pair_ex.iter().enumerate() {
7608            by_expert[expert as usize].push(pair as i32);
7609        }
7610
7611        let pair_tok_d = e.htod_i32(&pair_tok)?;
7612        let pair_ex_d = e.htod_i32(&pair_ex)?;
7613        let pair_w_d = e.htod(w_all)?;
7614        let tok_off_d = e.htod_i32(&tok_off)?;
7615        let tok_ids_d = e.htod_i32(&tok_ids)?;
7616
7617        let matvec = |proj: i32,
7618                      pair_rows: &CudaSlice<i32>,
7619                      aq: &CudaSlice<i8>,
7620                      ad: &CudaSlice<f32>,
7621                      in_f: usize,
7622                      out_f: usize,
7623                      qtype: i32,
7624                      row_bytes: usize|
7625         -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7626            e.moe_pairs_matvec_q8(
7627                table, proj, pair_rows, &pair_ex_d, aq, ad, in_f, out_f, n_expert, n_pairs, qtype,
7628                row_bytes,
7629            )
7630        };
7631
7632        let (gate_row_bytes, up_row_bytes) = if gu_il {
7633            let combined = m.gate_exps.row_bytes + m.up_exps.row_bytes;
7634            (combined, combined)
7635        } else {
7636            (m.gate_exps.row_bytes, m.up_exps.row_bytes)
7637        };
7638        let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
7639        let gate = matvec(
7640            0,
7641            &pair_tok_d,
7642            &zq,
7643            &zd,
7644            n_embd,
7645            n_ff_exp,
7646            m.gate_exps.qtype,
7647            gate_row_bytes,
7648        )?;
7649        let up = matvec(
7650            1,
7651            &pair_tok_d,
7652            &zq,
7653            &zd,
7654            n_embd,
7655            n_ff_exp,
7656            m.up_exps.qtype,
7657            up_row_bytes,
7658        )?;
7659        let mut act = e.uninit(n_pairs * n_ff_exp)?;
7660        Self::ffn_act_lim(
7661            e,
7662            cfg,
7663            &gate,
7664            &up,
7665            1.0,
7666            1.0,
7667            cfg.clamp_exp_at(il as u32),
7668            &mut act,
7669            n_pairs * n_ff_exp,
7670        )?;
7671        let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
7672        let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
7673        let pair_self_d = e.htod_i32(&pair_self)?;
7674        let down = matvec(
7675            2,
7676            &pair_self_d,
7677            &aq2,
7678            &ad2,
7679            n_ff_exp,
7680            n_embd,
7681            m.down_exps.qtype,
7682            m.down_exps.row_bytes,
7683        )?;
7684        let mut moe_out = e.uninit(t * n_embd)?;
7685        e.moe_pairs_scatter(
7686            &down,
7687            &pair_w_d,
7688            &tok_off_d,
7689            &tok_ids_d,
7690            &mut moe_out,
7691            t,
7692            n_embd,
7693        )?;
7694
7695        if std::env::var("MEMRA_MOE_STATS").is_ok() {
7696            let mut sizes: Vec<usize> = by_expert
7697                .iter()
7698                .filter_map(|pairs| (!pairs.is_empty()).then_some(pairs.len()))
7699                .collect();
7700            sizes.sort_unstable();
7701            let mean = sizes.iter().sum::<usize>() as f64 / sizes.len().max(1) as f64;
7702            println!(
7703                "moe-grouped il={il} t={t} dispatch=resident-q8-clamped-pairs active={}/{} \
7704                 m_e: min={} median={} mean={mean:.1} max={}",
7705                sizes.len(),
7706                n_expert,
7707                sizes.first().copied().unwrap_or(0),
7708                sizes.get(sizes.len() / 2).copied().unwrap_or(0),
7709                sizes.last().copied().unwrap_or(0),
7710            );
7711        }
7712        Ok(moe_out)
7713    }
7714
7715    fn moe_ffn_grouped_add_shared(
7716        e: &Engine,
7717        m: &MoeWeights,
7718        z: &CudaSlice<f32>,
7719        t: usize,
7720        cfg: &ModelConfig,
7721        il: u16,
7722        moe_out: &mut CudaSlice<f32>,
7723    ) -> Result<(), Box<dyn std::error::Error>> {
7724        let n_embd = cfg.n_embd as usize;
7725        if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
7726            (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
7727        {
7728            let n_ff_sh = gate_shexp.out_features();
7729            let sg_gate = e.matmul(gate_shexp, z, t)?;
7730            let sg_up = e.matmul(up_shexp, z, t)?;
7731            let mut sa = e.uninit(t * n_ff_sh)?;
7732            Self::ffn_act_lim(
7733                e,
7734                cfg,
7735                &sg_gate,
7736                &sg_up,
7737                1.0,
7738                1.0,
7739                cfg.clamp_shexp_at(il as u32),
7740                &mut sa,
7741                t * n_ff_sh,
7742            )?;
7743            let sh = e.matmul(down_shexp, &sa, t)?;
7744            let gate = match &m.gate_inp_shexp {
7745                Some(gate_inp_shexp) => {
7746                    if t < PRIME_MIN_T || crate::router_prefill_exact_on() {
7747                        e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
7748                    } else {
7749                        let raw = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
7750                        let mut gate = e.uninit(t)?;
7751                        e.sigmoid(&raw, &mut gate, t)?;
7752                        gate
7753                    }
7754                }
7755                None => e.htod(&vec![1.0f32; t])?,
7756            };
7757            e.add_scaled_rows(&sh, &gate, moe_out, n_embd, t)?;
7758        }
7759        Ok(())
7760    }
7761
7762    /// A2 expert-grouped MoE FFN (prefill path, MEMRA_MOE_GROUPED=1). Same semantics as moe_ffn:
7763    /// z [T, n_embd] -> moe_out [T, n_embd]. BIT-IDENTICAL to moe_ffn when using the slot scheme.
7764    pub(crate) fn moe_ffn_grouped(
7765        e: &Engine,
7766        m: &MoeWeights,
7767        z: &CudaSlice<f32>,
7768        t: usize,
7769        cfg: &ModelConfig,
7770        il: u16,
7771        max_block: usize,
7772    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7773        let moe = cfg.moe.as_ref().unwrap();
7774        let n_embd = cfg.n_embd as usize;
7775        let n_expert = moe.expert_count as usize;
7776        let n_used = moe.expert_used_count as usize;
7777        let n_ff_exp = moe.expert_ff_length as usize;
7778        // step35 per-layer SwiGLU clamp; None on every other arch / unclamped layer.
7779        let lim_exp = cfg.clamp_exp_at(il as u32);
7780
7781        // 1. ROUTER: exactly the same m-invariant sigmoid selector as the sequential path
7782        // (device by default; MEMRA_SIG_ROUTER=0 is the host oracle). The grouped dispatch never
7783        // enters the softmax-only pairs/dev router.
7784        let logits = Self::moe_router_logits(e, m, z, t, cfg)?;
7785        if let Some(sig) = cfg.sigmoid_router() {
7786            Self::trace_sigmoid_router_logits(e, il, t, n_expert, n_used, &logits, m, sig)?;
7787        }
7788        let (sel_all, w_all) = if let Some(sig) = cfg.sigmoid_router() {
7789            Self::moe_route_sigmoid_cfg(e, &logits, t, n_expert, n_used, m, sig)?
7790        } else {
7791            Self::moe_route_cfg(e, &logits, t, n_expert, n_used, m.active_experts.as_deref())?
7792        };
7793        crate::moesd::record_host_routes(il, n_expert, n_used, &sel_all)?;
7794        Self::trace_moe_routes(il, t, &sel_all, &w_all)?;
7795        Self::trace_moe_input(e, il, t, n_embd, z)?;
7796
7797        // Fits-VRAM Step35 under the PP stage split: use the expert-major q8 arithmetic directly
7798        // from the owning device's uniform resident slab. This is not `moe_ffn_pairs`: routing
7799        // already happened above and the activation is clamp-aware. Mixed layouts, remote slabs,
7800        // macro-scaled experts, q8-disabled configs, and spill fall through to metadata-aware A2.
7801        let no_exp_macros = m.gate_exps.macros.is_none()
7802            && m.up_exps.macros.is_none()
7803            && m.down_exps.macros.is_none();
7804        let resident_q8 = m.dev_exps.as_ref().filter(|dev| {
7805            m.has_uniform_expert_layout()
7806                && no_exp_macros
7807                && moe_q8_enabled()
7808                && q8_expert_supported(m.gate_exps.qtype)
7809                && q8_expert_supported(m.up_exps.qtype)
7810                && q8_expert_supported(m.down_exps.qtype)
7811                && moe_slab_enabled()
7812                && dev.dev == e.ctx().ordinal()
7813        });
7814        if let Some(dev) = resident_q8 {
7815            let mut moe_out = Self::moe_ffn_grouped_resident_q8(
7816                e,
7817                m,
7818                z,
7819                t,
7820                cfg,
7821                il,
7822                &sel_all,
7823                &w_all,
7824                &dev.ptr_row,
7825                dev.gu_il,
7826            )?;
7827            Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut moe_out)?;
7828            return Ok(moe_out);
7829        }
7830
7831        // 2. BUILD PER-EXPERT TOKEN LISTS (host-side grouping).
7832        // For each expert e, we need: which tokens use it, their positions in z, their top-k
7833        // slot index (for bit-identical accumulation), and their weights.
7834        struct ExpertGroup {
7835            tok_indices: Vec<i32>,  // indices into z rows (0..T-1)
7836            slot_indices: Vec<i32>, // top-k slot (0..n_used-1) for that token-expert pair
7837            weights: Vec<f32>,      // renormalized weight for that token-expert pair
7838        }
7839        let mut groups: Vec<ExpertGroup> = (0..n_expert)
7840            .map(|_| ExpertGroup {
7841                tok_indices: Vec::new(),
7842                slot_indices: Vec::new(),
7843                weights: Vec::new(),
7844            })
7845            .collect();
7846
7847        for tok in 0..t {
7848            for j in 0..n_used {
7849                let ex = sel_all[tok * n_used + j] as usize;
7850                let w = w_all[tok * n_used + j];
7851                groups[ex].tok_indices.push(tok as i32);
7852                groups[ex].slot_indices.push(j as i32);
7853                groups[ex].weights.push(w);
7854            }
7855        }
7856
7857        // 3. ALLOCATE SLOT BUFFER: [T, n_used, n_embd] f32, zero-initialized.
7858        // Each token's 8 expert contributions land in their respective slots.
7859        let mut slot_buf = e.zeros(t * n_used * n_embd)?;
7860        let mut wbuf = e.zeros(t * n_used)?; // [T, n_used] weight buffer for FMA reduce
7861
7862        // Expert weight dimensions (used in both cache and staging paths).
7863        let g_len = m.gate_exps.max_expert_bytes();
7864        let u_len = m.up_exps.max_expert_bytes();
7865        let d_len = m.down_exps.max_expert_bytes();
7866        let moe_q8 = m.has_uniform_expert_layout()
7867            && moe_q8_enabled()
7868            && q8_expert_supported(m.gate_exps.qtype)
7869            && q8_expert_supported(m.up_exps.qtype)
7870            && q8_expert_supported(m.down_exps.qtype);
7871        // The per-expert slab view is legal only for the ordinary contiguous uniform layout.
7872        // Interleaved GU slabs require the pointer-table fast path above.
7873        let slab_local = m
7874            .dev_exps
7875            .as_ref()
7876            .filter(|dev| !dev.gu_il && moe_slab_enabled() && dev.dev == e.ctx().ordinal());
7877        let use_cache =
7878            slab_local.is_none() && Engine::moe_cache_enabled() && !e.moe_cache_frozen();
7879        // The sequential no-cache/frozen staging oracle is f32. Use q8 only where sequential
7880        // also does: a local resident slab or a live SLRU dispatch.
7881        let grouped_q8 = moe_q8 && (slab_local.is_some() || use_cache);
7882
7883        // GPU scratch for staging (only allocated without a local slab or cache).
7884        let (mut scratch_g, mut scratch_u, mut scratch_d) = if slab_local.is_none() && !use_cache {
7885            (
7886                Some(e.alloc_u8(g_len)?),
7887                Some(e.alloc_u8(u_len)?),
7888                Some(e.alloc_u8(d_len)?),
7889            )
7890        } else {
7891            (None, None, None)
7892        };
7893
7894        // 4. PER ACTIVE EXPERT: gather, compute, scatter.
7895        // Processing ORDER: DESCENDING m_e (biggest token batches first) — the concluded winner
7896        // (rig5090 2026-07-04, the ascending-id arm and its MEMRA_MOE_ORDER seam removed): desc is
7897        // a first-forward win at partial cache capacity — the hot (big-m_e) experts are admitted
7898        // to the SLRU before the small-m_e tail can pollute it, so residency converges in ONE
7899        // forward instead of several: auto-cache T=501 126.9 -> 169.9 tok/s (1.34x), cap512
7900        // 119.6 -> 160.8 (and kills the rep-to-rep bimodal); wash (<2%) at cap64 pure-spill and
7901        // at long prompts where every expert stages regardless. Order is FREE to change without
7902        // breaking the byte-identity gate: the slot scheme pins each token's accumulation order
7903        // regardless of expert processing order (the whole point of the slots).
7904        let mut order: Vec<usize> = (0..n_expert)
7905            .filter(|&ex| !groups[ex].tok_indices.is_empty())
7906            .collect();
7907        order.sort_by(|&a, &b| {
7908            groups[b]
7909                .tok_indices
7910                .len()
7911                .cmp(&groups[a].tok_indices.len())
7912                .then(a.cmp(&b))
7913        });
7914        let mut m_dist: Vec<usize> = Vec::new(); // for stats
7915        let page_window = moe_page_prefetch_window();
7916        let worker_disk_prefetch = use_cache && crate::spill_pread::worker_enabled();
7917        if worker_disk_prefetch {
7918            if let Some(first) = grouped_worker_prefetch_position(order.len(), None) {
7919                Self::moe_prefetch_disk_expert(e, il, order[first], m, max_block, &[])?;
7920            }
7921        }
7922        for (order_pos, &ex) in order.iter().enumerate() {
7923            for next in page_prefetch_positions(order_pos, order.len(), page_window) {
7924                Self::moe_prefetch_host_expert(order[next], m);
7925            }
7926            if worker_disk_prefetch {
7927                if let Some(next) = grouped_worker_prefetch_position(order.len(), Some(order_pos)) {
7928                    use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
7929                    let keep = [
7930                        BlockId::new(il, PROJ_GATE, ex as u16),
7931                        BlockId::new(il, PROJ_UP, ex as u16),
7932                        BlockId::new(il, PROJ_DOWN, ex as u16),
7933                    ];
7934                    Self::moe_prefetch_disk_expert(e, il, order[next], m, max_block, &keep)?;
7935                }
7936            }
7937            let grp = &groups[ex];
7938            let m_e = grp.tok_indices.len();
7939            m_dist.push(m_e);
7940            let gl = m.gate_exps.expert_layout(ex);
7941            let ul = m.up_exps.expert_layout(ex);
7942            let dl = m.down_exps.expert_layout(ex);
7943
7944            // Upload index/weight arrays to device. The down-proj per-expert macro-scale
7945            // (ModelOpt weight_scale_2) folds into the scatter weights — post-matmul linear,
7946            // same fold as the sequential loop's `w[j] * macro_scale(ex)`. 1.0 for GGUF experts.
7947            let tok_idx_d = e.htod_i32(&grp.tok_indices)?;
7948            let slot_idx_d = e.htod_i32(&grp.slot_indices)?;
7949            let dmac = m.down_exps.macro_scale(ex);
7950            let weight_d = if dmac == 1.0 {
7951                e.htod(&grp.weights)?
7952            } else {
7953                let scaled: Vec<f32> = grp.weights.iter().map(|&w| w * dmac).collect();
7954                e.htod(&scaled)?
7955            };
7956
7957            // GATHER: collect m_e activation rows from z into a contiguous buffer.
7958            let mut gathered = e.zeros(m_e * n_embd)?;
7959            e.gather_rows(z, &tok_idx_d, &mut gathered, n_embd, m_e)?;
7960            let gv = gathered.slice(0..m_e * n_embd);
7961
7962            // Compute gate/up/down from the local slab, metadata-aware cache, or staging. The q8
7963            // form keeps each gathered row in the sequential arithmetic class at `m=m_e`.
7964            let y = if let Some(dev) = slab_local {
7965                let gate_start = ex * m.gate_exps.expert_stride;
7966                let up_start = ex * m.up_exps.expert_stride;
7967                let down_start = ex * m.down_exps.expert_stride;
7968                if grouped_q8 {
7969                    let (zq, zd) = e.quantize_q8_1(&gathered, m_e, n_embd)?;
7970                    let gate = e.qmatvec_expert_q8(
7971                        &dev.gate,
7972                        gate_start..gate_start + gl.len,
7973                        &zq,
7974                        &zd,
7975                        m_e,
7976                        m.gate_exps.in_f,
7977                        m.gate_exps.out_f,
7978                        gl.qtype,
7979                        gl.row_bytes,
7980                    )?;
7981                    let up = e.qmatvec_expert_q8(
7982                        &dev.up,
7983                        up_start..up_start + ul.len,
7984                        &zq,
7985                        &zd,
7986                        m_e,
7987                        m.up_exps.in_f,
7988                        m.up_exps.out_f,
7989                        ul.qtype,
7990                        ul.row_bytes,
7991                    )?;
7992                    let mut act = e.uninit(m_e * n_ff_exp)?;
7993                    Self::ffn_act_lim(
7994                        e,
7995                        cfg,
7996                        &gate,
7997                        &up,
7998                        m.gate_exps.macro_scale(ex),
7999                        m.up_exps.macro_scale(ex),
8000                        lim_exp,
8001                        &mut act,
8002                        m_e * n_ff_exp,
8003                    )?;
8004                    let (aq2, ad2) = e.quantize_q8_1(&act, m_e, n_ff_exp)?;
8005                    e.qmatvec_expert_q8(
8006                        &dev.down,
8007                        down_start..down_start + dl.len,
8008                        &aq2,
8009                        &ad2,
8010                        m_e,
8011                        m.down_exps.in_f,
8012                        m.down_exps.out_f,
8013                        dl.qtype,
8014                        dl.row_bytes,
8015                    )?
8016                } else {
8017                    let gate = e.qmatvec_view(
8018                        &dev.gate,
8019                        gate_start..gate_start + gl.len,
8020                        &gv,
8021                        m_e,
8022                        m.gate_exps.in_f,
8023                        m.gate_exps.out_f,
8024                        gl.qtype,
8025                        gl.row_bytes,
8026                    )?;
8027                    let up = e.qmatvec_view(
8028                        &dev.up,
8029                        up_start..up_start + ul.len,
8030                        &gv,
8031                        m_e,
8032                        m.up_exps.in_f,
8033                        m.up_exps.out_f,
8034                        ul.qtype,
8035                        ul.row_bytes,
8036                    )?;
8037                    let mut act = e.uninit(m_e * n_ff_exp)?;
8038                    Self::ffn_act_lim(
8039                        e,
8040                        cfg,
8041                        &gate,
8042                        &up,
8043                        m.gate_exps.macro_scale(ex),
8044                        m.up_exps.macro_scale(ex),
8045                        lim_exp,
8046                        &mut act,
8047                        m_e * n_ff_exp,
8048                    )?;
8049                    let actv = act.slice(0..m_e * n_ff_exp);
8050                    e.qmatvec_view(
8051                        &dev.down,
8052                        down_start..down_start + dl.len,
8053                        &actv,
8054                        m_e,
8055                        m.down_exps.in_f,
8056                        m.down_exps.out_f,
8057                        dl.qtype,
8058                        dl.row_bytes,
8059                    )?
8060                }
8061            } else if use_cache {
8062                use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
8063                if grouped_q8 {
8064                    let (zq, zd) = e.quantize_q8_1(&gathered, m_e, n_embd)?;
8065                    let gate = e.with_moe_cache(max_block, |cache, eng| {
8066                        let id = BlockId::new(il, PROJ_GATE, ex as u16);
8067                        let slot = cache.dispatch_source(id, m.gate_exps.expert_source(ex), eng)?;
8068                        eng.qmatvec_expert_q8(
8069                            cache.buf(slot),
8070                            0..gl.len,
8071                            &zq,
8072                            &zd,
8073                            m_e,
8074                            m.gate_exps.in_f,
8075                            m.gate_exps.out_f,
8076                            gl.qtype,
8077                            gl.row_bytes,
8078                        )
8079                    })?;
8080                    let up = e.with_moe_cache(max_block, |cache, eng| {
8081                        let id = BlockId::new(il, PROJ_UP, ex as u16);
8082                        let slot = cache.dispatch_source(id, m.up_exps.expert_source(ex), eng)?;
8083                        eng.qmatvec_expert_q8(
8084                            cache.buf(slot),
8085                            0..ul.len,
8086                            &zq,
8087                            &zd,
8088                            m_e,
8089                            m.up_exps.in_f,
8090                            m.up_exps.out_f,
8091                            ul.qtype,
8092                            ul.row_bytes,
8093                        )
8094                    })?;
8095                    let mut act = e.uninit(m_e * n_ff_exp)?;
8096                    Self::ffn_act_lim(
8097                        e,
8098                        cfg,
8099                        &gate,
8100                        &up,
8101                        m.gate_exps.macro_scale(ex),
8102                        m.up_exps.macro_scale(ex),
8103                        lim_exp,
8104                        &mut act,
8105                        m_e * n_ff_exp,
8106                    )?;
8107                    let (aq2, ad2) = e.quantize_q8_1(&act, m_e, n_ff_exp)?;
8108                    e.with_moe_cache(max_block, |cache, eng| {
8109                        let id = BlockId::new(il, PROJ_DOWN, ex as u16);
8110                        let slot = cache.dispatch_source(id, m.down_exps.expert_source(ex), eng)?;
8111                        eng.qmatvec_expert_q8(
8112                            cache.buf(slot),
8113                            0..dl.len,
8114                            &aq2,
8115                            &ad2,
8116                            m_e,
8117                            m.down_exps.in_f,
8118                            m.down_exps.out_f,
8119                            dl.qtype,
8120                            dl.row_bytes,
8121                        )
8122                    })?
8123                } else {
8124                    let gate = e.with_moe_cache(max_block, |cache, eng| {
8125                        let id = BlockId::new(il, PROJ_GATE, ex as u16);
8126                        let slot = cache.dispatch_source(id, m.gate_exps.expert_source(ex), eng)?;
8127                        eng.qmatvec_view(
8128                            cache.buf(slot),
8129                            0..gl.len,
8130                            &gv,
8131                            m_e,
8132                            m.gate_exps.in_f,
8133                            m.gate_exps.out_f,
8134                            gl.qtype,
8135                            gl.row_bytes,
8136                        )
8137                    })?;
8138                    let up = e.with_moe_cache(max_block, |cache, eng| {
8139                        let id = BlockId::new(il, PROJ_UP, ex as u16);
8140                        let slot = cache.dispatch_source(id, m.up_exps.expert_source(ex), eng)?;
8141                        eng.qmatvec_view(
8142                            cache.buf(slot),
8143                            0..ul.len,
8144                            &gv,
8145                            m_e,
8146                            m.up_exps.in_f,
8147                            m.up_exps.out_f,
8148                            ul.qtype,
8149                            ul.row_bytes,
8150                        )
8151                    })?;
8152                    let mut act = e.uninit(m_e * n_ff_exp)?;
8153                    Self::ffn_act_lim(
8154                        e,
8155                        cfg,
8156                        &gate,
8157                        &up,
8158                        m.gate_exps.macro_scale(ex),
8159                        m.up_exps.macro_scale(ex),
8160                        lim_exp,
8161                        &mut act,
8162                        m_e * n_ff_exp,
8163                    )?;
8164                    let actv = act.slice(0..m_e * n_ff_exp);
8165                    e.with_moe_cache(max_block, |cache, eng| {
8166                        let id = BlockId::new(il, PROJ_DOWN, ex as u16);
8167                        let slot = cache.dispatch_source(id, m.down_exps.expert_source(ex), eng)?;
8168                        eng.qmatvec_view(
8169                            cache.buf(slot),
8170                            0..dl.len,
8171                            &actv,
8172                            m_e,
8173                            m.down_exps.in_f,
8174                            m.down_exps.out_f,
8175                            dl.qtype,
8176                            dl.row_bytes,
8177                        )
8178                    })?
8179                }
8180            } else {
8181                let sg = scratch_g.as_mut().unwrap();
8182                let su = scratch_u.as_mut().unwrap();
8183                let sd = scratch_d.as_mut().unwrap();
8184                e.stage_expert(m.gate_exps.expert_bytes(ex), sg, 0)?;
8185                e.stage_expert(m.up_exps.expert_bytes(ex), su, 0)?;
8186                e.stage_expert(m.down_exps.expert_bytes(ex), sd, 0)?;
8187                if grouped_q8 {
8188                    let (zq, zd) = e.quantize_q8_1(&gathered, m_e, n_embd)?;
8189                    let gate = e.qmatvec_expert_q8(
8190                        sg,
8191                        0..gl.len,
8192                        &zq,
8193                        &zd,
8194                        m_e,
8195                        m.gate_exps.in_f,
8196                        m.gate_exps.out_f,
8197                        gl.qtype,
8198                        gl.row_bytes,
8199                    )?;
8200                    let up = e.qmatvec_expert_q8(
8201                        su,
8202                        0..ul.len,
8203                        &zq,
8204                        &zd,
8205                        m_e,
8206                        m.up_exps.in_f,
8207                        m.up_exps.out_f,
8208                        ul.qtype,
8209                        ul.row_bytes,
8210                    )?;
8211                    let mut act = e.uninit(m_e * n_ff_exp)?;
8212                    Self::ffn_act_lim(
8213                        e,
8214                        cfg,
8215                        &gate,
8216                        &up,
8217                        m.gate_exps.macro_scale(ex),
8218                        m.up_exps.macro_scale(ex),
8219                        lim_exp,
8220                        &mut act,
8221                        m_e * n_ff_exp,
8222                    )?;
8223                    let (aq2, ad2) = e.quantize_q8_1(&act, m_e, n_ff_exp)?;
8224                    e.qmatvec_expert_q8(
8225                        sd,
8226                        0..dl.len,
8227                        &aq2,
8228                        &ad2,
8229                        m_e,
8230                        m.down_exps.in_f,
8231                        m.down_exps.out_f,
8232                        dl.qtype,
8233                        dl.row_bytes,
8234                    )?
8235                } else {
8236                    let gate = e.qmatvec_view(
8237                        sg,
8238                        0..gl.len,
8239                        &gv,
8240                        m_e,
8241                        m.gate_exps.in_f,
8242                        m.gate_exps.out_f,
8243                        gl.qtype,
8244                        gl.row_bytes,
8245                    )?;
8246                    let up = e.qmatvec_view(
8247                        su,
8248                        0..ul.len,
8249                        &gv,
8250                        m_e,
8251                        m.up_exps.in_f,
8252                        m.up_exps.out_f,
8253                        ul.qtype,
8254                        ul.row_bytes,
8255                    )?;
8256                    let mut act = e.uninit(m_e * n_ff_exp)?;
8257                    Self::ffn_act_lim(
8258                        e,
8259                        cfg,
8260                        &gate,
8261                        &up,
8262                        m.gate_exps.macro_scale(ex),
8263                        m.up_exps.macro_scale(ex),
8264                        lim_exp,
8265                        &mut act,
8266                        m_e * n_ff_exp,
8267                    )?;
8268                    let actv = act.slice(0..m_e * n_ff_exp);
8269                    e.qmatvec_view(
8270                        sd,
8271                        0..dl.len,
8272                        &actv,
8273                        m_e,
8274                        m.down_exps.in_f,
8275                        m.down_exps.out_f,
8276                        dl.qtype,
8277                        dl.row_bytes,
8278                    )?
8279                }
8280            };
8281
8282            // SCATTER into slot buffer: each row goes to slot_buf[tok, slot, :].
8283            e.scatter_slot(
8284                &y,
8285                &tok_idx_d,
8286                &slot_idx_d,
8287                &weight_d,
8288                &mut slot_buf,
8289                &mut wbuf,
8290                n_embd,
8291                n_used,
8292                m_e,
8293            )?;
8294        }
8295
8296        // 5. REDUCE SLOTS: sum the 8 slots per token into the final moe_out.
8297        let mut moe_out = e.zeros(t * n_embd)?;
8298        e.reduce_slots(&slot_buf, &wbuf, &mut moe_out, n_embd, n_used, t)?;
8299
8300        // STATS: print m-distribution when MEMRA_MOE_STATS is set.
8301        if std::env::var("MEMRA_MOE_STATS").is_ok() && !m_dist.is_empty() {
8302            m_dist.sort_unstable();
8303            let active = m_dist.len();
8304            let mean = m_dist.iter().sum::<usize>() as f64 / active as f64;
8305            let median = m_dist[active / 2];
8306            let max_m = *m_dist.last().unwrap();
8307            let min_m = m_dist[0];
8308            let above16 = m_dist.iter().filter(|&&x| x >= 16).count();
8309            println!(
8310                "moe-grouped il={il} t={t} active={active}/{n_expert} \
8311                      m_e: min={min_m} median={median} mean={mean:.1} max={max_m} \
8312                      above_gemm_threshold(>=16)={above16}/{active}"
8313            );
8314        }
8315
8316        Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut moe_out)?;
8317        Ok(moe_out)
8318    }
8319
8320    /// Lane-3 M2: cross-stream MoE for lockstep decode. Routes all m stream rows in one
8321    /// batch, executes fully-HBM-resident experts through the grouped gather/GEMM/scatter
8322    /// machinery at m_e>1 (weight reads amortized across streams), and assigns any expert
8323    /// with a missing projection to that row's CPU companion call (whole-expert granularity,
8324    /// same rule as the sequential frozen path). Slot-pinned accumulation keeps each row's
8325    /// expert-sum order identical to the sequential path.
8326    pub(crate) fn moe_ffn_lockstep(
8327        &self,
8328        e: &Engine,
8329        m: &MoeWeights,
8330        zbatch: &CudaSlice<f32>,
8331        mrows: usize,
8332        il: u16,
8333        max_block: usize,
8334    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8335        use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
8336        let cfg = &self.cfg;
8337        let moe = cfg.moe.as_ref().unwrap();
8338        let n_embd = cfg.n_embd as usize;
8339        let n_expert = moe.expert_count as usize;
8340        let n_used = moe.expert_used_count as usize;
8341        let n_ff_exp = moe.expert_ff_length as usize;
8342        // step35 per-layer SwiGLU clamp; None on every other arch / unclamped layer.
8343        let lim_exp = cfg.clamp_exp_at(il as u32);
8344        let lim_shexp = cfg.clamp_shexp_at(il as u32);
8345
8346        let logits = e.matmul(&m.gate_inp, zbatch, mrows)?;
8347        if let Some(sig) = cfg.sigmoid_router() {
8348            Self::trace_sigmoid_router_logits(e, il, mrows, n_expert, n_used, &logits, m, sig)?;
8349        }
8350        let (sel_all, w_all) = if let Some(sig) = cfg.sigmoid_router() {
8351            Self::moe_route_sigmoid_cfg(e, &logits, mrows, n_expert, n_used, m, sig)?
8352        } else {
8353            Self::moe_route_cfg(
8354                e,
8355                &logits,
8356                mrows,
8357                n_expert,
8358                n_used,
8359                m.active_experts.as_deref(),
8360            )?
8361        };
8362        Self::trace_moe_routes(il, mrows, &sel_all, &w_all)?;
8363
8364        // Residency split at whole-expert granularity against the (frozen) cache.
8365        let resident_expert: Vec<bool> = e.with_moe_cache(max_block, |c, _| {
8366            Ok((0..n_expert)
8367                .map(|ex| {
8368                    [PROJ_GATE, PROJ_UP, PROJ_DOWN]
8369                        .into_iter()
8370                        .all(|p| c.resident(BlockId::new(il, p, ex as u16)).is_some())
8371                })
8372                .collect())
8373        })?;
8374
8375        struct Group {
8376            rows: Vec<i32>,
8377            slots: Vec<i32>,
8378            weights: Vec<f32>,
8379        }
8380        let mut groups: std::collections::HashMap<usize, Group> = Default::default();
8381        let mut cpu_rows: Vec<Vec<(usize, f32)>> = vec![Vec::new(); mrows];
8382        let mut cpu_by_expert: std::collections::HashMap<usize, Vec<(usize, f32)>> =
8383            Default::default();
8384        for row in 0..mrows {
8385            for j in 0..n_used {
8386                let ex = sel_all[row * n_used + j] as usize;
8387                let w = w_all[row * n_used + j];
8388                if resident_expert[ex] {
8389                    let group = groups.entry(ex).or_insert_with(|| Group {
8390                        rows: Vec::new(),
8391                        slots: Vec::new(),
8392                        weights: Vec::new(),
8393                    });
8394                    group.rows.push(row as i32);
8395                    group.slots.push(j as i32);
8396                    group.weights.push(w);
8397                } else {
8398                    crate::cpu_experts::record_incomplete_gpu_residency(0);
8399                    cpu_rows[row].push((ex, w));
8400                    cpu_by_expert.entry(ex).or_default().push((row, w));
8401                }
8402            }
8403        }
8404
8405        // CPU tickets first: reads/compute overlap the GPU grouped work below. Experts routed
8406        // by >=2 streams go through the multi-row ABI (weight decode amortized across rows);
8407        // each row's remaining experts stay one ordinary per-row call. Contribution FP-sum
8408        // order per row differs from the sequential single-call chunk — part of the
8409        // documented lockstep numeric class.
8410        let host_rows = e.dtoh(zbatch)?;
8411        let rows_ok = crate::cpu_experts::rows_supported();
8412        enum CpuPart {
8413            Single { row: usize },
8414            Rows { rows: Vec<usize> },
8415        }
8416        let mut tickets: Vec<(CpuPart, crate::cpu_experts::CpuExpertTicket)> = Vec::new();
8417        let mut rows_served: std::collections::HashSet<(usize, usize)> = Default::default();
8418        if rows_ok {
8419            let mut shared: Vec<(usize, Vec<(usize, f32)>)> = cpu_by_expert
8420                .into_iter()
8421                .filter(|(_, rows)| rows.len() >= 2)
8422                .collect();
8423            shared.sort_by_key(|(ex, _)| *ex);
8424            for (ex, mut row_weights) in shared {
8425                row_weights.sort_by_key(|(row, _)| *row);
8426                let inputs: Vec<(&[f32], f32)> = row_weights
8427                    .iter()
8428                    .map(|&(row, w)| (&host_rows[row * n_embd..(row + 1) * n_embd], w))
8429                    .collect();
8430                let job = crate::cpu_experts::prepare_rows_job(m, ex, &inputs)
8431                    .map_err(std::io::Error::other)?;
8432                for &(row, _) in &row_weights {
8433                    rows_served.insert((row, ex));
8434                }
8435                tickets.push((
8436                    CpuPart::Rows {
8437                        rows: row_weights.iter().map(|&(row, _)| row).collect(),
8438                    },
8439                    crate::cpu_experts::submit_rows(job).map_err(std::io::Error::other)?,
8440                ));
8441            }
8442        }
8443        for (row, selected) in cpu_rows.iter().enumerate() {
8444            let leftover: Vec<(usize, f32)> = selected
8445                .iter()
8446                .copied()
8447                .filter(|&(ex, _)| !rows_served.contains(&(row, ex)))
8448                .collect();
8449            if leftover.is_empty() {
8450                continue;
8451            }
8452            let host_row = &host_rows[row * n_embd..(row + 1) * n_embd];
8453            let job = crate::cpu_experts::prepare_job(m, il, &leftover, host_row)
8454                .map_err(std::io::Error::other)?;
8455            tickets.push((
8456                CpuPart::Single { row },
8457                crate::cpu_experts::submit(job).map_err(std::io::Error::other)?,
8458            ));
8459        }
8460
8461        let mut slot_buf = e.zeros(mrows * n_used * n_embd)?;
8462        let mut wbuf = e.zeros(mrows * n_used)?;
8463        let mut order: Vec<usize> = groups.keys().copied().collect();
8464        order.sort_by(|&a, &b| {
8465            groups[&b]
8466                .rows
8467                .len()
8468                .cmp(&groups[&a].rows.len())
8469                .then(a.cmp(&b))
8470        });
8471        for &ex in &order {
8472            let group = &groups[&ex];
8473            let m_e = group.rows.len();
8474            let gl = m.gate_exps.expert_layout(ex);
8475            let ul = m.up_exps.expert_layout(ex);
8476            let dl = m.down_exps.expert_layout(ex);
8477            let row_idx_d = e.htod_i32(&group.rows)?;
8478            let slot_idx_d = e.htod_i32(&group.slots)?;
8479            let dmac = m.down_exps.macro_scale(ex);
8480            let weight_d = if dmac == 1.0 {
8481                e.htod(&group.weights)?
8482            } else {
8483                let scaled: Vec<f32> = group.weights.iter().map(|&w| w * dmac).collect();
8484                e.htod(&scaled)?
8485            };
8486            let mut gathered = e.zeros(m_e * n_embd)?;
8487            e.gather_rows(zbatch, &row_idx_d, &mut gathered, n_embd, m_e)?;
8488            let gv = gathered.slice(0..m_e * n_embd);
8489            let gate = e.with_moe_cache(max_block, |c, eng| {
8490                let slot = c
8491                    .resident(BlockId::new(il, PROJ_GATE, ex as u16))
8492                    .ok_or("lockstep resident expert vanished (cache not frozen?)")?;
8493                eng.qmatvec_view(
8494                    c.buf(crate::moe_cache::DispatchSlot::Resident(slot)),
8495                    0..gl.len,
8496                    &gv,
8497                    m_e,
8498                    m.gate_exps.in_f,
8499                    m.gate_exps.out_f,
8500                    gl.qtype,
8501                    gl.row_bytes,
8502                )
8503            })?;
8504            let up = e.with_moe_cache(max_block, |c, eng| {
8505                let slot = c
8506                    .resident(BlockId::new(il, PROJ_UP, ex as u16))
8507                    .ok_or("lockstep resident expert vanished (cache not frozen?)")?;
8508                eng.qmatvec_view(
8509                    c.buf(crate::moe_cache::DispatchSlot::Resident(slot)),
8510                    0..ul.len,
8511                    &gv,
8512                    m_e,
8513                    m.up_exps.in_f,
8514                    m.up_exps.out_f,
8515                    ul.qtype,
8516                    ul.row_bytes,
8517                )
8518            })?;
8519            let mut act = e.zeros(m_e * n_ff_exp)?;
8520            Self::ffn_act_lim(
8521                e,
8522                cfg,
8523                &gate,
8524                &up,
8525                m.gate_exps.macro_scale(ex),
8526                m.up_exps.macro_scale(ex),
8527                lim_exp,
8528                &mut act,
8529                m_e * n_ff_exp,
8530            )?;
8531            let actv = act.slice(0..m_e * n_ff_exp);
8532            let y = e.with_moe_cache(max_block, |c, eng| {
8533                let slot = c
8534                    .resident(BlockId::new(il, PROJ_DOWN, ex as u16))
8535                    .ok_or("lockstep resident expert vanished (cache not frozen?)")?;
8536                eng.qmatvec_view(
8537                    c.buf(crate::moe_cache::DispatchSlot::Resident(slot)),
8538                    0..dl.len,
8539                    &actv,
8540                    m_e,
8541                    m.down_exps.in_f,
8542                    m.down_exps.out_f,
8543                    dl.qtype,
8544                    dl.row_bytes,
8545                )
8546            })?;
8547            e.scatter_slot(
8548                &y,
8549                &row_idx_d,
8550                &slot_idx_d,
8551                &weight_d,
8552                &mut slot_buf,
8553                &mut wbuf,
8554                n_embd,
8555                n_used,
8556                m_e,
8557            )?;
8558        }
8559        let mut moe_out = e.zeros(mrows * n_embd)?;
8560        e.reduce_slots(&slot_buf, &wbuf, &mut moe_out, n_embd, n_used, mrows)?;
8561
8562        // CPU contributions join BEFORE the shared expert (the sequential path's placement).
8563        let mut row_sums: Vec<Option<Vec<f32>>> = vec![None; mrows];
8564        for (part, ticket) in tickets {
8565            let cpu_output = ticket.wait().map_err(std::io::Error::other)?;
8566            let mut add_row = |row: usize, chunk: &[f32]| {
8567                let sum = row_sums[row].get_or_insert_with(|| vec![0.0f32; n_embd]);
8568                for (accumulator, value) in sum.iter_mut().zip(chunk) {
8569                    *accumulator += value;
8570                }
8571            };
8572            match part {
8573                CpuPart::Single { row } => add_row(row, &cpu_output),
8574                CpuPart::Rows { rows } => {
8575                    for (slot, row) in rows.into_iter().enumerate() {
8576                        add_row(row, &cpu_output[slot * n_embd..(slot + 1) * n_embd]);
8577                    }
8578                }
8579            }
8580        }
8581        for (row, sum) in row_sums.into_iter().enumerate() {
8582            let Some(sum) = sum else { continue };
8583            let cpu_output = e.htod(&sum)?;
8584            let mut dst = moe_out.slice_mut(row * n_embd..(row + 1) * n_embd);
8585            e.axpy_into(&cpu_output, 1.0, &mut dst, n_embd)?;
8586        }
8587
8588        if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
8589            (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
8590        {
8591            let n_ff_sh = gate_shexp.out_features();
8592            let sg_gate = e.matmul(gate_shexp, zbatch, mrows)?;
8593            let sg_up = e.matmul(up_shexp, zbatch, mrows)?;
8594            let mut sa = e.zeros(mrows * n_ff_sh)?;
8595            Self::ffn_act_lim(
8596                e,
8597                cfg,
8598                &sg_gate,
8599                &sg_up,
8600                1.0,
8601                1.0,
8602                lim_shexp,
8603                &mut sa,
8604                mrows * n_ff_sh,
8605            )?;
8606            let sh = e.matmul(down_shexp, &sa, mrows)?;
8607            // lockstep rows ARE decode tokens: fused sigmoid-dot per row so batched serving
8608            // decode matches the single-sequence decode chain bit-for-bit.
8609            let g = match &m.gate_inp_shexp {
8610                Some(gate_inp_shexp) => {
8611                    e.sigmoid_dot_rows(zbatch, gate_inp_shexp.float_data(), n_embd, mrows)?
8612                }
8613                None => e.htod(&vec![1.0f32; mrows])?,
8614            };
8615            e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, mrows)?;
8616        }
8617
8618        Ok(moe_out)
8619    }
8620}
8621
8622// ============================ gemma4 (R8 verified wiring) ==================================
8623// Node-for-node vs llama.cpp src/models/gemma4.cpp:180-405 (HANDOVER "R8 VERIFIED WIRING").
8624// v0 bring-up: full attention everywhere (exact for prompts < sliding_window 1024 — R6 masking
8625// later), sdpa_naive (hd 512 has no FA stamp), sequential host-staged MoE (the perf arms grow
8626// gemma variants after the correctness gate).
8627impl HybridModel {
8628    /// Per-layer attention geometry (R5): (head_dim, n_kv, n_head, rope_base, scale, is_swa).
8629    /// Rotary width the gemma-4 config DECLARES for layer `il`, per SWA/global class — the
8630    /// `n_rot` argument of the fused `rms_norm_qkv_rope*` kernels
8631    /// (`Engine::full_width_rope_only`, lane/graph-s-key-exactness-20260819).
8632    ///
8633    /// It exists so the fusion's full-width assumption is checked against the CONFIG instead of
8634    /// against itself: passing `hd` twice would guard nothing. On every gemma-4 artifact both
8635    /// routes set `rope_dims_{global,swa} == key_length_{global,swa}` (GGUF reads
8636    /// `rope.dimension_count(_swa)`, the safetensors route derives them from
8637    /// `global_head_dim`/`head_dim`), which is exactly WHY the fusion is legal here — gemma's
8638    /// "partial" rotary is per-dimension frequency scaling through the `rope_freqs` factors
8639    /// tensor, not an n_rot truncation. An arch that truncates instead (qwen3.5: 64 of 256;
8640    /// step35 full-attn: 64 of 128) must not reach these kernels, and now cannot silently.
8641    pub(crate) fn gemma4_rope_dims(&self, il: usize) -> usize {
8642        let g = self
8643            .cfg
8644            .gemma4
8645            .as_ref()
8646            .expect("gemma4_rope_dims on a non-gemma4 config");
8647        if g.swa_pattern[il] {
8648            g.rope_dims_swa as usize
8649        } else {
8650            g.rope_dims_global as usize
8651        }
8652    }
8653
8654    pub(crate) fn gemma4_geom(&self, il: usize) -> (usize, usize, usize, f32, f32, bool) {
8655        let g = self.cfg.gemma4.as_ref().unwrap();
8656        let swa = g.swa_pattern[il];
8657        let hd = if swa {
8658            g.key_length_swa
8659        } else {
8660            g.key_length_global
8661        } as usize;
8662        // attention scale = 1.0 (llama gemma4.cpp:11 "Gemma4 uses self.scaling = 1.0" — q/k are
8663        // per-head rms-normed; NOT the 1/sqrt(hd) default. Bring-up bug: 1/sqrt(hd) left token-0
8664        // rows exact (softmax over one element) while every later position drifted).
8665        (
8666            hd,
8667            g.head_count_kv[il] as usize,
8668            self.cfg.n_head as usize,
8669            if swa {
8670                g.rope_base_swa
8671            } else {
8672                g.rope_base_global
8673            },
8674            1.0,
8675            swa,
8676        )
8677    }
8678
8679    /// Suppress-token mask over t logits rows (tokenizer.ggml.suppress_tokens; no-op when the
8680    /// model ships none). NOT monotonic like softcap — must run before every argmax/sample, so
8681    /// every gemma4 logits tail (forward/prime/decode/dc/verify/e4b) calls this on device ld.
8682    pub(crate) fn gemma4_suppress(
8683        &self,
8684        e: &Engine,
8685        ld: &mut CudaSlice<f32>,
8686        t: usize,
8687    ) -> Result<(), Box<dyn std::error::Error>> {
8688        if let Some((ids, n)) = self.gemma4_aux.as_ref().and_then(|a| a.suppress_d.as_ref()) {
8689            // suppress_d is primary-owned by the audited invariant (serving picks the last/head
8690            // stage as primary, and this tail runs only after the last stage). The assert turns
8691            // that argued invariant into a checked one: any topology violating primary==head
8692            // trips here in debug instead of silently peer-reading a device-0 buffer.
8693            #[cfg(debug_assertions)]
8694            crate::debug_assert_tensor_stream_device(
8695                ids,
8696                &e.stream(),
8697                "gemma4_suppress.suppress_d",
8698            );
8699            e.mask_ids_rows(ld, ids, *n, self.output.out_features(), t)?;
8700        }
8701        Ok(())
8702    }
8703
8704    /// gemma4 attention (R5 geometry, R7 weightless V-norm on the RAW K projection, R9 dual rope).
8705    /// `cache`: Some => PRIME mode — append the T post-rope K / normed V rows into the quantized
8706    /// KV cache (same per-row quantize math as the decode append) and advance len. Fresh-prompt
8707    /// only (v0): attends within `tokens` via the f32 sdpa.
8708    #[allow(clippy::too_many_arguments)]
8709    /// ONE-PROGRAM seam (lane/gemma-restore-exactness-20260819, DEFAULT OFF): route dense
8710    /// gemma SWA prefill attention through the WINDOWED arm at every prompt length instead of
8711    /// switching program at `t > sliding_window`. The door is the measured cause of the
8712    /// prefix-restore exactness violation: below it the SWA layers attend unwindowed on f32
8713    /// operands, above it windowed on bf16 operands, so rows[0..932) of a 932-row prime and of
8714    /// a 1048-row prime differ in 59 of 60 layers. `=1` collapses it to one program so the
8715    /// published prefix KV stops depending on the total prompt length. Off by default because
8716    /// collapsing changes the bytes short prompts get today; the flip is an owner call.
8717    fn gemma_fa_one_program() -> bool {
8718        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8719        *ON.get_or_init(|| std::env::var("MEMRA_GEMMA_FA_ONE_PROGRAM").as_deref() == Ok("1"))
8720    }
8721
8722    fn gemma4_attn_prime(
8723        &self,
8724        e: &Engine,
8725        fa: &crate::hybrid::FullAttnLayer,
8726        il: usize,
8727        h: &CudaSlice<f32>,
8728        pos_d: &CudaSlice<i32>,
8729        t: usize,
8730        cache: Option<&mut Cache>,
8731        island: Option<&CudaSlice<i32>>,
8732    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8733        let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
8734        let eps = self.cfg.rms_eps;
8735        let aux = self.gemma4_aux.as_ref().unwrap();
8736        let ones = aux.ones(e);
8737        #[cfg(debug_assertions)]
8738        crate::debug_assert_tensor_stream_device(ones, &e.stream(), "gemma4_attn_prime.ones");
8739
8740        // quantize-once window: q/k/v share `h` — the MMQ D4 activation quantizes once
8741        // (h stays borrowed across the triple, so the cache key can't go stale).
8742        e.mmq_act_begin();
8743        let q0 = e.matmul(&fa.wq, h, t)?; // [t, nh*hd]
8744        if il == 0 && std::env::var("MEMRA_G4_PRIME_TRACE").as_deref() == Ok("1") {
8745            let v = e.dtoh(&q0)?;
8746            let nan = v.iter().filter(|x| x.is_nan()).count();
8747            let amax = v.iter().fold(0f32, |a, x| a.max(x.abs()));
8748            eprintln!(
8749                "[g4-prime-trace] L0 q0: nan={nan}/{} amax={amax:.3}",
8750                v.len()
8751            );
8752        }
8753        let k0 = e.matmul(&fa.wk, h, t)?; // [t, nkv*hd]
8754        // globals ship no v_proj (wv := wk at load) — V input is the SAME projection output;
8755        // reuse k0 instead of re-running the identical matmul (K=V dedup, 5 layers).
8756        let v0 = if swa {
8757            e.matmul(&fa.wv, h, t)?
8758        } else {
8759            e.clone_dtod(&k0)?
8760        };
8761        if il == 0 && std::env::var("MEMRA_G4_PRIME_TRACE").as_deref() == Ok("1") {
8762            for (tag, buf) in [("k0", &k0), ("v0", &v0)] {
8763                let v = e.dtoh(buf)?;
8764                let nan = v.iter().filter(|x| x.is_nan()).count();
8765                let amax = v.iter().fold(0f32, |a, x| a.max(x.abs()));
8766                eprintln!(
8767                    "[g4-prime-trace] L0 {tag}: nan={nan}/{} amax={amax:.3}",
8768                    v.len()
8769                );
8770            }
8771        }
8772
8773        let mut q = e.uninit(t * nh * hd)?;
8774        let mut k = e.uninit(t * nkv * hd)?;
8775        // R7: V = weightless rms_norm of the raw projection; NEVER roped.
8776        let mut v = e.uninit(t * nkv * hd)?;
8777        // 31B glue lane: producers emit the bf16 FA operands (norm emits vb; rope emits qb/kb
8778        // post-rope) — kills 3 f32->bf16 converts + re-reads per layer. Bit-identical operands
8779        // (same __float2bfloat16); MEMRA_FA_EMIT=0 reverts to the convert-in-FA path.
8780        static EMIT: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8781        // Island primes take the mask-capable naive kernel below; keep the operands f32
8782        // (the bf16 FA emit path has no island consumer). Text-only keeps emit unchanged.
8783        let emit = island.is_none()
8784            && t >= 16
8785            && crate::Engine::qkvnorm_w_on_prefill(nh * t + 2 * nkv * t, hd)
8786            && *EMIT.get_or_init(|| {
8787                std::env::var("MEMRA_FA_EMIT")
8788                    .map(|s| s != "0")
8789                    .unwrap_or(true)
8790            });
8791        let mut qb = e.alloc_uninit::<u8>(if emit { t * nh * hd * 2 } else { 1 })?;
8792        let mut kb = e.alloc_uninit::<u8>(if emit { t * nkv * hd * 2 } else { 1 })?;
8793        let mut vb = e.alloc_uninit::<u8>(if emit { t * nkv * hd * 2 } else { 1 })?;
8794        // f16-P/V door: emit V as f16 straight from the norm when this layer's FA consumer
8795        // reads f16 — kills the per-layer bf16->f16 re-encode (1.4GB/pass on 31B SWA).
8796        let v_f16 = emit
8797            && crate::fa_f16pv_on()
8798            && match hd {
8799                512 => true,
8800                256 => swa && crate::faw_hp_on() && nh % 2 == 0 && (nh / nkv) % 2 == 0,
8801                _ => false,
8802            };
8803        if emit {
8804            e.rms_norm_qkv_w4b(
8805                &q0,
8806                &k0,
8807                &v0,
8808                fa.q_norm.float_data(),
8809                fa.k_norm.float_data(),
8810                ones,
8811                &mut q,
8812                &mut k,
8813                &mut v,
8814                &mut vb,
8815                hd,
8816                nh * t,
8817                nkv * t,
8818                eps,
8819                v_f16,
8820            )?;
8821        } else {
8822            e.rms_norm_qkv(
8823                &q0,
8824                &k0,
8825                &v0,
8826                fa.q_norm.float_data(),
8827                fa.k_norm.float_data(),
8828                ones,
8829                &mut q,
8830                &mut k,
8831                &mut v,
8832                hd,
8833                nh * t,
8834                nkv * t,
8835                eps,
8836            )?;
8837        }
8838
8839        let ff = if swa {
8840            None
8841        } else {
8842            Some(
8843                aux.rope_freqs(e)
8844                    .expect("gemma4 global rope needs rope_freqs.weight"),
8845            )
8846        };
8847        #[cfg(debug_assertions)]
8848        if let Some(ff) = ff {
8849            crate::debug_assert_tensor_stream_device(
8850                ff,
8851                &e.stream(),
8852                "gemma4_attn_prime.rope_freqs",
8853            );
8854        }
8855        if emit {
8856            e.rope_neox2_bf16e(
8857                &mut q, &mut k, &mut qb, &mut kb, pos_d, hd, hd, nh, nkv, t, base, 1.0, ff,
8858            )?;
8859        } else {
8860            e.rope_neox2(&mut q, &mut k, pos_d, hd, hd, nh, nkv, t, base, 1.0, ff)?;
8861        }
8862
8863        if let Some(cache) = cache {
8864            let kvl = cache.kv[il].as_mut().unwrap();
8865            assert_eq!(kvl.len, 0, "gemma4 prime is fresh-prompt only (v0)");
8866            e.append_kv_quantized_rows(
8867                &k,
8868                &v,
8869                &mut kvl.k,
8870                &mut kvl.v,
8871                kvl.len,
8872                t,
8873                kvl.kv_dim_k,
8874                kvl.kv_dim_v,
8875                kvl.k_tok_bytes,
8876                kvl.v_tok_bytes,
8877                (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on()),
8878            )?;
8879            kvl.len += t;
8880        }
8881        let mut attn = e.zeros(t * nh * hd)?;
8882        // R6: SWA layers mask keys older than sliding_window once the prompt exceeds it
8883        // (windowed naive twin; fa windowed stamps later). Under the window, full attention
8884        // is exact — SWA rides fa_prefill (hd-256 stamp), the hd-512 globals stay naive.
8885        let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
8886        if let Some(span) = island {
8887            // Masked-prefill arm: every layer routes through the island-aware naive
8888            // kernel (correctness-first, same posture as the vision tower v1). The
8889            // window argument keeps the R6 shortcut: 0 while the prompt fits the
8890            // window, the real window beyond it.
8891            let w = if swa && t > win { win } else { 0 };
8892            e.sdpa_naive_island(&q, &k, &v, &mut attn, span, hd, nh, nkv, t, t, scale, w)?;
8893        } else if swa && (t > win || Self::gemma_fa_one_program()) {
8894            if hd == 256 && std::env::var("MEMRA_NOFA").is_err() {
8895                if emit {
8896                    e.fa_prefill_w_pre(
8897                        &qb, &kb, &vb, &mut attn, hd, nh, nkv, t, t, scale, true, win, v_f16,
8898                    )?;
8899                } else {
8900                    e.fa_prefill_w(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
8901                }
8902            } else {
8903                e.sdpa_naive_w(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
8904            }
8905        } else if hd == 256 && std::env::var("MEMRA_NOFA").is_err() {
8906            e.fa_prefill(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
8907        } else if hd == 512 && std::env::var("MEMRA_NOFA").is_err() {
8908            if emit {
8909                e.fa_prefill_hd512_pre(
8910                    &qb, &kb, &vb, &mut attn, hd, nh, nkv, t, t, scale, true, v_f16,
8911                )?;
8912            } else {
8913                e.fa_prefill_hd512(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
8914            }
8915        } else {
8916            e.sdpa_naive(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
8917        }
8918        Ok(e.matmul(&fa.wo, &attn, t)?)
8919    }
8920
8921    /// Back-compat wrapper (pure prefill, no cache).
8922    fn gemma4_attn(
8923        &self,
8924        e: &Engine,
8925        fa: &crate::hybrid::FullAttnLayer,
8926        il: usize,
8927        h: &CudaSlice<f32>,
8928        pos_d: &CudaSlice<i32>,
8929        t: usize,
8930    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8931        self.gemma4_attn_prime(e, fa, il, h, pos_d, t, None, None)
8932    }
8933
8934    /// gemma4 MoE with the expert input PRE-QUANTIZED (the q8z tail fusion). Caller guarantees
8935    /// the fast-arm conditions (resident dev slabs + dp4a qtypes) — decode t=1 and verify rows
8936    /// arms only, per-token kernel chains identical to the f32-input path (same quantize bytes:
8937    /// the q8z epilogue is quantize_q8_1 verbatim).
8938    fn gemma4_moe_q8(
8939        &self,
8940        e: &Engine,
8941        m: &crate::hybrid::MoeWeights,
8942        bits: &crate::hybrid::Gemma4MoeBits,
8943        mq: &(CudaSlice<i8>, CudaSlice<f32>),
8944        router_in: &CudaSlice<f32>,
8945        t: usize,
8946    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8947        let cfg = &self.cfg;
8948        let moe = cfg.moe.as_ref().unwrap();
8949        let n_embd = cfg.n_embd as usize;
8950        let n_expert = moe.expert_count as usize;
8951        let n_used = moe.expert_used_count as usize;
8952        let n_ff_exp = moe.expert_ff_length as usize;
8953        // router stays the two-launch pair: THREE fuse variants measured worse (serial-dot
8954        // -50%, 1024-thread warp-parallel -12% on topk sync overhead — jsonl 2026-07-11);
8955        // the pair's 12us is kernel time, not launch gaps.
8956        let logits = if crate::router_kernel_on() {
8957            e.router_gemv(m.gate_inp.float_data(), router_in, n_embd, n_expert, t)?
8958        } else {
8959            e.matmul(&m.gate_inp, router_in, t)?
8960        };
8961        let dev = m.dev_exps.as_ref().unwrap();
8962        let (sel_d, w_d) =
8963            e.moe_router_topk_scaled(&logits, t, n_expert, n_used, &bits.per_expert_scale_d)?;
8964        let (zq, zd) = mq;
8965        if t == 1 {
8966            let selv = sel_d.slice(0..n_used);
8967            let wv = w_d.slice(0..n_used);
8968            let act = e.moe_gate_up_gelu8_dev_q8(
8969                &dev.ptr_row,
8970                &selv,
8971                zq,
8972                zd,
8973                n_embd,
8974                n_ff_exp,
8975                n_used,
8976                n_expert,
8977                m.gate_exps.qtype,
8978                m.up_exps.qtype,
8979                m.gate_exps.row_bytes,
8980                m.up_exps.row_bytes,
8981            )?;
8982            let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
8983            let mut moe_out = e.uninit(n_embd)?;
8984            e.moe_down8_fma_dev_q8(
8985                &dev.ptr_row,
8986                &selv,
8987                &wv,
8988                &aq2,
8989                &ad2,
8990                &mut moe_out.slice_mut(0..n_embd),
8991                n_ff_exp,
8992                n_embd,
8993                n_used,
8994                n_expert,
8995                m.down_exps.qtype,
8996                m.down_exps.row_bytes,
8997            )?;
8998            return Ok(moe_out);
8999        }
9000        let csr = t <= 10 && std::env::var("MEMRA_GEMMA_CSR").as_deref() != Ok("0");
9001        let act = if csr {
9002            e.moe_gate_up_gelu8_dev_q8_csr(
9003                &dev.ptr_row,
9004                &sel_d,
9005                zq,
9006                zd,
9007                t * n_used,
9008                n_embd,
9009                n_ff_exp,
9010                n_used,
9011                n_expert,
9012                m.gate_exps.qtype,
9013                m.up_exps.qtype,
9014                m.gate_exps.row_bytes,
9015                m.up_exps.row_bytes,
9016            )?
9017        } else {
9018            e.moe_gate_up_gelu8_dev_q8_rows(
9019                &dev.ptr_row,
9020                &sel_d,
9021                zq,
9022                zd,
9023                t,
9024                n_embd,
9025                n_ff_exp,
9026                n_used,
9027                n_expert,
9028                m.gate_exps.qtype,
9029                m.up_exps.qtype,
9030                m.gate_exps.row_bytes,
9031                m.up_exps.row_bytes,
9032            )?
9033        };
9034        let (aq2, ad2) = e.quantize_q8_1(&act, t * n_used, n_ff_exp)?;
9035        let mut moe_out = e.uninit(t * n_embd)?;
9036        // down stays rows_g: the CSR dedup twin measured NEGATIVE at nsb=22 too (189.4 vs
9037        // 207.0 depth spec, bitwise-exact — jsonl 2026-07-10; qwen nsb=16 same verdict).
9038        e.moe_down8_fma_dev_q8_rows_g(
9039            &dev.ptr_row,
9040            &sel_d,
9041            &w_d,
9042            &aq2,
9043            &ad2,
9044            &mut moe_out,
9045            t,
9046            n_ff_exp,
9047            n_embd,
9048            n_used,
9049            n_expert,
9050            m.down_exps.qtype,
9051            m.down_exps.row_bytes,
9052        )?;
9053        Ok(moe_out)
9054    }
9055
9056    /// gemma4 MoE (R2 router prologue input supplied by caller, R3 per-expert output scale).
9057    /// Sequential host-staged v0 — softmax gating + renorm (moe_route, the qwen recipe), GELU
9058    /// experts, scale folded into the accumulate weight (post-matmul linear scale, exact fold).
9059    fn gemma4_moe(
9060        &self,
9061        e: &Engine,
9062        m: &crate::hybrid::MoeWeights,
9063        bits: &crate::hybrid::Gemma4MoeBits,
9064        moe_in: &CudaSlice<f32>,
9065        router_in: &CudaSlice<f32>,
9066        t: usize,
9067    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9068        let cfg = &self.cfg;
9069        let moe = cfg.moe.as_ref().unwrap();
9070        let n_embd = cfg.n_embd as usize;
9071        let n_expert = moe.expert_count as usize;
9072        let n_used = moe.expert_used_count as usize;
9073        let n_ff_exp = moe.expert_ff_length as usize;
9074
9075        // Router: the in-house GEMV for ALL small t (decode AND verify ride the same per-column
9076        // kernel — the cuBLASLt n-dependence flipped top-k at verify t on the 27B, d994271);
9077        // batched matmul only at real prefill.
9078        let logits = if t < PRIME_MIN_T && crate::router_kernel_on() {
9079            e.router_gemv(m.gate_inp.float_data(), router_in, n_embd, n_expert, t)?
9080        } else {
9081            e.matmul(&m.gate_inp, router_in, t)?
9082        };
9083
9084        // FAST SMALL-T ARM (decode t=1 AND spec verify t=2..15): device softmax-topk router,
9085        // then PER TOKEN the same fused gate_up GELU + down8 FMA launch pair over the resident
9086        // dev slabs — verify rides the EXACT decode kernel chain per token (dispatch-parity
9087        // law; the qwen "verify must be kernel-dispatch-identical to decode" lesson).
9088        if t < PRIME_MIN_T
9089            && m.dev_exps.as_ref().is_some_and(|d| !d.gu_il)
9090            && expert_dp4a_supported(m.gate_exps.qtype)
9091            && expert_dp4a_supported(m.up_exps.qtype)
9092            && expert_dp4a_supported(m.down_exps.qtype)
9093            && std::env::var("MEMRA_GEMMA_MOE_FAST").as_deref() != Ok("0")
9094        {
9095            let dev = m.dev_exps.as_ref().unwrap();
9096            let (sel_d, w_d) =
9097                e.moe_router_topk_scaled(&logits, t, n_expert, n_used, &bits.per_expert_scale_d)?;
9098            if t == 1 {
9099                let (zq, zd) = e.quantize_q8_1(moe_in, 1, n_embd)?;
9100                let selv = sel_d.slice(0..n_used);
9101                let wv = w_d.slice(0..n_used);
9102                let act = e.moe_gate_up_gelu8_dev_q8(
9103                    &dev.ptr_row,
9104                    &selv,
9105                    &zq,
9106                    &zd,
9107                    n_embd,
9108                    n_ff_exp,
9109                    n_used,
9110                    n_expert,
9111                    m.gate_exps.qtype,
9112                    m.up_exps.qtype,
9113                    m.gate_exps.row_bytes,
9114                    m.up_exps.row_bytes,
9115                )?;
9116                let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
9117                let mut moe_out = e.uninit(n_embd)?;
9118                e.moe_down8_fma_dev_q8(
9119                    &dev.ptr_row,
9120                    &selv,
9121                    &wv,
9122                    &aq2,
9123                    &ad2,
9124                    &mut moe_out.slice_mut(0..n_embd),
9125                    n_ff_exp,
9126                    n_embd,
9127                    n_used,
9128                    n_expert,
9129                    m.down_exps.qtype,
9130                    m.down_exps.row_bytes,
9131                )?;
9132                return Ok(moe_out);
9133            }
9134            // VERIFY ROWS TWINS (t=2..15): ONE launch pair for all tokens; per (token,row,slot)
9135            // bodies are the t=1 kernels VERBATIM (bit-identical to the per-token loop).
9136            // gate_up rides the CSR owner-scan dedup when t <= 10 (each duplicated expert's
9137            // weight stream decoded once; the qwen CSR class). MEMRA_GEMMA_CSR=0 -> rows.
9138            let (zq, zd) = e.quantize_q8_1(moe_in, t, n_embd)?;
9139            let csr = t <= 10 && std::env::var("MEMRA_GEMMA_CSR").as_deref() != Ok("0");
9140            let act = if csr {
9141                e.moe_gate_up_gelu8_dev_q8_csr(
9142                    &dev.ptr_row,
9143                    &sel_d,
9144                    &zq,
9145                    &zd,
9146                    t * n_used,
9147                    n_embd,
9148                    n_ff_exp,
9149                    n_used,
9150                    n_expert,
9151                    m.gate_exps.qtype,
9152                    m.up_exps.qtype,
9153                    m.gate_exps.row_bytes,
9154                    m.up_exps.row_bytes,
9155                )?
9156            } else {
9157                e.moe_gate_up_gelu8_dev_q8_rows(
9158                    &dev.ptr_row,
9159                    &sel_d,
9160                    &zq,
9161                    &zd,
9162                    t,
9163                    n_embd,
9164                    n_ff_exp,
9165                    n_used,
9166                    n_expert,
9167                    m.gate_exps.qtype,
9168                    m.up_exps.qtype,
9169                    m.gate_exps.row_bytes,
9170                    m.up_exps.row_bytes,
9171                )?
9172            };
9173            let (aq2, ad2) = e.quantize_q8_1(&act, t * n_used, n_ff_exp)?;
9174            let mut moe_out = e.uninit(t * n_embd)?;
9175            e.moe_down8_fma_dev_q8_rows_g(
9176                &dev.ptr_row,
9177                &sel_d,
9178                &w_d,
9179                &aq2,
9180                &ad2,
9181                &mut moe_out,
9182                t,
9183                n_ff_exp,
9184                n_embd,
9185                n_used,
9186                n_expert,
9187                m.down_exps.qtype,
9188                m.down_exps.row_bytes,
9189            )?;
9190            return Ok(moe_out);
9191        }
9192
9193        let (sel_all, mut w_all) = Self::moe_route(e, &logits, t, n_expert, n_used)?;
9194        for (i, &sx) in sel_all.iter().enumerate() {
9195            w_all[i] *= bits.per_expert_scale[sx as usize];
9196        }
9197
9198        // PREFILL PAIRS ARM (t >= 16): expert-major CSR over the resident slabs — ONE launch per
9199        // projection covers ALL (token,expert) pairs (the qwen pairs recipe, _em dot = expert_dot_g
9200        // which has the Q4_0 body; GELU pairs epilogue; R3 scale folded into pair_w).
9201        if t >= PRIME_MIN_T
9202            && m.dev_exps.as_ref().is_some_and(|d| !d.gu_il)
9203            && expert_dp4a_supported(m.gate_exps.qtype)
9204            && expert_dp4a_supported(m.up_exps.qtype)
9205            && expert_dp4a_supported(m.down_exps.qtype)
9206            && std::env::var("MEMRA_GEMMA_MOE_PAIRS").as_deref() != Ok("0")
9207        {
9208            let dev = m.dev_exps.as_ref().unwrap();
9209            let n_pairs = t * n_used;
9210            let pair_ex: Vec<i32> = sel_all.iter().map(|&x| x as i32).collect();
9211            let pair_tok: Vec<i32> = (0..n_pairs).map(|p| (p / n_used) as i32).collect();
9212            let tok_off: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
9213            let tok_ids: Vec<i32> = (0..n_pairs as i32).collect();
9214            let pt = e.htod_i32(&pair_tok)?;
9215            let pw = e.htod(&w_all)?;
9216            let toff = e.htod_i32(&tok_off)?;
9217            let tids = e.htod_i32(&tok_ids)?;
9218            let mut by_ex: Vec<Vec<i32>> = vec![Vec::new(); n_expert];
9219            for p in 0..n_pairs {
9220                by_ex[pair_ex[p] as usize].push(p as i32);
9221            }
9222            let mut ex_ids: Vec<i32> = Vec::new();
9223            let mut ex_off: Vec<i32> = vec![0];
9224            let mut ex_pairs: Vec<i32> = Vec::with_capacity(n_pairs);
9225            for (ex, list) in by_ex.iter().enumerate() {
9226                if list.is_empty() {
9227                    continue;
9228                }
9229                ex_ids.push(ex as i32);
9230                ex_pairs.extend_from_slice(list);
9231                ex_off.push(ex_pairs.len() as i32);
9232            }
9233            let n_active = ex_ids.len();
9234            let exi = e.htod_i32(&ex_ids)?;
9235            let exo = e.htod_i32(&ex_off)?;
9236            let exp_d = e.htod_i32(&ex_pairs)?;
9237            // GROUPED f16 LANE (MEMRA_MOE_F16G=1, round 46 arc 2): dequant active experts to
9238            // f16 once per projection + one grouped f16 GEMM over the CSR groups; CSR order
9239            // end-to-end (gelu is elementwise), one row permute before the scatter. The
9240            // ragged down k (704) needs no padding here — cublas takes any k.
9241            // f16-mirror numeric class, argmax/spec gated. PER-MODEL default OFF (round 50):
9242            // the gelu class regressed g26 board-2048 prefill -8.3% under the round-49
9243            // Hopper default — see moe_f16g_gemma_on.
9244            if crate::moe_f16g_gemma_on()
9245                && f16g_proj_ok(m.gate_exps.qtype, n_embd)
9246                && f16g_proj_ok(m.up_exps.qtype, n_embd)
9247                && f16g_proj_ok(m.down_exps.qtype, n_ff_exp)
9248            {
9249                let csr_tok: Vec<i32> = ex_pairs.iter().map(|&p| p / n_used as i32).collect();
9250                let csr_tok_d = e.htod_i32(&csr_tok)?;
9251                let (z_f16, z_s) = e.moe_f16g_act(moe_in, Some(&csr_tok_d), n_embd, n_pairs)?;
9252                let g_csr = e.moe_f16_grouped(
9253                    &dev.ptr_row,
9254                    0,
9255                    n_expert,
9256                    &exi,
9257                    &ex_off,
9258                    &exo,
9259                    &z_f16,
9260                    &z_s,
9261                    n_embd,
9262                    n_ff_exp,
9263                    n_active,
9264                    n_pairs,
9265                    m.gate_exps.qtype,
9266                    m.gate_exps.row_bytes,
9267                )?;
9268                let u_csr = e.moe_f16_grouped(
9269                    &dev.ptr_row,
9270                    1,
9271                    n_expert,
9272                    &exi,
9273                    &ex_off,
9274                    &exo,
9275                    &z_f16,
9276                    &z_s,
9277                    n_embd,
9278                    n_ff_exp,
9279                    n_active,
9280                    n_pairs,
9281                    m.up_exps.qtype,
9282                    m.up_exps.row_bytes,
9283                )?;
9284                let act_csr = e.moe_pairs_gelu_mul(&g_csr, &u_csr, n_pairs * n_ff_exp)?;
9285                let (a_f16, a_s) = e.moe_f16g_act(&act_csr, None, n_ff_exp, n_pairs)?;
9286                let d_csr = e.moe_f16_grouped(
9287                    &dev.ptr_row,
9288                    2,
9289                    n_expert,
9290                    &exi,
9291                    &ex_off,
9292                    &exo,
9293                    &a_f16,
9294                    &a_s,
9295                    n_ff_exp,
9296                    n_embd,
9297                    n_active,
9298                    n_pairs,
9299                    m.down_exps.qtype,
9300                    m.down_exps.row_bytes,
9301                )?;
9302                let y_down = e.rows_permute(&d_csr, &exp_d, n_pairs, n_embd)?;
9303                let mut moe_out = e.uninit(t * n_embd)?;
9304                e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
9305                if std::env::var("MEMRA_F16G_DEBUG").is_ok() {
9306                    let scan = |v: &[f32]| v.iter().filter(|x| !x.is_finite()).count();
9307                    let (yd, mo) = (e.dtoh(&y_down)?, e.dtoh(&moe_out)?);
9308                    eprintln!(
9309                        "[f16g-debug] post-permute bad={} post-scatter bad={}",
9310                        scan(&yd),
9311                        scan(&mo)
9312                    );
9313                }
9314                return Ok(moe_out);
9315            }
9316            // gate/up: int8-MMA expert GEMM (in_f = n_embd 2816 = 11x256 tiles ok); down keeps
9317            // the decode-once dp4a (in_f 704 fails the 256-superblock tiling).
9318            let mma =
9319                n_embd % 256 == 0 && std::env::var("MEMRA_GEMMA_MOE_MMA").as_deref() != Ok("0");
9320            let (gate, up) = if mma {
9321                let z_scr = e.mmq_iq_quantize_act(moe_in, n_embd, t)?;
9322                (
9323                    e.mmq_iq_experts(
9324                        &dev.ptr_row,
9325                        0,
9326                        n_expert,
9327                        &exi,
9328                        &exo,
9329                        &exp_d,
9330                        &pt,
9331                        &z_scr,
9332                        n_embd,
9333                        n_ff_exp,
9334                        n_active,
9335                        n_pairs,
9336                        t,
9337                        m.gate_exps.qtype,
9338                        m.gate_exps.row_bytes,
9339                    )?,
9340                    e.mmq_iq_experts(
9341                        &dev.ptr_row,
9342                        1,
9343                        n_expert,
9344                        &exi,
9345                        &exo,
9346                        &exp_d,
9347                        &pt,
9348                        &z_scr,
9349                        n_embd,
9350                        n_ff_exp,
9351                        n_active,
9352                        n_pairs,
9353                        t,
9354                        m.up_exps.qtype,
9355                        m.up_exps.row_bytes,
9356                    )?,
9357                )
9358            } else {
9359                let (zq, zd) = e.quantize_q8_1(moe_in, t, n_embd)?;
9360                (
9361                    e.moe_pairs_matvec_q8_dec(
9362                        &dev.ptr_row,
9363                        0,
9364                        &exi,
9365                        &exo,
9366                        &exp_d,
9367                        &pt,
9368                        &zq,
9369                        &zd,
9370                        n_embd,
9371                        n_ff_exp,
9372                        n_expert,
9373                        n_active,
9374                        n_pairs,
9375                        m.gate_exps.qtype,
9376                        m.gate_exps.row_bytes,
9377                    )?,
9378                    e.moe_pairs_matvec_q8_dec(
9379                        &dev.ptr_row,
9380                        1,
9381                        &exi,
9382                        &exo,
9383                        &exp_d,
9384                        &pt,
9385                        &zq,
9386                        &zd,
9387                        n_embd,
9388                        n_ff_exp,
9389                        n_expert,
9390                        n_active,
9391                        n_pairs,
9392                        m.up_exps.qtype,
9393                        m.up_exps.row_bytes,
9394                    )?,
9395                )
9396            };
9397            let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
9398            let pself = e.htod_i32(&pair_self)?;
9399            // DOWN through the int8-MMA expert GEMM (2026-07-31, g26 prefill lever): the
9400            // ragged k (n_ff_exp=704 on the 26B) rides a PADDED k-walk — in_f rounds up
9401            // to the 256-val superblock (768) while the act quantizer's zero padding
9402            // makes every padded-k product exactly zero (weight overread bytes multiply
9403            // zero int8 act values; the dev slab carries 144B tail slack for the OOB).
9404            // The old dp4a matvec was 11.3ms/call at m=T (the 0.07x prefill wall).
9405            // MEMRA_GEMMA_MOE_MMA=0 reverts down together with gate/up.
9406            // FUSED ACT-EPILOGUE (default on): gelu_tanh(gate)*up + D4 quantize in one
9407            // launch — no f32 act buffer (the fused kernel zero-pads the ragged tail
9408            // exactly like the two-pass quantizer). MEMRA_MOE_FUSE_ACTQ=0 rollback.
9409            // Scratch bytes are BYTE-IDENTICAL (kernel-check gated).
9410            let y_down = if mma {
9411                let in_pad = n_ff_exp.div_ceil(256) * 256;
9412                let a_scr = if crate::moe_fuse_actq_on() {
9413                    e.mmq_iq_fused_act_quant(&gate, &up, n_ff_exp, n_pairs, 1)?
9414                } else {
9415                    let act = e.moe_pairs_gelu_mul(&gate, &up, n_pairs * n_ff_exp)?;
9416                    e.mmq_iq_quantize_act(&act, n_ff_exp, n_pairs)?
9417                };
9418                e.mmq_iq_experts(
9419                    &dev.ptr_row,
9420                    2,
9421                    n_expert,
9422                    &exi,
9423                    &exo,
9424                    &exp_d,
9425                    &pself,
9426                    &a_scr,
9427                    in_pad,
9428                    n_embd,
9429                    n_active,
9430                    n_pairs,
9431                    n_pairs,
9432                    m.down_exps.qtype,
9433                    m.down_exps.row_bytes,
9434                )?
9435            } else {
9436                let act = e.moe_pairs_gelu_mul(&gate, &up, n_pairs * n_ff_exp)?;
9437                let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
9438                e.moe_pairs_matvec_q8_dec(
9439                    &dev.ptr_row,
9440                    2,
9441                    &exi,
9442                    &exo,
9443                    &exp_d,
9444                    &pself,
9445                    &aq2,
9446                    &ad2,
9447                    n_ff_exp,
9448                    n_embd,
9449                    n_expert,
9450                    n_active,
9451                    n_pairs,
9452                    m.down_exps.qtype,
9453                    m.down_exps.row_bytes,
9454                )?
9455            };
9456            let mut moe_out = e.uninit(t * n_embd)?;
9457            e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
9458            return Ok(moe_out);
9459        }
9460
9461        let g_len = m.gate_exps.expert_stride;
9462        let u_len = m.up_exps.expert_stride;
9463        let d_len = m.down_exps.expert_stride;
9464        // Resident dev slabs (fits-VRAM regime): read each expert straight from the device slab
9465        // at ex*stride — zero H2D, SAME qmatvec_view kernel/bytes as the staged path. Staging is
9466        // the spill fallback.
9467        let dev = m.dev_exps.as_ref().filter(|d| !d.gu_il);
9468        let (mut sg, mut su, mut sd) = if dev.is_some() {
9469            (None, None, None)
9470        } else {
9471            (
9472                Some(e.alloc_u8_uninit(g_len)?),
9473                Some(e.alloc_u8_uninit(u_len)?),
9474                Some(e.alloc_u8_uninit(d_len)?),
9475            )
9476        };
9477        let mut moe_out = e.zeros(t * n_embd)?;
9478        for tok in 0..t {
9479            let sel = &sel_all[tok * n_used..(tok + 1) * n_used];
9480            let w = &w_all[tok * n_used..(tok + 1) * n_used];
9481            let zt = moe_in.slice(tok * n_embd..(tok + 1) * n_embd);
9482            for (j, &ex) in sel.iter().enumerate() {
9483                let ex = ex as usize;
9484                let gate = match dev {
9485                    Some(d) => e.qmatvec_view(
9486                        &d.gate,
9487                        ex * g_len..(ex + 1) * g_len,
9488                        &zt,
9489                        1,
9490                        m.gate_exps.in_f,
9491                        m.gate_exps.out_f,
9492                        m.gate_exps.qtype,
9493                        m.gate_exps.row_bytes,
9494                    )?,
9495                    None => {
9496                        let sg = sg.as_mut().unwrap();
9497                        e.stage_expert(m.gate_exps.expert_bytes(ex), sg, 0)?;
9498                        e.qmatvec_view(
9499                            sg,
9500                            0..g_len,
9501                            &zt,
9502                            1,
9503                            m.gate_exps.in_f,
9504                            m.gate_exps.out_f,
9505                            m.gate_exps.qtype,
9506                            m.gate_exps.row_bytes,
9507                        )?
9508                    }
9509                };
9510                let up = match dev {
9511                    Some(d) => e.qmatvec_view(
9512                        &d.up,
9513                        ex * u_len..(ex + 1) * u_len,
9514                        &zt,
9515                        1,
9516                        m.up_exps.in_f,
9517                        m.up_exps.out_f,
9518                        m.up_exps.qtype,
9519                        m.up_exps.row_bytes,
9520                    )?,
9521                    None => {
9522                        let su = su.as_mut().unwrap();
9523                        e.stage_expert(m.up_exps.expert_bytes(ex), su, 0)?;
9524                        e.qmatvec_view(
9525                            su,
9526                            0..u_len,
9527                            &zt,
9528                            1,
9529                            m.up_exps.in_f,
9530                            m.up_exps.out_f,
9531                            m.up_exps.qtype,
9532                            m.up_exps.row_bytes,
9533                        )?
9534                    }
9535                };
9536                let mut act = e.uninit(n_ff_exp)?;
9537                e.gelu_tanh_mul(&gate, &up, &mut act, n_ff_exp)?;
9538                let actv = act.slice(0..n_ff_exp);
9539                let y = match dev {
9540                    Some(d) => e.qmatvec_view(
9541                        &d.down,
9542                        ex * d_len..(ex + 1) * d_len,
9543                        &actv,
9544                        1,
9545                        m.down_exps.in_f,
9546                        m.down_exps.out_f,
9547                        m.down_exps.qtype,
9548                        m.down_exps.row_bytes,
9549                    )?,
9550                    None => {
9551                        let sd = sd.as_mut().unwrap();
9552                        e.stage_expert(m.down_exps.expert_bytes(ex), sd, 0)?;
9553                        e.qmatvec_view(
9554                            sd,
9555                            0..d_len,
9556                            &actv,
9557                            1,
9558                            m.down_exps.in_f,
9559                            m.down_exps.out_f,
9560                            m.down_exps.qtype,
9561                            m.down_exps.row_bytes,
9562                        )?
9563                    }
9564                };
9565                let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
9566                e.axpy_into(&y, w[j], &mut dst, n_embd)?;
9567            }
9568        }
9569        Ok(moe_out)
9570    }
9571
9572    /// One gemma4 trunk layer (R8): x -> x_next.
9573    fn gemma4_layer(
9574        &self,
9575        e: &Engine,
9576        il: usize,
9577        layer: &crate::hybrid::HybridLayer,
9578        x: &CudaSlice<f32>,
9579        pos_d: &CudaSlice<i32>,
9580        t: usize,
9581    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9582        let n_embd = self.cfg.n_embd as usize;
9583        let eps = self.cfg.rms_eps;
9584
9585        let mut h = e.zeros(t * n_embd)?;
9586        e.rms_norm(x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
9587        let Mixer::Full(fa) = &layer.mixer else {
9588            panic!("gemma4 layer {il} not full-attn")
9589        };
9590        let o = self.gemma4_attn(e, fa, il, &h, pos_d, t)?;
9591        // gemma order: post_attention_norm applies to the ATTENTION OUTPUT, then the residual.
9592        let mut cur = e.zeros(t * n_embd)?;
9593        e.rms_norm(
9594            &o,
9595            layer.post_attn_norm.float_data(),
9596            &mut cur,
9597            n_embd,
9598            t,
9599            eps,
9600        )?;
9601        self.gemma4_layer_tail_add(e, layer, &cur, x, t)
9602    }
9603
9604    /// Everything after the attention output in a gemma4 layer: the residual add (cur + x ->
9605    /// attn_out) FUSED with the three attn_out norms, then shared FFN + router + MoE + combine +
9606    /// layer scale — shared verbatim by the prefill, decode and verify paths.
9607    fn gemma4_layer_tail_add(
9608        &self,
9609        e: &Engine,
9610        layer: &crate::hybrid::HybridLayer,
9611        cur: &CudaSlice<f32>,
9612        x: &CudaSlice<f32>,
9613        t: usize,
9614    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9615        Ok(self.gemma4_layer_tail_add_n(e, layer, cur, x, t, None)?.0)
9616    }
9617
9618    /// tail_add with the NEXT layer's attn_norm fused into the closing add+scale (one launch
9619    /// produces both x_next and h_next — the cross-layer fusion). `next_norm` None = last layer.
9620    fn gemma4_layer_tail_add_n(
9621        &self,
9622        e: &Engine,
9623        layer: &crate::hybrid::HybridLayer,
9624        cur: &CudaSlice<f32>,
9625        x: &CudaSlice<f32>,
9626        t: usize,
9627        next_norm: Option<&CudaSlice<f32>>,
9628    ) -> Result<(CudaSlice<f32>, Option<CudaSlice<f32>>), Box<dyn std::error::Error>> {
9629        let n_embd = self.cfg.n_embd as usize;
9630        let bits = layer.gemma4.as_ref().unwrap();
9631        let (sn, attn_out) = self.gemma4_layer_tail_core(e, layer, cur, x, t)?;
9632        let mut xn = e.uninit(t * n_embd)?;
9633        match next_norm {
9634            Some(w) => {
9635                let mut hn = e.uninit(t * n_embd)?;
9636                e.add_scale_rms_norm(
9637                    &sn,
9638                    &attn_out,
9639                    bits.layer_scale,
9640                    w,
9641                    &mut xn,
9642                    &mut hn,
9643                    n_embd,
9644                    t,
9645                    self.cfg.rms_eps,
9646                )?;
9647                Ok((xn, Some(hn)))
9648            }
9649            None => {
9650                e.add_scale(&sn, &attn_out, bits.layer_scale, &mut xn, t * n_embd)?;
9651                Ok((xn, None))
9652            }
9653        }
9654    }
9655
9656    /// Tail core: attn_out = cur+x (fused with the 3 norms), both FFN branches, the combine
9657    /// norm — returns (sn, attn_out) for the closing add+scale variants.
9658    fn gemma4_layer_tail_core(
9659        &self,
9660        e: &Engine,
9661        layer: &crate::hybrid::HybridLayer,
9662        cur: &CudaSlice<f32>,
9663        x: &CudaSlice<f32>,
9664        t: usize,
9665    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9666        self.gemma4_layer_tail_core_pn(e, layer, cur, x, t, None, false)
9667    }
9668
9669    /// tail_core with an optional PRE-NORM fold (glue-fusion lane): `pre_norm = Some(wa)`
9670    /// means `cur` is the RAW attention output and the dense entry runs
9671    /// rms(cur, wa) + residual-add + ffn_norm as ONE launch (E4B's post-attn norm fold).
9672    /// `defer_post_norm`: dense-arm exit returns RAW f0 (ffn_down output) instead of
9673    /// sn = rms(f0, post_ffw) — the caller fuses the post-norm into its residual emit
9674    /// (rms_pre_add_q8_1, E4B glue wave 5). MoE arm ignores it.
9675    fn gemma4_layer_tail_core_pn(
9676        &self,
9677        e: &Engine,
9678        layer: &crate::hybrid::HybridLayer,
9679        cur: &CudaSlice<f32>,
9680        x: &CudaSlice<f32>,
9681        t: usize,
9682        pre_norm: Option<&CudaSlice<f32>>,
9683        defer_post_norm: bool,
9684    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9685        let n_embd = self.cfg.n_embd as usize;
9686        let eps = self.cfg.rms_eps;
9687        let bits = layer.gemma4.as_ref().unwrap();
9688
9689        // DENSE gemma4 variants (31B/E4B — no MoE, no parallel branch): attn_out = cur + x
9690        // fused with the single ffn_norm; GELU_PAR ffn; post_ffw_norm.
9691        let Some(mbits) = bits.moe_bits.as_ref() else {
9692            let crate::hybrid::Ffn::Dense {
9693                ffn_gate,
9694                ffn_up,
9695                ffn_down,
9696            } = &layer.ffn
9697            else {
9698                panic!("gemma4 dense layer without Dense ffn")
9699            };
9700            let mut attn_out = e.uninit(t * n_embd)?;
9701            let mut zsh = e.uninit(t * n_embd)?;
9702            // wave-2: with the pre-norm fold active the entry ALSO emits zsh q8_1 — the
9703            // t=1 fused2 gate/up consume the pair with no standalone quantize launch.
9704            let mut zpair: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
9705            match pre_norm {
9706                Some(wa) if t == 1 => {
9707                    zpair = Some(e.rms_pre_add_rms_norm_q8z(
9708                        cur,
9709                        wa,
9710                        x,
9711                        bits.ffn_norm.float_data(),
9712                        &mut attn_out,
9713                        &mut zsh,
9714                        n_embd,
9715                        t,
9716                        eps,
9717                    )?);
9718                }
9719                Some(wa) => e.rms_pre_add_rms_norm(
9720                    cur,
9721                    wa,
9722                    x,
9723                    bits.ffn_norm.float_data(),
9724                    &mut attn_out,
9725                    &mut zsh,
9726                    n_embd,
9727                    t,
9728                    eps,
9729                )?,
9730                None => e.add_rms_norm(
9731                    cur,
9732                    x,
9733                    bits.ffn_norm.float_data(),
9734                    &mut attn_out,
9735                    &mut zsh,
9736                    n_embd,
9737                    t,
9738                    eps,
9739                )?,
9740            }
9741            let n_ff = ffn_gate.out_features();
9742            // FFN persistent slab (counter-barrier form) FALSIFIED here 2026-07-14
9743            // (falsification #7, jsonl row): PDL glue already hides the launch boundaries
9744            // it fused; barrier + worst-segment occupancy net −0.3% (31B depth) / −2.3%
9745            // (E4B spec). Down's act dependency is all-to-all, so sentinel sync cannot
9746            // rescue segment C — the megakernel front is closed for the dense tail.
9747            let (gate, up) = if t == 1 {
9748                let (zq, zd) = match zpair {
9749                    Some(p) => p,
9750                    None => e.quantize_q8_1(&zsh, 1, n_embd)?,
9751                };
9752                match e.matmul_q4_fused2(ffn_gate, ffn_up, &zq, &zd)? {
9753                    Some(p) => p,
9754                    // NVFP4mix dense trunk: gate/up are both NVFP4 (down stays Q8_0).
9755                    None => match e.matmul_nvfp4_fused2(ffn_gate, ffn_up, &zq, &zd, 1)? {
9756                        Some(p) => p,
9757                        None => (
9758                            e.matmul_pre(ffn_gate, &zq, &zd, &zsh, 1)?,
9759                            e.matmul_pre(ffn_up, &zq, &zd, &zsh, 1)?,
9760                        ),
9761                    },
9762                }
9763            } else {
9764                // BATCHED FUSED2 (DEFAULT ON 2026-07-13, MEMRA_F2B=0 seam): one segmented
9765                // launch for the verify's gate+up — the up segment's blocks fill SMs as
9766                // the gate segment drains (the launch-tail mechanism behind the b-tier
9767                // plateau; first positive after six falsified in-kernel variants).
9768                static F2B: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9769                let f2b = *F2B.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0"));
9770                let fused = if f2b {
9771                    let (zq, zd) = e.quantize_q8_1(&zsh, t, n_embd)?;
9772                    e.matmul_q4_fused2_batched(ffn_gate, ffn_up, &zq, &zd, t)?
9773                } else {
9774                    None
9775                };
9776                match fused {
9777                    Some(p) => p,
9778                    None => {
9779                        // quantize-once window: gate/up share `zsh` (borrowed across the pair).
9780                        e.mmq_act_begin();
9781                        (e.matmul(ffn_gate, &zsh, t)?, e.matmul(ffn_up, &zsh, t)?)
9782                    }
9783                }
9784            };
9785            let mut act = e.uninit(t * n_ff)?;
9786            // act quantize folds into the GELU epilogue (bit-identical q8_1 rounding);
9787            // ffn_down rides matmul_pre — one quantize launch fewer per layer.
9788            let f0 = if e.uses_q8_1_fast(ffn_down) {
9789                let upv = e.view(&up, t * n_ff);
9790                let up_all = upv.slice(0..t * n_ff);
9791                let (aq, ad) = e.gelu_tanh_mul_q8_1(&gate, &up_all, &mut act, n_ff, t)?;
9792                e.matmul_pre(ffn_down, &aq, &ad, &act, t)?
9793            } else {
9794                e.gelu_tanh_mul(&gate, &up, &mut act, t * n_ff)?;
9795                e.matmul(ffn_down, &act, t)?
9796            };
9797            if defer_post_norm {
9798                return Ok((f0, attn_out));
9799            }
9800            let mut sn = e.uninit(t * n_embd)?;
9801            e.rms_norm(
9802                &f0,
9803                bits.post_ffw_norm.float_data(),
9804                &mut sn,
9805                n_embd,
9806                t,
9807                eps,
9808            )?;
9809            return Ok((sn, attn_out));
9810        };
9811
9812        assert!(pre_norm.is_none(), "pre-norm fold is dense-entry only");
9813        // MoE variant (26B): attn_out = cur + x fused with the three attn_out norms
9814        // (ffn_norm + router-scale + pre_ffw_norm_2): ONE launch, chains verbatim. At small t
9815        // (decode + verify) the zsh and moe_in outputs are EMITTED q8_1 (both consumers are
9816        // quantized matmuls — two quantize launches + two f32 round-trips fold away).
9817        let mut attn_out = e.uninit(t * n_embd)?;
9818        let mut router_in = e.uninit(t * n_embd)?;
9819        let fast_moe = match &layer.ffn {
9820            crate::hybrid::Ffn::Moe(m) => {
9821                m.dev_exps.as_ref().is_some_and(|d| !d.gu_il)
9822                    && expert_dp4a_supported(m.gate_exps.qtype)
9823                    && expert_dp4a_supported(m.up_exps.qtype)
9824                    && expert_dp4a_supported(m.down_exps.qtype)
9825                    && std::env::var("MEMRA_GEMMA_MOE_FAST").as_deref() != Ok("0")
9826            }
9827            _ => false,
9828        };
9829        let q8z = t < PRIME_MIN_T && fast_moe;
9830        let (zsh_f32, zsh_q8, moe_q8) = if q8z {
9831            let (z0, m2) = e.add_rms_norm3_q8z(
9832                cur,
9833                x,
9834                bits.ffn_norm.float_data(),
9835                &mbits.router_scale_pre,
9836                mbits.pre_ffw_norm_2.float_data(),
9837                &mut attn_out,
9838                &mut router_in,
9839                n_embd,
9840                t,
9841                eps,
9842            )?;
9843            (None, Some(z0), Some(m2))
9844        } else {
9845            let mut zsh = e.uninit(t * n_embd)?;
9846            let mut moe_in = e.uninit(t * n_embd)?;
9847            e.add_rms_norm3(
9848                cur,
9849                x,
9850                bits.ffn_norm.float_data(),
9851                &mbits.router_scale_pre,
9852                mbits.pre_ffw_norm_2.float_data(),
9853                &mut attn_out,
9854                &mut zsh,
9855                &mut router_in,
9856                &mut moe_in,
9857                n_embd,
9858                t,
9859                eps,
9860            )?;
9861            (Some((zsh, moe_in)), None, None)
9862        };
9863        let attn_out2 = attn_out;
9864        #[allow(unused_variables)]
9865        let attn_out = &attn_out2;
9866        let n_ff = mbits.shared_gate.out_features();
9867        let (gate, up) = if let Some((zq, zd)) = zsh_q8.as_ref() {
9868            if t == 1 {
9869                match e.matmul_q4_fused2(&mbits.shared_gate, &mbits.shared_up, zq, zd)? {
9870                    Some(p) => p,
9871                    None => match e.matmul_nvfp4_fused2(
9872                        &mbits.shared_gate,
9873                        &mbits.shared_up,
9874                        zq,
9875                        zd,
9876                        1,
9877                    )? {
9878                        Some(p) => p,
9879                        None => {
9880                            let h0 = e.zeros(0)?;
9881                            (
9882                                e.matmul_pre(&mbits.shared_gate, zq, zd, &h0, 1)?,
9883                                e.matmul_pre(&mbits.shared_up, zq, zd, &h0, 1)?,
9884                            )
9885                        }
9886                    },
9887                }
9888            } else {
9889                // verify t 2..15: the batched mmvq twins consume the pre-quantized pair.
9890                let h0 = e.zeros(0)?;
9891                (
9892                    e.matmul_pre(&mbits.shared_gate, zq, zd, &h0, t)?,
9893                    e.matmul_pre(&mbits.shared_up, zq, zd, &h0, t)?,
9894                )
9895            }
9896        } else {
9897            let (zsh, _) = zsh_f32.as_ref().unwrap();
9898            (
9899                e.matmul(&mbits.shared_gate, zsh, t)?,
9900                e.matmul(&mbits.shared_up, zsh, t)?,
9901            )
9902        };
9903        let mut act = e.uninit(t * n_ff)?;
9904        e.gelu_tanh_mul(&gate, &up, &mut act, t * n_ff)?;
9905        let mlp0 = e.matmul(&mbits.shared_down, &act, t)?;
9906        let crate::hybrid::Ffn::Moe(m) = &layer.ffn else {
9907            panic!("gemma4 layer not MoE")
9908        };
9909        let moe0 = match (&moe_q8, &zsh_f32) {
9910            (Some(mq), _) => self.gemma4_moe_q8(e, m, mbits, mq, &router_in, t)?,
9911            (None, Some((_, moe_in))) => self.gemma4_moe(e, m, mbits, moe_in, &router_in, t)?,
9912            _ => unreachable!(),
9913        };
9914        // post_ffw_norm_1(mlp0) + post_ffw_norm_2(moe0): one fused launch, per-row verbatim.
9915        let mut mlp = e.uninit(t * n_embd)?;
9916        let mut moe = e.uninit(t * n_embd)?;
9917        e.rms_norm2x(
9918            &mlp0,
9919            &moe0,
9920            mbits.post_ffw_norm_1.float_data(),
9921            mbits.post_ffw_norm_2.float_data(),
9922            &mut mlp,
9923            &mut moe,
9924            n_embd,
9925            t,
9926            eps,
9927        )?;
9928
9929        // combine: rms_norm(mlp + moe, post_ffw_norm) + attn_out, then the layer output scalar.
9930        // add+norm fused (add_rms_norm == add then rms_norm, kernel-check-pinned identity).
9931        let mut sum = e.uninit(t * n_embd)?;
9932        let mut sn = e.uninit(t * n_embd)?;
9933        e.add_rms_norm(
9934            &mlp,
9935            &moe,
9936            bits.post_ffw_norm.float_data(),
9937            &mut sum,
9938            &mut sn,
9939            n_embd,
9940            t,
9941            eps,
9942        )?;
9943        Ok((sn, attn_out2))
9944    }
9945
9946    /// tail_add with the next attn_norm emitted PRE-QUANTIZED q8_1 (decode/verify loops).
9947    /// pn-fold front of `gemma4_layer_tail_add_nq` (MEMRA_G4_PNFOLD, GAP-DIAGNOSIS
9948    /// verdict 7): takes the RAW attention output `o` and, on the dense trunk, folds
9949    /// post_attn_norm into the tail entry and post_ffw_norm into the residual exit —
9950    /// the E4B glue chain (13594) backported to the decode/verify trio. Kills the
9951    /// standalone rms_norm(o) + rms_norm(f0) (+ffn quantize) launches per layer.
9952    /// BITS-CHANGING vs the two-launch chain; every caller rides this ONE front so
9953    /// decode == verify == graph parity holds by construction at either seam value.
9954    /// Non-dense (26B MoE) and seam-off fall to the unfused chain unchanged.
9955    pub(crate) fn gemma4_layer_tail_add_nq_pn(
9956        &self,
9957        e: &Engine,
9958        layer: &crate::hybrid::HybridLayer,
9959        o: &CudaSlice<f32>,
9960        x: &CudaSlice<f32>,
9961        t: usize,
9962        next_norm: Option<&CudaSlice<f32>>,
9963    ) -> Result<(CudaSlice<f32>, Option<(CudaSlice<i8>, CudaSlice<f32>)>), Box<dyn std::error::Error>>
9964    {
9965        let n_embd = self.cfg.n_embd as usize;
9966        let eps = self.cfg.rms_eps;
9967        let bits = layer.gemma4.as_ref().unwrap();
9968        if Engine::g4_pnfold_on() && matches!(layer.ffn, crate::hybrid::Ffn::Dense { .. }) {
9969            let (f0, attn_out) = self.gemma4_layer_tail_core_pn(
9970                e,
9971                layer,
9972                o,
9973                x,
9974                t,
9975                Some(layer.post_attn_norm.float_data()),
9976                true,
9977            )?;
9978            let mut xn = e.uninit(t * n_embd)?;
9979            return match next_norm {
9980                Some(w) => {
9981                    let pair = e.rms_pre_add_scale_rms_norm_q8_1(
9982                        &f0,
9983                        bits.post_ffw_norm.float_data(),
9984                        &attn_out,
9985                        bits.layer_scale,
9986                        w,
9987                        &mut xn,
9988                        n_embd,
9989                        t,
9990                        eps,
9991                    )?;
9992                    Ok((xn, Some(pair)))
9993                }
9994                None => {
9995                    let mut sn = e.uninit(t * n_embd)?;
9996                    e.rms_norm(
9997                        &f0,
9998                        bits.post_ffw_norm.float_data(),
9999                        &mut sn,
10000                        n_embd,
10001                        t,
10002                        eps,
10003                    )?;
10004                    e.add_scale(&sn, &attn_out, bits.layer_scale, &mut xn, t * n_embd)?;
10005                    Ok((xn, None))
10006                }
10007            };
10008        }
10009        let mut cur = e.uninit(t * n_embd)?;
10010        e.rms_norm(
10011            o,
10012            layer.post_attn_norm.float_data(),
10013            &mut cur,
10014            n_embd,
10015            t,
10016            eps,
10017        )?;
10018        self.gemma4_layer_tail_add_nq(e, layer, &cur, x, t, next_norm)
10019    }
10020
10021    pub(crate) fn gemma4_layer_tail_add_nq(
10022        &self,
10023        e: &Engine,
10024        layer: &crate::hybrid::HybridLayer,
10025        cur: &CudaSlice<f32>,
10026        x: &CudaSlice<f32>,
10027        t: usize,
10028        next_norm: Option<&CudaSlice<f32>>,
10029    ) -> Result<(CudaSlice<f32>, Option<(CudaSlice<i8>, CudaSlice<f32>)>), Box<dyn std::error::Error>>
10030    {
10031        let n_embd = self.cfg.n_embd as usize;
10032        let bits = layer.gemma4.as_ref().unwrap();
10033        let (sn, attn_out) = self.gemma4_layer_tail_core(e, layer, cur, x, t)?;
10034        let mut xn = e.uninit(t * n_embd)?;
10035        match next_norm {
10036            Some(w) => {
10037                let pair = e.add_scale_rms_norm_q8_1(
10038                    &sn,
10039                    &attn_out,
10040                    bits.layer_scale,
10041                    w,
10042                    &mut xn,
10043                    n_embd,
10044                    t,
10045                    self.cfg.rms_eps,
10046                )?;
10047                Ok((xn, Some(pair)))
10048            }
10049            None => {
10050                e.add_scale(&sn, &attn_out, bits.layer_scale, &mut xn, t * n_embd)?;
10051                Ok((xn, None))
10052            }
10053        }
10054    }
10055
10056    /// gemma4 prefill: `last_only` = forward_last semantics (lm_head on the final row only).
10057    /// R4: final logits softcapped 30*tanh(l/30) on host (monotonic — argmax unaffected).
10058    fn gemma4_forward(
10059        &self,
10060        e: &Engine,
10061        tokens: &[u32],
10062        last_only: bool,
10063    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
10064        // E4B routes to its own forward regardless of the caller's entry point (forward /
10065        // forward_last / prime paths all funnel here for gemma4).
10066        if self.is_gemma4_e4b() {
10067            return self.gemma4_e4b_forward(e, tokens, last_only);
10068        }
10069        let n_embd = self.cfg.n_embd as usize;
10070        let t = tokens.len();
10071        let pos: Vec<i32> = (0..t as i32).collect();
10072        let pos_d = e.htod_i32(&pos)?;
10073
10074        let mut x = self.embed(e, tokens)?;
10075        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
10076        // MEMRA_GEMMA_PROBE=1: per-layer trunk stats (host rms + max of the LAST token row) —
10077        // the bring-up bisect vs llama-eval-callback node stats.
10078        let probe = std::env::var("MEMRA_GEMMA_PROBE").is_ok();
10079        let stat =
10080            |e: &Engine, x: &CudaSlice<f32>, tag: &str| -> Result<(), Box<dyn std::error::Error>> {
10081                let h = e.dtoh(x)?;
10082                let bad = h.iter().filter(|v| !v.is_finite()).count();
10083                let mx = h
10084                    .iter()
10085                    .filter(|v| v.is_finite())
10086                    .fold(0.0f32, |m, v| m.max(v.abs()));
10087                eprintln!(
10088                    "[gemma-probe] {tag}: tok0_first3={:?} bad={bad} max={mx:.3e}",
10089                    &h[..3]
10090                );
10091                Ok(())
10092            };
10093        if probe {
10094            stat(e, &x, "embed")?;
10095        }
10096        for (il, layer) in self.layers.iter().enumerate() {
10097            x = self.gemma4_layer(e, il, layer, &x, &pos_d, t)?;
10098            if probe {
10099                stat(e, &x, &format!("L{il}"))?;
10100            }
10101        }
10102        let mut hn = e.zeros(t * n_embd)?;
10103        e.rms_norm(
10104            &x,
10105            self.output_norm.float_data(),
10106            &mut hn,
10107            n_embd,
10108            t,
10109            self.cfg.rms_eps,
10110        )?;
10111        let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
10112        let n_vocab = self.output.out_features();
10113        let logits = if last_only {
10114            let hv = e.view(&hn, t * n_embd);
10115            let last_row = hv.slice((t - 1) * n_embd..t * n_embd);
10116            let mut hlast = e.zeros(n_embd)?;
10117            e.copy_view_into(&mut hlast, 0, &last_row, n_embd)?;
10118            let mut ld = e.matmul(&self.output, &hlast, 1)?;
10119            e.softcap(&mut ld, cap, n_vocab)?;
10120            self.gemma4_suppress(e, &mut ld, 1)?;
10121            e.dtoh(&ld)?
10122        } else {
10123            let mut ld = e.matmul(&self.output, &hn, t)?;
10124            e.softcap(&mut ld, cap, t * n_vocab)?;
10125            self.gemma4_suppress(e, &mut ld, t)?;
10126            e.dtoh(&ld)?
10127        };
10128        Ok(logits)
10129    }
10130
10131    /// gemma4 BATCHED PROMPT PRIME v0 (fresh cache only): the prefill graph over the whole
10132    /// prompt with each layer's post-rope K / weightless-normed V appended into the quantized
10133    /// KV cache (decode-append row math). Returns (last-row softcapped logits, h_seed = last
10134    /// pre-output_norm hidden, hiddens = full pre-output_norm stack [T, n_embd]).
10135    pub(crate) fn gemma4_prime(
10136        &self,
10137        e: &Engine,
10138        tokens: &[u32],
10139        cache: &mut Cache,
10140        overlay: Option<&crate::vision::EmbedOverlay>,
10141    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10142        // Err, not assert (2026-08-07, lane/gemma4-serve-gaps): a served gemma4 prompt longer
10143        // than the worker's prefill tick used to chunk here, and chunk 2 (pos > 0) killed the
10144        // whole worker process on this line. The worker now primes gemma4 monolithically and
10145        // routes continuation suffixes tokenwise; this is the per-request backstop.
10146        if cache.pos != 0 {
10147            return Err(
10148                "gemma4 prime v0 is fresh-prompt only (no continuation/chunked prime) \
10149                        — prime the full prompt in one call or decode tokenwise"
10150                    .into(),
10151            );
10152        }
10153        let n_embd = self.cfg.n_embd as usize;
10154        let eps = self.cfg.rms_eps;
10155        let t = tokens.len();
10156        let pos: Vec<i32> = (0..t as i32).collect();
10157        let pos_d = e.htod_i32(&pos)?;
10158        let mut x = self.embed(e, tokens)?;
10159        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
10160        // Masked-prefill arm (lane/gemma-vision): splice the tower rows AFTER the
10161        // sqrt(n_embd) text scale — the reference scales token batches only
10162        // (llama.cpp gemma4.cpp:182 `ubatch.token ? sqrtf(n_embd) : 1.0f`), so image
10163        // embeddings enter the trunk raw. Each span becomes an attention ISLAND:
10164        // bidirectional within itself, causal+SWA everywhere else, matching the
10165        // reference's llama_set_causal_attn(false) image batch exactly.
10166        let island: Option<CudaSlice<i32>> = match overlay {
10167            Some(ov) => {
10168                let mut span_id = vec![-1i32; t];
10169                for (i, &(pos, row_off, n_rows)) in ov.spans.iter().enumerate() {
10170                    if pos + n_rows > t {
10171                        return Err(format!(
10172                            "gemma4 overlay span {i} [{pos}, {}) exceeds the prompt ({t})",
10173                            pos + n_rows
10174                        )
10175                        .into());
10176                    }
10177                    let view = ov.rows.slice(row_off * n_embd..(row_off + n_rows) * n_embd);
10178                    e.copy_view_into(&mut x, pos * n_embd, &view, n_rows * n_embd)?;
10179                    for s in span_id.iter_mut().skip(pos).take(n_rows) {
10180                        *s = i as i32;
10181                    }
10182                }
10183                // MEMRA_GV_FORCE_CAUSAL=1: WRONG-ARM probe seam — splice the rows but
10184                // keep the plain causal mask. Exists only so the decisive probe can show
10185                // the island mask itself changes the answer; never on in serving.
10186                if std::env::var("MEMRA_GV_FORCE_CAUSAL").as_deref() == Ok("1") {
10187                    None
10188                } else {
10189                    Some(e.htod_i32(&span_id)?)
10190                }
10191            }
10192            None => None,
10193        };
10194        for (il, layer) in self.layers.iter().enumerate() {
10195            let mut h = e.zeros(t * n_embd)?;
10196            e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
10197            let Mixer::Full(fa) = &layer.mixer else {
10198                panic!("gemma4 layer not full-attn")
10199            };
10200            let trace = il == 0 && std::env::var("MEMRA_G4_PRIME_TRACE").as_deref() == Ok("1");
10201            if trace {
10202                let v = e.dtoh(&h)?;
10203                let nan = v.iter().filter(|x| x.is_nan()).count();
10204                eprintln!("[g4-prime-trace] L0 post-attn_norm: nan={nan}/{}", v.len());
10205            }
10206            let o =
10207                self.gemma4_attn_prime(e, fa, il, &h, &pos_d, t, Some(cache), island.as_ref())?;
10208            if trace {
10209                let v = e.dtoh(&o)?;
10210                let nan = v.iter().filter(|x| x.is_nan()).count();
10211                eprintln!("[g4-prime-trace] L0 post-attn: nan={nan}/{}", v.len());
10212            }
10213            let mut cur = e.zeros(t * n_embd)?;
10214            e.rms_norm(
10215                &o,
10216                layer.post_attn_norm.float_data(),
10217                &mut cur,
10218                n_embd,
10219                t,
10220                eps,
10221            )?;
10222            x = self.gemma4_layer_tail_add(e, layer, &cur, &x, t)?;
10223            self.dflash_tap(e, cache, il, &x, t)?;
10224            // MEMRA_G4_PRIME_TRACE=1: per-layer NaN hunt (recipe-probe diagnostic; off = zero cost)
10225            if std::env::var("MEMRA_G4_PRIME_TRACE").as_deref() == Ok("1") {
10226                let h = e.dtoh(&x)?;
10227                let nan = h.iter().filter(|v| v.is_nan()).count();
10228                let amax = h.iter().fold(0f32, |a, v| a.max(v.abs()));
10229                eprintln!(
10230                    "[g4-prime-trace] layer {il}: nan={nan}/{} amax={amax:.3}",
10231                    h.len()
10232                );
10233                if nan > 0 {
10234                    return Err(format!("g4-prime-trace: first NaN at layer {il}").into());
10235                }
10236            }
10237        }
10238        cache.pos += t;
10239        let hiddens = e.clone_dtod(&x)?;
10240        let xv = e.view(&x, t * n_embd);
10241        let last_row = xv.slice((t - 1) * n_embd..t * n_embd);
10242        let mut h_seed = e.zeros(n_embd)?;
10243        e.copy_view_into(&mut h_seed, 0, &last_row, n_embd)?;
10244        let mut hn = e.uninit(n_embd)?;
10245        e.rms_norm(
10246            &h_seed,
10247            self.output_norm.float_data(),
10248            &mut hn,
10249            n_embd,
10250            1,
10251            eps,
10252        )?;
10253        let mut ld = e.matmul(&self.output, &hn, 1)?;
10254        let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
10255        e.softcap(&mut ld, cap, self.output.out_features())?;
10256        self.gemma4_suppress(e, &mut ld, 1)?;
10257        let logits = e.dtoh(&ld)?;
10258        Ok((logits, h_seed, hiddens))
10259    }
10260
10261    /// gemma4 T=1 decode attention: per-layer geometry, quantized-KV append + fa_decode
10262    /// (vec kernels at hd 256, generic scalar at the globals' hd 512), weightless V-norm,
10263    /// dual rope, scale 1.0. Takes the attn-normed input PRE-QUANTIZED (the cross-layer
10264    /// fused norm emits q8 directly — the f32 h never materializes).
10265    fn gemma4_decode_attn(
10266        &self,
10267        e: &Engine,
10268        fa: &crate::hybrid::FullAttnLayer,
10269        il: usize,
10270        hq: &CudaSlice<i8>,
10271        hdq: &CudaSlice<f32>,
10272        pos_d: &CudaSlice<i32>,
10273        cache: &mut Cache,
10274    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10275        let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
10276        let eps = self.cfg.rms_eps;
10277        let aux = self.gemma4_aux.as_ref().unwrap();
10278        let ones = aux.ones(e);
10279        #[cfg(debug_assertions)]
10280        crate::debug_assert_tensor_stream_device(ones, &e.stream(), "gemma4_decode_attn.ones");
10281        let (hq, hdq) = (hq, hdq);
10282        let h0 = e.zeros(0)?;
10283        let h = &h0;
10284        let (q0, k0, v0) = if swa {
10285            match e.matmul_q4_fused3(&fa.wq, &fa.wk, &fa.wv, &hq, &hdq)? {
10286                Some(t3) => t3,
10287                // NVFP4mix trio is MIXED-type (wv stays Q8_0), so fused3 can never
10288                // match — fuse the uniform (q,k) pair and take v as its own single.
10289                None => match e.matmul_nvfp4_fused2(&fa.wq, &fa.wk, &hq, &hdq, 1)? {
10290                    Some((q0, k0)) => {
10291                        let v0 = e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?;
10292                        (q0, k0, v0)
10293                    }
10294                    None => (
10295                        e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
10296                        e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
10297                        e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?,
10298                    ),
10299                },
10300            }
10301        } else {
10302            let (q0, k0) = match e.matmul_q4_fused2(&fa.wq, &fa.wk, &hq, &hdq)? {
10303                Some(p) => p,
10304                None => match e.matmul_nvfp4_fused2(&fa.wq, &fa.wk, &hq, &hdq, 1)? {
10305                    Some(p) => p,
10306                    None => (
10307                        e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
10308                        e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
10309                    ),
10310                },
10311            };
10312            let v0 = e.clone_dtod(&k0)?;
10313            (q0, k0, v0)
10314        };
10315        let mut q = e.uninit(nh * hd)?;
10316        let mut k = e.uninit(nkv * hd)?;
10317        let mut v = e.uninit(nkv * hd)?;
10318        // E4B wave-3 fold, m=1 completion (2026-07-23): the rows arms took this fold at
10319        // 550fcfa5; the decode-step trio kept 2 launches/layer it doesn't need.
10320        let ff = if swa {
10321            None
10322        } else {
10323            Some(
10324                aux.rope_freqs(e)
10325                    .expect("gemma4 global rope needs rope_freqs.weight"),
10326            )
10327        };
10328        #[cfg(debug_assertions)]
10329        if let Some(ff) = ff {
10330            crate::debug_assert_tensor_stream_device(
10331                ff,
10332                &e.stream(),
10333                "gemma4_decode_attn.rope_freqs",
10334            );
10335        }
10336        let kvl = cache.kv[il].as_mut().unwrap();
10337        let kv_fp8 = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
10338        if crate::Engine::qkv_append_on() {
10339            // append fold, eager completion (zoo-fusion arc): the dc arms took this fold
10340            // 2026-07-23; the eager step kept the 2-launch pair it doesn't need. Host-len
10341            // twin of the dc fold — bit-identical bodies, one launch per layer.
10342            e.rms_norm_qkv_rope_append(
10343                &q0,
10344                &k0,
10345                &v0,
10346                fa.q_norm.float_data(),
10347                fa.k_norm.float_data(),
10348                ones,
10349                &mut q,
10350                &mut k,
10351                &mut v,
10352                hd,
10353                self.gemma4_rope_dims(il),
10354                nh,
10355                nkv,
10356                pos_d,
10357                nh,
10358                nkv,
10359                base,
10360                1.0,
10361                ff,
10362                eps,
10363                &mut kvl.k,
10364                &mut kvl.v,
10365                kvl.len,
10366                kvl.k_tok_bytes,
10367                kvl.v_tok_bytes,
10368                kv_fp8,
10369            )?;
10370        } else {
10371            e.rms_norm_qkv_rope(
10372                &q0,
10373                &k0,
10374                &v0,
10375                fa.q_norm.float_data(),
10376                fa.k_norm.float_data(),
10377                ones,
10378                &mut q,
10379                &mut k,
10380                &mut v,
10381                hd,
10382                self.gemma4_rope_dims(il),
10383                nh,
10384                nkv,
10385                pos_d,
10386                nh,
10387                nkv,
10388                base,
10389                1.0,
10390                ff,
10391                eps,
10392            )?;
10393            e.append_kv_quantized(
10394                &k,
10395                &v,
10396                &mut kvl.k,
10397                &mut kvl.v,
10398                kvl.len,
10399                kvl.kv_dim_k,
10400                kvl.kv_dim_v,
10401                kvl.k_tok_bytes,
10402                kvl.v_tok_bytes,
10403                kv_fp8,
10404            )?;
10405        }
10406        kvl.len += 1;
10407        // R6 decode: SWA layers attend only the last `sliding_window` keys — a token-aligned
10408        // VIEW OFFSET into the quantized cache (keys carry absolute rope; the mask is purely
10409        // positional). Globals attend the full history.
10410        let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
10411        let mut attn = e.uninit(nh * hd)?;
10412        // global layers: SAME hd512 rows twin as verify with t=1 (parity law).
10413        if !swa
10414            && hd == 512
10415            && kvl.len >= crate::fa512_min_tkv()
10416            && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
10417        {
10418            let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
10419            let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
10420            // device-len: eager syncs the counter (async arg-store), dc keeps it live.
10421            let base = kvl.len as i32;
10422            e.i32_set_k(&mut kvl.len_d, base)?;
10423            e.fa_decode_rows(
10424                &q,
10425                &kp,
10426                &vp,
10427                &mut attn,
10428                hd,
10429                nh,
10430                nkv,
10431                kvl.len - 1,
10432                1,
10433                scale,
10434                kvl.k_tok_bytes,
10435                kvl.v_tok_bytes,
10436                Some((&kvl.len_d, -1)),
10437                false,
10438                false,
10439                None,
10440            )?;
10441            return Ok(e.matmul(&fa.wo, &attn, 1)?);
10442        }
10443        // windowed regime: SAME rows_w kernel as verify with t=1 (parity law — see verify_attn).
10444        if swa
10445            && kvl.len > win
10446            && hd == 256
10447            && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
10448        {
10449            let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
10450            let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
10451            let base = kvl.len as i32;
10452            e.i32_set_k(&mut kvl.len_d, base)?;
10453            e.fa_decode_rows_w(
10454                &q,
10455                &kp,
10456                &vp,
10457                &mut attn,
10458                hd,
10459                nh,
10460                nkv,
10461                &kvl.len_d,
10462                -1,
10463                1,
10464                scale,
10465                win,
10466                kvl.k_tok_bytes,
10467                kvl.v_tok_bytes,
10468                None,
10469            )?;
10470            return Ok(e.matmul(&fa.wo, &attn, 1)?);
10471        }
10472        let (off_tok, t_kv) = if swa && kvl.len > win {
10473            (kvl.len - win, win)
10474        } else {
10475            (0, kvl.len)
10476        };
10477        let k_view = e.view_u8_range(
10478            &kvl.k,
10479            off_tok * kvl.k_tok_bytes,
10480            (off_tok + t_kv) * kvl.k_tok_bytes,
10481        );
10482        let v_view = e.view_u8_range(
10483            &kvl.v,
10484            off_tok * kvl.v_tok_bytes,
10485            (off_tok + t_kv) * kvl.v_tok_bytes,
10486        );
10487        e.fa_decode_kvmod(
10488            &q,
10489            &k_view,
10490            &v_view,
10491            &mut attn,
10492            hd,
10493            nh,
10494            nkv,
10495            t_kv,
10496            scale,
10497            kvl.k_tok_bytes,
10498            kvl.v_tok_bytes,
10499            swa && crate::Engine::wkv_on(),
10500        )?;
10501        Ok(e.matmul(&fa.wo, &attn, 1)?)
10502    }
10503
10504    /// gemma4 DEVICE-COUNTER decode step (graph arc): token id + rope pos + KV lengths live in
10505    /// device counters; ZERO varying host kernel args. `cap_bucket_max` = Some(bucket) for graph
10506    /// capture (host mirrors untouched, n_splits from bucket, full-buffer KV views) / None for
10507    /// the eager-dc gate path (host mirrors advanced, live geometry — bit-identical target =
10508    /// gemma4_decode_step_h's token stream). V1 scope: t_kv <= sliding_window (no window views
10509    /// in-graph; the driver gates).
10510    #[allow(clippy::too_many_arguments)]
10511    pub fn gemma4_decode_step_dc(
10512        &self,
10513        e: &Engine,
10514        token_d: &CudaSlice<u32>,
10515        pos_d: &mut CudaSlice<i32>,
10516        embd_gpu: &CudaSlice<u8>,
10517        embd_qt: i32,
10518        embd_rb: usize,
10519        cache: &mut Cache,
10520        n_vocab: usize,
10521        cap_bucket_max: Option<(usize, usize)>,
10522    ) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
10523        let mut tok_out = e.stream().alloc_zeros::<u32>(1)?;
10524        self.gemma4_decode_step_dc_into(
10525            e,
10526            token_d,
10527            pos_d,
10528            embd_gpu,
10529            embd_qt,
10530            embd_rb,
10531            cache,
10532            n_vocab,
10533            cap_bucket_max,
10534            &mut tok_out,
10535        )?;
10536        Ok(tok_out)
10537    }
10538
10539    /// CAPTURE body: argmax lands in the PERSISTENT `tok_out` (same buffer = same address on
10540    /// every replay; pass `token_d` itself for the self-feeding graph loop).
10541    #[allow(clippy::too_many_arguments)]
10542    pub fn gemma4_decode_step_dc_into(
10543        &self,
10544        e: &Engine,
10545        token_d: &CudaSlice<u32>,
10546        pos_d: &mut CudaSlice<i32>,
10547        embd_gpu: &CudaSlice<u8>,
10548        embd_qt: i32,
10549        embd_rb: usize,
10550        cache: &mut Cache,
10551        n_vocab: usize,
10552        cap_bucket_max: Option<(usize, usize)>,
10553        tok_out: &mut CudaSlice<u32>,
10554    ) -> Result<(), Box<dyn std::error::Error>> {
10555        let n_embd = self.cfg.n_embd as usize;
10556        let eps = self.cfg.rms_eps;
10557        let mut x = e.embed_gather_device(embd_gpu, token_d, n_embd, embd_qt, embd_rb)?;
10558        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
10559        let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
10560        let n_layers = self.layers.len();
10561        for (il, layer) in self.layers.iter().enumerate() {
10562            let (hq, hdq) = match h_carry.take() {
10563                Some(p) => p,
10564                None => {
10565                    e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, 1, eps)?
10566                }
10567            };
10568            let Mixer::Full(fa) = &layer.mixer else {
10569                panic!("gemma4 layer {il} not full-attn")
10570            };
10571            let o =
10572                self.gemma4_decode_attn_dc(e, fa, il, &hq, &hdq, pos_d, cache, cap_bucket_max)?;
10573            let next_norm = if il + 1 < n_layers {
10574                Some(self.layers[il + 1].attn_norm.float_data())
10575            } else {
10576                None
10577            };
10578            let (xn, hn) = self.gemma4_layer_tail_add_nq_pn(e, layer, &o, &x, 1, next_norm)?;
10579            x = xn;
10580            h_carry = hn;
10581        }
10582        let mut hn = e.uninit(n_embd)?;
10583        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
10584        let mut logits = e.matmul(&self.output, &hn, 1)?;
10585        self.gemma4_suppress(e, &mut logits, 1)?; // cap skipped (monotonic); the mask is not
10586        e.argmax_token_device_into(&logits, tok_out, n_vocab)?;
10587        e.inc_seqlen(pos_d)?;
10588        if cap_bucket_max.is_none() {
10589            cache.pos += 1;
10590        }
10591        Ok(())
10592    }
10593
10594    /// Persistent transient slots for the ALLOC-FREE captured dc step (the graph door):
10595    /// every buffer the step produces per token lives here, allocated ONCE pre-capture, so
10596    /// the captured graph carries zero cuMemAllocAsync/Free nodes (the 226us/launch tax,
10597    /// osrt 2026-07-23). Sized for the model's max per-layer shapes.
10598
10599    /// Build the slot set (call OUTSIDE any capture).
10600    pub fn g4_dc_slots(&self, e: &Engine) -> Result<G4DcSlots, Box<dyn std::error::Error>> {
10601        let n_embd = self.cfg.n_embd as usize;
10602        let n_vocab = self.output.out_features();
10603        let n_layers = self.layers.len();
10604        let (mut qmax, mut kvmax, mut ffmax) = (0usize, 0usize, 0usize);
10605        for il in 0..n_layers {
10606            let (hd, nkv, nh, _b, _s, _w) = self.gemma4_geom(il);
10607            qmax = qmax.max(nh * hd);
10608            kvmax = kvmax.max(nkv * hd);
10609            if let crate::hybrid::Ffn::Dense { ffn_gate, .. } = &self.layers[il].ffn {
10610                ffmax = ffmax.max(ffn_gate.out_features());
10611            }
10612        }
10613        Ok(G4DcSlots {
10614            x: e.uninit(n_embd)?,
10615            xn: e.uninit(n_embd)?,
10616            cur: e.uninit(n_embd)?,
10617            hq: e.alloc_i8_uninit(n_embd)?,
10618            hd_: e.uninit(n_embd / 32)?,
10619            q0: e.uninit(qmax)?,
10620            k0: e.uninit(kvmax)?,
10621            v0: e.uninit(kvmax)?,
10622            q: e.uninit(qmax)?,
10623            k: e.uninit(kvmax)?,
10624            v: e.uninit(kvmax)?,
10625            attn: e.uninit(qmax)?,
10626            o: e.uninit(n_embd)?,
10627            attn_out: e.uninit(n_embd)?,
10628            zsh: e.uninit(n_embd)?,
10629            // zq/zd feed the wo matvec (nh*hd rows — 4096 on hd512 globals > n_embd),
10630            // the ffn entry (n_embd) and the lm_head (n_embd): size for the max.
10631            zq: e.alloc_i8_uninit(n_embd.max(qmax))?,
10632            zd: e.uninit(n_embd.max(qmax) / 32)?,
10633            gate: e.uninit(ffmax)?,
10634            up: e.uninit(ffmax)?,
10635            act: e.uninit(ffmax)?,
10636            actq: e.alloc_i8_uninit(ffmax)?,
10637            actd: e.uninit(ffmax / 32)?,
10638            f0: e.uninit(n_embd)?,
10639            sn: e.uninit(n_embd)?,
10640            hn: e.uninit(n_embd)?,
10641            logits: e.uninit(n_vocab)?,
10642        })
10643    }
10644
10645    /// m=1 pre-quantized matvec into a slot — mirrors matmul_pre's m=1 mmvq route exactly
10646    /// (rp4-mirror bytes; mmvq_supports guaranteed for gemma4 q4_0/q6_K).
10647    fn g4_matvec_m1_into(
10648        &self,
10649        e: &Engine,
10650        w: &crate::model::GpuTensor,
10651        aq: &CudaSlice<i8>,
10652        ad: &CudaSlice<f32>,
10653        y: &mut CudaSlice<f32>,
10654    ) -> Result<(), Box<dyn std::error::Error>> {
10655        use crate::model::GpuTensor;
10656        let (bytes, qtype, row_bytes, scale, rp) = match w {
10657            GpuTensor::Quant {
10658                bytes,
10659                qtype,
10660                row_bytes,
10661                scale,
10662                rp,
10663                ..
10664            } => (bytes, *qtype, *row_bytes, *scale, *rp),
10665            _ => return Err("g4_matvec_m1_into: non-quant tensor".into()),
10666        };
10667        let (mbytes, mrp) = match w {
10668            GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
10669            _ => (bytes, rp),
10670        };
10671        e.qmatvec_mmvq_into(
10672            mbytes,
10673            aq,
10674            ad,
10675            1,
10676            w.in_features(),
10677            w.out_features(),
10678            qtype,
10679            row_bytes,
10680            scale,
10681            mrp,
10682            y,
10683        )
10684    }
10685
10686    /// ALLOC-FREE dc step (capture body): kernel-for-kernel mirror of
10687    /// `gemma4_decode_step_dc_into` at t=1 with every transient slot-fed. Dense gemma4 only
10688    /// (12B/31B; uniform q4_0 trunk guarantees the fused2/3 arms).
10689    #[allow(clippy::too_many_arguments)]
10690    pub fn gemma4_decode_step_dc_slotted(
10691        &self,
10692        e: &Engine,
10693        token_d: &CudaSlice<u32>,
10694        pos_d: &mut CudaSlice<i32>,
10695        embd_gpu: &CudaSlice<u8>,
10696        embd_qt: i32,
10697        embd_rb: usize,
10698        cache: &mut Cache,
10699        n_vocab: usize,
10700        cap_bucket_max: Option<(usize, usize)>,
10701        sl: &mut G4DcSlots,
10702        tok_out: &mut CudaSlice<u32>,
10703        ring: Option<(&mut CudaSlice<u32>, usize)>,
10704    ) -> Result<(), Box<dyn std::error::Error>> {
10705        let n_embd = self.cfg.n_embd as usize;
10706        let eps = self.cfg.rms_eps;
10707        e.embed_gather_device_into(embd_gpu, token_d, &mut sl.x, n_embd, embd_qt, embd_rb)?;
10708        e.scale_inplace(&mut sl.x, (n_embd as f32).sqrt(), n_embd)?;
10709        let n_layers = self.layers.len();
10710        let mut has_carry = false;
10711        for il in 0..n_layers {
10712            if !has_carry {
10713                e.rms_norm_q8_1_into(
10714                    &sl.x,
10715                    self.layers[il].attn_norm.float_data(),
10716                    n_embd,
10717                    1,
10718                    eps,
10719                    &mut sl.hq,
10720                    &mut sl.hd_,
10721                )?;
10722            }
10723            has_carry = true;
10724            let layer = &self.layers[il];
10725            let Mixer::Full(fa) = &layer.mixer else {
10726                panic!("gemma4 layer {il} not full-attn")
10727            };
10728            self.gemma4_decode_attn_dc_slotted(e, fa, il, pos_d, cache, cap_bucket_max, sl)?;
10729            // pn-fold: the tail's entry launch consumes RAW sl.o (post_attn folded there);
10730            // the standalone norm only survives on the unfused seam arm.
10731            if !Engine::g4_pnfold_on() {
10732                e.rms_norm(
10733                    &sl.o,
10734                    layer.post_attn_norm.float_data(),
10735                    &mut sl.cur,
10736                    n_embd,
10737                    1,
10738                    eps,
10739                )?;
10740            }
10741            let next_norm = if il + 1 < n_layers {
10742                Some(self.layers[il + 1].attn_norm.float_data())
10743            } else {
10744                None
10745            };
10746            self.gemma4_layer_tail_slotted(e, layer, next_norm, sl)?;
10747            std::mem::swap(&mut sl.x, &mut sl.xn);
10748        }
10749        e.rms_norm(
10750            &sl.x,
10751            self.output_norm.float_data(),
10752            &mut sl.hn,
10753            n_embd,
10754            1,
10755            eps,
10756        )?;
10757        e.quantize_q8_1_into(&sl.hn, 1, n_embd, &mut sl.zq, &mut sl.zd)?;
10758        // lm_head via the same m=1 mmvq route as matmul (q6_K on the gemma4 family).
10759        {
10760            let (zq, zd) = (&sl.zq, &sl.zd);
10761            let zq = unsafe { &*(zq as *const CudaSlice<i8>) };
10762            let zd = unsafe { &*(zd as *const CudaSlice<f32>) };
10763            self.g4_matvec_m1_into(e, &self.output, zq, zd, &mut sl.logits)?;
10764        }
10765        self.gemma4_suppress(e, &mut sl.logits, 1)?;
10766        e.argmax_token_device_into(&sl.logits, tok_out, n_vocab)?;
10767        if let Some((ring, base)) = ring {
10768            // in-graph token parking: slot = (pos - base) % ring.len(); the door drains
10769            // with ONE sync per chunk instead of a per-token dtoh (the launch-serialization
10770            // fix — llama's 885us/launch overlaps GPU work because nothing waits per token).
10771            e.plain_tok_ring(tok_out, pos_d, base, ring)?;
10772        }
10773        e.inc_seqlen(pos_d)?;
10774        if cap_bucket_max.is_none() {
10775            cache.pos += 1;
10776        }
10777        Ok(())
10778    }
10779
10780    /// Slot-fed dc attention: mirrors gemma4_decode_attn_dc's CAPTURE arm kernel-for-kernel.
10781    #[allow(clippy::too_many_arguments)]
10782    fn gemma4_decode_attn_dc_slotted(
10783        &self,
10784        e: &Engine,
10785        fa: &crate::hybrid::FullAttnLayer,
10786        il: usize,
10787        pos_d: &CudaSlice<i32>,
10788        cache: &mut Cache,
10789        cap_bucket_max: Option<(usize, usize)>,
10790        sl: &mut G4DcSlots,
10791    ) -> Result<(), Box<dyn std::error::Error>> {
10792        let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
10793        let eps = self.cfg.rms_eps;
10794        let aux = self.gemma4_aux.as_ref().unwrap();
10795        let ones = aux.ones(e);
10796        #[cfg(debug_assertions)]
10797        crate::debug_assert_tensor_stream_device(
10798            ones,
10799            &e.stream(),
10800            "gemma4_decode_attn_dc_slotted.ones",
10801        );
10802        {
10803            let hq = unsafe { &*(&sl.hq as *const CudaSlice<i8>) };
10804            let hdq = unsafe { &*(&sl.hd_ as *const CudaSlice<f32>) };
10805            if swa {
10806                if !e.matmul_q4_fused3_into(
10807                    &fa.wq, &fa.wk, &fa.wv, hq, hdq, &mut sl.q0, &mut sl.k0, &mut sl.v0,
10808                )? {
10809                    // NVFP4mix trio is MIXED-type (wv stays Q8_0): fuse the uniform
10810                    // (q,k) pair, v through the generic m1 slot matvec — the same two
10811                    // bodies the eager arm runs, slot-fed for the alloc-free capture.
10812                    if e.matmul_nvfp4_fused2_into(&fa.wq, &fa.wk, hq, hdq, &mut sl.q0, &mut sl.k0)?
10813                    {
10814                        self.g4_matvec_m1_into(e, &fa.wv, hq, hdq, &mut sl.v0)?;
10815                    } else {
10816                        return Err("slotted step: fused3 unavailable (non-uniform trunk)".into());
10817                    }
10818                }
10819            } else {
10820                if !e.matmul_q4_fused2_into(&fa.wq, &fa.wk, hq, hdq, &mut sl.q0, &mut sl.k0)?
10821                    && !e
10822                        .matmul_nvfp4_fused2_into(&fa.wq, &fa.wk, hq, hdq, &mut sl.q0, &mut sl.k0)?
10823                {
10824                    return Err("slotted step: fused2 unavailable".into());
10825                }
10826                let k0r = unsafe { &*(&sl.k0 as *const CudaSlice<f32>) };
10827                e.copy_into(&mut sl.v0, 0, k0r, nkv * hd)?;
10828            }
10829        }
10830        // E4B wave-3 fold, m=1 completion (2026-07-23) — MUST mirror the dc_into arm
10831        // kernel-for-kernel (graph stream-identity gate).
10832        let ff = if swa {
10833            None
10834        } else {
10835            Some(
10836                aux.rope_freqs(e)
10837                    .expect("gemma4 global rope needs rope_freqs.weight"),
10838            )
10839        };
10840        #[cfg(debug_assertions)]
10841        if let Some(ff) = ff {
10842            crate::debug_assert_tensor_stream_device(
10843                ff,
10844                &e.stream(),
10845                "gemma4_decode_attn_dc_slotted.rope_freqs",
10846            );
10847        }
10848        let kvl = cache.kv[il].as_mut().unwrap();
10849        let kv_fp8 = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
10850        if crate::Engine::qkv_append_on() {
10851            // append fold (2026-07-23): mirrors dc_into.
10852            e.rms_norm_qkv_rope_append_dc(
10853                &sl.q0,
10854                &sl.k0,
10855                &sl.v0,
10856                fa.q_norm.float_data(),
10857                fa.k_norm.float_data(),
10858                ones,
10859                &mut sl.q,
10860                &mut sl.k,
10861                &mut sl.v,
10862                hd,
10863                self.gemma4_rope_dims(il),
10864                nh,
10865                nkv,
10866                pos_d,
10867                nh,
10868                nkv,
10869                base,
10870                1.0,
10871                ff,
10872                eps,
10873                &mut kvl.k,
10874                &mut kvl.v,
10875                &kvl.len_d,
10876                kvl.k_tok_bytes,
10877                kvl.v_tok_bytes,
10878                kv_fp8,
10879            )?;
10880        } else {
10881            e.rms_norm_qkv_rope(
10882                &sl.q0,
10883                &sl.k0,
10884                &sl.v0,
10885                fa.q_norm.float_data(),
10886                fa.k_norm.float_data(),
10887                ones,
10888                &mut sl.q,
10889                &mut sl.k,
10890                &mut sl.v,
10891                hd,
10892                self.gemma4_rope_dims(il),
10893                nh,
10894                nkv,
10895                pos_d,
10896                nh,
10897                nkv,
10898                base,
10899                1.0,
10900                ff,
10901                eps,
10902            )?;
10903            e.append_kv_quantized_dc(
10904                &sl.k,
10905                &sl.v,
10906                &mut kvl.k,
10907                &mut kvl.v,
10908                &kvl.len_d,
10909                kvl.kv_dim_k,
10910                kvl.kv_dim_v,
10911                kvl.k_tok_bytes,
10912                kvl.v_tok_bytes,
10913                kv_fp8,
10914            )?;
10915        }
10916        e.inc_seqlen(&mut kvl.len_d)?;
10917        let (b_swa, b_glob) = cap_bucket_max.expect("slotted step is capture-only");
10918        let k_view = e.view_u8(&kvl.k, kvl.k.len());
10919        let v_view = e.view_u8(&kvl.v, kvl.v.len());
10920        let rows_on = std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0");
10921        let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
10922        // combine-q8 emit (wave-5b m=1 port): the rows arms quantize inside the combine —
10923        // the standalone quantize launch runs only on the non-rows fallback. MUST mirror
10924        // the dc_into arm branch-for-branch (stream gate).
10925        let mut fa_q8 = false;
10926        if !swa && hd == 512 && b_glob >= crate::fa512_min_tkv() && rows_on {
10927            e.fa_decode_rows(
10928                &sl.q,
10929                &k_view,
10930                &v_view,
10931                &mut sl.attn,
10932                hd,
10933                nh,
10934                nkv,
10935                b_glob - 1,
10936                1,
10937                scale,
10938                kvl.k_tok_bytes,
10939                kvl.v_tok_bytes,
10940                Some((&kvl.len_d, -1)),
10941                false,
10942                false,
10943                Some((&mut sl.zq, &mut sl.zd)),
10944            )?;
10945            fa_q8 = true;
10946        } else if swa && b_swa > win && hd == 256 && rows_on {
10947            e.fa_decode_rows_w(
10948                &sl.q,
10949                &k_view,
10950                &v_view,
10951                &mut sl.attn,
10952                hd,
10953                nh,
10954                nkv,
10955                &kvl.len_d,
10956                -1,
10957                1,
10958                scale,
10959                win,
10960                kvl.k_tok_bytes,
10961                kvl.v_tok_bytes,
10962                Some((&mut sl.zq, &mut sl.zd)),
10963            )?;
10964            fa_q8 = true;
10965        } else {
10966            let b = if swa { b_swa } else { b_glob };
10967            e.fa_decode_dc(
10968                &sl.q,
10969                &k_view,
10970                &v_view,
10971                &mut sl.attn,
10972                hd,
10973                nh,
10974                nkv,
10975                &kvl.len_d,
10976                b,
10977                scale,
10978                kvl.k_tok_bytes,
10979                kvl.v_tok_bytes,
10980                swa && crate::Engine::wkv_on(),
10981            )?;
10982        }
10983        if !fa_q8 {
10984            let aq = unsafe { &*(&sl.attn as *const CudaSlice<f32>) };
10985            e.quantize_q8_1_into(aq, 1, nh * hd, &mut sl.zq, &mut sl.zd)?;
10986        }
10987        {
10988            let zq = unsafe { &*(&sl.zq as *const CudaSlice<i8>) };
10989            let zd = unsafe { &*(&sl.zd as *const CudaSlice<f32>) };
10990            self.g4_matvec_m1_into(e, &fa.wo, zq, zd, &mut sl.o)?;
10991        }
10992        Ok(())
10993    }
10994
10995    /// Slot-fed dense layer tail: mirrors gemma4_layer_tail_core (t=1, no pre-norm fold) +
10996    /// tail_add_nq kernel-for-kernel; the next layer's (hq, hd_) carry lands in the slots.
10997    fn gemma4_layer_tail_slotted(
10998        &self,
10999        e: &Engine,
11000        layer: &crate::hybrid::HybridLayer,
11001        next_norm: Option<&CudaSlice<f32>>,
11002        sl: &mut G4DcSlots,
11003    ) -> Result<(), Box<dyn std::error::Error>> {
11004        let n_embd = self.cfg.n_embd as usize;
11005        let eps = self.cfg.rms_eps;
11006        let bits = layer.gemma4.as_ref().unwrap();
11007        let crate::hybrid::Ffn::Dense {
11008            ffn_gate,
11009            ffn_up,
11010            ffn_down,
11011        } = &layer.ffn
11012        else {
11013            return Err("slotted tail: dense ffn only".into());
11014        };
11015        let pnfold = Engine::g4_pnfold_on();
11016        if pnfold {
11017            // pn-fold entry (E4B glue, slot-fed): rms(o, post_attn) + residual + ffn_norm
11018            // + q8 emit in ONE launch — replaces rms_norm + add_rms_norm + quantize.
11019            let or = unsafe { &*(&sl.o as *const CudaSlice<f32>) };
11020            let xr = unsafe { &*(&sl.x as *const CudaSlice<f32>) };
11021            e.rms_pre_add_rms_norm_q8z_into(
11022                or,
11023                layer.post_attn_norm.float_data(),
11024                xr,
11025                bits.ffn_norm.float_data(),
11026                &mut sl.attn_out,
11027                &mut sl.zsh,
11028                n_embd,
11029                1,
11030                eps,
11031                &mut sl.zq,
11032                &mut sl.zd,
11033            )?;
11034        } else {
11035            e.add_rms_norm(
11036                &sl.cur,
11037                &sl.x,
11038                bits.ffn_norm.float_data(),
11039                &mut sl.attn_out,
11040                &mut sl.zsh,
11041                n_embd,
11042                1,
11043                eps,
11044            )?;
11045        }
11046        let n_ff = ffn_gate.out_features();
11047        if !pnfold {
11048            let zshr = unsafe { &*(&sl.zsh as *const CudaSlice<f32>) };
11049            e.quantize_q8_1_into(zshr, 1, n_embd, &mut sl.zq, &mut sl.zd)?;
11050        }
11051        {
11052            let zq = unsafe { &*(&sl.zq as *const CudaSlice<i8>) };
11053            let zd = unsafe { &*(&sl.zd as *const CudaSlice<f32>) };
11054            if !e.matmul_q4_fused2_into(ffn_gate, ffn_up, zq, zd, &mut sl.gate, &mut sl.up)?
11055                && !e.matmul_nvfp4_fused2_into(
11056                    ffn_gate,
11057                    ffn_up,
11058                    zq,
11059                    zd,
11060                    &mut sl.gate,
11061                    &mut sl.up,
11062                )?
11063            {
11064                return Err("slotted tail: ffn fused2 unavailable".into());
11065            }
11066        }
11067        debug_assert!(e.uses_q8_1_fast(ffn_down));
11068        {
11069            let upr = unsafe { &*(&sl.up as *const CudaSlice<f32>) };
11070            let upv = e.view(upr, n_ff);
11071            let up_all = upv.slice(0..n_ff);
11072            let gr = unsafe { &*(&sl.gate as *const CudaSlice<f32>) };
11073            e.gelu_tanh_mul_q8_1_into(
11074                gr,
11075                &up_all,
11076                &mut sl.act,
11077                n_ff,
11078                1,
11079                &mut sl.actq,
11080                &mut sl.actd,
11081            )?;
11082        }
11083        {
11084            let aq = unsafe { &*(&sl.actq as *const CudaSlice<i8>) };
11085            let ad = unsafe { &*(&sl.actd as *const CudaSlice<f32>) };
11086            self.g4_matvec_m1_into(e, ffn_down, aq, ad, &mut sl.f0)?;
11087        }
11088        if pnfold {
11089            // pn-fold exit: rms(f0, post_ffw) + scaled residual + next attn_norm + q8
11090            // emit in ONE launch — replaces rms_norm + add_scale_rms_norm_q8_1.
11091            if let Some(w) = next_norm {
11092                let f0r = unsafe { &*(&sl.f0 as *const CudaSlice<f32>) };
11093                let aor = unsafe { &*(&sl.attn_out as *const CudaSlice<f32>) };
11094                e.rms_pre_add_scale_rms_norm_q8_1_into(
11095                    f0r,
11096                    bits.post_ffw_norm.float_data(),
11097                    aor,
11098                    bits.layer_scale,
11099                    w,
11100                    &mut sl.xn,
11101                    n_embd,
11102                    1,
11103                    eps,
11104                    &mut sl.hq,
11105                    &mut sl.hd_,
11106                )?;
11107                return Ok(());
11108            }
11109        }
11110        e.rms_norm(
11111            &sl.f0,
11112            bits.post_ffw_norm.float_data(),
11113            &mut sl.sn,
11114            n_embd,
11115            1,
11116            eps,
11117        )?;
11118        match next_norm {
11119            Some(w) => {
11120                e.add_scale_rms_norm_q8_1_into(
11121                    &sl.sn,
11122                    &sl.attn_out,
11123                    bits.layer_scale,
11124                    w,
11125                    &mut sl.xn,
11126                    n_embd,
11127                    1,
11128                    eps,
11129                    &mut sl.hq,
11130                    &mut sl.hd_,
11131                )?;
11132            }
11133            None => {
11134                e.add_scale(&sl.sn, &sl.attn_out, bits.layer_scale, &mut sl.xn, n_embd)?;
11135            }
11136        }
11137        Ok(())
11138    }
11139
11140    /// dc attention: same math as gemma4_decode_attn, KV slot/lengths from device counters.
11141    #[allow(clippy::too_many_arguments)]
11142    fn gemma4_decode_attn_dc(
11143        &self,
11144        e: &Engine,
11145        fa: &crate::hybrid::FullAttnLayer,
11146        il: usize,
11147        hq: &CudaSlice<i8>,
11148        hdq: &CudaSlice<f32>,
11149        pos_d: &CudaSlice<i32>,
11150        cache: &mut Cache,
11151        cap_bucket_max: Option<(usize, usize)>,
11152    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11153        let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
11154        let eps = self.cfg.rms_eps;
11155        let aux = self.gemma4_aux.as_ref().unwrap();
11156        let ones = aux.ones(e);
11157        #[cfg(debug_assertions)]
11158        crate::debug_assert_tensor_stream_device(ones, &e.stream(), "gemma4_decode_attn_dc.ones");
11159        let (q0, k0, v0) = if swa {
11160            match e.matmul_q4_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hdq)? {
11161                Some(t3) => t3,
11162                // NVFP4mix: wv is Q8_0 (mixed trio) — fuse the (q,k) pair, v single.
11163                None => match e.matmul_nvfp4_fused2(&fa.wq, &fa.wk, hq, hdq, 1)? {
11164                    Some((q0, k0)) => {
11165                        let h0 = e.zeros(0)?;
11166                        let v0 = e.matmul_pre(&fa.wv, hq, hdq, &h0, 1)?;
11167                        (q0, k0, v0)
11168                    }
11169                    None => {
11170                        let h0 = e.zeros(0)?;
11171                        (
11172                            e.matmul_pre(&fa.wq, hq, hdq, &h0, 1)?,
11173                            e.matmul_pre(&fa.wk, hq, hdq, &h0, 1)?,
11174                            e.matmul_pre(&fa.wv, hq, hdq, &h0, 1)?,
11175                        )
11176                    }
11177                },
11178            }
11179        } else {
11180            let (q0, k0) = match e.matmul_q4_fused2(&fa.wq, &fa.wk, hq, hdq)? {
11181                Some(p) => p,
11182                None => match e.matmul_nvfp4_fused2(&fa.wq, &fa.wk, hq, hdq, 1)? {
11183                    Some(p) => p,
11184                    None => {
11185                        let h0 = e.zeros(0)?;
11186                        (
11187                            e.matmul_pre(&fa.wq, hq, hdq, &h0, 1)?,
11188                            e.matmul_pre(&fa.wk, hq, hdq, &h0, 1)?,
11189                        )
11190                    }
11191                },
11192            };
11193            let v0 = e.clone_dtod(&k0)?;
11194            (q0, k0, v0)
11195        };
11196        let mut q = e.uninit(nh * hd)?;
11197        let mut k = e.uninit(nkv * hd)?;
11198        let mut v = e.uninit(nkv * hd)?;
11199        // E4B wave-3 fold, m=1 completion (2026-07-23): mirrored by the slotted arm.
11200        let ff = if swa {
11201            None
11202        } else {
11203            Some(
11204                aux.rope_freqs(e)
11205                    .expect("gemma4 global rope needs rope_freqs.weight"),
11206            )
11207        };
11208        #[cfg(debug_assertions)]
11209        if let Some(ff) = ff {
11210            crate::debug_assert_tensor_stream_device(
11211                ff,
11212                &e.stream(),
11213                "gemma4_decode_attn_dc.rope_freqs",
11214            );
11215        }
11216        let kvl = cache.kv[il].as_mut().unwrap();
11217        let kv_fp8 = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
11218        if crate::Engine::qkv_append_on() {
11219            // append fold (2026-07-23): norm+rope+cache-append in ONE launch.
11220            e.rms_norm_qkv_rope_append_dc(
11221                &q0,
11222                &k0,
11223                &v0,
11224                fa.q_norm.float_data(),
11225                fa.k_norm.float_data(),
11226                ones,
11227                &mut q,
11228                &mut k,
11229                &mut v,
11230                hd,
11231                self.gemma4_rope_dims(il),
11232                nh,
11233                nkv,
11234                pos_d,
11235                nh,
11236                nkv,
11237                base,
11238                1.0,
11239                ff,
11240                eps,
11241                &mut kvl.k,
11242                &mut kvl.v,
11243                &kvl.len_d,
11244                kvl.k_tok_bytes,
11245                kvl.v_tok_bytes,
11246                kv_fp8,
11247            )?;
11248        } else {
11249            e.rms_norm_qkv_rope(
11250                &q0,
11251                &k0,
11252                &v0,
11253                fa.q_norm.float_data(),
11254                fa.k_norm.float_data(),
11255                ones,
11256                &mut q,
11257                &mut k,
11258                &mut v,
11259                hd,
11260                self.gemma4_rope_dims(il),
11261                nh,
11262                nkv,
11263                pos_d,
11264                nh,
11265                nkv,
11266                base,
11267                1.0,
11268                ff,
11269                eps,
11270            )?;
11271            e.append_kv_quantized_dc(
11272                &k,
11273                &v,
11274                &mut kvl.k,
11275                &mut kvl.v,
11276                &kvl.len_d,
11277                kvl.kv_dim_k,
11278                kvl.kv_dim_v,
11279                kvl.k_tok_bytes,
11280                kvl.v_tok_bytes,
11281                kv_fp8,
11282            )?;
11283        }
11284        e.inc_seqlen(&mut kvl.len_d)?;
11285        let mut attn = e.uninit(nh * hd)?;
11286        // combine-q8 emit carrier (wave-5b m=1 port): rows arms fill this; the tail then
11287        // rides g4_matvec_m1_into instead of matmul's internal quantize.
11288        let mut fa_q8: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
11289        // Weight prefetch NOT wired here (26B/31B/qwen probes 2026-07-13: 26B flat
11290        // 196.6/196.0 vs 196.4/196.1 — MoE ffn dilutes wo; 31B −0.2% — dense decode sits
11291        // at the DRAM wall, no idle window to front-load into). E4B keeps the arm
11292        // (gemma4_e4b_attn, +0.65% valid window).
11293        match cap_bucket_max {
11294            None => {
11295                // dc-EAGER: host knows the length — R6 window views EXACTLY like the eager
11296                // decode (SWA layers attend the last `sliding_window` keys); the device
11297                // counters carry only the append slot + the graph seam.
11298                kvl.len += 1;
11299                let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
11300                if !swa
11301                    && hd == 512
11302                    && kvl.len >= crate::fa512_min_tkv()
11303                    && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
11304                {
11305                    // global layers: SAME hd512 rows twin as verify, t=1 (parity law).
11306                    // dc: len_d is live (inc_seqlen) — device-len rides it, plus=-1.
11307                    let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
11308                    let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
11309                    let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
11310                    e.fa_decode_rows(
11311                        &q,
11312                        &kp,
11313                        &vp,
11314                        &mut attn,
11315                        hd,
11316                        nh,
11317                        nkv,
11318                        kvl.len - 1,
11319                        1,
11320                        scale,
11321                        kvl.k_tok_bytes,
11322                        kvl.v_tok_bytes,
11323                        Some((&kvl.len_d, -1)),
11324                        false,
11325                        false,
11326                        Some((&mut aq8, &mut ad8)),
11327                    )?;
11328                    fa_q8 = Some((aq8, ad8));
11329                } else if swa
11330                    && kvl.len > win
11331                    && hd == 256
11332                    && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
11333                {
11334                    // windowed regime: SAME rows_w kernel as verify, t=1 (parity law).
11335                    let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
11336                    let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
11337                    let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
11338                    e.fa_decode_rows_w(
11339                        &q,
11340                        &kp,
11341                        &vp,
11342                        &mut attn,
11343                        hd,
11344                        nh,
11345                        nkv,
11346                        &kvl.len_d,
11347                        -1,
11348                        1,
11349                        scale,
11350                        win,
11351                        kvl.k_tok_bytes,
11352                        kvl.v_tok_bytes,
11353                        Some((&mut aq8, &mut ad8)),
11354                    )?;
11355                    fa_q8 = Some((aq8, ad8));
11356                } else {
11357                    let (off_tok, t_kv) = if swa && kvl.len > win {
11358                        (kvl.len - win, win)
11359                    } else {
11360                        (0, kvl.len)
11361                    };
11362                    let k_view = e.view_u8_range(
11363                        &kvl.k,
11364                        off_tok * kvl.k_tok_bytes,
11365                        (off_tok + t_kv) * kvl.k_tok_bytes,
11366                    );
11367                    let v_view = e.view_u8_range(
11368                        &kvl.v,
11369                        off_tok * kvl.v_tok_bytes,
11370                        (off_tok + t_kv) * kvl.v_tok_bytes,
11371                    );
11372                    e.fa_decode_kvmod(
11373                        &q,
11374                        &k_view,
11375                        &v_view,
11376                        &mut attn,
11377                        hd,
11378                        nh,
11379                        nkv,
11380                        t_kv,
11381                        scale,
11382                        kvl.k_tok_bytes,
11383                        kvl.v_tok_bytes,
11384                        swa && crate::Engine::wkv_on(),
11385                    )?;
11386                }
11387            }
11388            Some((b_swa, b_glob)) => {
11389                // capture: full-buffer views + device length. PARITY LAW port (graph arc step
11390                // 3): the SAME regime branches as dc-eager, pinned per arm — b_swa is the
11391                // exact-key t_kv for the dc family (n_splits constant per bucket), b_glob is
11392                // the RUNG max for the rows family (kernels derive per-replay splits from
11393                // kvl.len_d; grid/partials sized for the rung end, excess splits exit).
11394                let k_view = e.view_u8(&kvl.k, kvl.k.len());
11395                let v_view = e.view_u8(&kvl.v, kvl.v.len());
11396                let rows_on = std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0");
11397                let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
11398                if !swa && hd == 512 && b_glob >= crate::fa512_min_tkv() && rows_on {
11399                    let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
11400                    e.fa_decode_rows(
11401                        &q,
11402                        &k_view,
11403                        &v_view,
11404                        &mut attn,
11405                        hd,
11406                        nh,
11407                        nkv,
11408                        b_glob - 1,
11409                        1,
11410                        scale,
11411                        kvl.k_tok_bytes,
11412                        kvl.v_tok_bytes,
11413                        Some((&kvl.len_d, -1)),
11414                        false,
11415                        false,
11416                        Some((&mut aq8, &mut ad8)),
11417                    )?;
11418                    fa_q8 = Some((aq8, ad8));
11419                } else if swa && b_swa > win && hd == 256 && rows_on {
11420                    let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
11421                    e.fa_decode_rows_w(
11422                        &q,
11423                        &k_view,
11424                        &v_view,
11425                        &mut attn,
11426                        hd,
11427                        nh,
11428                        nkv,
11429                        &kvl.len_d,
11430                        -1,
11431                        1,
11432                        scale,
11433                        win,
11434                        kvl.k_tok_bytes,
11435                        kvl.v_tok_bytes,
11436                        Some((&mut aq8, &mut ad8)),
11437                    )?;
11438                    fa_q8 = Some((aq8, ad8));
11439                } else {
11440                    let b = if swa { b_swa } else { b_glob };
11441                    e.fa_decode_dc(
11442                        &q,
11443                        &k_view,
11444                        &v_view,
11445                        &mut attn,
11446                        hd,
11447                        nh,
11448                        nkv,
11449                        &kvl.len_d,
11450                        b,
11451                        scale,
11452                        kvl.k_tok_bytes,
11453                        kvl.v_tok_bytes,
11454                        swa && crate::Engine::wkv_on(),
11455                    )?;
11456                }
11457            }
11458        }
11459        // combine-q8 emit (wave-5b m=1 port): the rows arms produced the wo activation pair
11460        // in-combine — ride the slotted arm's exact matvec route (parity by construction).
11461        if let Some((aq8, ad8)) = fa_q8 {
11462            let mut y = e.uninit(fa.wo.out_features())?;
11463            self.g4_matvec_m1_into(e, &fa.wo, &aq8, &ad8, &mut y)?;
11464            return Ok(y);
11465        }
11466        Ok(e.matmul(&fa.wo, &attn, 1)?)
11467    }
11468
11469    /// gemma4 GRAPH-REPLAY greedy loop: per (swa-key, global-key) fa bucket, capture ONE full
11470    /// dc step (self-feeding: argmax writes token_d in-graph) and replay it — one graph launch
11471    /// per token, one 4B dtoh. V1 scope: whole generation under the sliding window (no window
11472    /// views in-graph); caller gates and falls back to the dc-eager loop.
11473    pub fn gemma4_generate_graph(
11474        &self,
11475        e: &Engine,
11476        prompt_pos: usize,
11477        first_token: u32,
11478        cache: &mut Cache,
11479        max_new: usize,
11480        eos: &[u32],
11481        mut on_token: impl FnMut(u32) -> bool,
11482    ) -> Result<(Vec<u32>, crate::decode::StopReason), Box<dyn std::error::Error>> {
11483        if self.is_gemma4_e4b() {
11484            return Err(
11485                "E4B graph serving is unwired (HANDOVER-E4B.md) — dc-eager is the serving arm"
11486                    .into(),
11487            );
11488        }
11489        use crate::decode::StopReason;
11490        let n_vocab = self.output.out_features();
11491        let n_embd = self.cfg.n_embd as usize;
11492        let embd_gpu = self
11493            .embd_gpu
11494            .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload"));
11495        let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
11496        for kvl in cache.kv.iter_mut().flatten() {
11497            e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
11498        }
11499        let mut token_d = e.stream().clone_htod(&[first_token])?;
11500        let mut pos_d = e.htod_i32(&[prompt_pos as i32])?;
11501        let g4 = self.cfg.gemma4.as_ref().unwrap();
11502        let (hd_s, hd_g) = (g4.key_length_swa as usize, g4.key_length_global as usize);
11503        // per-layer nkv: swa vs global counts from the pattern (uniform within class).
11504        let nkv_s = g4
11505            .head_count_kv
11506            .iter()
11507            .zip(g4.swa_pattern.iter())
11508            .find(|p| *p.1)
11509            .map(|p| *p.0 as usize)
11510            .unwrap_or(8);
11511        let nkv_g = g4
11512            .head_count_kv
11513            .iter()
11514            .zip(g4.swa_pattern.iter())
11515            .find(|p| !*p.1)
11516            .map(|p| *p.0 as usize)
11517            .unwrap_or(2);
11518        let mut graphs: std::collections::HashMap<
11519            ((bool, usize), (bool, usize), bool, bool),
11520            (
11521                cudarc::driver::CudaGraph,
11522                Vec<Box<dyn std::any::Any + Send>>,
11523            ),
11524        > = Default::default();
11525        // ALLOC-FREE capture: persistent transient slots — the captured graph carries zero
11526        // mem nodes (the 226us/launch tax). Slots must outlive every cached graph.
11527        let mut slots = self.g4_dc_slots(e)?;
11528        // Chunked replay ring: tokens park on-device; ONE drain sync per chunk. ring_base is
11529        // baked at the door entry (the modulo keeps every capture valid indefinitely).
11530        const RING: usize = 64;
11531        // DRAIN pinned at 1 (2026-07-23): relaunching the SAME graph exec before its prior
11532        // launch completes is ILLEGAL (chunk=4 -> ILLEGAL_ADDRESS; chunk=1 clean). Pipelining
11533        // needs alternating exec instances — and the measured payoff was ~0 (the ~200us
11534        // cuGraphLaunch host cost already overlaps its own launch's GPU work; llama pays
11535        // 885us/launch the same way). MEMRA_GRAPH_DRAIN raises it only for experiments.
11536        const DRAIN: usize = 1;
11537        let mut ring = e.stream().alloc_zeros::<u32>(RING)?;
11538        let ring_base = prompt_pos;
11539        let mut out = Vec::with_capacity(max_new);
11540        let mut reason = StopReason::MaxNew;
11541        let mut next = first_token;
11542        let mut captures = 0usize;
11543        for _ in 0..max_new {
11544            out.push(next);
11545            if eos.contains(&next) {
11546                reason = StopReason::Eos;
11547                break;
11548            }
11549            if !on_token(next) {
11550                reason = StopReason::Callback;
11551                break;
11552            }
11553            let t_kv = cache.pos + 1;
11554            // Bucket key per ARM (graph arc step 3):
11555            //  - swa component: dc family under the window (exact (fa_vec, n_splits) key —
11556            //    n_splits must be constant per bucket); rows_w above it (window-constant, so
11557            //    the component collapses to a single marker).
11558            //  - global component: dc family under the fa512 floor (exact key); rows_dpl16
11559            //    at/above it — the kernel derives splits from len_d per replay, so buckets
11560            //    are power-of-2 RUNGS (one capture per doubling; grid sized for the rung end).
11561            let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
11562            let f512 = crate::fa512_min_tkv();
11563            let key_s = if t_kv > win {
11564                (true, usize::MAX)
11565            } else {
11566                e.fa_bucket_key(t_kv, hd_s, nkv_s, crate::Engine::wkv_on())
11567            };
11568            let (key_g, rung_end) = if t_kv >= f512 {
11569                // strict upper bound: the rung must ROLL at exact powers (t_kv==1024 starts
11570                // the [1024,2048) bucket) — sizing covers every replayed T_kv < end.
11571                let end = (t_kv + 1).next_power_of_two().max(f512 * 2);
11572                ((true, end), end)
11573            } else {
11574                (e.fa_bucket_key(t_kv, hd_g, nkv_g, false), t_kv)
11575            };
11576            let key = (key_s, key_g, t_kv >= f512, t_kv > win);
11577            if !graphs.contains_key(&key) {
11578                let bucket_max = (t_kv, rung_end);
11579                // snapshot device+host state (the 3 capture-warmup runs leave no residue).
11580                let snap = cache.snapshot(e)?;
11581                let pos_save = e.dtoh_i32_one(&pos_d)?;
11582                let len_save: Vec<Option<i32>> = cache
11583                    .kv
11584                    .iter()
11585                    .map(|k| k.as_ref().map(|kvl| e.dtoh_i32_one(&kvl.len_d).unwrap()))
11586                    .collect();
11587                let tok_save = e.dtoh_u32_one(&token_d)?;
11588                // RETAINED capture (2026-07-23): the plain capture_graph left pool-transient
11589                // clones as dead COPY NODES replayed every launch — the E4B 0.74ms/token
11590                // regression class, and this door's measured -8.8%. The keeper pins warmup
11591                // transients so the captured graph holds kernel nodes only.
11592                let graph = {
11593                    let tok_ref = &mut token_d;
11594                    let pos_ref = &mut pos_d;
11595                    let cache_ref = &mut *cache;
11596                    let slots_ref = &mut slots;
11597                    let ring_ref = &mut ring;
11598                    e.capture_graph_retained_flags(
11599                        cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY,
11600                        |e| {
11601                        // self-feeding: the argmax writes token_d itself.
11602                        let tok_in = unsafe { &*(tok_ref as *const CudaSlice<u32>) };
11603                        let sl = unsafe { &mut *(slots_ref as *mut G4DcSlots) };
11604                        let rg = unsafe { &mut *(ring_ref as *mut CudaSlice<u32>) };
11605                        self.gemma4_decode_step_dc_slotted(e, tok_in, pos_ref, embd_gpu, qt, rb,
11606                                                           cache_ref, n_vocab, Some(bucket_max),
11607                                                           sl, tok_ref, Some((rg, ring_base)))
11608                    })?
11609                };
11610                cache.rollback(e, &snap, 0)?;
11611                e.set_i32_one(&mut pos_d, pos_save)?;
11612                for (il, ls) in len_save.iter().enumerate() {
11613                    if let (Some(kvl), Some(v)) = (cache.kv[il].as_mut(), ls) {
11614                        e.set_i32_one(&mut kvl.len_d, *v)?;
11615                    }
11616                }
11617                e.set_u32_one(&mut token_d, tok_save)?;
11618                if std::env::var("MEMRA_GRAPH_CENSUS").as_deref() == Ok("1") {
11619                    if let Ok(c) = crate::graph_update::node_census(&graph.0) {
11620                        eprintln!("[graph-census] {c:?}");
11621                    }
11622                }
11623                graphs.insert(key, graph);
11624                captures += 1;
11625            }
11626            // CHUNKED REPLAY: enqueue up to DRAIN launches back-to-back (the ~200us host
11627            // cost of each cuGraphLaunch overlaps the PREVIOUS launch's ~11ms GPU work),
11628            // then ONE sync + ring drain. The chunk must stay inside this bucket key and
11629            // the budget; capture warmups already emitted their tokens through the ring.
11630            let mut chunk = 1usize;
11631            let drain_cap: usize = std::env::var("MEMRA_GRAPH_DRAIN")
11632                .ok()
11633                .and_then(|v| v.parse().ok())
11634                .unwrap_or(DRAIN);
11635            while chunk < drain_cap && out.len() + chunk < max_new {
11636                let t_next = cache.pos + 1 + chunk;
11637                let key_s2 = if t_next > win {
11638                    (true, usize::MAX)
11639                } else {
11640                    e.fa_bucket_key(t_next, hd_s, nkv_s, crate::Engine::wkv_on())
11641                };
11642                let key_g2 = if t_next >= f512 {
11643                    (true, (t_next + 1).next_power_of_two().max(f512 * 2))
11644                } else {
11645                    e.fa_bucket_key(t_next, hd_g, nkv_g, false)
11646                };
11647                if (key_s2, key_g2, t_next >= f512, t_next > win) != key {
11648                    break;
11649                }
11650                chunk += 1;
11651            }
11652            let g = &graphs.get(&key).unwrap().0;
11653            for _ in 0..chunk {
11654                g.launch()?;
11655            }
11656            e.stream().synchronize()?;
11657            let ringh = e.dtoh_u32(&ring)?;
11658            for j in 0..chunk {
11659                let pos_j = cache.pos + j;
11660                let tok_j = ringh[(pos_j - ring_base) % RING];
11661                cache.pos += 0; // advanced below in one shot
11662                if j + 1 == chunk {
11663                    next = tok_j;
11664                } else {
11665                    out.push(tok_j);
11666                    if eos.contains(&tok_j) || !on_token(tok_j) {
11667                        reason = if eos.contains(&tok_j) {
11668                            StopReason::Eos
11669                        } else {
11670                            StopReason::Callback
11671                        };
11672                        // roll device/host state back to the stop point.
11673                        let keep = cache.pos + j + 1;
11674                        e.set_i32_one(&mut pos_d, keep as i32)?;
11675                        for kvl in cache.kv.iter_mut().filter_map(|k| k.as_mut()) {
11676                            e.set_i32_one(&mut kvl.len_d, keep as i32)?;
11677                            kvl.len = keep;
11678                        }
11679                        cache.pos = keep;
11680                        if std::env::var("MEMRA_GRAPH_STATS").is_ok() {
11681                            eprintln!("[gemma-graph] captures={captures} buckets={}", graphs.len());
11682                        }
11683                        return Ok((out, reason));
11684                    }
11685                }
11686            }
11687            cache.pos += chunk;
11688            for kvl in cache.kv.iter_mut().filter_map(|k| k.as_mut()) {
11689                kvl.len += chunk;
11690            }
11691        }
11692        if std::env::var("MEMRA_GRAPH_STATS").is_ok() {
11693            eprintln!("[gemma-graph] captures={captures} buckets={}", graphs.len());
11694        }
11695        Ok((out, reason))
11696    }
11697
11698    /// gemma4 VERIFY step (spec decode): t tokens batched through the trunk at positions
11699    /// pos0..pos0+t-1, K/V rows appended to the quantized cache (caller rolls back rejected
11700    /// rows via kvl.len), per-token causal windowed attend (fa_decode per token — the same
11701    /// kernel family as T=1 decode at each token's t_kv). Returns [t, n_vocab] softcapped
11702    /// logits (host) + advances cache.pos by t.
11703    pub(crate) fn gemma4_decode_step_t(
11704        &self,
11705        e: &Engine,
11706        tokens: &[u32],
11707        pos0: usize,
11708        cache: &mut Cache,
11709    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
11710        Ok(self.gemma4_decode_step_t_h(e, tokens, pos0, cache)?.0)
11711    }
11712
11713    /// GREEDY verify: per-row DEVICE argmax (t x 4B host traffic instead of the t x 1MB logits
11714    /// stack — softcap skipped: tanh is monotonic, per-row argmax unaffected). Returns
11715    /// (argmax ids [t], post-output_norm hidden stack [t, n_embd]).
11716    pub(crate) fn gemma4_decode_step_t_am(
11717        &self,
11718        e: &Engine,
11719        tokens: &[u32],
11720        pos0: usize,
11721        cache: &mut Cache,
11722    ) -> Result<(Vec<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11723        let (ld, hn) = self.gemma4_verify_trunk(e, tokens, pos0, cache, None)?;
11724        let t = tokens.len();
11725        let n_vocab = self.output.out_features();
11726        let mut toks = e.stream().alloc_zeros::<u32>(t)?;
11727        for i in 0..t {
11728            e.argmax_token_device_col(&ld, i, n_vocab, &mut toks, i)?;
11729        }
11730        Ok((e.dtoh_u32(&toks)?, hn))
11731    }
11732
11733    /// Device-token verify (async spec round): tokens live in tok_d[0..t]; per-row argmax
11734    /// lands in a DEVICE buffer (no host logits). Returns (vam_d [t] u32 device, hn stack).
11735    pub(crate) fn gemma4_decode_step_t_am_dev(
11736        &self,
11737        e: &Engine,
11738        tok_d: &CudaSlice<u32>,
11739        t: usize,
11740        pos0: usize,
11741        cache: &mut Cache,
11742    ) -> Result<(CudaSlice<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11743        let (ld, hn) = self.gemma4_verify_trunk(e, &vec![0u32; t], pos0, cache, Some(tok_d))?;
11744        let n_vocab = self.output.out_features();
11745        let mut vam = e.stream().alloc_zeros::<u32>(t)?;
11746        for i in 0..t {
11747            e.argmax_token_device_col(&ld, i, n_vocab, &mut vam, i)?;
11748        }
11749        Ok((vam, hn))
11750    }
11751
11752    /// gemma4 verify + the POST-output_norm hidden stack [t, n_embd] (the drafter's h input —
11753    /// llama's h_nextn convention).
11754    pub(crate) fn gemma4_decode_step_t_h(
11755        &self,
11756        e: &Engine,
11757        tokens: &[u32],
11758        pos0: usize,
11759        cache: &mut Cache,
11760    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11761        let (mut ld, hn) = self.gemma4_verify_trunk(e, tokens, pos0, cache, None)?;
11762        let t = tokens.len();
11763        let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
11764        e.softcap(&mut ld, cap, t * self.output.out_features())?;
11765        Ok((e.dtoh(&ld)?, hn))
11766    }
11767
11768    /// Persistent device scratch for the BURST verify stream (one per generation): rope pos
11769    /// rows [cap] + per-row t_kv counters (device-filled from the round counter, zero H2D).
11770    pub(crate) fn verify_stream_scratch(
11771        &self,
11772        e: &Engine,
11773        cap: usize,
11774    ) -> Result<VerifyStreamScratch, Box<dyn std::error::Error>> {
11775        Ok(VerifyStreamScratch {
11776            pos_d: e.htod_i32(&vec![0i32; cap])?,
11777            row_ctrs: (0..cap)
11778                .map(|_| e.htod_i32(&[0]))
11779                .collect::<Result<_, _>>()?,
11780        })
11781    }
11782
11783    /// BURST verify trunk (device-slot twin): tokens from a device buffer, rope positions
11784    /// iota'd from `ctr`, appends/attention at the layers' len_d counters (verify_attn_stream),
11785    /// NO host cache.pos/len advance. Returns (per-row device argmaxes [t], post-norm hidden
11786    /// stack [t, n_embd]) — the burst's accept/seed inputs, zero host readbacks.
11787    /// `scr` = PERSISTENT per-gen scratch (pos rows + per-row t_kv counters): the first cut
11788    /// htod_i32-allocated these per call — 9 pageable H2D copies per round, each a stream
11789    /// sync, exactly the turnaround the burst exists to remove.
11790    pub(crate) fn gemma4_verify_t_am_stream(
11791        &self,
11792        e: &Engine,
11793        tok_d: &CudaSlice<u32>,
11794        t: usize,
11795        ctr: &CudaSlice<i32>,
11796        hint: usize,
11797        cache: &mut Cache,
11798        scr: &mut VerifyStreamScratch,
11799    ) -> Result<(CudaSlice<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11800        let n_embd = self.cfg.n_embd as usize;
11801        let eps = self.cfg.rms_eps;
11802        assert!(t <= scr.row_ctrs.len() && t <= 64);
11803        e.i32_iota_from(ctr, &mut scr.pos_d, t)?;
11804        for i in 0..t {
11805            e.i32_copy_add(ctr, &mut scr.row_ctrs[i], (i + 1) as i32)?;
11806        }
11807        let (pos_d, row_ctrs) = (&scr.pos_d, &scr.row_ctrs);
11808        let embd_gpu = self
11809            .embd_gpu
11810            .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload"));
11811        let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
11812        let mut x = e.embed_gather_device_td(embd_gpu, tok_d, t, n_embd, qt, rb)?;
11813        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
11814        let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
11815        let n_layers = self.layers.len();
11816        for (il, layer) in self.layers.iter().enumerate() {
11817            let (hq, hdq) = match h_carry.take() {
11818                Some(p) => p,
11819                None => {
11820                    e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, t, eps)?
11821                }
11822            };
11823            let Mixer::Full(fa) = &layer.mixer else {
11824                panic!("gemma4 layer {il} not full-attn")
11825            };
11826            let o = self
11827                .gemma4_verify_attn_stream(e, fa, il, &hq, &hdq, pos_d, t, cache, hint, row_ctrs)?;
11828            let next_norm = if il + 1 < n_layers {
11829                Some(self.layers[il + 1].attn_norm.float_data())
11830            } else {
11831                None
11832            };
11833            let (xn, hn) = self.gemma4_layer_tail_add_nq_pn(e, layer, &o, &x, t, next_norm)?;
11834            x = xn;
11835            h_carry = hn;
11836            self.dflash_tap(e, cache, il, &x, t)?;
11837        }
11838        let mut hn = e.uninit(t * n_embd)?;
11839        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
11840        let ld = e.matmul(&self.output, &hn, t)?;
11841        let n_vocab = self.output.out_features();
11842        let mut vam = e.stream().alloc_zeros::<u32>(t)?;
11843        for i in 0..t {
11844            e.argmax_token_device_col(&ld, i, n_vocab, &mut vam, i)?;
11845        }
11846        Ok((vam, hn))
11847    }
11848
11849    /// Verify trunk core: returns (UN-softcapped logits device [t, n_vocab], post-output_norm
11850    /// hidden stack [t, n_embd]); appends KV rows + advances cache.pos.
11851    /// DFlash tap write (dflash lane): copy the post-layer residual rows of `x` into the
11852    /// armed sink at the tap slot for `il` (row-major [t, n_taps*hidden]). Per-row D2D
11853    /// copies — t <= block_size on verify; prime pays t*n_taps once per prompt (dedicated
11854    /// kernel later if it shows in the profile).
11855    pub(crate) fn dflash_tap(
11856        &self,
11857        e: &Engine,
11858        cache: &mut Cache,
11859        il: usize,
11860        x: &CudaSlice<f32>,
11861        t: usize,
11862    ) -> Result<(), Box<dyn std::error::Error>> {
11863        let Some(taps) = cache.dflash_taps.as_mut() else {
11864            return Ok(());
11865        };
11866        let Some(slot) = taps.layer_ids.iter().position(|&l| l == il) else {
11867            return Ok(());
11868        };
11869        let h = taps.hidden;
11870        let n_taps = taps.layer_ids.len();
11871        let base = taps.base;
11872        debug_assert!(
11873            base + t <= taps.t,
11874            "tap window {base}+{t} exceeds sink {}",
11875            taps.t
11876        );
11877        let xv = e.view(x, t * h);
11878        for r in 0..t {
11879            let row = xv.slice(r * h..(r + 1) * h);
11880            e.copy_view_into(&mut taps.buf, (base + r) * n_taps * h + slot * h, &row, h)?;
11881        }
11882        Ok(())
11883    }
11884
11885    fn gemma4_verify_trunk(
11886        &self,
11887        e: &Engine,
11888        tokens: &[u32],
11889        pos0: usize,
11890        cache: &mut Cache,
11891        tok_dev: Option<&CudaSlice<u32>>,
11892    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11893        let n_embd = self.cfg.n_embd as usize;
11894        let eps = self.cfg.rms_eps;
11895        let t = tokens.len();
11896        let pos: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
11897        let pos_d = e.htod_i32(&pos)?;
11898        let mut x = match tok_dev {
11899            Some(td) => {
11900                let embd_gpu = self
11901                    .embd_gpu
11902                    .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload"));
11903                let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
11904                e.embed_gather_device_td(embd_gpu, td, t, n_embd, qt, rb)?
11905            }
11906            None => e.htod(&self.embd.gather(n_embd, tokens))?,
11907        };
11908        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
11909        let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
11910        let n_layers = self.layers.len();
11911        for (il, layer) in self.layers.iter().enumerate() {
11912            let (hq, hdq) = match h_carry.take() {
11913                Some(p) => p,
11914                None => {
11915                    e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, t, eps)?
11916                }
11917            };
11918            let Mixer::Full(fa) = &layer.mixer else {
11919                panic!("gemma4 layer {il} not full-attn")
11920            };
11921            let o = self.gemma4_verify_attn(e, fa, il, &hq, &hdq, &pos_d, t, cache)?;
11922            let next_norm = if il + 1 < n_layers {
11923                Some(self.layers[il + 1].attn_norm.float_data())
11924            } else {
11925                None
11926            };
11927            let (xn, hn) = self.gemma4_layer_tail_add_nq_pn(e, layer, &o, &x, t, next_norm)?;
11928            x = xn;
11929            h_carry = hn;
11930            self.dflash_tap(e, cache, il, &x, t)?;
11931        }
11932        let mut hn = e.uninit(t * n_embd)?;
11933        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
11934        let mut ld = e.matmul(&self.output, &hn, t)?;
11935        self.gemma4_suppress(e, &mut ld, t)?; // before the per-row argmax consumers
11936        cache.pos += t;
11937        Ok((ld, hn))
11938    }
11939
11940    /// Verify attention: project/norm/rope t rows, append them to the cache, then attend each
11941    /// row causally over [win_off_i .. base+i] via the SAME fa_decode dispatch as T=1 decode.
11942    /// BURST verify attention (device-slot twin of `gemma4_verify_attn`): the append lands
11943    /// at the layer's len_d counter (rows_dc), attention bases ride the counter (rows /
11944    /// rows_w with base_dev), and NO host len is read or bumped — `hint` is a host UPPER
11945    /// bound on base_len used only for split sizing and arm gating (the burst loop passes
11946    /// pos0 + burst slack). Host len mirrors re-sync at the burst drain.
11947    #[allow(clippy::too_many_arguments)]
11948    fn gemma4_verify_attn_stream(
11949        &self,
11950        e: &Engine,
11951        fa: &crate::hybrid::FullAttnLayer,
11952        il: usize,
11953        hq: &CudaSlice<i8>,
11954        hdq: &CudaSlice<f32>,
11955        pos_d: &CudaSlice<i32>,
11956        t: usize,
11957        cache: &mut Cache,
11958        hint: usize,
11959        row_ctrs: &[CudaSlice<i32>],
11960    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11961        let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
11962        let eps = self.cfg.rms_eps;
11963        let aux = self.gemma4_aux.as_ref().unwrap();
11964        let ones = aux.ones(e);
11965        #[cfg(debug_assertions)]
11966        crate::debug_assert_tensor_stream_device(
11967            ones,
11968            &e.stream(),
11969            "gemma4_verify_attn_stream.ones",
11970        );
11971        let h0 = e.zeros(0)?;
11972        let h = &h0;
11973        // BATCHED FUSED qkv (MEMRA_F2B=1, megakernel microcosm): swa layers fuse all three,
11974        // globals fuse q,k (v := k clone). Bit-identical per row; segments tail-fill.
11975        static F2B_QKV: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11976        let f2b = *F2B_QKV.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0"));
11977        let fused_qkv = if f2b {
11978            if swa {
11979                e.matmul_q4_fused3_batched(&fa.wq, &fa.wk, &fa.wv, hq, hdq, t)?
11980                    .map(|(a, b, c)| (a, b, Some(c)))
11981            } else {
11982                e.matmul_q4_fused2_batched(&fa.wq, &fa.wk, hq, hdq, t)?
11983                    .map(|(a, b)| (a, b, None))
11984            }
11985        } else {
11986            None
11987        };
11988        let (q0, k0, v0) = match fused_qkv {
11989            Some((a, b, cv)) => {
11990                let v = match cv {
11991                    Some(c) => c,
11992                    None => e.clone_dtod(&b)?,
11993                };
11994                (a, b, v)
11995            }
11996            None => {
11997                let q0 = e.matmul_pre(&fa.wq, hq, hdq, h, t)?;
11998                let k0 = e.matmul_pre(&fa.wk, hq, hdq, h, t)?;
11999                let v0 = if swa {
12000                    e.matmul_pre(&fa.wv, hq, hdq, h, t)?
12001                } else {
12002                    e.clone_dtod(&k0)?
12003                };
12004                (q0, k0, v0)
12005            }
12006        };
12007        let mut q = e.uninit(t * nh * hd)?;
12008        let mut k = e.uninit(t * nkv * hd)?;
12009        let mut v = e.uninit(t * nkv * hd)?;
12010        // E4B wave-3 fold (2026-07-23 decode-dust port): q/k/v norms + q/k rope in ONE
12011        // launch — rope math verbatim on the normed rows; V ones-rms, never roped.
12012        let ff = if swa {
12013            None
12014        } else {
12015            Some(
12016                aux.rope_freqs(e)
12017                    .expect("gemma4 global rope needs rope_freqs.weight"),
12018            )
12019        };
12020        #[cfg(debug_assertions)]
12021        if let Some(ff) = ff {
12022            crate::debug_assert_tensor_stream_device(
12023                ff,
12024                &e.stream(),
12025                "gemma4_verify_attn_stream.rope_freqs",
12026            );
12027        }
12028        e.rms_norm_qkv_rope(
12029            &q0,
12030            &k0,
12031            &v0,
12032            fa.q_norm.float_data(),
12033            fa.k_norm.float_data(),
12034            ones,
12035            &mut q,
12036            &mut k,
12037            &mut v,
12038            hd,
12039            self.gemma4_rope_dims(il),
12040            nh * t,
12041            nkv * t,
12042            pos_d,
12043            nh,
12044            nkv,
12045            base,
12046            1.0,
12047            ff,
12048            eps,
12049        )?;
12050        let kvl = cache.kv[il].as_mut().unwrap();
12051        // append at the DEVICE slot; the counter advances by t on-device.
12052        e.append_kv_quantized_rows_dc(
12053            &k,
12054            &v,
12055            &mut kvl.k,
12056            &mut kvl.v,
12057            &kvl.len_d,
12058            t,
12059            kvl.kv_dim_k,
12060            kvl.kv_dim_v,
12061            kvl.k_tok_bytes,
12062            kvl.v_tok_bytes,
12063            (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on()),
12064        )?;
12065        // len_d is NOT advanced here: the burst's device rollback (spec_rollback_stream) is
12066        // the sole len writer after this round's attention (base stays = old len, plus = 0).
12067        let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
12068        let mut attn = e.uninit(t * nh * hd)?;
12069        let k_view = e.view_u8(&kvl.k, kvl.k.len());
12070        let v_view = e.view_u8(&kvl.v, kvl.v.len());
12071        // arm gating on the host HINT (upper bound; burst entry guards hint >= the vec floor
12072        // and a stable window regime — the same rung/regime keys as the draft graph).
12073        if swa && hint + 1 >= win {
12074            // fully-windowed rows: per-row window geometry from the counter (plus = 0: len_d
12075            // still holds the pre-append len; row r's T_kv = ctr + r + 1).
12076            e.fa_decode_rows_w(
12077                &q,
12078                &k_view,
12079                &v_view,
12080                &mut attn,
12081                hd,
12082                nh,
12083                nkv,
12084                &kvl.len_d,
12085                0,
12086                t,
12087                scale,
12088                win,
12089                kvl.k_tok_bytes,
12090                kvl.v_tok_bytes,
12091                None,
12092            )?;
12093        } else if hd == 512 && hint + t < crate::fa512_min_tkv() {
12094            // globals UNDER the fa512 crossover: eager runs the per-row fa_decode_kvmod
12095            // fallback there (parity with t=1 decode) — mirror it with per-row fa_decode_dc
12096            // (bit-correct for any t_kv <= bucket; the drafter dc arc proved the pairing).
12097            // Burst entry gates the horizon onto one side of the crossover, so hint decides
12098            // for every row.
12099            // NO .max(512): fa_decode_dc gates its hd512 dpl16-vs-scalar pick on bucket_max
12100            // (mirroring eager's fa512 floor) — forcing 512 here flipped the arm to dpl16
12101            // while eager ran scalar (il=5 KV drift, the burst's 4/128). SAME LAW capped
12102            // from above (2026-07-13): a regime-pinned hint near the floor (round-graph
12103            // captures use f512-1) pow2-rounds PAST it — clamp under the floor, or the arm
12104            // re-flips to dpl16 (the round-graph 4/64). The scalar unified self-splits, so
12105            // any bucket >= the live length is exact.
12106            let bucket = (hint + t + 2)
12107                .next_power_of_two()
12108                .min(crate::fa512_min_tkv().saturating_sub(1));
12109            let qv = e.view(&q, t * nh * hd);
12110            for i in 0..t {
12111                let q_row = qv.slice(i * nh * hd..(i + 1) * nh * hd);
12112                let mut q_one = e.uninit(nh * hd)?;
12113                e.copy_view_into(&mut q_one, 0, &q_row, nh * hd)?;
12114                let mut a_one = e.uninit(nh * hd)?;
12115                e.fa_decode_dc(
12116                    &q_one,
12117                    &k_view,
12118                    &v_view,
12119                    &mut a_one,
12120                    hd,
12121                    nh,
12122                    nkv,
12123                    &row_ctrs[i],
12124                    bucket,
12125                    scale,
12126                    kvl.k_tok_bytes,
12127                    kvl.v_tok_bytes,
12128                    false,
12129                )?;
12130                e.copy_into(&mut attn, i * nh * hd, &a_one, nh * hd)?;
12131            }
12132        } else if hd == 512 {
12133            // globals past the crossover: dpl16 dc twin via the shared rows wrapper (hint
12134            // sizes splits — upper bound; splits beyond the device len exit in-kernel).
12135            e.fa_decode_rows(
12136                &q,
12137                &k_view,
12138                &v_view,
12139                &mut attn,
12140                hd,
12141                nh,
12142                nkv,
12143                hint,
12144                t,
12145                scale,
12146                kvl.k_tok_bytes,
12147                kvl.v_tok_bytes,
12148                Some((&kvl.len_d, 0)),
12149                false,
12150                false,
12151                None,
12152            )?;
12153        } else {
12154            // hd256 under-window: v4 device-len rows twin.
12155            e.fa_decode_rows_dc(
12156                &q,
12157                &k_view,
12158                &v_view,
12159                &mut attn,
12160                hd,
12161                nh,
12162                nkv,
12163                &kvl.len_d,
12164                hint + t,
12165                t,
12166                scale,
12167                kvl.k_tok_bytes,
12168                kvl.v_tok_bytes,
12169                0,
12170                swa && crate::Engine::wkv_on(),
12171            )?;
12172        }
12173        Ok(e.matmul(&fa.wo, &attn, t)?)
12174    }
12175
12176    fn gemma4_verify_attn(
12177        &self,
12178        e: &Engine,
12179        fa: &crate::hybrid::FullAttnLayer,
12180        il: usize,
12181        hq: &CudaSlice<i8>,
12182        hdq: &CudaSlice<f32>,
12183        pos_d: &CudaSlice<i32>,
12184        t: usize,
12185        cache: &mut Cache,
12186    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12187        let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
12188        let eps = self.cfg.rms_eps;
12189        let aux = self.gemma4_aux.as_ref().unwrap();
12190        let ones = aux.ones(e);
12191        #[cfg(debug_assertions)]
12192        crate::debug_assert_tensor_stream_device(ones, &e.stream(), "gemma4_verify_attn.ones");
12193        let n_embd = self.cfg.n_embd as usize;
12194        let _ = n_embd;
12195
12196        let h0 = e.zeros(0)?;
12197        let h = &h0;
12198        // BATCHED FUSED qkv (MEMRA_F2B=1, megakernel microcosm): swa layers fuse all three,
12199        // globals fuse q,k (v := k clone). Bit-identical per row; segments tail-fill.
12200        static F2B_QKV: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12201        let f2b = *F2B_QKV.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0"));
12202        let fused_qkv = if f2b {
12203            if swa {
12204                e.matmul_q4_fused3_batched(&fa.wq, &fa.wk, &fa.wv, hq, hdq, t)?
12205                    .map(|(a, b, c)| (a, b, Some(c)))
12206            } else {
12207                e.matmul_q4_fused2_batched(&fa.wq, &fa.wk, hq, hdq, t)?
12208                    .map(|(a, b)| (a, b, None))
12209            }
12210        } else {
12211            None
12212        };
12213        let (q0, k0, v0) = match fused_qkv {
12214            Some((a, b, cv)) => {
12215                let v = match cv {
12216                    Some(c) => c,
12217                    None => e.clone_dtod(&b)?,
12218                };
12219                (a, b, v)
12220            }
12221            None => {
12222                let q0 = e.matmul_pre(&fa.wq, hq, hdq, h, t)?;
12223                let k0 = e.matmul_pre(&fa.wk, hq, hdq, h, t)?;
12224                let v0 = if swa {
12225                    e.matmul_pre(&fa.wv, hq, hdq, h, t)?
12226                } else {
12227                    e.clone_dtod(&k0)?
12228                };
12229                (q0, k0, v0)
12230            }
12231        };
12232        let mut q = e.uninit(t * nh * hd)?;
12233        let mut k = e.uninit(t * nkv * hd)?;
12234        let mut v = e.uninit(t * nkv * hd)?;
12235        // E4B wave-3 fold (2026-07-23 decode-dust port): q/k/v norms + q/k rope in ONE
12236        // launch — rope math verbatim on the normed rows; V ones-rms, never roped.
12237        let ff = if swa {
12238            None
12239        } else {
12240            Some(
12241                aux.rope_freqs(e)
12242                    .expect("gemma4 global rope needs rope_freqs.weight"),
12243            )
12244        };
12245        #[cfg(debug_assertions)]
12246        if let Some(ff) = ff {
12247            crate::debug_assert_tensor_stream_device(
12248                ff,
12249                &e.stream(),
12250                "gemma4_verify_attn.rope_freqs",
12251            );
12252        }
12253        e.rms_norm_qkv_rope(
12254            &q0,
12255            &k0,
12256            &v0,
12257            fa.q_norm.float_data(),
12258            fa.k_norm.float_data(),
12259            ones,
12260            &mut q,
12261            &mut k,
12262            &mut v,
12263            hd,
12264            self.gemma4_rope_dims(il),
12265            nh * t,
12266            nkv * t,
12267            pos_d,
12268            nh,
12269            nkv,
12270            base,
12271            1.0,
12272            ff,
12273            eps,
12274        )?;
12275        let kvl = cache.kv[il].as_mut().unwrap();
12276        let base_len = kvl.len;
12277        e.append_kv_quantized_rows(
12278            &k,
12279            &v,
12280            &mut kvl.k,
12281            &mut kvl.v,
12282            base_len,
12283            t,
12284            kvl.kv_dim_k,
12285            kvl.kv_dim_v,
12286            kvl.k_tok_bytes,
12287            kvl.v_tok_bytes,
12288            (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on()),
12289        )?;
12290        kvl.len += t;
12291        let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
12292        let mut attn = e.uninit(t * nh * hd)?;
12293        // ROWS twin when no window offset is in play (globals are hd512-ineligible; SWA rows
12294        // under the window need no offset): ONE launch, per-row causal == the per-token loop.
12295        let rows_ok = (hd == 256 && base_len + 1 >= crate::fa_vec_min_tkv())
12296            // gemma globals: hd512 rows twin in the dpl16 vec regime (row 0 gates the batch);
12297            // decode rides the SAME symbol at t=1 (parity law).
12298            || (hd == 512 && !swa && base_len + 1 >= crate::fa512_min_tkv());
12299        if rows_ok && (!swa || base_len + t <= win) {
12300            let k_view = e.view_u8(&kvl.k, (base_len + t) * kvl.k_tok_bytes);
12301            let v_view = e.view_u8(&kvl.v, (base_len + t) * kvl.v_tok_bytes);
12302            if hd == 512 {
12303                // device-len twin: sync the counter to the verify base (async arg-store).
12304                e.i32_set_k(&mut kvl.len_d, base_len as i32)?;
12305                e.fa_decode_rows(
12306                    &q,
12307                    &k_view,
12308                    &v_view,
12309                    &mut attn,
12310                    hd,
12311                    nh,
12312                    nkv,
12313                    base_len,
12314                    t,
12315                    scale,
12316                    kvl.k_tok_bytes,
12317                    kvl.v_tok_bytes,
12318                    Some((&kvl.len_d, 0)),
12319                    false,
12320                    swa && crate::Engine::wkv_on(),
12321                    None,
12322                )?;
12323            } else {
12324                // hd256: the SAME v4_dc symbol the burst verify launches (PARITY LAW — the
12325                // host-base rows_v4 twin compiles apart from v4_dc and the two-symbol split
12326                // drifted the burst's persisted KV at il=8; one symbol, counter arg-stored).
12327                e.i32_set_k(&mut kvl.len_d, base_len as i32)?;
12328                e.fa_decode_rows_dc(
12329                    &q,
12330                    &k_view,
12331                    &v_view,
12332                    &mut attn,
12333                    hd,
12334                    nh,
12335                    nkv,
12336                    &kvl.len_d,
12337                    base_len + t,
12338                    t,
12339                    scale,
12340                    kvl.k_tok_bytes,
12341                    kvl.v_tok_bytes,
12342                    0,
12343                    swa && crate::Engine::wkv_on(),
12344                )?;
12345            }
12346            return Ok(e.matmul(&fa.wo, &attn, t)?);
12347        }
12348        // WINDOWED rows twin (deep ctx, every row fully windowed): one launch, ABSOLUTE-index
12349        // per-row geometry over the prefix view. PARITY LAW (2026-07-10 root-cause): textually
12350        // identical kernels do NOT compile bit-identically (nvcc unrolls fa_decode_vec_q 2x vs
12351        // its rows_w clone — SASS-proven; the unpinned score `+=` chain then rounds apart), so
12352        // decode-vs-verify parity comes from BOTH sides launching THIS SAME rows_w kernel
12353        // (decode passes t=1) — bitwise equal per position by symbol identity, any lane.
12354        // MEMRA_GEMMA_ROWS_W=0 -> per-token loop (decode falls back to fa_decode views too).
12355        if hd == 256
12356            && swa
12357            && base_len + 1 >= win
12358            && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
12359        {
12360            let k_view = e.view_u8(&kvl.k, (base_len + t) * kvl.k_tok_bytes);
12361            let v_view = e.view_u8(&kvl.v, (base_len + t) * kvl.v_tok_bytes);
12362            e.i32_set_k(&mut kvl.len_d, base_len as i32)?;
12363            e.fa_decode_rows_w(
12364                &q,
12365                &k_view,
12366                &v_view,
12367                &mut attn,
12368                hd,
12369                nh,
12370                nkv,
12371                &kvl.len_d,
12372                0,
12373                t,
12374                scale,
12375                win,
12376                kvl.k_tok_bytes,
12377                kvl.v_tok_bytes,
12378                None,
12379            )?;
12380            return Ok(e.matmul(&fa.wo, &attn, t)?);
12381        }
12382        for i in 0..t {
12383            let avail = base_len + i + 1;
12384            let (off_tok, t_kv) = if swa && avail > win {
12385                (avail - win, win)
12386            } else {
12387                (0, avail)
12388            };
12389            let k_view = e.view_u8_range(
12390                &kvl.k,
12391                off_tok * kvl.k_tok_bytes,
12392                (off_tok + t_kv) * kvl.k_tok_bytes,
12393            );
12394            let v_view = e.view_u8_range(
12395                &kvl.v,
12396                off_tok * kvl.v_tok_bytes,
12397                (off_tok + t_kv) * kvl.v_tok_bytes,
12398            );
12399            let qi = e.view(&q, t * nh * hd);
12400            let q_row = qi.slice(i * nh * hd..(i + 1) * nh * hd);
12401            let mut q_one = e.uninit(nh * hd)?;
12402            e.copy_view_into(&mut q_one, 0, &q_row, nh * hd)?;
12403            let mut a_one = e.uninit(nh * hd)?;
12404            // straddle rounds: windowed rows must ride the SAME rows_w kernel decode uses
12405            // (parity law above); global hd512 rows past the fa512 floor likewise ride the
12406            // rows_dpl16 twin; remaining rows keep the gated fa_decode pair.
12407            if swa
12408                && avail > win
12409                && hd == 256
12410                && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
12411            {
12412                let kp = e.view_u8(&kvl.k, avail * kvl.k_tok_bytes);
12413                let vp = e.view_u8(&kvl.v, avail * kvl.v_tok_bytes);
12414                e.i32_set_k(&mut kvl.len_d, (avail - 1) as i32)?;
12415                e.fa_decode_rows_w(
12416                    &q_one,
12417                    &kp,
12418                    &vp,
12419                    &mut a_one,
12420                    hd,
12421                    nh,
12422                    nkv,
12423                    &kvl.len_d,
12424                    0,
12425                    1,
12426                    scale,
12427                    win,
12428                    kvl.k_tok_bytes,
12429                    kvl.v_tok_bytes,
12430                    None,
12431                )?;
12432            } else if !swa
12433                && hd == 512
12434                && avail >= crate::fa512_min_tkv()
12435                && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
12436            {
12437                let kp = e.view_u8(&kvl.k, avail * kvl.k_tok_bytes);
12438                let vp = e.view_u8(&kvl.v, avail * kvl.v_tok_bytes);
12439                e.i32_set_k(&mut kvl.len_d, (avail - 1) as i32)?;
12440                e.fa_decode_rows(
12441                    &q_one,
12442                    &kp,
12443                    &vp,
12444                    &mut a_one,
12445                    hd,
12446                    nh,
12447                    nkv,
12448                    avail - 1,
12449                    1,
12450                    scale,
12451                    kvl.k_tok_bytes,
12452                    kvl.v_tok_bytes,
12453                    Some((&kvl.len_d, 0)),
12454                    false,
12455                    false,
12456                    None,
12457                )?;
12458            } else {
12459                e.fa_decode_kvmod(
12460                    &q_one,
12461                    &k_view,
12462                    &v_view,
12463                    &mut a_one,
12464                    hd,
12465                    nh,
12466                    nkv,
12467                    t_kv,
12468                    scale,
12469                    kvl.k_tok_bytes,
12470                    kvl.v_tok_bytes,
12471                    swa && crate::Engine::wkv_on(),
12472                )?;
12473            }
12474            e.copy_into(&mut attn, i * nh * hd, &a_one, nh * hd)?;
12475        }
12476        Ok(e.matmul(&fa.wo, &attn, t)?)
12477    }
12478
12479    /// gemma4 T=1 decode step: R8 layer graph over the cache; returns (softcapped logits host,
12480    /// h_seed = pre-output_norm hidden). Advances cache.pos.
12481    pub(crate) fn gemma4_decode_step_h(
12482        &self,
12483        e: &Engine,
12484        token: u32,
12485        cache: &mut Cache,
12486    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12487        // M1-PP2 door (crate::pp): 2-stage split with an explicit activation handoff.
12488        // Default OFF — unset env means this branch is never taken. The gemma4 arm stays
12489        // 2-stage in M2 (the N-stage gate model is the generic-arm 9B); N>2 warns + runs
12490        // unsplit rather than guessing a fence.
12491        if let Some(split) = crate::pp::pp2_split(self.layers.len()) {
12492            return self.gemma4_decode_step_h_pp2(e, token, cache, split);
12493        }
12494        if crate::pp::pp_cuts(self.layers.len()).is_some() {
12495            crate::pp::warn_unwired_once("gemma4 eager decode (N>2)");
12496        }
12497        let n_embd = self.cfg.n_embd as usize;
12498        let eps = self.cfg.rms_eps;
12499        let pos_d = e.htod_i32(&[cache.pos as i32])?;
12500        let mut x = e.htod(&self.embd.gather(n_embd, &[token]))?;
12501        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
12502        // cross-layer fusion: each tail's closing add+scale also EMITS the next layer's
12503        // attn-normed input pre-quantized q8_1 (the mixer consumes only quantized matmuls).
12504        let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
12505        let n_layers = self.layers.len();
12506        for (il, layer) in self.layers.iter().enumerate() {
12507            let (hq, hdq) = match h_carry.take() {
12508                Some(p) => p,
12509                None => {
12510                    e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, 1, eps)?
12511                }
12512            };
12513            let Mixer::Full(fa) = &layer.mixer else {
12514                panic!("gemma4 layer {il} not full-attn")
12515            };
12516            let o = self.gemma4_decode_attn(e, fa, il, &hq, &hdq, &pos_d, cache)?;
12517            let next_norm = if il + 1 < n_layers {
12518                Some(self.layers[il + 1].attn_norm.float_data())
12519            } else {
12520                None
12521            };
12522            let (xn, hn) = self.gemma4_layer_tail_add_nq_pn(e, layer, &o, &x, 1, next_norm)?;
12523            x = xn;
12524            h_carry = hn;
12525        }
12526        let mut hn = e.uninit(n_embd)?;
12527        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
12528        let h_seed = e.clone_dtod(&x)?;
12529        let mut ld = e.matmul(&self.output, &hn, 1)?;
12530        let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
12531        e.softcap(&mut ld, cap, self.output.out_features())?; // R4 on device (262k host tanh ~ms/step)
12532        self.gemma4_suppress(e, &mut ld, 1)?;
12533        let logits = e.dtoh(&ld)?;
12534        cache.pos += 1;
12535        Ok((logits, h_seed))
12536    }
12537
12538    /// M1-PP2 stage subgraph (gemma4 arm): layers [lo, hi) of the gemma4 T=1 walk with the
12539    /// pre-quantized next-norm carry LOCAL to the range. Enters with a materialized residual
12540    /// `x` (the range head runs its own `rms_norm_q8_1` against ITS layer's attn_norm — for
12541    /// lo == 0 that is exactly the unsplit loop's il==0 arm); exits with the residual
12542    /// materialized (range tail passes next_norm = None, exactly the unsplit last-layer arm).
12543    /// Bit-identity of the cut relies on the kernel-check-pinned `add_scale_rms_norm_q8_1 ==
12544    /// add_scale then rms_norm_q8_1` identity (`pp2-gate` verifies end-to-end).
12545    fn gemma4_decode_layers(
12546        &self,
12547        e: &Engine,
12548        mut x: CudaSlice<f32>,
12549        lo: usize,
12550        hi: usize,
12551        pos_d: &CudaSlice<i32>,
12552        cache: &mut Cache,
12553    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12554        let n_embd = self.cfg.n_embd as usize;
12555        let eps = self.cfg.rms_eps;
12556        let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
12557        for il in lo..hi {
12558            let layer = &self.layers[il];
12559            let (hq, hdq) = match h_carry.take() {
12560                Some(p) => p,
12561                // range head: il == lo — norm against THIS layer's attn_norm.
12562                None => {
12563                    e.rms_norm_q8_1(&x, self.layers[il].attn_norm.float_data(), n_embd, 1, eps)?
12564                }
12565            };
12566            let Mixer::Full(fa) = &layer.mixer else {
12567                panic!("gemma4 layer {il} not full-attn")
12568            };
12569            let o = self.gemma4_decode_attn(e, fa, il, &hq, &hdq, pos_d, cache)?;
12570            let next_norm = if il + 1 < hi {
12571                Some(self.layers[il + 1].attn_norm.float_data())
12572            } else {
12573                None
12574            };
12575            let (xn, hn) = self.gemma4_layer_tail_add_nq_pn(e, layer, &o, &x, 1, next_norm)?;
12576            x = xn;
12577            h_carry = hn;
12578        }
12579        Ok(x)
12580    }
12581
12582    /// M1-PP2 (increment 2): `gemma4_decode_step_h` as TWO stage subgraphs, each on its
12583    /// own stream (and device, under MEMRA_PP_DEVICES) with the transport-selected
12584    /// boundary handoff — same choreography as the generic arm (decode.rs), same
12585    /// ownership contract (crate::pp). Stage 0 = embed+scale + layers [0, split);
12586    /// stage 1 = layers [split, n) + output_norm + softcapped head.
12587    /// Each stage uploads its own copy of the step's position scalar on its own stream.
12588    /// MEMRA_PP_STREAMS=0 = the increment-1 same-stream seam. Gate: `pp2-gate`.
12589    fn gemma4_decode_step_h_pp2(
12590        &self,
12591        e: &Engine,
12592        token: u32,
12593        cache: &mut Cache,
12594        split: usize,
12595    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12596        if crate::pp::pp2_streams_off() {
12597            return self.gemma4_decode_step_h_pp2_samestream(e, token, cache, split);
12598        }
12599        let rt = crate::pp::Pp2Rt::get(e)?;
12600        let e0 = rt.engine(0, e);
12601        let e1 = rt.engine(1, e);
12602        let n_embd = self.cfg.n_embd as usize;
12603        let eps = self.cfg.rms_eps;
12604        let pos = cache.pos as i32;
12605
12606        // ---- STAGE 0 (its own stream): embed + sqrt(n_embd) scale + layers [0, split) ----
12607        let slot = {
12608            let _st0 = rt.enter(0);
12609            let pos_d = e0.htod_i32(&[pos])?;
12610            #[cfg(debug_assertions)]
12611            crate::debug_assert_tensor_stream_device(
12612                &pos_d,
12613                &e0.stream(),
12614                "gemma4_decode_step_h_pp2.stage0.pos_d",
12615            );
12616            let mut x = e0.htod(&self.embd.gather(n_embd, &[token]))?;
12617            e0.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
12618            let x = self.gemma4_decode_layers(e0, x, 0, split, &pos_d, cache)?;
12619            rt.tx(0, &x, n_embd)?
12620        };
12621
12622        // ---- STAGE 1 (its own stream): RX + layers [split, n) + softcapped head ----
12623        let _st1 = rt.enter(1);
12624        let pos_d = e1.htod_i32(&[pos])?;
12625        #[cfg(debug_assertions)]
12626        crate::debug_assert_tensor_stream_device(
12627            &pos_d,
12628            &e1.stream(),
12629            "gemma4_decode_step_h_pp2.stage1.pos_d",
12630        );
12631        let x = rt.rx(0, slot, n_embd)?;
12632        let x = self.gemma4_decode_layers(e1, x, split, self.layers.len(), &pos_d, cache)?;
12633
12634        let mut hn = e1.uninit(n_embd)?;
12635        e1.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
12636        let h_seed = e1.clone_dtod(&x)?;
12637        let mut ld = e1.matmul(&self.output, &hn, 1)?;
12638        let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
12639        e1.softcap(&mut ld, cap, self.output.out_features())?;
12640        self.gemma4_suppress(e1, &mut ld, 1)?;
12641        let logits = e1.dtoh(&ld)?;
12642        cache.pos += 1;
12643        Ok((logits, h_seed))
12644    }
12645
12646    /// MEMRA_PP_STREAMS=0 rollback seam: the increment-1 gemma4 pp2 body verbatim — both
12647    /// stage subgraphs on the ambient compute stream, boundary = two plain dtod copies.
12648    fn gemma4_decode_step_h_pp2_samestream(
12649        &self,
12650        e: &Engine,
12651        token: u32,
12652        cache: &mut Cache,
12653        split: usize,
12654    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12655        let n_embd = self.cfg.n_embd as usize;
12656        let eps = self.cfg.rms_eps;
12657        let pos_d = e.htod_i32(&[cache.pos as i32])?;
12658
12659        // ---- STAGE 0: embed + sqrt(n_embd) scale + layers [0, split) ----
12660        let mut x = e.htod(&self.embd.gather(n_embd, &[token]))?;
12661        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
12662        let x = self.gemma4_decode_layers(e, x, 0, split, &pos_d, cache)?;
12663
12664        // ---- STAGE BOUNDARY: explicit [n_embd] activation handoff (TX copy, RX copy) ----
12665        let boundary_tx = e.clone_dtod(&x)?;
12666        let boundary_rx = e.clone_dtod(&boundary_tx)?;
12667
12668        // ---- STAGE 1: layers [split, n) + output_norm + softcapped head ----
12669        let x =
12670            self.gemma4_decode_layers(e, boundary_rx, split, self.layers.len(), &pos_d, cache)?;
12671
12672        let mut hn = e.uninit(n_embd)?;
12673        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
12674        let h_seed = e.clone_dtod(&x)?;
12675        let mut ld = e.matmul(&self.output, &hn, 1)?;
12676        let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
12677        e.softcap(&mut ld, cap, self.output.out_features())?;
12678        self.gemma4_suppress(e, &mut ld, 1)?;
12679        let logits = e.dtoh(&ld)?;
12680        cache.pos += 1;
12681        Ok((logits, h_seed))
12682    }
12683}
12684
12685// ============================ step35 (Step-3.7-Flash) ==================================
12686// Node-for-node vs llama.cpp src/models/step35.cpp:216-300. WHY THIS IS A DEDICATED MIXER
12687// FAMILY and not a few branches inside the generic `full_attn*` chain:
12688//
12689//   1. `n_head` IS PER LAYER (`step35.attention.head_count` is an ARRAY: 64 on full-attn
12690//      layers, 96 on SWA). Every generic site reads the `cfg.n_head` scalar, so wq/wo/attn_gate
12691//      shapes and the FA head counts would be wrong on 33 of 45 layers.
12692//   2. `cfg.rope_dim_count` is 128 for this arch, but FULL-attn layers rotate only 64 dims
12693//      (upstream halves `n_rot_full` after the generic loader defaults it). A generic mixer
12694//      would rotate 128 dims on the full layers — silently wrong, plausible-looking logits.
12695//   3. Dual rope base (5e6 full / 1e4 SWA) + `rope_freqs.weight` on FULL layers ONLY.
12696//   4. SWA 3:1 (`sliding_window` 512, pattern [false,true,true,true]).
12697//   5. The head-wise gate is a SEPARATE `attn_gate.weight [n_embd, n_head_l]` (one pre-sigmoid
12698//      scalar per head, broadcast over head_dim) applied to the attention output BEFORE wo —
12699//      not the qwen35 fused-in-wq per-DIM gate that `attn_out_gate()`/`q_gate_split` handle.
12700//      Its input is the POST-attn_norm hidden (`cur`), not the residual.
12701//
12702// Attention scale is the DEFAULT 1/sqrt(n_embd_head_k) (step35.cpp:255) — NOT gemma4's 1.0.
12703impl HybridModel {
12704    /// Step35's declarative geometry row. Missing rows are artifact/config errors; never
12705    /// synthesize a drafter or trunk layer from a neighboring class.
12706    pub(crate) fn step35_geom(&self, il: usize) -> memra_gguf::config::LayerGeometry {
12707        let geometry = self
12708            .cfg
12709            .layer_geometry(il as u32)
12710            .unwrap_or_else(|| panic!("step35 layer {il} has no geometry-table row"));
12711        debug_assert_eq!(
12712            geometry.attention_gate,
12713            memra_gguf::config::AttentionGateKind::SeparateHead
12714        );
12715        geometry
12716    }
12717
12718    /// step35 attention core: everything from the q/k/v projection outputs through the head-wise
12719    /// gate, returning the GATED attention output PRE-`wo` (the caller runs wo, so this composes
12720    /// with both the plain `matmul(&fa.wo, ..)` sites and the f16/into-slab GEMM sites).
12721    ///
12722    /// `hg` = the POST-attn_norm hidden state (upstream's `cur`) — the gate projection's input
12723    /// on single-sequence paths. `gt_pre` supplies the already-projected per-head gate for a
12724    /// cross-request batch; exactly one of `hg` or `gt_pre` must be present.
12725    /// `cache`:
12726    ///   * `Some` => PRIME mode: append this chunk's post-rope K / raw V rows into the resident
12727    ///     quantized cache and attend THROUGH the cache view, exactly like the generic
12728    ///     `full_attn_prime_fa_dispatch` (quantize-then-attend on chunk 0 too, so chunk size
12729    ///     cannot decide where a precision edge falls — the grain-free chunk-invariance
12730    ///     contract, lane/chunkinv-flip).
12731    ///   * `None` => pure prefill (forward/forward_last/t2probe): attend over this batch's f32
12732    ///     q/k/v, no cache side effect.
12733    ///
12734    /// SWA WINDOW, and why prefill needs `sdpa_naive_w_quantized_view`: the view is trimmed to
12735    /// the oldest key ANY query in this chunk can reach (`off = base_len - (win-1)`), which
12736    /// bounds the view at `win-1+t` rows, but the mask is still required — inside the chunk,
12737    /// query `qt` may only see view keys `[qt, qt+win-1]`, so the earlier keys the trimmed view
12738    /// still contains must be masked per query. memra's window convention
12739    /// (`sdpa_naive_w_f32`: mask `t < q_pos-(window-1)`, `q_pos = (T_kv-T)+qt`) is bit-for-bit
12740    /// upstream's `LLAMA_SWA_TYPE_STANDARD` (`llama-hparams.h:359`: mask `p1-p0 >= n_swa`), and
12741    /// step35.cpp:6 sets exactly that swa_type — verified verbatim on both sides.
12742    ///
12743    /// Note the f32 `sdpa_naive_w` floor's shared memory is `t_kv*4` bytes, so the SWA arm needs
12744    /// `t_kv = win-1+t <= 12287` — i.e. it relies on chunked prefill (MEMRA_PRIME_CHUNK, default
12745    /// 4096). A monolithic 32k prime would exceed the 48 KiB dynamic-smem default.
12746    ///
12747    /// SWA ARM SELECTION IS KEYED ON `seq_end`, NOT ON THE CHUNK'S OWN `t_kv` (lane/step35-chunkfix,
12748    /// 2026-08-07). `seq_end` = the ABSOLUTE end position of the whole prime call
12749    /// (`cache.pos + prompt_len` at entry; `t` on the cacheless prefill path) — a property of the
12750    /// REQUEST, identical at every MEMRA_PRIME_CHUNK value. The old predicate `swa && t_kv > win`
12751    /// read the chunk's own extent, which made kernel selection — and therefore the logits, the
12752    /// hidden rows, and the generated text — a function of the chunk size:
12753    ///   a chunk [b,e) has off = max(0, b-(win-1)), t_kv = e-off, so b <= win-1 => off=0 => FA iff
12754    ///   e <= win, while every later chunk has t_kv = t+(win-1) > win. The FA rows were therefore a
12755    ///   contiguous PREFIX [0,P) with P = c*floor(win/c) for c <= win (else 0), and the output
12756    ///   depended only on P. Measured (research/step37-p2-20260806, closed-form receipt
12757    ///   `raw/chunkinv-step35-GAP2-CONFIRMED-20260807.txt`): at T=4883, c=512 and c=64 (P=512) are
12758    ///   byte-identical to each other but DIFFER from the c=4096 default (P=0) by maxdiff 1.813e0
12759    ///   with greedy text diverging at step 6; c=513 (P=0) is bit-identical to the default. A
12760    ///   one-token change in a documented machine-config knob changed the answer.
12761    /// So the pre-fix comment "same cache bytes, same numeric class" was false as written: same
12762    /// bytes, DIFFERENT numeric class — swapping `fa_prefill_view_ws` for the f32 windowed floor on
12763    /// the same rows moves the logits by ~1.8.
12764    ///
12765    /// Keying on `seq_end` makes P identically 0 for every chunk size, so MEMRA_PRIME_CHUNK is a
12766    /// pure memory/transient knob again (research/step35-chunkfix-20260807). It is
12767    /// correct-by-construction rather than a tolerance argument — the windowed kernel computes the
12768    /// same masked attention the FA arm computed (for e <= win the window mask is a no-op under
12769    /// causal, which is exactly why FA was legal there), it is now simply the ONLY arm used once
12770    /// the request passes the window. It also cannot move the shipped default: at chunk=4096 every
12771    /// chunk of a `seq_end > win` prime ALREADY had t_kv > win (chunk 0 has t_kv = min(chunk,
12772    /// seq_end) > win; every later chunk has t_kv >= t+win-1 > win since `PRIME_MIN_T` keeps
12773    /// t >= 16), and a `seq_end <= win` prime keeps every chunk on FA exactly as before — so this
12774    /// is a no-op on both default regimes and only removes the FA prefix at small chunk values.
12775    /// The `t_kv <= 12287` smem ceiling is respected: the chunks whose arm changes are precisely
12776    /// those with t_kv <= win = 512.
12777    #[allow(clippy::too_many_arguments)]
12778    fn step35_attn_pre_wo(
12779        &self,
12780        e: &Engine,
12781        fa: &FullAttnLayer,
12782        mut g3: Vec<CudaSlice<f32>>,
12783        hg: Option<&CudaSlice<f32>>,
12784        gt_pre: Option<&CudaSlice<f32>>,
12785        pos_d: &CudaSlice<i32>,
12786        t: usize,
12787        cache: Option<&mut Cache>,
12788        il: usize,
12789        seq_end: usize,
12790    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12791        let geometry = self.step35_geom(il);
12792        let hd = geometry.head_dim_k as usize;
12793        let nkv = geometry.n_head_kv as usize;
12794        let nh = geometry.n_head as usize;
12795        let rbase = geometry.rope_base;
12796        let scale = geometry.attention_scale();
12797        let swa = geometry.window.is_some();
12798        let eps = self.cfg.rms_eps;
12799        let win = geometry.window.unwrap_or(0) as usize;
12800        let n_rot = geometry.n_rot as usize;
12801
12802        let v = g3.pop().unwrap();
12803        let k0 = g3.pop().unwrap();
12804        let q0 = g3.pop().unwrap();
12805
12806        // q/k RMSNorm over head_dim rows (attn_q_norm/attn_k_norm [128] F32), then the
12807        // per-layer PARTIAL NEOX rope: n_rot = 64 on full layers (dims 64..127 pass through
12808        // unrotated), 128 on SWA. rope_freqs (llama3-style per-dim factors) on FULL only.
12809        let mut q = e.uninit(t * nh * hd)?;
12810        e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, nh * t, eps)?;
12811        let mut k = e.uninit(t * nkv * hd)?;
12812        e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, nkv * t, eps)?;
12813        let ff = if geometry.rope_factors {
12814            self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
12815        } else {
12816            None
12817        };
12818        #[cfg(debug_assertions)]
12819        if let Some(ff) = ff {
12820            crate::debug_assert_tensor_stream_device(
12821                ff,
12822                &e.stream(),
12823                "step35_attn_pre_wo.rope_freqs",
12824            );
12825        }
12826        e.rope_neox2(&mut q, &mut k, pos_d, hd, n_rot, nh, nkv, t, rbase, 1.0, ff)?;
12827
12828        let mut attn = e.uninit(t * nh * hd)?;
12829        match cache {
12830            Some(cache) => {
12831                let base_len = cache.kv[il].as_ref().unwrap().len;
12832                // Read per layer call, never in a measured default.
12833                let legacy_tkv = std::env::var("MEMRA_STEP35_SWA_TKV").as_deref() == Ok("1");
12834                let legacy_calllocal = std::env::var("MEMRA_PRIME_CALLLOCAL").as_deref() == Ok("1");
12835                let off = if swa {
12836                    let raw = base_len.saturating_sub(win - 1);
12837                    if legacy_tkv || legacy_calllocal {
12838                        raw
12839                    } else {
12840                        raw & !31usize
12841                    }
12842                } else {
12843                    0
12844                };
12845                {
12846                    let kvl = cache.kv[il].as_mut().unwrap();
12847                    assert!(kvl.len + t <= cache.max_ctx, "step35 prime: KV overflow");
12848                    let write_row = e.prepare_kv_append(kvl, off, t)?;
12849                    e.append_kv_quantized_rows(
12850                        &k,
12851                        &v,
12852                        &mut kvl.k,
12853                        &mut kvl.v,
12854                        write_row,
12855                        t,
12856                        kvl.kv_dim_k,
12857                        kvl.kv_dim_v,
12858                        kvl.k_tok_bytes,
12859                        kvl.v_tok_bytes,
12860                        crate::Engine::kv_fp8_on(),
12861                    )?;
12862                    kvl.len += t;
12863                    let new_len = kvl.len as i32;
12864                    e.set_i32_one(&mut kvl.len_d, new_len)?;
12865                }
12866                let kvl = cache.kv[il].as_ref().unwrap();
12867                // Both segmentation seams restore the FULL pre-fix arithmetic: their legacy
12868                // arm predicate (`t_kv` below, per-call `seq_end` in prime_cache) AND the
12869                // unaligned view offset here. Both halves are load-bearing for the canaries:
12870                // on the FA default the predicate arms agree bitwise wherever they can differ
12871                // (a t_kv<=win view has no maskable key, so windowed==unwindowed FA bit-for-bit).
12872                // That made predicate-only seams INERT under Lever A — first chunkinv35c, then
12873                // tickinv35c. The raw offset restores the live segmentation-variant mechanism
12874                // on the current FA path: its tile grid starts at the chunk/call boundary.
12875                // SWA: trim the view to the oldest key any query in this chunk can reach —
12876                // ALIGNED DOWN to the FA tile size (BK=32). The raw trim is chunk-DEPENDENT
12877                // (off = base_len-(win-1), and base_len is a chunk boundary), and the FA
12878                // kernel's online-softmax recurrence groups keys into BK tiles relative to
12879                // the VIEW START — so an unaligned off regroups the same absolute keys into
12880                // different tiles at different chunk sizes = different (m,l) rounding =
12881                // chunk-dependent bits. chunkinv35 caught exactly this on the first FA-arm
12882                // battery (first_div == chunk size; raw/leverA-gates-20260807T135541Z.log).
12883                // Aligning off to 32 pins tiles to ABSOLUTE key positions for every chunk
12884                // size; the <=31 extra leading keys are older than EVERY query's window
12885                // (all queries sit at >= base_len, so keys < base_len-(win-1) are masked
12886                // for all of them) and a fully-masked key is an exact-0.0 no-op in both
12887                // kernels (NEG_INF -> p=0.0; l+=0.0 and O+=0.0 are bitwise identity), so
12888                // the floor arm's bits do not move either (gated: G2f, battery 2).
12889                let t_kv = base_len + t - off;
12890                let physical = kvl.physical_rows(off, off + t_kv)?;
12891                let k_view = e.view_u8_range(
12892                    &kvl.k,
12893                    physical.start * kvl.k_tok_bytes,
12894                    physical.end * kvl.k_tok_bytes,
12895                );
12896                let v_view = e.view_u8_range(
12897                    &kvl.v,
12898                    physical.start * kvl.v_tok_bytes,
12899                    physical.end * kvl.v_tok_bytes,
12900                );
12901                // CHUNK-INVARIANT PREDICATE: `seq_end` (the request's absolute end position), not
12902                // this chunk's `t_kv`. See the doc note — keying on t_kv made P = c*floor(win/c)
12903                // rows take FA and the output a function of MEMRA_PRIME_CHUNK.
12904                // MEMRA_STEP35_SWA_TKV=1 is the ROLLBACK SEAM to the pre-fix arithmetic — BOTH
12905                // halves: this predicate AND the unaligned view offset above (`legacy_tkv`).
12906                // It is what gives the step35 chunkinv gate its canary teeth: chunk-VARIANT by
12907                // construction, so the invariance assertion MUST break under it (the seam whose
12908                // absence made the original canary inert — GAP 1 in research/step37-p2-20260806;
12909                // and a predicate-only seam went inert AGAIN under the FA default, battery 2's
12910                // CANARY UNEXPECTEDLY MATCHED — see the offset comment). Read per call, not
12911                // cached (probes flip it in-process). Never on in a measured default run.
12912                let swa_naive = if legacy_tkv {
12913                    t_kv > win
12914                } else {
12915                    seq_end > win
12916                };
12917                if swa && swa_naive {
12918                    // Windowed mask needed (see the doc note). DEFAULT since lane/pp-prefill
12919                    // 2026-08-07: the windowed hd128 FA stamp (`fa_prefill_view_ws_w_hd128`) —
12920                    // the anatomy profile measured the f32 floor at 565 ms/layer on a pp4096
12921                    // (41% of the whole prime) while the unwindowed hd128 FA family did the
12922                    // strictly harder causal-4096 in 3.3 ms. NOTE t_kv can be <= win here
12923                    // (a chunk-0 view on a small chunk); the windowed kernel handles that
12924                    // identically to the unwindowed one modulo the mask, which is the point.
12925                    // NEW NUMERIC CLASS vs the floor (bf16-MMA online softmax vs f32 serial),
12926                    // selected on `seq_end` like every arm here, so the class is uniform for
12927                    // the whole request at every MEMRA_PRIME_CHUNK — chunkinv holds by the
12928                    // same construction as the chunkfix. MEMRA_STEP35_SWA_FA=0 = rollback to
12929                    // the f32 floor (the previous numeric config, kept as the A/B seam).
12930                    if std::env::var("MEMRA_STEP35_SWA_FA").as_deref() == Ok("0") {
12931                        e.sdpa_naive_w_quantized_view(
12932                            &q,
12933                            &k_view,
12934                            &v_view,
12935                            &mut attn,
12936                            hd,
12937                            nh,
12938                            nkv,
12939                            t,
12940                            t_kv,
12941                            scale,
12942                            true,
12943                            win,
12944                            kvl.k_tok_bytes,
12945                            kvl.v_tok_bytes,
12946                        )?;
12947                    } else {
12948                        e.fa_prefill_view_ws_w_hd128(
12949                            &q,
12950                            &k_view,
12951                            &v_view,
12952                            &mut attn,
12953                            hd,
12954                            nh,
12955                            nkv,
12956                            t,
12957                            t_kv,
12958                            scale,
12959                            true,
12960                            win,
12961                            kvl.k_tok_bytes,
12962                            kvl.v_tok_bytes,
12963                        )?;
12964                    }
12965                } else if std::env::var("MEMRA_NOFA").is_ok() {
12966                    e.sdpa_naive_quantized_view(
12967                        &q,
12968                        &k_view,
12969                        &v_view,
12970                        &mut attn,
12971                        hd,
12972                        nh,
12973                        nkv,
12974                        t,
12975                        t_kv,
12976                        scale,
12977                        true,
12978                        kvl.k_tok_bytes,
12979                        kvl.v_tok_bytes,
12980                    )?;
12981                } else {
12982                    // seq_end <= win (or a full-attn layer): no query in the WHOLE request can
12983                    // reach past the window, so the window mask is a no-op under causal and every
12984                    // chunk rides the hd128 dequant-once FA prefill — one arm for the whole
12985                    // request either way, which is what makes the chunk size arithmetic-free.
12986                    e.fa_prefill_view_ws(
12987                        &q,
12988                        &k_view,
12989                        &v_view,
12990                        &mut attn,
12991                        hd,
12992                        nh,
12993                        nkv,
12994                        t,
12995                        t_kv,
12996                        scale,
12997                        true,
12998                        kvl.k_tok_bytes,
12999                        kvl.v_tok_bytes,
13000                        crate::Engine::kv_fp8_on(),
13001                    )?;
13002                }
13003            }
13004            None => {
13005                // Cacheless prefill (forward/forward_last/t2probe) is MONOLITHIC — there is no
13006                // chunk loop on this path, so seq_end == t and the predicate is unchanged. The
13007                // assert pins that: if a chunked cacheless prefill is ever added, it must thread
13008                // seq_end here too or it re-opens the same door.
13009                debug_assert_eq!(
13010                    seq_end, t,
13011                    "step35 cacheless prefill is monolithic (seq_end == t)"
13012                );
13013                if swa && seq_end > win {
13014                    e.sdpa_naive_w(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
13015                } else if std::env::var("MEMRA_NOFA").is_ok() {
13016                    e.sdpa_naive(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
13017                } else {
13018                    e.fa_prefill(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
13019                }
13020            }
13021        }
13022
13023        // HEAD-WISE GATE (step35.cpp:267-285): gate = attn_gate(cur) -> [t, n_head_l];
13024        // attn *= sigmoid(gate) broadcast over head_dim. BEFORE wo.
13025        let gw = fa
13026            .attn_gate
13027            .as_ref()
13028            .ok_or("step35 layer is missing attn_gate.weight (head-wise attention gate)")?;
13029        let gt_owned = if gt_pre.is_none() {
13030            Some(e.matmul(
13031                gw,
13032                hg.ok_or("step35 attention needs hg when gt_pre is absent")?,
13033                t,
13034            )?)
13035        } else {
13036            None
13037        };
13038        let gt = gt_pre.or(gt_owned.as_ref()).unwrap();
13039        let mut ag = e.uninit(t * nh * hd)?;
13040        e.attn_head_gate(&attn, gt, &mut ag, None, hd, nh, t)?;
13041        Ok(ag)
13042    }
13043
13044    /// step35 PREFILL mixer, no cache side effect (the `full_attn` contract: `forward`,
13045    /// `forward_last`, t2probe). Post-`wo`.
13046    pub(crate) fn step35_attn(
13047        &self,
13048        e: &Engine,
13049        fa: &FullAttnLayer,
13050        h: &CudaSlice<f32>,
13051        pos_d: &CudaSlice<i32>,
13052        t: usize,
13053        il: usize,
13054    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13055        let g3 = e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?;
13056        // Monolithic path: the request ends at t (no chunk loop). See step35_attn_pre_wo's note.
13057        let ag = self.step35_attn_pre_wo(e, fa, g3, Some(h), None, pos_d, t, None, il, t)?;
13058        Ok(e.matmul(&fa.wo, &ag, t)?)
13059    }
13060
13061    /// step35 PRIME mixer (the `full_attn_prime` contract: append this chunk's K/V into the
13062    /// resident quantized cache, attend through the cache view). Post-`wo`.
13063    ///
13064    /// `seq_end` = the ABSOLUTE end position of the whole prime request (chunk-invariant); it
13065    /// selects the SWA arm. See `step35_attn_pre_wo`'s doc note for why it cannot be this chunk's
13066    /// own extent.
13067    #[allow(clippy::too_many_arguments)]
13068    pub(crate) fn step35_attn_prime(
13069        &self,
13070        e: &Engine,
13071        fa: &FullAttnLayer,
13072        h: &CudaSlice<f32>,
13073        hx: Option<&CudaSlice<u8>>,
13074        pos_d: &CudaSlice<i32>,
13075        t: usize,
13076        cache: &mut Cache,
13077        il: usize,
13078        seq_end: usize,
13079    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13080        let g3 = match hx {
13081            Some(xh) => e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], h, xh, t)?,
13082            None => e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?,
13083        };
13084        let ag =
13085            self.step35_attn_pre_wo(e, fa, g3, Some(h), None, pos_d, t, Some(cache), il, seq_end)?;
13086        Ok(e.matmul(&fa.wo, &ag, t)?)
13087    }
13088
13089    /// step35 T=1 decode attention (post-`wo`, matching `full_attn_decode_pre`'s contract).
13090    /// `pre_q` = the attn-input norm's q8_1 pair when the caller took the norm-fusion lever
13091    /// (then `h` is a zero-length placeholder and EVERY projection here — including the gate —
13092    /// must be on the q8_1 fast path; `mixer_in_q8_1_fast` enforces that for step35 by also
13093    /// requiring `attn_gate`).
13094    ///
13095    /// SWA decode is a token-aligned VIEW OFFSET into the quantized cache (the gemma4 R6
13096    /// pattern): keys carry absolute rope and the mask is purely positional, so the single
13097    /// query at `len-1` attending the last `win` rows IS the windowed result — no mask kernel.
13098    #[allow(clippy::too_many_arguments)]
13099    pub(crate) fn step35_decode_attn(
13100        &self,
13101        e: &Engine,
13102        fa: &FullAttnLayer,
13103        il: usize,
13104        h: &CudaSlice<f32>,
13105        pre_q: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
13106        pos_d: &CudaSlice<i32>,
13107        cache: &mut Cache,
13108    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13109        let geometry = self.step35_geom(il);
13110        let hd = geometry.head_dim_k as usize;
13111        let nkv = geometry.n_head_kv as usize;
13112        let nh = geometry.n_head as usize;
13113        let rbase = geometry.rope_base;
13114        let scale = geometry.attention_scale();
13115        let swa = geometry.window.is_some();
13116        let eps = self.cfg.rms_eps;
13117        let win = geometry.window.unwrap_or(0) as usize;
13118        let n_rot = geometry.n_rot as usize;
13119        let n_embd = self.cfg.n_embd as usize;
13120        let gw = fa
13121            .attn_gate
13122            .as_ref()
13123            .ok_or("step35 layer is missing attn_gate.weight (head-wise attention gate)")?;
13124
13125        let (q0, k0, v0, gt) = match pre_q {
13126            Some((hq, hdq)) => {
13127                debug_assert!(
13128                    e.uses_q8_1_fast(gw),
13129                    "step35 pre-quantized decode requires attn_gate on the q8_1 fast path \
13130                     (h is a zero-length placeholder here) — see mixer_in_q8_1_fast"
13131                );
13132                let (a, b, c) = match e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hdq)? {
13133                    Some(t3) => t3,
13134                    None => (
13135                        e.matmul_pre(&fa.wq, hq, hdq, h, 1)?,
13136                        e.matmul_pre(&fa.wk, hq, hdq, h, 1)?,
13137                        e.matmul_pre(&fa.wv, hq, hdq, h, 1)?,
13138                    ),
13139                };
13140                let gt = e.matmul_pre(gw, hq, hdq, h, 1)?;
13141                (a, b, c, gt)
13142            }
13143            None => {
13144                if e.uses_q8_1_fast(&fa.wq)
13145                    && e.uses_q8_1_fast(&fa.wk)
13146                    && e.uses_q8_1_fast(&fa.wv)
13147                    && e.uses_q8_1_fast(gw)
13148                {
13149                    let (hq, hdq) = e.quantize_q8_1(h, 1, n_embd)?;
13150                    let (a, b, c) = match e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, &hq, &hdq)? {
13151                        Some(t3) => t3,
13152                        None => (
13153                            e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
13154                            e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
13155                            e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?,
13156                        ),
13157                    };
13158                    let gt = e.matmul_pre(gw, &hq, &hdq, h, 1)?;
13159                    (a, b, c, gt)
13160                } else {
13161                    (
13162                        e.matmul(&fa.wq, h, 1)?,
13163                        e.matmul(&fa.wk, h, 1)?,
13164                        e.matmul(&fa.wv, h, 1)?,
13165                        e.matmul(gw, h, 1)?,
13166                    )
13167                }
13168            }
13169        };
13170
13171        let mut q = e.uninit(nh * hd)?;
13172        e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, nh, eps)?;
13173        let mut k = e.uninit(nkv * hd)?;
13174        e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, nkv, eps)?;
13175        let ff = if swa {
13176            None
13177        } else {
13178            self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
13179        };
13180        #[cfg(debug_assertions)]
13181        if let Some(ff) = ff {
13182            crate::debug_assert_tensor_stream_device(
13183                ff,
13184                &e.stream(),
13185                "step35_decode_attn.rope_freqs",
13186            );
13187        }
13188        e.rope_neox2(&mut q, &mut k, pos_d, hd, n_rot, nh, nkv, 1, rbase, 1.0, ff)?;
13189
13190        if std::env::var("MEMRA_NOFA").is_ok() {
13191            return Err(
13192                "MEMRA_NOFA (naive f32 SDPA) is incompatible with the quantized KV \
13193                        cache; unset MEMRA_NOFA to use fa_decode"
13194                    .into(),
13195            );
13196        }
13197        let kvl = cache.kv[il].as_mut().unwrap();
13198        let next_len = kvl.len + 1;
13199        let (off, t_kv) = if swa && next_len > win {
13200            (next_len - win, win)
13201        } else {
13202            (0, next_len)
13203        };
13204        let write_row = e.prepare_kv_append(kvl, off & !31usize, 1)?;
13205        e.append_kv_quantized(
13206            &k,
13207            &v0,
13208            &mut kvl.k,
13209            &mut kvl.v,
13210            write_row,
13211            kvl.kv_dim_k,
13212            kvl.kv_dim_v,
13213            kvl.k_tok_bytes,
13214            kvl.v_tok_bytes,
13215            crate::Engine::kv_fp8_on(),
13216        )?;
13217        kvl.len = next_len;
13218        let physical = kvl.physical_rows(off, off + t_kv)?;
13219        let k_view = e.view_u8_range(
13220            &kvl.k,
13221            physical.start * kvl.k_tok_bytes,
13222            physical.end * kvl.k_tok_bytes,
13223        );
13224        let v_view = e.view_u8_range(
13225            &kvl.v,
13226            physical.start * kvl.v_tok_bytes,
13227            physical.end * kvl.v_tok_bytes,
13228        );
13229        let mut attn = e.uninit(nh * hd)?;
13230        e.fa_decode_kvmod(
13231            &q,
13232            &k_view,
13233            &v_view,
13234            &mut attn,
13235            hd,
13236            nh,
13237            nkv,
13238            t_kv,
13239            scale,
13240            kvl.k_tok_bytes,
13241            kvl.v_tok_bytes,
13242            crate::Engine::kv_fp8_on(),
13243        )?;
13244
13245        let mut ag = e.uninit(nh * hd)?;
13246        e.attn_head_gate(&attn, &gt, &mut ag, None, hd, nh, 1)?;
13247        Ok(e.matmul(&fa.wo, &ag, 1)?)
13248    }
13249}
13250
13251// ===================================================================================== //
13252//  gemma-4 E4B (per-layer embeddings + KV-sharing) — FIRST-LIGHT forward.               //
13253//  Dedicated simple path (Stage-B matmuls, per-row causal attention over the quantized  //
13254//  cache) so the tuned 26B/31B paths stay untouched. Arch: research/gemma4-bringup/     //
13255//  e4b-arch-map.md; llama reference: src/models/gemma4.cpp (E4B arms). Wired: forward   //
13256//  (prefill logits), prime (tokenwise-equivalent batched trunk), eager decode_step_h.   //
13257//  NOT wired: dc/graph serving arms, verify/spec, chunked prime (HANDOVER-E4B.md).      //
13258// ===================================================================================== //
13259impl HybridModel {
13260    pub fn is_gemma4_e4b(&self) -> bool {
13261        self.gemma4_aux.as_ref().is_some_and(|a| a.e4b.is_some())
13262    }
13263
13264    /// E4B per-layer geometry: nh/nkv derive from the layer's OWN tensor shapes (the E4B GGUF
13265    /// ships a scalar head_count_kv; swa q 8x256 / kv 2x256, global q 4x512 / kv 1x512).
13266    /// KV-shared layers report the SHARE TARGET's kv count (their wk IS the target's tensor).
13267    fn gemma4_e4b_geom(&self, il: usize) -> (usize, usize, usize, f32, f32, bool) {
13268        let g = self.cfg.gemma4.as_ref().unwrap();
13269        let swa = g.swa_pattern[il];
13270        let hd = if swa {
13271            g.key_length_swa
13272        } else {
13273            g.key_length_global
13274        } as usize;
13275        let Mixer::Full(fa) = &self.layers[il].mixer else {
13276            panic!("e4b layer {il} not full-attn")
13277        };
13278        let nh = fa.wq.out_features() / hd;
13279        let nkv = fa.wk.out_features() / hd;
13280        (
13281            hd,
13282            nkv,
13283            nh,
13284            if swa {
13285                g.rope_base_swa
13286            } else {
13287                g.rope_base_global
13288            },
13289            1.0,
13290            swa,
13291        )
13292    }
13293
13294    /// E4B KV-share target: Some(target) for the trailing shared layers, None = own cache.
13295    fn gemma4_e4b_kv_target(&self, il: usize) -> Option<usize> {
13296        self.layers[il]
13297            .gemma4
13298            .as_ref()
13299            .and_then(|b| b.e4b.as_ref())
13300            .and_then(|e4| e4.kv_share.map(|t| t as usize))
13301    }
13302
13303    /// E4B prologue: inp_pl [t][n_layer][n_epl] =
13304    ///   ( gather(per_layer_tok_embd, tok)*sqrt(n_epl)
13305    ///   + rms_norm(model_proj . x_scaled * 1/sqrt(n_embd), proj_norm) ) * 1/sqrt(2)
13306    /// (llama gemma4.cpp build_inp_per_layer + project_per_layer_inputs, exact order).
13307    fn gemma4_e4b_inp_pl(
13308        &self,
13309        e: &Engine,
13310        tokens: &[u32],
13311        x_scaled: &CudaSlice<f32>,
13312        t: usize,
13313    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13314        let tok_d = e.stream().clone_htod(&tokens.to_vec())?;
13315        self.gemma4_e4b_inp_pl_dev(e, &tok_d, x_scaled, t)
13316    }
13317
13318    /// Device-token prologue core (dc arm shares it: token ids never touch the host).
13319    fn gemma4_e4b_inp_pl_dev(
13320        &self,
13321        e: &Engine,
13322        tok_d: &CudaSlice<u32>,
13323        x_scaled: &CudaSlice<f32>,
13324        t: usize,
13325    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13326        let aux = self.gemma4_aux.as_ref().unwrap();
13327        let m = aux.e4b.as_ref().unwrap();
13328        let n_embd = self.cfg.n_embd as usize;
13329        let n_layer = self.layers.len();
13330        let width = m.n_epl * n_layer;
13331        let tbl = m.tok_tbl_gpu.get_or_init(|| {
13332            e.upload_u8(&m.tok_embd_bytes)
13333                .expect("e4b per-layer token table upload")
13334        });
13335        let mut a =
13336            e.embed_gather_device_td(tbl, tok_d, t, width, m.tok_embd_qt, m.tok_embd_row_bytes)?;
13337        e.scale_inplace(&mut a, (m.n_epl as f32).sqrt(), t * width)?;
13338        let mut p = e.matmul(&m.model_proj, x_scaled, t)?;
13339        e.scale_inplace(&mut p, 1.0 / (n_embd as f32).sqrt(), t * width)?;
13340        let mut pn = e.uninit(t * width)?;
13341        e.rms_norm(
13342            &p,
13343            m.proj_norm.float_data(),
13344            &mut pn,
13345            m.n_epl,
13346            t * n_layer,
13347            self.cfg.rms_eps,
13348        )?;
13349        let mut out = e.uninit(t * width)?;
13350        e.add_scale(&a, &pn, 1.0 / 2f32.sqrt(), &mut out, t * width)?;
13351        Ok(out)
13352    }
13353
13354    /// E4B attention (t-wide, causal, per-row fa over the QUANTIZED cache — first-light
13355    /// correctness path; the fa rows arms come later). Own-KV layers project+norm+rope k/v
13356    /// and append t rows; KV-shared layers are Q-only over the target layer's cache (which
13357    /// already holds this forward's rows — the target runs earlier in the stack).
13358    #[allow(clippy::too_many_arguments)]
13359    fn gemma4_e4b_attn(
13360        &self,
13361        e: &Engine,
13362        il: usize,
13363        hq: &CudaSlice<i8>,
13364        hdq: &CudaSlice<f32>,
13365        pos_d: &CudaSlice<i32>,
13366        t: usize,
13367        cache: &mut Cache,
13368        dc_bucket: Option<usize>,
13369    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13370        let (hd, nkv, nh, base, scale, swa) = self.gemma4_e4b_geom(il);
13371        let eps = self.cfg.rms_eps;
13372        let aux = self.gemma4_aux.as_ref().unwrap();
13373        let ones = aux.ones(e);
13374        #[cfg(debug_assertions)]
13375        crate::debug_assert_tensor_stream_device(ones, &e.stream(), "gemma4_e4b_attn.ones");
13376        let Mixer::Full(fa) = &self.layers[il].mixer else {
13377            unreachable!()
13378        };
13379        // pre-quantized layer input (E4B fusion port, 2026-07-12): ONE norm+quant feeds
13380        // wq/wk/wv via matmul_pre — the first-light path quantized the same h three times
13381        // (the profile's 342 quantize_q8_1/token; glue = 26% of the 6.1ms token).
13382        let h0 = e.zeros(0)?;
13383        let h = &h0;
13384
13385        let ff = if swa {
13386            None
13387        } else {
13388            Some(
13389                aux.rope_freqs(e)
13390                    .expect("e4b global rope needs rope_freqs.weight"),
13391            )
13392        };
13393        #[cfg(debug_assertions)]
13394        if let Some(ff) = ff {
13395            crate::debug_assert_tensor_stream_device(ff, &e.stream(), "gemma4_e4b_attn.rope_freqs");
13396        }
13397        let share = self.gemma4_e4b_kv_target(il);
13398        // Own-KV arms keep the f32 (k, v) alive for the prime fa arm below.
13399        let mut kv_f32: Option<(CudaSlice<f32>, CudaSlice<f32>)> = None;
13400        let mut q;
13401        if let Some(_tgt) = share {
13402            let q0 = e.matmul_pre(&fa.wq, hq, hdq, h, t)?;
13403            q = e.uninit(t * nh * hd)?;
13404            // Q-only through the same fused norm+rope kernel (rk = 0: the k/v segments are
13405            // empty; q0 stands in for the unused k/v pointers).
13406            let mut kdummy = e.uninit(1)?;
13407            let mut vdummy = e.uninit(1)?;
13408            e.rms_norm_qkv_rope(
13409                &q0,
13410                &q0,
13411                &q0,
13412                fa.q_norm.float_data(),
13413                fa.q_norm.float_data(),
13414                ones,
13415                &mut q,
13416                &mut kdummy,
13417                &mut vdummy,
13418                hd,
13419                self.gemma4_rope_dims(il),
13420                nh * t,
13421                0,
13422                pos_d,
13423                nh,
13424                1,
13425                base,
13426                1.0,
13427                ff,
13428                eps,
13429            )?;
13430        } else {
13431            // wave-4b: ONE concat matvec (wq|wk|wv) at t == 1 when the cat tensor exists;
13432            // else the fused3 grid launch; else per-matvec. The cat output is contiguous
13433            // q|k|v rows — the cat norm+rope twin consumes it directly.
13434            let e4bits = self.layers[il].gemma4.as_ref().and_then(|g| g.e4b.as_ref());
13435            let cat = e4bits.and_then(|e4| e4.qkv_cat.as_ref());
13436            q = e.uninit(t * nh * hd)?;
13437            let mut k = e.uninit(t * nkv * hd)?;
13438            let mut v = e.uninit(t * nkv * hd)?;
13439            if t == 1 && cat.is_some() {
13440                let qkv0 = e.matmul_pre(cat.unwrap(), hq, hdq, h, 1)?;
13441                e.rms_norm_qkv_rope_cat(
13442                    &qkv0,
13443                    fa.q_norm.float_data(),
13444                    fa.k_norm.float_data(),
13445                    ones,
13446                    &mut q,
13447                    &mut k,
13448                    &mut v,
13449                    hd,
13450                    self.gemma4_rope_dims(il),
13451                    nh,
13452                    nkv,
13453                    pos_d,
13454                    nh,
13455                    nkv,
13456                    base,
13457                    1.0,
13458                    ff,
13459                    eps,
13460                )?;
13461            } else {
13462                let (q0, k0, v0) = match if t == 1 {
13463                    e.matmul_q4_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hdq)?
13464                } else {
13465                    // E4B verify f3 port (2026-07-14): the 31B segmented-grid batched qkv
13466                    // on E4B's real-V triple — same MEMRA_F2B seam, bit-identical per row.
13467                    static F2B_QKV: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13468                    if *F2B_QKV.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0")) {
13469                        e.matmul_q4_fused3_batched(&fa.wq, &fa.wk, &fa.wv, hq, hdq, t)?
13470                    } else {
13471                        None
13472                    }
13473                } {
13474                    Some(triple) => triple,
13475                    None => (
13476                        e.matmul_pre(&fa.wq, hq, hdq, h, t)?,
13477                        e.matmul_pre(&fa.wk, hq, hdq, h, t)?,
13478                        e.matmul_pre(&fa.wv, hq, hdq, h, t)?,
13479                    ), // E4B: real v (K != V)
13480                };
13481                // wave-3 fold: q/k/v norms + q/k rope in ONE launch (rope math verbatim on
13482                // the normed rows; V ones-rms, never roped).
13483                e.rms_norm_qkv_rope(
13484                    &q0,
13485                    &k0,
13486                    &v0,
13487                    fa.q_norm.float_data(),
13488                    fa.k_norm.float_data(),
13489                    ones,
13490                    &mut q,
13491                    &mut k,
13492                    &mut v,
13493                    hd,
13494                    self.gemma4_rope_dims(il),
13495                    nh * t,
13496                    nkv * t,
13497                    pos_d,
13498                    nh,
13499                    nkv,
13500                    base,
13501                    1.0,
13502                    ff,
13503                    eps,
13504                )?;
13505            }
13506            let kvl = cache.kv[il].as_mut().unwrap();
13507            // class flag must match the cache dims (the g-threading sweep hardcoded `false`
13508            // here and the wkv default corrupted E4B: q8_0 bytes into an e4m3 cache — the
13509            // degenerate tok-0 stream, 2026-07-12).
13510            let cls = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
13511            if dc_bucket.is_some() {
13512                // DC arm (graph serving): append at the len_d slot, advance the counter
13513                // in-stream — replay-correct, no host len in the launch args. Host mirrors
13514                // are NOT touched here (the replay loop owns them; a bump at capture-record
13515                // time would double-count the capture iteration).
13516                debug_assert!(t == 1);
13517                // wave 5c: append + len_d inc fused (one launch; single-block ordering).
13518                e.append_kv_quantized_row_dc_inc(
13519                    &k,
13520                    &v,
13521                    &mut kvl.k,
13522                    &mut kvl.v,
13523                    &mut kvl.len_d,
13524                    kvl.kv_dim_k,
13525                    kvl.kv_dim_v,
13526                    kvl.k_tok_bytes,
13527                    kvl.v_tok_bytes,
13528                    cls,
13529                )?;
13530            } else {
13531                e.append_kv_quantized_rows(
13532                    &k,
13533                    &v,
13534                    &mut kvl.k,
13535                    &mut kvl.v,
13536                    kvl.len,
13537                    t,
13538                    kvl.kv_dim_k,
13539                    kvl.kv_dim_v,
13540                    kvl.k_tok_bytes,
13541                    kvl.v_tok_bytes,
13542                    cls,
13543                )?;
13544                kvl.len += t;
13545            }
13546            kv_f32 = Some((k, v));
13547        }
13548        // attention: per-row causal fa over the (own or target) quantized cache. The cache
13549        // already contains this forward's rows in both arms; row i attends [.., base+i].
13550        let kvl_idx = share.unwrap_or(il);
13551        let kvl = cache.kv[kvl_idx].as_ref().unwrap();
13552        let base_len = kvl.len - t; // pre-append length (target appended this forward too)
13553        let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
13554        let mut attn = e.uninit(t * nh * hd)?;
13555        // PRIME FA ARMS (perf lane 3, completed 2026-07-31 — the H100 board pinned E4B
13556        // prefill at 482 tok/s: the per-row loop ran the DECODE kernel once per token per
13557        // layer). Fresh-prompt (base_len == 0) batched attention, one launch per layer:
13558        //   - own-KV hd256 under the window: f32 fa_prefill (full causal exact — 26B pattern)
13559        //   - own-KV hd256 above the window (swa): fa_prefill_w windowed twin (12B pattern)
13560        //   - own-KV hd512 globals: fa_prefill_hd512 (12B/31B globals pattern)
13561        //   - KV-shared layers (no f32 k/v): fa_prefill_view over the target's QUANTIZED
13562        //     rows (the T=K verify kernel; the target appended this forward's rows already).
13563        //     swa-shared above the window keeps the per-row loop (no windowed view twin).
13564        // Same numeric class split as the 26B prime (f32/batched prefill vs per-row
13565        // quantized decode); run-gen argmax + chat gates arbitrate. MEMRA_NOFA=1 reverts all.
13566        if t > 1 && base_len == 0 && std::env::var("MEMRA_NOFA").is_err() {
13567            if let Some((kf, vf)) = &kv_f32 {
13568                if hd == 256 && t <= win {
13569                    e.fa_prefill(&q, kf, vf, &mut attn, hd, nh, nkv, t, t, scale, true)?;
13570                    return Ok(e.matmul(&fa.wo, &attn, t)?);
13571                }
13572                if hd == 256 && swa && t > win {
13573                    e.fa_prefill_w(&q, kf, vf, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
13574                    return Ok(e.matmul(&fa.wo, &attn, t)?);
13575                }
13576                if hd == 512 && !swa {
13577                    e.fa_prefill_hd512(&q, kf, vf, &mut attn, hd, nh, nkv, t, t, scale, true)?;
13578                    return Ok(e.matmul(&fa.wo, &attn, t)?);
13579                }
13580            } else if share.is_some() {
13581                let g = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
13582                let k_view = e.view_u8(&kvl.k, kvl.k.len());
13583                let v_view = e.view_u8(&kvl.v, kvl.v.len());
13584                if hd == 256 && (!swa || t <= win) {
13585                    // inline-dequant quantized-view prefill (fa_prefill_q stamps 256/128)
13586                    e.fa_prefill_view(
13587                        &q,
13588                        &k_view,
13589                        &v_view,
13590                        &mut attn,
13591                        hd,
13592                        nh,
13593                        nkv,
13594                        t,
13595                        t,
13596                        scale,
13597                        true,
13598                        kvl.k_tok_bytes,
13599                        kvl.v_tok_bytes,
13600                        g,
13601                    )?;
13602                    return Ok(e.matmul(&fa.wo, &attn, t)?);
13603                }
13604                // remaining shared classes (swa above the window; hd512 globals): dequant
13605                // the target's t rows ONCE to f32, then the same f32 twins as own-KV.
13606                let kv_dim = nkv * hd;
13607                let mut kf = e.uninit(t * kv_dim)?;
13608                let mut vf = e.uninit(t * kv_dim)?;
13609                e.fa_dequant_kv_view_f32(
13610                    &k_view,
13611                    &v_view,
13612                    &mut kf,
13613                    &mut vf,
13614                    kv_dim,
13615                    kv_dim,
13616                    t,
13617                    kvl.k_tok_bytes,
13618                    kvl.v_tok_bytes,
13619                    g,
13620                )?;
13621                if hd == 512 {
13622                    e.fa_prefill_hd512(&q, &kf, &vf, &mut attn, hd, nh, nkv, t, t, scale, true)?;
13623                } else {
13624                    e.fa_prefill_w(&q, &kf, &vf, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
13625                }
13626                return Ok(e.matmul(&fa.wo, &attn, t)?);
13627            }
13628        }
13629        if let Some(bucket) = dc_bucket {
13630            // DC arm (t == 1, under-window regime — the generate gate enforces it): ONE
13631            // fa_decode_dc over the live counter. len_d already advanced past this token
13632            // (t_kv = len_d[0], the cache.rs contract). KV-shared layers read the target's
13633            // counter (advanced when the target ran earlier in the stack).
13634            assert!(t == 1);
13635            // hd512 globals: eager (kvmod) picks the SCALAR unified below the fa512 floor,
13636            // and under the window every live t_kv sits below it — cap the capture bucket
13637            // under the floor so fa_decode_dc bakes the same scalar symbol, or the graph
13638            // rides the dpl16 twin and its numeric class diverges from the dc-eager stream
13639            // (E4B-GRAPH-GATE 2/64, 2026-07-12).
13640            let bucket = if hd == 512 && win <= crate::fa512_min_tkv() {
13641                bucket.min(crate::fa512_min_tkv().saturating_sub(1))
13642            } else {
13643                bucket
13644            };
13645            let k_view = e.view_u8(&kvl.k, kvl.k.len());
13646            let v_view = e.view_u8(&kvl.v, kvl.v.len());
13647            let g = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
13648            // Weight prefetch (SOTA item 3, 2026-07-13): wo's decode plane prefetched into
13649            // L2 across the fa window — fa reads KV only, the weight DRAM lanes are idle
13650            // there (E4B valid window +0.65%: 196.8 vs 195.6). Value-free scheduling op,
13651            // captured into the dc graph like any other launch. Extending the cascade to
13652            // the ffn gate/up planes measured NEGATIVE (193.9 vs 195.8 — 29MB/layer floods
13653            // the fill path and evicts still-hot lines); wo-only is the shipped shape.
13654            // MEMRA_WPF=0 rollback seam.
13655            if crate::Engine::wpf_level() >= 1 {
13656                e.prefetch_weight_l2(&fa.wo)?;
13657            }
13658            // wave 5b: the combine emits the wo input q8 pair directly — the standalone
13659            // quantize launch + the f32 attn round-trip fold away (t=1 fast path only).
13660            if e.uses_q8_1_fast(&fa.wo) {
13661                let mut oq = e.alloc_i8_uninit(nh * hd)?;
13662                let mut od = e.zeros(nh * hd / 32)?;
13663                e.fa_decode_dc_q8(
13664                    &q,
13665                    &k_view,
13666                    &v_view,
13667                    &mut attn,
13668                    hd,
13669                    nh,
13670                    nkv,
13671                    &kvl.len_d,
13672                    bucket,
13673                    scale,
13674                    kvl.k_tok_bytes,
13675                    kvl.v_tok_bytes,
13676                    g,
13677                    Some((&mut oq, &mut od)),
13678                )?;
13679                return Ok(e.matmul_pre(&fa.wo, &oq, &od, &attn, t)?);
13680            }
13681            e.fa_decode_dc(
13682                &q,
13683                &k_view,
13684                &v_view,
13685                &mut attn,
13686                hd,
13687                nh,
13688                nkv,
13689                &kvl.len_d,
13690                bucket,
13691                scale,
13692                kvl.k_tok_bytes,
13693                kvl.v_tok_bytes,
13694                g,
13695            )?;
13696            return Ok(e.matmul(&fa.wo, &attn, t)?);
13697        }
13698        for i in 0..t {
13699            let avail = base_len + i + 1;
13700            let (off_tok, t_kv) = if swa && avail > win {
13701                (avail - win, win)
13702            } else {
13703                (0, avail)
13704            };
13705            let k_view = e.view_u8_range(
13706                &kvl.k,
13707                off_tok * kvl.k_tok_bytes,
13708                (off_tok + t_kv) * kvl.k_tok_bytes,
13709            );
13710            let v_view = e.view_u8_range(
13711                &kvl.v,
13712                off_tok * kvl.v_tok_bytes,
13713                (off_tok + t_kv) * kvl.v_tok_bytes,
13714            );
13715            let qv = e.view(&q, t * nh * hd);
13716            let q_row = qv.slice(i * nh * hd..(i + 1) * nh * hd);
13717            let mut q_one = e.uninit(nh * hd)?;
13718            e.copy_view_into(&mut q_one, 0, &q_row, nh * hd)?;
13719            let mut a_one = e.uninit(nh * hd)?;
13720            // read class MUST match the append class (globals are e4m3 under gkv): the
13721            // swa-only flag decoded global-layer e4m3 bytes as q8_0/q5_1 — attention over
13722            // garbage on EVERY E4B global layer, the cross-mode maxdiff-30 root (2026-07-12).
13723            e.fa_decode_kvmod(
13724                &q_one,
13725                &k_view,
13726                &v_view,
13727                &mut a_one,
13728                hd,
13729                nh,
13730                nkv,
13731                t_kv,
13732                scale,
13733                kvl.k_tok_bytes,
13734                kvl.v_tok_bytes,
13735                (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on()),
13736            )?;
13737            e.copy_into(&mut attn, i * nh * hd, &a_one, nh * hd)?;
13738        }
13739        Ok(e.matmul(&fa.wo, &attn, t)?)
13740    }
13741
13742    /// E4B trunk: embed -> prologue -> layers (attn + dense ffn via the 31B tail_core + the
13743    /// per-layer-embedding tail + layer scale) -> output_norm. Returns (softcapped logits
13744    /// device [t, n_vocab], pre-output_norm hidden [t, n_embd]). Appends t rows per own-KV
13745    /// layer; does NOT advance cache.pos (caller owns pos).
13746    fn gemma4_e4b_trunk(
13747        &self,
13748        e: &Engine,
13749        tokens: &[u32],
13750        pos0: usize,
13751        cache: &mut Cache,
13752        head_last: bool,
13753    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13754        let n_embd = self.cfg.n_embd as usize;
13755        let t = tokens.len();
13756        let pos: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
13757        let pos_d = e.htod_i32(&pos)?;
13758        let mut x = e.htod(&self.embd.gather(n_embd, tokens))?;
13759        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
13760        let inp_pl = self.gemma4_e4b_inp_pl(e, tokens, &x, t)?;
13761        self.gemma4_e4b_trunk_core(e, x, inp_pl, &pos_d, t, cache, None, true, head_last)
13762    }
13763
13764    /// Layer stack + head over prebuilt (x_scaled, inp_pl, device pos) — everything below
13765    /// here is device-driven, so the dc arm shares it verbatim (stream identity with the
13766    /// eager chain by construction: SAME functions, not twins).
13767    fn gemma4_e4b_trunk_core(
13768        &self,
13769        e: &Engine,
13770        x_in: CudaSlice<f32>,
13771        inp_pl: CudaSlice<f32>,
13772        pos_d: &CudaSlice<i32>,
13773        t: usize,
13774        cache: &mut Cache,
13775        dc_bucket: Option<usize>,
13776        cap_logits: bool,
13777        head_last: bool,
13778    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13779        let n_embd = self.cfg.n_embd as usize;
13780        let eps = self.cfg.rms_eps;
13781        let n_layer = self.layers.len();
13782        let mut x = x_in;
13783        let aux_e4b = self.gemma4_aux.as_ref().unwrap().e4b.as_ref().unwrap();
13784        let n_epl = aux_e4b.n_epl;
13785
13786        // cross-layer fusion (2026-07-12 port of the 26B/31B trunk structure): each layer's
13787        // closing add_scale also EMITS the next layer's attn-normed input pre-quantized
13788        // q8_1 (add_scale_rms_norm_q8_1); the LAST layer emits through output_norm so the
13789        // head rides matmul_pre too. First layer's pair comes from a standalone fused
13790        // norm+quant.
13791        let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
13792        for il in 0..n_layer {
13793            let layer = &self.layers[il];
13794            let (hq, hdq) = match h_carry.take() {
13795                Some(p) => p,
13796                None => e.rms_norm_q8_1(&x, layer.attn_norm.float_data(), n_embd, t, eps)?,
13797            };
13798            let o = self.gemma4_e4b_attn(e, il, &hq, &hdq, pos_d, t, cache, dc_bucket)?;
13799            // dense ffn tail with the post-attn norm FOLDED into its entry (one launch for
13800            // rms(o, post_attn_norm) + residual add + ffn_norm — glue-fusion lane).
13801            let bits = layer.gemma4.as_ref().unwrap();
13802            let e4b = bits.e4b.as_ref().expect("e4b layer bits");
13803            // glue wave 5: the tail DEFERS its post_ffw norm — the FFN exit fuses
13804            // rms(f0, post_ffw) + residual add + q8 emit into ONE launch (rms_pre_add_q8_1),
13805            // killing the rms_norm + add_q8_1 pair per layer. T-GENERIC since 2026-07-13:
13806            // the fused single-phase reduction is NOT FP-order-identical to the unfused
13807            // rms_norm+add pair (the original "bit-identical chain" claim was FALSE — the
13808            // t==1 gate left the batched VERIFY on the unfused chain and split E4B verify
13809            // from decode by logit maxdiff ~0.45, the greedy tie-flip at depth: gate 135/256.
13810            // NOFUSE bisect: disabling ONLY this fusion -> verify maxdiff 0.000e0). With the
13811            // gate dropped, decode AND verify ride the same fused chain — parity by
13812            // construction, VERIFY-GATE 0.000e0.
13813            let fuse_exit = e.uses_q8_1_fast(&e4b.inp_gate);
13814            let (sn, attn_out) = self.gemma4_layer_tail_core_pn(
13815                e,
13816                layer,
13817                &o,
13818                &x,
13819                t,
13820                Some(layer.post_attn_norm.float_data()),
13821                fuse_exit,
13822            )?;
13823            let mut resid = e.uninit(t * n_embd)?;
13824            // per-layer-embedding tail: resid += rms(proj . (gelu(inp_gate . resid) * inp_pl[il]))
13825            // wave-2: the residual add emits q8_1 alongside — inp_gate rides matmul_pre.
13826            // (PLE one-block mega-fusion PROBED NEGATIVE 2026-07-13: argmax-correct but
13827            // 126 vs 189 tok/s — one SM pulling 0.74MB of weights loses to the multi-block
13828            // launch chain it replaced; jsonl row. Kernel deleted per doctrine.)
13829            let g = if fuse_exit {
13830                // sn here = RAW f0 (post_ffw deferred).
13831                let (rq, rd) = e.rms_pre_add_q8_1(
13832                    &sn,
13833                    bits.post_ffw_norm.float_data(),
13834                    &attn_out,
13835                    &mut resid,
13836                    n_embd,
13837                    t,
13838                    self.cfg.rms_eps,
13839                )?;
13840                e.matmul_pre(&e4b.inp_gate, &rq, &rd, &resid, t)?
13841            } else {
13842                e.add(&sn, &attn_out, &mut resid, t * n_embd)?;
13843                e.matmul(&e4b.inp_gate, &resid, t)?
13844            };
13845            let mut act = e.uninit(t * n_epl)?;
13846            let y = if t == 1 && e.uses_q8_1_fast(&e4b.proj) {
13847                let ipv = e.view(&inp_pl, n_epl * n_layer);
13848                let row = ipv.slice(il * n_epl..(il + 1) * n_epl);
13849                let (aq, ad) = e.gelu_tanh_mul_q8_1(&g, &row, &mut act, n_epl, 1)?;
13850                e.matmul_pre(&e4b.proj, &aq, &ad, &act, t)?
13851            } else {
13852                let mut inp_this = e.uninit(t * n_epl)?;
13853                e.copy_rows_strided(
13854                    &inp_pl,
13855                    &mut inp_this,
13856                    n_epl,
13857                    t,
13858                    n_epl * n_layer,
13859                    il * n_epl,
13860                )?;
13861                e.gelu_tanh_mul(&g, &inp_this, &mut act, t * n_epl)?;
13862                e.matmul(&e4b.proj, &act, t)?
13863            };
13864            // rms(y, post_norm) + (yn + resid)*layer_scale + next-layer norm+quant emit,
13865            // ONE launch (glue-fusion lane; last layer emits through output_norm).
13866            let next_norm = if il + 1 < n_layer {
13867                self.layers[il + 1].attn_norm.float_data()
13868            } else {
13869                self.output_norm.float_data()
13870            };
13871            let mut xn = e.uninit(t * n_embd)?;
13872            let pair = e.rms_pre_add_scale_rms_norm_q8_1(
13873                &y,
13874                e4b.post_norm.float_data(),
13875                &resid,
13876                bits.layer_scale,
13877                next_norm,
13878                &mut xn,
13879                n_embd,
13880                t,
13881                eps,
13882            )?;
13883            h_carry = Some(pair);
13884            x = xn;
13885        }
13886        // the head consumes the last layer's fused (output_norm) emit. head_last callers
13887        // (prime, last_only forward) need only the final row's logits — the all-T head is
13888        // t*n_vocab of discarded work (E4B prime: ~134ms GEMM + a 2.26GB dtoh kept 1 row).
13889        let (oq, odq) = h_carry.take().unwrap();
13890        let h0 = e.zeros(0)?;
13891        let hm = if head_last { 1 } else { t };
13892        let (hq, hd) = if head_last && t > 1 {
13893            let mut q1 = e.uninit_i8(n_embd)?;
13894            e.dtod_copy_view_i8(&oq.slice((t - 1) * n_embd..t * n_embd), &mut q1)?;
13895            let nb = n_embd / 32;
13896            let mut d1 = e.uninit(nb)?;
13897            e.dtod_copy_view(&odq.slice((t - 1) * nb..t * nb), &mut d1)?;
13898            (q1, d1)
13899        } else {
13900            (oq, odq)
13901        };
13902        let mut ld = e.matmul_pre(&self.output, &hq, &hd, &h0, hm)?;
13903        // softcap is strictly monotonic — greedy (argmax-only) consumers skip it, matching
13904        // the 26B/31B dc precedent (their dc head goes matmul -> argmax with no cap).
13905        // Logit-returning callers (host logits / spec prime) keep the capped emit.
13906        if cap_logits {
13907            let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
13908            e.softcap(&mut ld, cap, hm * self.output.out_features())?;
13909        }
13910        self.gemma4_suppress(e, &mut ld, hm)?; // mask both capped and argmax-only consumers
13911        Ok((ld, x))
13912    }
13913
13914    /// E4B batched VERIFY (device tokens, the spec round's t=K+1 step): t rows through the
13915    /// e4b trunk (per-row causal attention; own-KV layers append t rows host-len, KV-shared
13916    /// layers ride their targets), per-row device argmax + the POST-output_norm hidden
13917    /// stack (the drafter's h convention). Advances cache.pos/kvl.len by t — the spec
13918    /// round rolls back rejected rows (shared layers have no KvLayer, so the plain rewind
13919    /// covers exactly the layers that appended).
13920    pub fn gemma4_e4b_decode_step_t_am_dev(
13921        &self,
13922        e: &Engine,
13923        tok_d: &CudaSlice<u32>,
13924        t: usize,
13925        pos0: usize,
13926        cache: &mut Cache,
13927    ) -> Result<(CudaSlice<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13928        let n_embd = self.cfg.n_embd as usize;
13929        let eps = self.cfg.rms_eps;
13930        let pos: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
13931        let pos_d = e.htod_i32(&pos)?;
13932        let embd_gpu = self
13933            .embd_gpu
13934            .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload"));
13935        let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
13936        let mut x = e.embed_gather_device_td(embd_gpu, tok_d, t, n_embd, qt, rb)?;
13937        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
13938        let inp_pl = self.gemma4_e4b_inp_pl_dev(e, tok_d, &x, t)?;
13939        let (ld, xp) =
13940            self.gemma4_e4b_trunk_core(e, x, inp_pl, &pos_d, t, cache, None, true, false)?;
13941        // softcap is monotonic — the per-row argmax is invariant to it (the trunk's head
13942        // emit is already capped, matching the eager chain bit-for-bit).
13943        let n_vocab = self.output.out_features();
13944        let mut vam = e.stream().alloc_zeros::<u32>(t)?;
13945        for i in 0..t {
13946            e.argmax_token_device_col(&ld, i, n_vocab, &mut vam, i)?;
13947        }
13948        let mut hn = e.uninit(t * n_embd)?;
13949        e.rms_norm(&xp, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
13950        cache.pos += t;
13951        Ok((vam, hn))
13952    }
13953
13954    /// E4B verify + host logits + POST-output_norm hidden stack (the short-prompt spec
13955    /// prime path — mirror of `gemma4_decode_step_t_h`).
13956    pub(crate) fn gemma4_e4b_decode_step_t_h(
13957        &self,
13958        e: &Engine,
13959        tokens: &[u32],
13960        pos0: usize,
13961        cache: &mut Cache,
13962    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13963        let n_embd = self.cfg.n_embd as usize;
13964        let eps = self.cfg.rms_eps;
13965        let t = tokens.len();
13966        let (ld, xp) = self.gemma4_e4b_trunk(e, tokens, pos0, cache, false)?;
13967        let mut hn = e.uninit(t * n_embd)?;
13968        e.rms_norm(&xp, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
13969        cache.pos += t;
13970        Ok((e.dtoh(&ld)?, hn))
13971    }
13972
13973    /// E4B GRAPH-CAPTURABLE dc step: same trunk as the dc step but token_d is updated IN
13974    /// PLACE (self-feeding replay) and every launch arg is a device counter — pos from
13975    /// pos_d (inc'd in-stream), KV slots from len_d (advanced in-stream), attention from
13976    /// fa_decode_dc at `bucket`. Host mirrors (cache.pos / kvl.len) advance in the caller's
13977    /// replay loop. UNDER-WINDOW regime only (the caller gates pos + budget < window).
13978    pub fn gemma4_e4b_decode_step_dcg(
13979        &self,
13980        e: &Engine,
13981        token_d: &mut CudaSlice<u32>,
13982        pos_d: &mut CudaSlice<i32>,
13983        embd_gpu: &CudaSlice<u8>,
13984        embd_qt: i32,
13985        embd_rb: usize,
13986        cache: &mut Cache,
13987        n_vocab: usize,
13988        bucket: usize,
13989    ) -> Result<(), Box<dyn std::error::Error>> {
13990        let n_embd = self.cfg.n_embd as usize;
13991        let mut x = e.embed_gather_device(embd_gpu, token_d, n_embd, embd_qt, embd_rb)?;
13992        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
13993        let inp_pl = self.gemma4_e4b_inp_pl_dev(e, token_d, &x, 1)?;
13994        let (ld, _x) =
13995            self.gemma4_e4b_trunk_core(e, x, inp_pl, pos_d, 1, cache, Some(bucket), false, false)?;
13996        e.argmax_token_device_into(&ld, token_d, n_vocab)?;
13997        e.inc_seqlen(pos_d)?;
13998        Ok(())
13999    }
14000
14001    /// E4B DEVICE-COUNTER decode step (the dc serving arm): token id rides `token_d`, the
14002    /// greedy argmax lands in the returned device buffer — 4B/token host traffic. The layer
14003    /// stack is gemma4_e4b_trunk_core, i.e. the SAME functions the eager chain runs (stream
14004    /// identity by construction, not by twin-kernel parity). Host KV mirrors advance like the
14005    /// 26B dc-eager arm (window views are host math); len_d stays synced by the caller's
14006    /// entry sync + the appends here don't read it. Graph capture is NOT wired (no
14007    /// cap_bucket_max) — the E4B graph arc comes after the perf gates.
14008    #[allow(clippy::too_many_arguments)]
14009    pub fn gemma4_e4b_decode_step_dc(
14010        &self,
14011        e: &Engine,
14012        token_d: &CudaSlice<u32>,
14013        pos_d: &mut CudaSlice<i32>,
14014        embd_gpu: &CudaSlice<u8>,
14015        embd_qt: i32,
14016        embd_rb: usize,
14017        cache: &mut Cache,
14018        n_vocab: usize,
14019    ) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
14020        let n_embd = self.cfg.n_embd as usize;
14021        let eps = self.cfg.rms_eps;
14022        let mut x = e.embed_gather_device(embd_gpu, token_d, n_embd, embd_qt, embd_rb)?;
14023        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
14024        let inp_pl = self.gemma4_e4b_inp_pl_dev(e, token_d, &x, 1)?;
14025        let (ld, _x) =
14026            self.gemma4_e4b_trunk_core(e, x, inp_pl, pos_d, 1, cache, None, false, false)?;
14027        let mut tok_out = e.stream().alloc_zeros::<u32>(1)?;
14028        e.argmax_token_device_into(&ld, &mut tok_out, n_vocab)?;
14029        e.inc_seqlen(pos_d)?;
14030        cache.pos += 1;
14031        let _ = eps;
14032        Ok(tok_out)
14033    }
14034
14035    /// E4B eager T=1 decode step (decode_step_h contract): returns (softcapped logits host,
14036    /// pre-output_norm hidden). Advances cache.pos.
14037    pub(crate) fn gemma4_e4b_decode_step_h(
14038        &self,
14039        e: &Engine,
14040        token: u32,
14041        cache: &mut Cache,
14042    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14043        let (ld, x) = self.gemma4_e4b_trunk(e, &[token], cache.pos, cache, false)?;
14044        let logits = e.dtoh(&ld)?;
14045        cache.pos += 1;
14046        Ok((logits, x))
14047    }
14048
14049    /// E4B batched prime (prime_cache contract: (last-row logits host, h_seed pre-norm last
14050    /// row, hidden stack)). First-light: the t-wide trunk (per-row attention) — correct, not
14051    /// fast; the prefill fa arms come later.
14052    pub(crate) fn gemma4_e4b_prime(
14053        &self,
14054        e: &Engine,
14055        tokens: &[u32],
14056        cache: &mut Cache,
14057    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14058        // Err, not assert (2026-08-07, lane/gemma4-serve-gaps): same served-chunked-prompt
14059        // process-kill as gemma4_prime — refuse per-request.
14060        if cache.pos != 0 {
14061            return Err(
14062                "e4b prime is fresh-prompt only (v0) — prime the full prompt in one \
14063                        call or decode tokenwise"
14064                    .into(),
14065            );
14066        }
14067        let n_embd = self.cfg.n_embd as usize;
14068        let t = tokens.len();
14069        let (ld, x) = self.gemma4_e4b_trunk(e, tokens, 0, cache, true)?;
14070        cache.pos += t;
14071        let last = e.dtoh(&ld)?; // head_last: ld is already the final row only
14072        let xv = e.view(&x, t * n_embd);
14073        let row = xv.slice((t - 1) * n_embd..t * n_embd);
14074        let mut h_seed = e.uninit(n_embd)?;
14075        e.copy_view_into(&mut h_seed, 0, &row, n_embd)?;
14076        Ok((last, h_seed, x))
14077    }
14078
14079    /// E4B prefill logits (forward contract — no persistent cache; scratch cache internally).
14080    pub(crate) fn gemma4_e4b_forward(
14081        &self,
14082        e: &Engine,
14083        tokens: &[u32],
14084        last_only: bool,
14085    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
14086        let mut cache = Cache::new(e, &self.cfg, tokens.len() + 8)?;
14087        let (ld, _x) = self.gemma4_e4b_trunk(e, tokens, 0, &mut cache, last_only)?;
14088        Ok(e.dtoh(&ld)?) // head_last already reduced to the final row when last_only
14089    }
14090}
14091
14092#[cfg(test)]
14093mod prime_chunk_schedule_tests {
14094    use super::{
14095        PRIME_MIN_T, PRIME_PIPE_MIN_CHUNK, dynamic_prime_chunk_ranges, fixed_prime_chunk_ranges,
14096        fixed_prime_chunk_ranges_for_ring,
14097    };
14098
14099    fn sizes(ranges: &[(usize, usize)]) -> Vec<usize> {
14100        ranges.iter().map(|(start, end)| end - start).collect()
14101    }
14102
14103    fn auto_chunk(t: usize) -> usize {
14104        t.div_ceil(8).max(PRIME_PIPE_MIN_CHUNK).min(4096)
14105    }
14106
14107    #[test]
14108    fn fixed_schedule_retains_measured_geometry() {
14109        assert_eq!(
14110            sizes(&fixed_prime_chunk_ranges(461, 128)),
14111            vec![128, 128, 128, 77]
14112        );
14113        assert_eq!(
14114            sizes(&fixed_prime_chunk_ranges(1833, 230)),
14115            vec![230, 230, 230, 230, 230, 230, 230, 223]
14116        );
14117        assert_eq!(sizes(&fixed_prime_chunk_ranges(4096, 512)), vec![512; 8]);
14118        let capped = sizes(&fixed_prime_chunk_ranges_for_ring(8200, 4096, true));
14119        assert_eq!(capped, vec![4096, 4088, 16]);
14120        assert!(capped.iter().all(|&rows| rows <= 4096));
14121        assert_eq!(
14122            sizes(&fixed_prime_chunk_ranges_for_ring(4100, 4096, false)),
14123            vec![4100],
14124            "flag-off schedule remains byte-for-byte the legacy monolithic tail",
14125        );
14126    }
14127
14128    #[test]
14129    fn dynamic_schedule_matches_registered_shapes() {
14130        let cases = [
14131            (461, vec![64, 141, 132, 124]),
14132            (1833, vec![115, 269, 260, 252, 244, 237, 231, 225]),
14133            (4096, vec![256, 602, 580, 563, 545, 531, 516, 503]),
14134        ];
14135        for (t, expected) in cases {
14136            let chunk = auto_chunk(t);
14137            let fixed = fixed_prime_chunk_ranges(t, chunk);
14138            assert_eq!(
14139                sizes(&dynamic_prime_chunk_ranges(t, chunk, &fixed)),
14140                expected
14141            );
14142        }
14143    }
14144
14145    #[test]
14146    fn dynamic_schedule_covers_exactly_and_shrinks_after_fill() {
14147        for t in 256..=8192 {
14148            let chunk = auto_chunk(t);
14149            let fixed = fixed_prime_chunk_ranges(t, chunk);
14150            let dynamic = dynamic_prime_chunk_ranges(t, chunk, &fixed);
14151            assert_eq!(dynamic.len(), fixed.len(), "T={t}");
14152            assert_eq!(dynamic.first().unwrap().0, 0, "T={t}");
14153            assert_eq!(dynamic.last().unwrap().1, t, "T={t}");
14154            for pair in dynamic.windows(2) {
14155                assert_eq!(pair[0].1, pair[1].0, "T={t}");
14156            }
14157            assert!(
14158                dynamic
14159                    .iter()
14160                    .all(|(start, end)| end - start >= PRIME_MIN_T),
14161                "T={t} sizes={:?}",
14162                sizes(&dynamic)
14163            );
14164            if dynamic.len() >= 3 {
14165                let chunk_sizes = sizes(&dynamic);
14166                assert!(
14167                    chunk_sizes[0] < chunk_sizes[1],
14168                    "T={t} sizes={chunk_sizes:?}"
14169                );
14170                assert!(
14171                    chunk_sizes[1..].windows(2).all(|pair| pair[0] >= pair[1]),
14172                    "T={t} sizes={chunk_sizes:?}"
14173                );
14174            }
14175        }
14176    }
14177}
14178
14179#[cfg(test)]
14180mod page_prefetch_tests {
14181    use super::{
14182        grouped_worker_prefetch_position, page_prefetch_positions,
14183        page_prefetch_window_from_values, worker_prefetch_positions,
14184    };
14185
14186    #[test]
14187    fn page_prefetch_window_keeps_existing_opt_in_default() {
14188        assert_eq!(page_prefetch_window_from_values(false, None), 0);
14189        assert_eq!(page_prefetch_window_from_values(false, Some("8")), 0);
14190        assert_eq!(page_prefetch_window_from_values(true, None), 1);
14191        assert_eq!(page_prefetch_window_from_values(true, Some("bad")), 1);
14192        assert_eq!(page_prefetch_window_from_values(true, Some("0")), 0);
14193        assert_eq!(page_prefetch_window_from_values(true, Some("8")), 8);
14194    }
14195
14196    #[test]
14197    fn rolling_page_prefetch_advises_each_future_expert_once() {
14198        let advised: Vec<_> = (0..7)
14199            .flat_map(|position| page_prefetch_positions(position, 7, 3))
14200            .collect();
14201        assert_eq!(advised, vec![1, 2, 3, 4, 5, 6]);
14202
14203        let one_ahead: Vec<_> = (0..4)
14204            .flat_map(|position| page_prefetch_positions(position, 4, 1))
14205            .collect();
14206        assert_eq!(one_ahead, vec![1, 2, 3]);
14207        assert!(page_prefetch_positions(0, 4, 0).is_empty());
14208    }
14209
14210    #[test]
14211    fn grouped_worker_prefetch_primes_first_then_each_known_next_once() {
14212        assert_eq!(grouped_worker_prefetch_position(0, None), None);
14213        let positions: Vec<_> = std::iter::once(grouped_worker_prefetch_position(4, None).unwrap())
14214            .chain(
14215                (0..4).filter_map(|position| grouped_worker_prefetch_position(4, Some(position))),
14216            )
14217            .collect();
14218        assert_eq!(positions, vec![0, 1, 2, 3]);
14219        assert_eq!(grouped_worker_prefetch_position(1, Some(0)), None);
14220    }
14221
14222    #[test]
14223    fn rolling_worker_prefetch_primes_current_and_each_future_expert_once() {
14224        let queued: Vec<_> = (0..8)
14225            .flat_map(|position| worker_prefetch_positions(position, 8, 5))
14226            .collect();
14227        assert_eq!(queued, (0..8).collect::<Vec<_>>());
14228
14229        let one_at_a_time: Vec<_> = (0..4)
14230            .flat_map(|position| worker_prefetch_positions(position, 4, 1))
14231            .collect();
14232        assert_eq!(one_at_a_time, vec![0, 1, 2, 3]);
14233        assert!(worker_prefetch_positions(0, 4, 0).is_empty());
14234    }
14235}
14236
14237pub struct G4DcSlots {
14238    x: CudaSlice<f32>,
14239    xn: CudaSlice<f32>,
14240    cur: CudaSlice<f32>,
14241    hq: CudaSlice<i8>,
14242    hd_: CudaSlice<f32>,
14243    q0: CudaSlice<f32>,
14244    k0: CudaSlice<f32>,
14245    v0: CudaSlice<f32>,
14246    q: CudaSlice<f32>,
14247    k: CudaSlice<f32>,
14248    v: CudaSlice<f32>,
14249    attn: CudaSlice<f32>,
14250    o: CudaSlice<f32>,
14251    attn_out: CudaSlice<f32>,
14252    zsh: CudaSlice<f32>,
14253    zq: CudaSlice<i8>,
14254    zd: CudaSlice<f32>,
14255    gate: CudaSlice<f32>,
14256    up: CudaSlice<f32>,
14257    act: CudaSlice<f32>,
14258    actq: CudaSlice<i8>,
14259    actd: CudaSlice<f32>,
14260    f0: CudaSlice<f32>,
14261    sn: CudaSlice<f32>,
14262    hn: CudaSlice<f32>,
14263    logits: CudaSlice<f32>,
14264}