Skip to main content

memra_engine/
lib.rs

1//! memra engine: Stage-1 correctness-first forward-pass kernels + ops, on sm_120 via cudarc.
2
3use cudarc::driver::{
4    CudaContext, CudaFunction, CudaModule, CudaSlice, CudaStream, LaunchConfig, PushKernelArg,
5};
6use cudarc::nvrtc::Ptx;
7use std::sync::{Arc, Mutex};
8
9#[cfg(debug_assertions)]
10pub(crate) fn debug_assert_tensor_stream_device<T>(
11    tensor: &CudaSlice<T>,
12    stream: &CudaStream,
13    site: &str,
14) {
15    let tensor_dev = tensor.ordinal();
16    let stream_dev = stream.context().ordinal();
17    assert_eq!(
18        tensor_dev, stream_dev,
19        "PP cross-device tensor read at {site}: tensor on dev{tensor_dev}, stream on dev{stream_dev}"
20    );
21}
22
23pub use memra_gguf;
24pub use memra_runtime;
25
26pub mod forward;
27pub mod hybrid;
28pub mod hybrid_forward;
29pub mod model;
30pub mod sigrouter_contract;
31pub mod vision;
32pub mod vision_gemma;
33pub mod vision_pre;
34/// The dual cache lives in the shared `memra-kv` crate (Phase D extraction); this
35/// re-export keeps every `crate::cache::` / `memra_engine::cache::` path unchanged.
36pub mod cache {
37    pub use memra_kv::*;
38}
39pub mod decode;
40pub mod decode_batch;
41pub mod dflash;
42pub mod eagle;
43pub mod gemma_spec;
44pub mod graph_update;
45/// MLA (multi-head latent attention) CPU f32 reference — GLM-5.2 bring-up lane increment 1.
46/// Naive vs absorbed decode forms + NORM/NEOX rope permutation, unit-tested; the permanent
47/// oracle for the MLA kernel family (`research/mla-bringup-20260801/DESIGN.md`). No CUDA deps.
48pub mod mla;
49pub mod moesd;
50pub mod parallel;
51pub mod pp;
52pub mod round_stream;
53pub mod spec;
54pub use memra_sampling as sampler;
55
56/// In-house MoE router GEMV on the spec-verify small-t path (DEFAULT ON since 2026-07-10:
57/// battery green on 35B p2/p3 K=1..8, acceptance bit-identical, +2-4% spec e2e — replaces
58/// ~240 per-column cuBLAS gemv launches/round). MEMRA_ROUTER_KERNEL=0 is the rollback seam.
59/// MoE grouped f16 GEMM door (experimental until gated), f16-mirror numeric class:
60/// per-layer expert dequant to f16 + one grouped f16 GEMM over the CSR groups.
61///   MEMRA_MOE_F16G=1  cublasGemmGroupedBatchedEx (round 46 arc 2). The grouped API issues
62///                     through cublas-internal streams NOT ordered with ours — v1 pays a full
63///                     stream sync per projection (round-47 ledgered defect).
64///   MEMRA_MOE_F16G=2  single-kernel grouped GEMM on the engine stream (round 49): ordered by
65///                     construction, zero syncs, f32 C with the act row-scale folded in.
66/// DEFAULT (2026-08-01, round 49 promotion): mode 1 on the Hopper lane — with the 41/41
67/// dequant coverage fix the q35 board-2048 prime measured 5490 (MMQ) / 8380 (mode 1,
68/// +53%) / 7990 (mode 2) x3 interleaved on the H100, argmax MATCH — the last board loss
69/// flips. The 5090 measured FLAT (858GB/s makes the dequant-workspace traffic cancel the
70/// GEMM win) — but that verdict is for expert banks the int8-MMA MMQ arm can take
71/// (IQ3_S/IQ4_XS/Q4_0). MEMRA_MOE_F16G=0 kills anywhere.
72///
73/// HOPPER RE-VERDICT (2026-08-02, lane/h100-flip-full): mode 2 with full direct coverage
74/// (Q4_K/Q6_K/IQ4_XS/IQ3_S tile loaders, lane/iq-direct-loaders) + the deep tail
75/// (lane/sk-tail-form) FLIPS past cublas mode 1 on the H100 — q35 board-2048 prime
76/// 13163.6 (mode 2, cross=32) vs 8626.5 (mode 1) vs 8073.4 (round-51 sk form), +52.6%,
77/// interleaved x5 zero overlap, argmax MATCH 30/30. The round-54 NO-FLIP (8547 vs 8112)
78/// was coverage-priced at 5.2% direct; ~100% coverage kills the workspace pass and the
79/// verdict inverts. Hopper naked default -> mode 2 (this arm); the gemma (gelu) site
80/// stays env-explicit-only via moe_f16g_gemma_on (Err => closed, unaffected by this arm).
81///
82/// MODE-2 DEFAULT (sm_120a naked, 2026-08-02, lane/f16g-default-rearb): with the direct
83/// tile loaders covering Q4_K/Q6_K/IQ4_XS/IQ3_S, the sk visitor beats the int8-MMA MMQ
84/// tiles on the IQ-bank models too (q35 board-2048 +33.9%, KAT pp512 +46.7% / pp2048
85/// +30.6% — research/iq-direct-loaders-20260802 §3-5, confirmed + full battery in
86/// research/f16g-default-rearb-20260802/), so every f16g-admitted expert layer rides
87/// mode 2 naked. Decode/verify stay on dp4a (t >= 16 floor). f16-mirror numeric class
88/// for naked q35/KAT prefill+prime — new token-sha anchors stamped in the rearb lane.
89///
90/// AUTO-KQUANT (mode 3, 2026-08-02, lane/q4k-expert-prefill): the previous sm_120a
91/// default, kept reachable via MEMRA_MOE_F16G=3. The mode-2 sk form is admitted ONLY for
92/// layers the MMA MMQ arm rejects (k-quant expert projections — Q3_K/Q4_K/Q6_K), i.e.
93/// exactly where the baseline is the per-pair moe_pairs_matvec_q8_em fallback with zero
94/// token reuse (Ornith-35B Q4_K_M board-2048 1098.2 -> 3453.7, 3.14x,
95/// research/q4k-expert-prefill-20260802/). Its "IQ banks keep their measured-faster MMQ
96/// tiles" ruling was priced BEFORE the IQ direct loaders and is refuted on the 5090 —
97/// the k-quant-only admission survives as the rollback seam, not the default.
98/// The gemma (gelu) site stays env-explicit-only (moe_f16g_gemma_on).
99pub fn moe_f16g_mode() -> u8 {
100    static M: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
101    *M.get_or_init(|| match std::env::var("MEMRA_MOE_F16G").as_deref() {
102        Ok("0") => 0,
103        Ok("2") => 2,
104        Ok("3") => 3,
105        Ok(_) => 1,
106        // Both arches independently re-arbitrated to mode 2 on 2026-08-02
107        // (5090: lane/f16g-default-rearb; H100: lane/h100-flip-full) — unset = 2 everywhere.
108        Err(_) => 2,
109    })
110}
111/// Mode-2 sk kernel form policy (round 51, lane/sk-bm128): the single-kernel grouped GEMM runs
112/// as a persistent problem-visitor over the real CSR tiles with two tile forms. Returns
113/// (shape_sel, cross) for the FFI:
114///   MEMRA_F16G_SK=0    -> (-1, _): the round-49 grid-scan kernel (rollback seam).
115///   MEMRA_F16G_SK=32   -> all groups on the 32x64x32 2-stage form (cross = i32::MAX).
116///   MEMRA_F16G_SK=128  -> all groups on the 128x64x64 3-stage form (cross = 1; groups fall
117///                         back to 32x64 in-launcher when the device/in_f can't take it).
118///   unset              -> hybrid split: groups with m_e >= MEMRA_F16G_SK_CROSS ride the 128
119///                         form. Default cross = 64 (5090 sweep 2026-08-01, receipts
120///                         research/sk-bm128-20260801/; H100 re-swept on the direct+tail
121///                         form 2026-08-02, lane/h100-flip-full: {16,32,64} ->
122///                         12868/13192/13225 — 64 wins there too, the pre-direct 32
123///                         verdict was stale).
124pub fn moe_f16g_sk_params() -> (i32, i32) {
125    static P: std::sync::OnceLock<(i32, i32)> = std::sync::OnceLock::new();
126    *P.get_or_init(|| match std::env::var("MEMRA_F16G_SK").as_deref() {
127        Ok("0") => (-1, 0),
128        Ok("32") => (0, i32::MAX),
129        Ok("128") => (0, 1),
130        _ => {
131            let cross = std::env::var("MEMRA_F16G_SK_CROSS")
132                .ok()
133                .and_then(|v| v.parse().ok())
134                .unwrap_or(64);
135            (0, cross)
136        }
137    })
138}
139/// DIRECT-FROM-QUANT sk tile loaders (lane/kquant-tile-loaders, 2026-08-02; IQ classes added
140/// by lane/iq-direct-loaders): Q4_K/Q6_K/IQ4_XS/IQ3_S expert projections on the mode-2/3 sk
141/// visitor forms dequant their weight tiles in-register from the quant superblocks instead of
142/// running the per-(layer,projection) dequant pass into an f16 workspace (41.8% of Ornith-35B
143/// t=512 kernel time — the pp512 wall, research/q4k-expert-prefill-20260802 §5; the IQ classes
144/// are 94.8% of q35's bank bytes — the h100-sk-direct coverage pricing). Bit-identical to the
145/// workspace path by construction (kernel-check "f16g-kq-direct" gates it bitwise) — a
146/// data-movement change, not a numeric-class change. Default ON; MEMRA_F16G_DIRECT=0 reverts
147/// to the workspace path everywhere; MEMRA_F16G_DIRECT=kq keeps the k-quant loaders and
148/// reverts only the IQ classes (the iq-direct-loaders A/B seam — the pre-lane shipped config).
149pub fn moe_f16g_direct_on(qtype: i32) -> bool {
150    static M: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
151    let m = *M.get_or_init(|| match std::env::var("MEMRA_F16G_DIRECT").as_deref() {
152        Ok("0") => 0,
153        Ok("kq") => 1,
154        _ => 2,
155    });
156    match m {
157        0 => false,
158        1 => qtype == QT_Q4_K || qtype == QT_Q6_K,
159        _ => true,
160    }
161}
162/// DEEP-TAIL sk form (lane/sk-tail-form, 2026-08-02): groups below the visitor crossover ride
163/// a 32x64x64 3-STAGE cp.async tile instead of the round-51 32x64x32 2-stage — the same 32-row
164/// tile (zero extra padding), 2 k-blocks in flight instead of 1 and half the syncs per k. The
165/// H100 ncu pricing (research/sk-bm128-20260801) put the 2-stage tail at 31% of the sk GEMM
166/// stage under q35's routing skew. Bit-identical to every other sk form by construction
167/// (kernel-check "f16g-sk" gates all tail arms maxdiff==0); exists in both the workspace-f16
168/// and direct-from-quant variants. Default ON; MEMRA_F16G_TAIL=0 = rollback to the 2-stage
169/// tail. in_f % 64 != 0 falls back in-launcher.
170pub fn moe_f16g_tail_on() -> bool {
171    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
172    *ON.get_or_init(|| std::env::var("MEMRA_F16G_TAIL").as_deref() != Ok("0"))
173}
174
175/// Per-model door for the gemma-MoE (gelu) grouped path: round 49's Hopper default
176/// REGRESSED g26 board-2048 prefill -8.3% interleaved x5 on-box (def median 10380,
177/// wild 8.9k-11.7k spread; off 11317, ±0.13%) — the +6-15% probe verdict didn't
178/// survive the board workload (stale-verdict law, round 50). The silu/qwen class
179/// keeps the round-49 default (q35 +53% board-2048). Explicit MEMRA_MOE_F16G=1/2
180/// still opens this door for A/B.
181pub fn moe_f16g_gemma_on() -> bool {
182    static M: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
183    *M.get_or_init(|| !matches!(std::env::var("MEMRA_MOE_F16G").as_deref(), Ok("0") | Err(_)))
184}
185
186/// Fused act-epilogue (silu/gelu-mul + q8_1_mmq quantize in one launch) for the MoE prefill
187/// MMA arms. Byte-identical to the two-pass path (kernel-check gated) — default ON.
188/// MEMRA_MOE_FUSE_ACTQ=0 is the rollback/A-B seam.
189pub fn moe_fuse_actq_on() -> bool {
190    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
191    *ON.get_or_init(|| std::env::var("MEMRA_MOE_FUSE_ACTQ").as_deref() != Ok("0"))
192}
193
194/// PREFILL router m-invariance (lane/concat-prime-exact, 2026-08-02). The batched cuBLASLt
195/// router GEMM changes a row's logits when OTHER rows join the call (probed: first change at
196/// m=65 on the Ornith-35B router, 3.9e-3 — while the MMQ/f16 trunk GEMMs are bit-identical
197/// across m). Feeding a top-k discontinuity, that made a served request's expert selection a
198/// function of its CO-ARRIVALS under cross-request prime batching. The in-house router GEMV
199/// is m-invariant, so prefill uses it too and routing depends on a session's own tokens only.
200/// DEFAULT ON: it is the serving isolation contract, and it is the same kernel decode and spec
201/// verify already use (dispatch parity, one router kernel for every t).
202/// MEMRA_ROUTER_PREFILL_EXACT=0 reverts to the batched GEMM.
203pub fn router_prefill_exact_on() -> bool {
204    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
205    *ON.get_or_init(|| std::env::var("MEMRA_ROUTER_PREFILL_EXACT").as_deref() != Ok("0"))
206}
207
208pub fn router_kernel_on() -> bool {
209    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
210    *ON.get_or_init(|| {
211        let on = std::env::var("MEMRA_ROUTER_KERNEL").as_deref() != Ok("0");
212        if !on {
213            eprintln!("[memra] router kernel OFF (rollback: per-column cuBLAS gemv)");
214        }
215        on
216    })
217}
218
219/// FAST-ROUTER batch twin (lane/fast-router, 2026-08-02). The concat-prime exactness fix
220/// (router_prefill_exact_on) routes prefill through router_gemv — m-invariant, but a
221/// per-(expert,token) GEMV program with zero operand reuse, so q35 board-2048 prefill paid
222/// -10% on the 5090. router_gemv_f32_w8_batch register-tiles (8x8 expert-x-token) the same
223/// per-row FP chains (BIT-IDENTICAL per row — kernel-check sweeps m=1..2048 on real router
224/// weights), so the t crossover below is pure perf, not a numeric config. Swept on-box
225/// (research/fast-router-20260802/crossover-router*.jsonl): plain wins t<=4, batch +7-9%
226/// at t=8, 1.9x at t=16 rising to 3.45x at t=2048 — MIN_T=8. Decode t=1 and spec verify
227/// t<8 keep the plain w8 form. MEMRA_ROUTER_BATCH=0 forces plain at every t (rollback
228/// seam, perf-only: bits are equal by the kernel-check gate).
229/// Killed arms (same sweep, JSONL is the record): the 8x16 tile lost to 8x8 at every t
230/// (128-accumulator register pressure beats the halved w-traffic), and the same-shape
231/// sigmoid_dot_rows twin (out_f=1) measured 0.62-0.89x at every prefill t
232/// (launch-latency-bound, ~7us/layer at m=2048) — both bit-identity-PASSED before dying.
233pub const ROUTER_BATCH_MIN_T: usize = 8;
234pub fn router_batch_on() -> bool {
235    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
236    *ON.get_or_init(|| std::env::var("MEMRA_ROUTER_BATCH").as_deref() != Ok("0"))
237}
238mod cpu_experts;
239#[cfg(memra_cutlass)]
240pub mod cutlass_ffi;
241pub mod dsv4_ffi;
242pub mod dsv4_gpu;
243pub mod f16_ffi;
244pub mod fp8_ffi;
245pub mod mmq_ffi;
246pub mod moe_cache;
247pub mod prime_graph;
248pub mod spill;
249mod spill_pread;
250
251// Fatbins are EMBEDDED (crates-release lane, 2026-08-04): build.rs still writes them to
252// OUT_DIR, but the bytes ship inside the binary via include_bytes! and load through
253// cuModuleLoadData. Distribution contract: a prebuilt or cargo-installed binary must be
254// self-contained — the old baked OUT_DIR *paths* pointed at the builder's temp dir and
255// broke every machine that wasn't the build machine. Same bytes, same module image;
256// the runtime MEMRA_GEMM_FATBIN tune-seam override below is preserved.
257const FATBIN: &[u8] = include_bytes!(env!("MEMRA_ENGINE_FATBIN"));
258const HYBRID_FATBIN: &[u8] = include_bytes!(env!("MEMRA_HYBRID_FATBIN"));
259const QMATVEC_FATBIN: &[u8] = include_bytes!(env!("MEMRA_QMATVEC_FATBIN"));
260const FLASH_FATBIN: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN"));
261const GEMM_FATBIN: &[u8] = include_bytes!(env!("MEMRA_GEMM_FATBIN"));
262const ROUTER_FATBIN: &[u8] = include_bytes!(env!("MEMRA_ROUTER_FATBIN"));
263/// spec_sample.cu: sampled-spec primitives (Philox Gumbel-max / softmax gather / residual sampler).
264const SAMPLE_FATBIN: &[u8] = include_bytes!(env!("MEMRA_SAMPLE_FATBIN"));
265
266/// TUNE SEAM (tools/sweep): a RUNTIME `MEMRA_GEMM_FATBIN=<path>` overrides the baked-in
267/// qmatvec_gemm.cu fatbin path (build.rs bakes the same name at COMPILE time via
268/// cargo:rustc-env — that constant is the default). Lets the sweep harness swap in a
269/// `-D`-tuned fatbin per process with NO rust rebuild. Unset at runtime => the
270/// compile-time default (zero behavior change).
271fn gemm_fatbin_bytes() -> std::borrow::Cow<'static, [u8]> {
272    assert!(
273        !(portable_mma_gated() && std::env::var_os("MEMRA_GEMM_FATBIN").is_some()),
274        "MEMRA_GEMM_FATBIN overrides are not allowed in the portable CUDA lane"
275    );
276    match std::env::var("MEMRA_GEMM_FATBIN") {
277        Ok(path) => std::borrow::Cow::Owned(
278            std::fs::read(&path).unwrap_or_else(|e| panic!("MEMRA_GEMM_FATBIN read {path}: {e}")),
279        ),
280        Err(_) => std::borrow::Cow::Borrowed(GEMM_FATBIN),
281    }
282}
283
284/// Phase A (ARCHITECTURE-H100.md): sm_90a re-enables the portable-PTX tensor-core paths
285/// (int8 mma.m16n8k32/k16.s8, bf16 m16n8k16, ldmatrix, cp.async — all sm_80-class, native
286/// on Hopper) that the portable boot lane gates off. Dispatch guards that used to test
287/// `cfg!(memra_portable_cuda)` test this instead; sm_89 keeps the pure-portable behavior.
288/// The sm_120a/sm_100a-only MMA kinds (mxf4nvf4, kind::f8f6f4) are NOT covered — their
289/// launchers stay fail-closed stubs on 90a and their dispatch arms stay arch-gated.
290pub(crate) const fn portable_mma_gated() -> bool {
291    cfg!(memra_portable_cuda) && !cfg!(memra_hopper_mma)
292}
293
294/// The GDN K4/K5 mma pair's UNSET-env default — ONE definition for the three read sites
295/// (gdn_mma_enabled, the k123 pre-work, gdn_scan_chunked's dispatch). They read the env
296/// per call ON PURPOSE (kernel-check toggles it to pin both configs), so the shared part
297/// is this compile-time constant: ON for Hopper-MMA builds (the original 90a promotion)
298/// and for sm_120a builds (lane/moeprime-nvfp4-direct, 2026-08-21 — measured on one RTX
299/// PRO 6000 ornith15 pp14715 +6-8% and the local 5090 q38-27b +1-2%, both orders both
300/// rigs). A site defaulting differently from its peers arms the mma pre-work while the
301/// scan takes the scalar route — measured as a 0.8% LOSS, the drift this helper kills.
302pub(crate) const fn gdn_mma_default_on() -> bool {
303    cfg!(memra_hopper_mma) || konst_eq(env!("MEMRA_BUILT_CUDA_ARCH"), "120a")
304}
305
306/// const str-eq (std `==` on &str is not const-stable on this toolchain floor).
307const fn konst_eq(a: &str, b: &str) -> bool {
308    let (a, b) = (a.as_bytes(), b.as_bytes());
309    if a.len() != b.len() {
310        return false;
311    }
312    let mut i = 0;
313    while i < a.len() {
314        if a[i] != b[i] {
315            return false;
316        }
317        i += 1;
318    }
319    true
320}
321
322/// The legacy quantized prefill GEMMs are tuned and validated for sm_120a; sm_90a re-admits
323/// them through the Hopper-MMA lane (int8 m16n8k32.s8 is sm_80-class PTX).  Keep the policy
324/// in a pure helper so the dispatch guard can be regression-tested without constructing an
325/// Engine or allocating a GPU tensor.
326const fn legacy_quant_gemm_allowed(portable_cuda: bool, hopper_mma: bool, no_gemm: bool) -> bool {
327    (!portable_cuda || hopper_mma) && !no_gemm
328}
329
330// ---- KV-cache format selection (kvbytes lane, 2026-07-08; default OFF = daily config) ----
331// `MEMRA_KV_K` = q8_0 (default, 34 B/32elem) | fp8 (raw e4m3, 32 B — the -6% K-bytes arm)
332// `MEMRA_KV_V` = q5_1 (default, 24 B/32elem) | q4_0 (18 B, -25% V bytes) | fp8 (32 B, +33%)
333// A non-default format is a NEW NUMERIC CONFIG: its own run-gen argmax baseline is legal,
334// but the gate battery (kernel-check, run-spec self-consistency) must pass WITHIN it and
335// the choice is explicit env, never silent. flash_attn.cu is compiled once per format pair
336// (build.rs); the kernels keep their names — Engine::new just loads the matching fatbin.
337const FLASH_FATBIN_VQ4: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_VQ4"));
338const FLASH_FATBIN_VF8: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_VF8"));
339const FLASH_FATBIN_KF8: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_KF8"));
340const FLASH_FATBIN_KF8VQ4: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_KF8VQ4"));
341const FLASH_FATBIN_KF8VF8: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_KF8VF8"));
342
343/// KV format policy moved to the shared `memra-kv` crate (Phase D); re-exported so the
344/// fatbin router below and every existing `crate::kv_blk_bytes()` call site is unchanged.
345pub use memra_kv::{kv_blk_bytes, kv_cache_formats};
346
347/// The flash_attn fatbin matching the selected KV formats.
348fn flash_fatbin_bytes() -> &'static [u8] {
349    match kv_cache_formats() {
350        ("q8_0", "q5_1") => FLASH_FATBIN,
351        ("q8_0", "q4_0") => FLASH_FATBIN_VQ4,
352        ("q8_0", "fp8") => FLASH_FATBIN_VF8,
353        ("fp8", "q5_1") => FLASH_FATBIN_KF8,
354        ("fp8", "q4_0") => FLASH_FATBIN_KF8VQ4,
355        ("fp8", "fp8") => FLASH_FATBIN_KF8VF8,
356        other => unreachable!("kv_cache_formats returned {other:?}"),
357    }
358}
359
360/// TUNE SEAM (tools/sweep): kernel1 (Q8_0/Q4_K/Q5_K) launch-tile override,
361/// `MEMRA_GEMM_K1_LAUNCH="BM,BN,NWARP"`. MUST match the `-D K1_BM/K1_BN/NWARP` the swept
362/// fatbin was compiled with (the .cu tile and the host launch grid/block have to agree —
363/// the hardcoded (128,128,8) in qmatvec_gemm/qmatvec_gemm_raw is the shipped default).
364/// Kernel2 (Q6_K/NVFP4) launch is untouched. Unset or malformed => None => shipped
365/// defaults (zero behavior change).
366fn k1_launch_override() -> Option<(u32, u32, u32)> {
367    static K1: std::sync::OnceLock<Option<(u32, u32, u32)>> = std::sync::OnceLock::new();
368    *K1.get_or_init(|| {
369        let v = std::env::var("MEMRA_GEMM_K1_LAUNCH").ok()?;
370        let p: Vec<u32> = v.split(',').filter_map(|s| s.trim().parse().ok()).collect();
371        match p.as_slice() {
372            [bm, bn, w] => Some((*bm, *bn, *w)),
373            _ => None,
374        }
375    })
376}
377
378/// H100 wgmma prefill-GEMM seam (task 8, ARCHITECTURE-H100.md): OPT-IN (MEMRA_WGMMA=1).
379/// v0 verdict (2026-07-26, N=5 pp512 9B-Q8_0): wgmma 3845 tok/s vs MMQ 8692 — the
380/// standalone harness's "688us MMQ ref" was a pp2048-shape figure, so v0 (unpipelined,
381/// 64x64 tile, wait_group<0> every 32-K step) is ~3x SLOWER per launch at m=512 model
382/// shapes. Default stays MMQ until the pipelined version beats it N=5 (repo law).
383/// Correctness stays pinned regardless: kernel-check's wgmma case is cfg-gated, not env-gated.
384pub(crate) fn wgmma_gemm_enabled() -> bool {
385    static V: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
386    *V.get_or_init(|| std::env::var("MEMRA_WGMMA").as_deref() == Ok("1"))
387}
388
389/// TUNE SEAM: keys per FA-decode split (`MEMRA_FA_SPLIT` forces a fixed size; default 64). Smaller
390/// splits raise grid.y so grid = n_head_kv * n_splits fills the 82 SMs at short/mid ctx (vec path
391/// launches only n_head_kv=8 CTAs per split). Swept clock-locked 2026-07-03 (graph tg128): 32 beat
392/// 64 at ctx 128/512 (+0.5/+1.2%) and lost at 2048 (-3%) — BUT the adaptive 32/64 default BROKE the
393/// MTP spec-decode exact-match gate (run-spec K=1/2 self-consistency FAIL with 32; PASS with 64):
394/// the split count changes the combine's FP summation order, and the spec verify's batched forward
395/// only argmax-matches single-step decode under the 64-split order on real prompts. Spec exactness
396/// (the bigger lever) outranks a <=1.2% decode win -> default stays FIXED 64; sweeps use the env.
397/// Takes t_kv so eager, _dc capture, and fa_geom_eager stay signature-compatible for future
398/// adaptive retries (any retry MUST pass run-spec self-consistency first).
399/// Minimum t_kv for the warp-per-token vec FA path (below it the scalar path's 4x-more-blocks
400/// hides latency better — measured crossover, see `fa_decode`). Shared by fa_decode / fa_decode_dc /
401/// fa_geom_eager / fa_decode_rows-eligibility (spec verify) so the kernel pick NEVER diverges
402/// between eager decode and the verify (the spec-exactness law).
403pub const FA_VEC_MIN_TKV: usize = 96;
404/// Env-overridable crossover (MEMRA_FA_VEC_MIN, default FA_VEC_MIN_TKV). The 96 floor was
405/// measured on the qwen geometry (nkv=2); gemma4 SWA layers run nkv=8 = 4x the vec grid,
406/// which moves the crossover — sweep per model, adopt per the battery.
407pub fn fa_vec_min_tkv() -> usize {
408    static V: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
409    *V.get_or_init(|| {
410        std::env::var("MEMRA_FA_VEC_MIN")
411            .ok()
412            .and_then(|v| v.parse().ok())
413            .unwrap_or_else(|| FA_VEC_MIN_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
414    })
415}
416
417/// f16-P/V class (DEFAULT since 2026-07-23 stamp v4; MEMRA_FA_F16PV=0 = f32-class rollback):
418/// llama-fa=1-style f16 P + f16 P@V accumulation on the hd512/SWA prefill stamps
419/// (KQ/softmax/normalize stay f32). Laptop stamp: 12B 1.045x, 31B 0.979x vs llama.
420///
421/// SPEC-SERVING FLIP (2026-07-26, the wkv acceptance-law pattern): with MEMRA_DRAFT set the
422/// default is OFF. f16 P/V shifts the PRIME's hidden states/KV in the sub-argmax logit
423/// space the drafter feeds on — argmax gates stay MATCH while depth acceptance falls off a
424/// cliff (26B d1736 0.883 -> 0.405, -40% e2e; f16pv-off alone restores 0.846/314 tok/s —
425/// the perf-ci acceptance battery is the only gate that sees this class). Explicit
426/// MEMRA_FA_F16PV always wins; plain serving keeps the f16 prefill win.
427pub fn fa_f16pv_on() -> bool {
428    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
429    *ON.get_or_init(|| {
430        std::env::var("MEMRA_FA_F16PV")
431            .map(|v| v != "0")
432            .unwrap_or_else(|_| std::env::var("MEMRA_DRAFT").is_err())
433    })
434}
435
436/// hd512 head-pair arm (DEFAULT since stamp v4; MEMRA_FA512_HP=0 reverts to sp16): GQA
437/// ncols2=2 — 2 heads per CTA share each staged K/V tile, Q register-resident. Engages
438/// when n_head is even and the GQA group (n_head/n_head_kv) is even.
439pub fn fa512_hp_on() -> bool {
440    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
441    *ON.get_or_init(|| std::env::var("MEMRA_FA512_HP").as_deref() != Ok("0"))
442}
443
444/// SWA head-pair arm (DEFAULT since stamp v4; MEMRA_FAW_HP=0 reverts to p1): llama-class
445/// windowed geometry — 32 q-rows x 2 heads per CTA sharing staged K/V, f16 P@V
446/// accumulation. Even n_head and even GQA group required (guarded per call).
447pub fn faw_hp_on() -> bool {
448    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
449    *ON.get_or_init(|| std::env::var("MEMRA_FAW_HP").as_deref() != Ok("0"))
450}
451
452/// 4-warp sp16 experiment arm (MEMRA_FA512_W4=1, requires the f16pv door): GEMM0 split-K
453/// 4-way + GEMM1 4x128 O-dims. Own partial-sum order — oracle-band gated. Returns warp
454/// count (2 = base sp16). 8-warp arm measured NEGATIVE 2026-07-23 (jsonl) and removed.
455pub fn fa512_wide_warps() -> usize {
456    static N: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
457    *N.get_or_init(|| match std::env::var("MEMRA_FA512_W4").as_deref() {
458        Ok("1") => 4,
459        _ => 2,
460    })
461}
462
463/// hd-512 vec crossover floor (MEMRA_FA512_MIN, default 512) — shared by fa_decode dispatch
464/// and the gemma global-layer rows/parity call sites.
465pub fn fa512_min_tkv() -> usize {
466    static FA512_MIN: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
467    *FA512_MIN.get_or_init(|| {
468        std::env::var("MEMRA_FA512_MIN")
469            .ok()
470            .and_then(|v| v.parse().ok())
471            .unwrap_or(512)
472    })
473}
474/// Per-model crossover default, set at model load BEFORE the first decode (per-model
475/// numeric-config adoption law). qwen keeps the measured 96; gemma4 (nkv=8 SWA) measured
476/// vec-always fastest: 119.9 (96) / 130.0 (48) / 133.2 (1) tok/s tg128-regime, 2026-07-10.
477pub static FA_VEC_MIN_DEFAULT: std::sync::atomic::AtomicUsize =
478    std::sync::atomic::AtomicUsize::new(FA_VEC_MIN_TKV);
479/// Per-model windowed-split default (MEMRA_FA_SPW overrides): gemma MoE (26B, nkv=8) measured
480/// 32 (grid-limited t=1 under the raw-e4m3 sV ceiling, 2026-07-12); dense gemma (31B)
481/// measured 64 (37.13/37.12 vs 36.87/36.86 at 1.7k, N=2 — different attention geometry).
482pub static FA_SPW_DEFAULT: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(32);
483/// Per-model hd512 (gemma globals) split default (MEMRA_FA_SP512 overrides): 26B measured 16
484/// (2026-07-11 N=2), dense 31B measured 32 (36.86/36.93 vs 36.73/36.73 at 1.7k, 2026-07-12).
485/// fused t=1 q4_0 pair/triple row mapping: true = mr1 (one row/warp). Per-model default
486/// (dense gemma wins +1.1% short / +0.6% depth on the 31B; MoE 26B REGRESSES −1.2% —
487/// its shared-expert fused2 shapes lose to the finer grid). MEMRA_Q40_MR env still wins.
488pub static FUSED_MR1_DEFAULT: std::sync::atomic::AtomicBool =
489    std::sync::atomic::AtomicBool::new(false);
490/// Per-model router-GEMV form (2026-07-31): the 8-warp twin is +8.8% on the H100 q35
491/// decode step (router was 14.8% of it) with argmax + spec self-consistency green on
492/// qwen-class MoE both rigs. The gemma-4 26B knife-edge block (2026-07-31, single
493/// synthetic prompt) was RE-ARBITRATED 2026-08-01 on 6 real prompts — gate outcomes
494/// identical to the lone-warp arm, +13% g26 decode — so gemma4 rides the default too
495/// (research/g26-decode-20260801/). MEMRA_ROUTER_V2 env overrides either way.
496pub static ROUTER_W8_DEFAULT: std::sync::atomic::AtomicBool =
497    std::sync::atomic::AtomicBool::new(true);
498pub static FA_SP512_DEFAULT: std::sync::atomic::AtomicUsize =
499    std::sync::atomic::AtomicUsize::new(16);
500/// Per-model rms_norm block size (per-model numeric-config law: the per-thread partial-sum
501/// split changes with blockDim -> different FP order -> battery-arbitrated per model).
502/// qwen keeps the shipped 256; gemma4 adopts 1024 (single-row 2816-col norms are one-block
503/// latency-bound at 256 threads — 7us/launch measured).
504pub static RMS_BLOCK_DEFAULT: std::sync::atomic::AtomicU32 = std::sync::atomic::AtomicU32::new(256);
505/// gemma4 fa split ladder switch (set at model load; see fa_split_keys).
506pub static FA_SP_GEMMA: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
507/// Per-model stream-k entry override for SPEC serving (-1 = unset → env/default;
508/// 0 = force tiling; 1 = admit the deterministic form selector). The former timing
509/// selector made identical boots choose different fold orders; `MEMRA_MMQ_SK_FORM` is the
510/// explicit numerical-form seam. mmq_ffi reads this before the env.
511pub static MMQ_SK_FORCE: std::sync::atomic::AtomicI8 = std::sync::atomic::AtomicI8::new(-1);
512/// Per-model FP8-KV door — lives in memra-kv next to the format policy it drives
513/// (re-export keeps `crate::KV_FP8_FORCE` setters in model.rs/hybrid.rs working).
514pub use memra_kv::KV_FP8_FORCE;
515pub(crate) fn rms_block() -> u32 {
516    static V: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
517    *V.get_or_init(|| {
518        std::env::var("MEMRA_RMS_BLOCK")
519            .ok()
520            .and_then(|v| v.parse().ok())
521            .unwrap_or_else(|| RMS_BLOCK_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
522    })
523}
524
525pub(crate) fn fa_split_keys(t_kv: usize, n_head_kv: usize) -> usize {
526    static S: std::sync::OnceLock<Option<usize>> = std::sync::OnceLock::new();
527    if let Some(forced) = *S.get_or_init(|| {
528        std::env::var("MEMRA_FA_SPLIT")
529            .ok()
530            .and_then(|v| v.parse().ok())
531            .filter(|&s: &usize| s >= 8 && s % 8 == 0)
532    }) {
533        return forced;
534    }
535    // CTX-ADAPTIVE default (2026-07-05 40k sweep: sp32 24.5 vs sp128 26.0 tok/s = +5.8% — at
536    // deep ctx the n_splits count explodes (40k/32 = 1265 splits x 8 kv-heads) and the combine
537    // + partial-buffer cost dominates; at short ctx small splits fill the SMs). Exactness: split
538    // size only changes the PARTITION of keys; the rows/combine order per split is fixed and the
539    // gate battery (kernel-check + run-spec K=1..8) arbitrates every default change.
540    //
541    // SM-AWARE SHORT-CTX RUNG (2026-07-06 g7e): the 32-key rung was tuned on the 82-SM 5090.
542    // On 188 SMs the vec grid (n_head_kv x n_splits CTAs) starves at short ctx — the 35B has
543    // n_head_kv=2, so ctx128/split32 = 8 CTAs on 188 SMs. Measured on g7e (N=1 sweep + N=3
544    // interleaved confirm): 35B ctx128 sp16 179 vs sp32 161 (+11%), ctx512 178 vs 158, ctx2048
545    // flat, ctx>=4096 sp64 edges sp16 by ~3%; 27B ctx128 70.9 vs 66.3 (+7%); 9B 177 vs 163
546    // (+9%). Rigs <=100 SMs keep the validated 5090 ladder EXACTLY (default unchanged there —
547    // rig-divergence law: this branch is measured on 188 SMs only).
548    // gemma4 all-16 ladder probe REVERTED (2026-07-10): +1.3 plain at d1736 (157.5 vs 156.2)
549    // but depth VERIFY collapsed (spec 203.5 -> 169 — the windowed rows' per-row combine over
550    // 64 splits). The mixed default (swa nkv=8 -> 32, globals nkv=2 -> 8-ladder) stays; a
551    // caller-split policy would break row-vs-decode split parity. FA_SP_GEMMA kept as a seam.
552    if FA_SP_GEMMA.load(std::sync::atomic::Ordering::Relaxed)
553        && std::env::var("MEMRA_FA_SP16").as_deref() == Ok("1")
554    {
555        return if t_kv <= 8192 {
556            16
557        } else if t_kv <= 16384 {
558            64
559        } else {
560            128
561        };
562    }
563    let big_rig = fa_sm_count() >= 128;
564    if big_rig {
565        let _ = n_head_kv;
566        if t_kv <= 2048 {
567            16
568        } else if t_kv <= 16384 {
569            64
570        } else {
571            128
572        }
573    } else if n_head_kv <= 4 {
574        // KV-HEAD-AWARE RUNG (2026-07-08, 5090): the 8192->32 rung was validated on kv=8 models
575        // (27B/9B: 8 heads x n_splits fills 82 SMs). The 35B has n_head_kv=2 — at ctx512/sp32
576        // the vec grid is 2 x 20 = 40 CTAs on 82 SMs (half idle). Measured (35B, run-gen 128tok
577        // N=1 sweep + N=3 confirm): sp8 162.1 / sp16 161.3 / sp32 159.4 at short ctx.
578        // DEPTH TAPER (same day, the deep-ctx lesson re-learned on this rung): sp8 at d6257 =
579        // 782 splits -> combine + partial-buffer cost dominates (141.2 tok/s); the d6257 sweep
580        // says sp64 = 153.0 (sp16/32 147, sp96 147.6, sp128 141). Few-kv-head models need the
581        // taper EARLIER than kv=8 (per-split grid 4x thinner, same per-split combine cost).
582        // Crossover hunt: sp8 vs sp64 = 156.7/155.9 at d3072, 151.7/155.6 at d4096 -> boundary 3072.
583        // RUNG RE-SWEPT UNDER THE DEEP KERNEL (2026-08-02, lane/ladder-3072 — the stale-verdict
584        // law: the 3072 boundary was calibrated on the conflicted v4 core; the deep rewrite cut
585        // vec cost ~1.2-1.4x while combine scales with n_splits, so sp8's combine bill
586        // dominates far earlier). Kernel receipts (quiet-rig nsys, deep vec + combine us):
587        // d1024 sp8 17.1 vs sp64 10.6; d2048 31.0 vs 12.2; d3072 44.0 vs 18.3. e2e run-gen
588        // tg128 N=3 interleaved (KAT + q35, research/ladder-3072-20260802/): sp8 loses at
589        // EVERY depth >= 1024 (KAT d2048 182.6 vs 188.0 = -2.9%, d3072 175.9 vs 186.4 =
590        // -5.6%; q35 d4096 169.2 vs 182.6 = -7.4%); d512 flat (+-0.2%, inside noise). sp32
591        // ties sp64 within noise in the mid band and loses at d4096 -> no extra rung.
592        // Boundary 3072 -> 512: sp8 keeps only the short-ctx band it was validated on
593        // (ctx128-512); sp64 takes over where the deep kernel made combine the bill.
594        if t_kv <= 512 {
595            8
596        } else if t_kv <= 16384 {
597            64
598        } else {
599            128
600        }
601    } else {
602        if t_kv <= 8192 {
603            32
604        } else if t_kv <= 16384 {
605            64
606        } else {
607            128
608        }
609    }
610}
611
612/// SM count of device 0, cached (used by fa_split_keys' rig-size rung; primary-context query,
613/// same attribute Engine::batched_variant reads).
614fn fa_sm_count() -> i32 {
615    static N: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
616    *N.get_or_init(|| {
617        cudarc::driver::result::init().ok();
618        cudarc::driver::result::device::get(0)
619            .and_then(|d| unsafe { cudarc::driver::result::device::get_attribute(
620                d, cudarc::driver::sys::CUdevice_attribute_enum::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT) })
621            .unwrap_or(82)
622    })
623}
624
625/// FA-prefill kernel-name suffix for a head_dim (the template-stamped twins in flash_attn.cu):
626/// 256 = the original names (qwen35 class, dispatch unchanged), 128 = `_hd128` (MiniMax-M3).
627/// Any other dim errors — callers gate to sdpa_naive before dispatching FA.
628fn fa_hd_suffix(head_dim: usize) -> Result<&'static str, Box<dyn std::error::Error>> {
629    match head_dim {
630        256 => Ok(""),
631        128 => Ok("_hd128"),
632        d => Err(format!(
633            "fa_prefill: no kernel stamped for head_dim={d} (only 256/128); \
634                          callers must gate to sdpa_naive"
635        )
636        .into()),
637    }
638}
639
640/// Quant type codes matching qmatvec.cu QType enum.
641pub const QT_Q8_0: i32 = 0;
642pub const QT_Q4_K: i32 = 1;
643pub const QT_Q6_K: i32 = 2;
644pub const QT_Q5_K: i32 = 3;
645pub const QT_Q3_K: i32 = 4;
646pub const QT_IQ4_XS: i32 = 5;
647pub const QT_IQ3_S: i32 = 6;
648pub const QT_NVFP4: i32 = 7;
649/// Checkpoint-native FP8-E4M3 (MEMRA_ST_E4M3, lane e4m3dec): raw safetensors e4m3 weight bytes
650/// [out_f, in_f] row-major (row_bytes == in_f), per-tensor f32 weight_scale in GpuTensor `scale`
651/// (fused at the mmvq write / post-matmul scale_inplace). Decode = qmatvec_e4m3_mmvq (+ _b2/_b4/_b8
652/// batched twins); prefill (m>=16) = the cuBLASLt FP8 GEMM on the SAME resident bytes (fp8_ffi.rs)
653/// — ONE weight copy total, no Q8_0 re-encode duplicate.
654pub const QT_F8_E4M3: i32 = 10;
655/// Device-side tag for the A6 SPLIT-PLANE repacked NVFP4 layout (Stage-A generic kernel only;
656/// GpuTensor keeps qtype=QT_NVFP4 + an `rp` flag — this tag never lives in a GpuTensor).
657pub const QT_NVFP4_RP: i32 = 9;
658/// Unquantized f32 weight (safetensors MoE Path A: experts dequantized to f32 host-resident).
659pub const QT_F32: i32 = 8;
660pub const QT_BF16: i32 = 11;
661pub const QT_Q4_0: i32 = 12; // gemma-4 QAT GGUF weight format (18B/32: fp16 d + nibbles)
662/// GGUF Q2_K. Appended after the existing Q4_0 code so kernel ABI values do not move.
663/// Mixed-expert artifacts use the generic f32-dequant staged kernel until a target-rig-gated
664/// dp4a/MMQ implementation exists.
665pub const QT_Q2_K: i32 = 13;
666/// Checkpoint-native FP8-E4M3 with a BLOCK-128 weight-scale GRID (lane/fp8-blk128-decode,
667/// 2026-08-05) — the Qwen-official FP8 / DeepSeek-V3 scale class. Same raw e4m3 bytes as
668/// `QT_F8_E4M3` ([out_f, in_f] row-major, row_bytes == in_f), but the dequant scale is
669/// `GpuTensor::Quant.blk` (`Fp8BlockScales`, [ceil(out_f/128), ceil(in_f/128)] f32) and the
670/// scalar `scale` field is 1.0 by the layout contract.
671///
672/// WHY A DISTINCT CODE rather than `QT_F8_E4M3` + a `blk` flag: every existing QT_F8_E4M3
673/// consumer (qmatvec_e4m3_mmvq and its batched/fused twins, e4m3_fused_params,
674/// matmul_pre_dual_noscale's F8 arm, try_fp8_gemm) threads exactly ONE scalar weight scale. Under
675/// a shared code, any consumer that was not taught the grid would still MATCH and would dequant
676/// every tile at scale 1.0 — a silent numeric corruption. Under a distinct code every untaught
677/// consumer refuses loudly instead (`mmvq_supports`/`gemm_supports`/`mmq_supports` return false;
678/// the mmvq name match panics), so a missed dispatch site is a crash or a refusal receipt, never
679/// wrong numbers. Decode = `qmatvec_e4m3_blk_mmvq`; prefill (m>=16) = the per-block FP8 MMQ tile
680/// on the SAME resident bytes+grid (fp8_ffi::try_fp8_blk_mmq) — ONE weight copy total.
681pub const QT_F8_E4M3_BLK: i32 = 14;
682
683/// Engine device context: CUDA context, stream, loaded kernel modules, cuBLASLt (via runtime::Gpu).
684pub struct Engine {
685    pub gpu: memra_runtime::Gpu,
686    module: Arc<CudaModule>,
687    hybrid: Arc<CudaModule>,
688    qmatvec: Arc<CudaModule>,
689    flash: Arc<CudaModule>,
690    /// FP8-GLOBALS module (2026-07-11): the kf8vf8 fatbin loaded ALONGSIDE the default —
691    /// gemma GLOBAL layers (hd512) append + attend in e4m3 (dequant-latency arc, HANDOVER).
692    /// Lazy: loaded on first global-format use; None until then.
693    flash_g: std::sync::OnceLock<Arc<CudaModule>>,
694    gemm: Arc<CudaModule>,
695    router: Arc<CudaModule>,
696    /// Sampled-spec kernels (research/sampled-spec-impl-map.md piece A).
697    sample: Arc<CudaModule>,
698    /// EDGE-1 §B: one shared SLRU expert-residency cache, lazily built on first MoE dispatch under
699    /// MEMRA_MOE_CACHE. `Mutex` makes it multi-agent safe (§E.2); the lock covers only lookup/admit/
700    /// memcpy-issue (µs), NOT the GEMM, so streams still overlap. `None` => cache disabled.
701    moe_cache: Mutex<Option<crate::moe_cache::MoeSlotCache>>,
702    /// Exact retained expert-block lengths collected after model load. Mixed-layout models use
703    /// this inventory to preallocate fixed-address size classes instead of sizing every slot to
704    /// the single largest block. The cache still owns every address for its full lifetime.
705    moe_cache_layout: Mutex<Option<Vec<usize>>>,
706    /// CAPTURE-RETAIN mode (graph arc, 2026-07-12): while a graph capture (and its allocator
707    /// warmups) runs, every Engine allocation is ALSO kept alive here — a captured graph's
708    /// transient buffers must never return to the pool, or later allocations (e.g. the spec
709    /// verify between replays) reuse their addresses and the replay reads/writes live memory
710    /// (the draft-graph corruption root cause). Fast-path cost when off: one relaxed atomic.
711    capture_keep_on: std::sync::atomic::AtomicBool,
712    /// VERIFY-EXACT scope (dflash lane, 2026-07-13): when set, matmul/matmul_pre skip the
713    /// m>=16 prefill-GEMM branches so a t>=16 batched VERIFY rides the decode-exact b-tier
714    /// class (the parity law). The t=16 dflash verify tripped the GEMM threshold — 770us/
715    /// matmul (54% of the round) AND a different FP order than decode (issue-10 landmine).
716    verify_exact: std::sync::atomic::AtomicBool,
717    capture_keep: Mutex<Vec<Box<dyn std::any::Any + Send>>>,
718    /// EDGE-1 §C.2: dedicated H2D copy stream for async prefetch (event-synced to the compute stream).
719    pub copy_stream: Arc<CudaStream>,
720    /// Resident CUTLASS NVFP4 prefill scratch (workspace + a_packed + sfa_linear + sfa_sw + y + alpha),
721    /// allocated ONCE and grown to the largest prefill GEMM shape, then reused per-call. Removes the
722    /// 6 fresh allocations + alpha htod that `cutlass_fp4_gemm` did every prefill matmul (~200/prefill).
723    /// Safe as a single shared buffer because all GPU compute serializes on the one `gpu.stream` worker
724    /// thread (the server runs one GPU worker; no concurrent CUTLASS GEMMs share this scratch). `None`
725    /// until the first CUTLASS FP4 GEMM. Mutex guards lazy build/grow only (matches `moe_cache`).
726    #[cfg(memra_cutlass)]
727    cutlass_scratch: Mutex<Option<crate::cutlass_ffi::CutlassScratch>>,
728    /// FP8-ACT PREFILL scratch (MEMRA_PP_FP8): quantized-activation buffer + scale block + cuBLASLt
729    /// workspace, allocated once and grown to the largest prefill m*k (see fp8_ffi.rs). `None`
730    /// until the first FP8 prefill GEMM; Mutex guards lazy build/grow only (matches cutlass_scratch).
731    fp8_scratch: Mutex<Option<crate::fp8_ffi::Fp8Scratch>>,
732    /// f16-P/V door: pooled V re-encode buffer (bf16->f16) for the hd512 _pre path. Lazy-grow;
733    /// per-call cudaMalloc was a laptop-regression suspect (VRAM pressure, 31B nkv=4 = 4x bytes).
734    fa_vf16_scratch: Mutex<Option<CudaSlice<u8>>>,
735    /// Pooled fa-decode split partials (part_o, part_m, part_l): per-call zeros() was 3
736    /// alloc+memset pairs per fa launch (~144 mem nodes per decode token — the graph door's
737    /// residual launch tax) — lazy-grow, memset-prefix per use, stream-ordered reuse.
738    fa_part_pool: Mutex<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>>,
739    /// Retired fa-part pool generations (#68): old buffers whose addresses captured graphs may
740    /// have baked — kept alive for the Engine's lifetime instead of returning to the async pool
741    /// (see the RETIRE-ON-GROW comment at the realloc sites). Doubling growth bounds the total.
742    fa_part_retired: Mutex<Vec<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>>,
743    /// name -> resolved CudaFunction (capture-safe lookups; see `func`).
744    fn_cache: Mutex<std::collections::HashMap<String, CudaFunction>>,
745    f16_scratch: Mutex<Option<crate::f16_ffi::F16Scratch>>,
746    /// RANK1 LEVER (parallel argmax): resident pass-1 partials scratch (part_v[NB] f32, part_i[NB] i32),
747    /// allocated ONCE on first parallel-argmax call and reused. Stable pointers so the 2-pass argmax
748    /// is CUDA-graph-capturable (the buffer is referenced by both captured passes; lazy-allocated
749    /// before capture under the generate_graph tracking-off window so it carries no events).
750    argmax_partials: Mutex<Option<(CudaSlice<f32>, CudaSlice<i32>)>>,
751    /// ARC B (chunk-prime dequant-once): resident bf16 K/V workspace for `fa_prefill_view_ws`
752    /// ((K bytes, V bytes) u8 buffers holding [t_kv, kv_dim] bf16). Grown lazily to the largest
753    /// (t_kv, kv_dim) seen, REUSED across layers/chunks/calls (contents rewritten per launch —
754    /// safe because all compute serializes on the one gpu.stream). ~82MB at 40k ctx on the 27B.
755    prime_deqw_ws: Mutex<Option<(CudaSlice<u8>, CudaSlice<u8>)>>,
756    /// LAUNCH-STRUCTURE STAGE 1: persistent PINNED (cacheable, flags=0) host staging buffer for the
757    /// fused-router sel/w readback — one async DtoH pair + ONE sync instead of two synced dtohs.
758    /// Grown lazily; reused every MoE layer (single-threaded decode serializes on the sync).
759    router_stage: Mutex<Option<PinnedStage>>,
760}
761
762/// FAVENDOR lane env gate (2026-07-08): MEMRA_FA_V2=1 dispatches the llama-fattn-vec-mechanism
763/// decode kernels (fa_decode_vec_q_v2 / fa_decode_vec_q_rows_v2 / fa_decode_vec_q_v2_dc):
764/// tile-batched online softmax (one alpha rescale per 32-key tile instead of per key) + wide-load
765/// block dequant in the staging phase. NOTE rev2: llama's register streaming (no smem) was ALSO
766/// tried and measured 2x WORSE at depth in our gqa-warps frame — the smem KV-tile broadcast stays
767/// (see the kernel comment). NEW NUMERIC CONFIG (tile-level softmax regrouping changes FP order vs
768/// the per-key twins) — own argmax baseline; eager decode, the spec-verify rows path AND the
769/// graph _dc path switch TOGETHER (the spec-exactness law). Default OFF. Read per call (not
770/// OnceLock) so the gate battery can A/B within one process, matching the MEMRA_NO_FA_VEC pattern.
771fn fa_v2_on() -> bool {
772    // DEFAULT ON since 2026-07-08 (MEMRA_FA_V2=0 reverts): tile-batched online softmax, e2e
773    // measured across every model x depth — 35B 168.7->173.4 (d512) / 153.1->158.5 (d6257),
774    // 9B 131.2->132.7 / 108.4->124.5 (+15% — the engine-wide depth-slope fix), 27B 47.2->47.7 /
775    // 42.2->44.9. One-time numeric-config change; kernel-check + argmax + spec self-consistency
776    // + graph bit-identity green on all three models.
777    std::env::var("MEMRA_FA_V2")
778        .map(|v| v != "0")
779        .unwrap_or(true)
780}
781
782/// FA v3 gate (default ON since 2026-07-09; MEMRA_FA_V3=0 reverts to v2 — research/fa/fa_v3_design.md):
783/// HYBRID decode twins (fa_decode_vec_q_v3 / _rows_v3 / _v3_dc): llama's int8-dp4a K.Q with
784/// register-quantized Q (no K dequant, no K smem) + OUR CTA-shared staged bf16 V tile + OUR
785/// split partition/combine. NEW NUMERIC CONFIG (int8 Q quantization changes the K.Q accumulation
786/// vs the bf16-roundtrip FMA chain) — own argmax baseline; eager decode, the spec-verify rows
787/// path AND the graph _dc path switch TOGETHER (the spec-exactness law). Read per call so the
788/// gate battery can A/B within one process (the MEMRA_FA_V2 pattern).
789fn fa_v3_on() -> bool {
790    // DEFAULT ON since 2026-07-09 (MEMRA_FA_V3=0 reverts to v2): dp4a-K hybrid FA decode —
791    // fa kernel -21-23% at depth (micro), 35B spec p3 +5% (190->200, the last spec cell),
792    // d6257 +1.7%. Own numeric config; full battery green on 35B+9B incl graph bit-identity.
793    std::env::var("MEMRA_FA_V3")
794        .map(|v| v != "0")
795        .unwrap_or(true)
796}
797
798/// The v3 dp4a K path reads RAW q8_0 bytes (34B blocks) and stages q5_1 V verbatim — it is only
799/// correct on the DEFAULT KV formats — and needs dpl % 4 == 0 consecutive quants per lane
800/// (head_dim % 128 == 0; both daily models are hd256). All three dispatch sites share this
801/// predicate so the twins can never diverge.
802fn fa_v4_mode() -> &'static str {
803    static M: std::sync::OnceLock<String> = std::sync::OnceLock::new();
804    M.get_or_init(|| std::env::var("MEMRA_FA_V4").unwrap_or_default())
805}
806fn fa_v4_on() -> bool {
807    fa_v4_mode() != "0"
808} // DEFAULT ON 2026-07-10 (MEMRA_FA_V4=0 rollback)
809/// t_kv-conditional v4 pick (gemma depth lesson 2026-07-10: v4's key-per-lane pipeline starves
810/// at the 1024-window with short splits — MEMRA_FA_V4=0 measured depth plain 158.0 vs 156.7).
811/// Threshold MEMRA_FA_V4_MAX (default usize::MAX = unchanged behavior; gemma sets 1024 at load
812/// via FA_V4_MAX_DEFAULT). Applied at EVERY dispatch site (eager, rows, rows_w, dc) so verify
813/// stays kernel-family-identical to decode at the same t_kv.
814/// Per-model deep-ctx smem floor default (MEMRA_FA_SMEM_TKV env overrides): gemma pushes it
815/// above the 1024 window so the windowed decode + verify rows share the REGISTER family.
816pub static FA_SMEM_TKV_DEFAULT: std::sync::atomic::AtomicUsize =
817    std::sync::atomic::AtomicUsize::new(1024);
818pub static FA_V4_MAX_DEFAULT: std::sync::atomic::AtomicUsize =
819    std::sync::atomic::AtomicUsize::new(usize::MAX);
820pub fn fa_v4_at_pub(t_kv: usize) -> bool {
821    fa_v4_at(t_kv)
822}
823fn fa_v4_at(t_kv: usize) -> bool {
824    static M: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
825    let mx = *M.get_or_init(|| {
826        std::env::var("MEMRA_FA_V4_MAX")
827            .ok()
828            .and_then(|v| v.parse().ok())
829            .unwrap_or_else(|| FA_V4_MAX_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
830    });
831    fa_v4_on() && t_kv < mx
832}
833/// FA-DEEP gate (2026-08-02, lane fa-decode-deep): deep-ctx v4 twins
834/// (fa_decode_vec_q_v4_deep / _deep_dc) — the depth-decode lane's priced fix. Unlike
835/// v2/v3/v4 this is NOT a numeric config: the deep twins run the v4 program VERBATIM
836/// (same split partition, same softmax/accumulation order, same partials/combine) and only
837/// move the smem physical layout (bank de-conflict row pads) + the load schedule (next-tile
838/// L2 prefetch) — kernel-check pins bitdiff==0 vs the v4 twins across depths, so eager /
839/// rows-verify / graph / seqs stay mutually bit-identical wherever the threshold falls.
840/// Engages at t_kv >= MEMRA_FA_DEEP_MIN. The swept floor is 0 = ALWAYS ON where v4 ran
841/// (fa-deep-bench fine grid 96..6144, 2026-08-02: deep flat-or-better at EVERY depth,
842/// 1.01-1.26x, no losing cell — so there is no engagement boundary and no new
843/// capture-recapture edge; the env stays as a sweep/diagnostic seam only).
844/// MEMRA_FA_DEEP=0 is the rollback seam. Read per call so the battery + bench can A/B
845/// within one process (the v2/v3 pattern).
846pub const FA_DEEP_MIN_DEFAULT: usize = 0;
847fn fa_deep_at(t_kv: usize) -> bool {
848    if std::env::var("MEMRA_FA_DEEP").as_deref() == Ok("0") {
849        return false;
850    }
851    let min = std::env::var("MEMRA_FA_DEEP_MIN")
852        .ok()
853        .and_then(|v| v.parse().ok())
854        .unwrap_or(FA_DEEP_MIN_DEFAULT);
855    t_kv >= min
856}
857/// Public twin (kernel-check builds the deep-vs-v4 bit pin; bench sweeps the floor).
858pub fn fa_deep_at_pub(t_kv: usize) -> bool {
859    fa_deep_at(t_kv)
860}
861
862fn fa_v3_active(head_dim: usize) -> bool {
863    // v3's dp4a-K walk reads raw q8_0 K bytes — no e4m3 arm; the fp8-KV arm (MEMRA_KV_FP8)
864    // must fall back like any non-default KV format (the rows_dc stream path asserts on it).
865    fa_v3_on()
866        && head_dim % 128 == 0
867        && kv_cache_formats() == ("q8_0", "q5_1")
868        && !Engine::kv_fp8_on()
869}
870
871/// BATCHED-TICK increment 2 (2026-08-01): true iff a row at this t_kv would take the v4
872/// eager arm in `fa_decode_kvmod`'s dispatch — the exact precondition for the z-batched
873/// `fa_decode_vec_q_seqs_v4` twin to reproduce its per-seq program bit-identically.
874/// Mirrors the kvmod predicates: vec on + above the vec floor + hd256 + inside the v4
875/// window + the PRODUCTION v4 body (the noB3/stage phase probes are wrong-output) + the
876/// default flash module (no fp8-KV g-module). Callers must ALSO group rows on one
877/// `fa_split_keys` rung (the rows-twins' straddle law) before batching.
878pub fn fa_seqs_eligible(t_kv: usize, head_dim: usize) -> bool {
879    std::env::var("MEMRA_NO_FA_VEC").is_err()
880        && t_kv >= fa_vec_min_tkv()
881        && head_dim == 256
882        && fa_v4_at(t_kv)
883        && !matches!(fa_v4_mode(), "noB3" | "stage")
884        && !Engine::kv_fp8_on()
885}
886/// Public twin of the crate-private split ladder (kernel-check builds the seqs-vs-loop pin).
887pub fn fa_split_keys_pub(t_kv: usize, n_head_kv: usize) -> usize {
888    fa_split_keys(t_kv, n_head_kv)
889}
890
891/// A raw pinned (page-locked, CACHEABLE — flags=0, not write-combined) host allocation for
892/// DtoH staging. cudarc's `alloc_pinned` uses CU_MEMHOSTALLOC_WRITECOMBINED, which is right for
893/// HtoD streams but pathologically slow for host READS — the router readback is host-read-heavy,
894/// so we allocate through `result::malloc_host` with flags=0 directly.
895struct PinnedStage {
896    ptr: *mut u8,
897    cap: usize,
898}
899unsafe impl Send for PinnedStage {}
900impl PinnedStage {
901    fn new(cap: usize) -> Result<Self, Box<dyn std::error::Error>> {
902        let ptr = unsafe { cudarc::driver::result::malloc_host(cap, 0)? } as *mut u8;
903        Ok(PinnedStage { ptr, cap })
904    }
905}
906impl Drop for PinnedStage {
907    fn drop(&mut self) {
908        let _ = unsafe { cudarc::driver::result::free_host(self.ptr as _) };
909    }
910}
911
912/// Number of pass-1 blocks for the parallel argmax (fan-out across SMs to saturate HBM). 256 blocks
913/// x 256 threads = 65536 threads covering the 248K-vocab scan in ~4 strided loads/thread.
914pub const ARGMAX_NB: usize = 256;
915
916/// crate-visible alias for the batched FA3 shim entry (hybrid_forward's batch arm).
917pub(crate) use memra_fa3_vl as fa3_vl_raw;
918
919unsafe extern "C" {
920    /// FA3 v10 shim (cu/fa3_prefill.cu): TMA-swizzled wgmma FA, fresh causal hd256.
921    fn memra_fa3_prefill(
922        q16: *const core::ffi::c_void,
923        k16: *const core::ffi::c_void,
924        v16: *const core::ffi::c_void,
925        o: *mut f32,
926        t: i32,
927        h: i32,
928        hkv: i32,
929        d: i32,
930        scale: f32,
931        stream: *mut core::ffi::c_void,
932    ) -> i32;
933    /// batched varlen twin: host arrays of device pointers per seq (B <= 8).
934    pub(crate) fn memra_fa3_vl(
935        q16s: *const *const core::ffi::c_void,
936        k16s: *const *const core::ffi::c_void,
937        v16s: *const *const core::ffi::c_void,
938        os: *const *mut f32,
939        ts: *const i32,
940        b: i32,
941        h: i32,
942        hkv: i32,
943        d: i32,
944        scale: f32,
945        stream: *mut core::ffi::c_void,
946    ) -> i32;
947}
948
949/// STAGE-2 GROUPED DECODE: 8 expert weight-block device pointers passed BY VALUE as one kernel
950/// param (matches the CUDA `wptr8_t` struct: 8x 64-bit pointers, `#[repr(C)]` => identical
951/// layout). The pointers are SLRU cache-slot base addresses — fixed for the engine's lifetime
952/// (slots are never re-allocated), so passing raw values is stable across the launch.
953#[repr(C)]
954#[derive(Clone, Copy)]
955pub struct WPtr8(pub [u64; 8]);
956unsafe impl cudarc::driver::DeviceRepr for WPtr8 {}
957
958/// task #18 varlen GDN: per-seq args for gdn_chunk_{state,output}_mma_vl — one launch
959/// runs all B<=8 sequences' K4/K5 (CUDA `gdnseq_t`/`gdnvl_t`, layout-identical repr(C)).
960/// Raw addresses are valid for the launch: every referenced buffer outlives the call and
961/// all work is on the single compute stream (same discipline as the f16 GEMM FFI).
962#[repr(C)]
963#[derive(Clone, Copy, Default)]
964pub struct GdnSeqVl {
965    pub kb16: u64,
966    pub gcum: u64,
967    pub beta: u64,
968    pub u: u64,
969    pub wb16: u64,
970    pub y: u64,
971    pub ssnap: u64,
972    pub state_in: u64,
973    pub state_out: u64,
974    pub q: u64,
975    pub p: u64,
976    pub o: u64,
977    pub k: u64,
978    pub v: u64,
979    pub g: u64,
980    pub a: u64,
981    pub w: u64,
982    pub t: i32,
983    pub nc: i32,
984}
985unsafe impl cudarc::driver::DeviceRepr for GdnSeqVl {}
986#[repr(C)]
987#[derive(Clone, Copy)]
988pub struct GdnVl8(pub [GdnSeqVl; 8]);
989unsafe impl cudarc::driver::DeviceRepr for GdnVl8 {}
990
991/// task #22: per-seq wgmma-fused extras (CUDA `gdnw_t`/`gdnwvl_t`) — qb16 mirror +
992/// pre-masked Pb16, riding NEXT TO GdnSeqVl so the base struct stays untouched.
993#[repr(C)]
994#[derive(Clone, Copy, Default)]
995pub struct GdnWVl {
996    pub qb16: u64,
997    pub pb16: u64,
998}
999unsafe impl cudarc::driver::DeviceRepr for GdnWVl {}
1000#[repr(C)]
1001#[derive(Clone, Copy)]
1002pub struct GdnWVl8(pub [GdnWVl; 8]);
1003unsafe impl cudarc::driver::DeviceRepr for GdnWVl8 {}
1004
1005/// task #18 increment 3: per-seq PREP/TAIL args (CUDA `gdnprep_t`/`gdnprepvl_t`).
1006#[repr(C)]
1007#[derive(Clone, Copy, Default)]
1008pub struct GdnPrepVl {
1009    pub qkv: u64,
1010    pub conv_state: u64,
1011    pub conv_out: u64,
1012    pub q_g: u64,
1013    pub k_g: u64,
1014    pub v_g: u64,
1015    pub q_l2: u64,
1016    pub k_l2: u64,
1017    pub beta_raw: u64,
1018    pub alpha: u64,
1019    pub beta: u64,
1020    pub g_log: u64,
1021    pub o: u64,
1022    pub z: u64,
1023    pub gn: u64,
1024    pub gn16: u64,
1025    pub kb16: u64,
1026    pub qb16: u64,
1027    pub t: i32,
1028    pub pad: i32,
1029}
1030unsafe impl cudarc::driver::DeviceRepr for GdnPrepVl {}
1031#[repr(C)]
1032#[derive(Clone, Copy)]
1033pub struct GdnPrepVl8(pub [GdnPrepVl; 8]);
1034unsafe impl cudarc::driver::DeviceRepr for GdnPrepVl8 {}
1035
1036/// task #18 (attn side): per-seq varlen FA args (CUDA `faseq_t`/`favl_t`).
1037#[repr(C)]
1038#[derive(Clone, Copy, Default)]
1039pub struct FaSeqVl {
1040    pub q: u64,
1041    pub k16: u64,
1042    pub v16: u64,
1043    pub o: u64,
1044    pub kf: u64,
1045    pub vf: u64,
1046    pub t: i32,
1047    pub pad: i32,
1048}
1049unsafe impl cudarc::driver::DeviceRepr for FaSeqVl {}
1050#[repr(C)]
1051#[derive(Clone, Copy)]
1052pub struct FaVl8(pub [FaSeqVl; 8]);
1053unsafe impl cudarc::driver::DeviceRepr for FaVl8 {}
1054
1055/// task #18 (attn pre-FA): per-seq split/norm/rope/append args (CUDA `attnpre_t`).
1056#[repr(C)]
1057#[derive(Clone, Copy, Default)]
1058pub struct AttnPreVl {
1059    pub qf: u64,
1060    pub kf: u64,
1061    pub vf: u64,
1062    pub q: u64,
1063    pub gate: u64,
1064    pub qn: u64,
1065    pub kn: u64,
1066    pub kc: u64,
1067    pub vc: u64,
1068    pub t: i32,
1069    pub pad: i32,
1070}
1071unsafe impl cudarc::driver::DeviceRepr for AttnPreVl {}
1072#[repr(C)]
1073#[derive(Clone, Copy)]
1074pub struct AttnPreVl8(pub [AttnPreVl; 8]);
1075unsafe impl cudarc::driver::DeviceRepr for AttnPreVl8 {}
1076
1077/// task #18 increment 2: one sequence's FULL chunk-buffer set (alloc-only; the
1078/// varlen K1-K5 chain fills them).
1079pub struct GdnChunkBufs {
1080    pub gcum: CudaSlice<f32>,
1081    pub a: CudaSlice<f32>,
1082    pub p: CudaSlice<f32>,
1083    pub u: CudaSlice<f32>,
1084    pub w: CudaSlice<f32>,
1085    pub kb16: CudaSlice<u8>,
1086    pub wb16: CudaSlice<u8>,
1087    pub y16: CudaSlice<u8>,
1088    pub ssnap16: CudaSlice<u8>,
1089    pub qb16: CudaSlice<u8>,
1090    pub pb16: CudaSlice<u8>,
1091    pub o: CudaSlice<f32>,
1092    pub t: usize,
1093    pub nc: usize,
1094}
1095
1096/// STAGE-2 GROUPED DECODE: the 8 routed-expert weights by value (CUDA `f32x8_t`).
1097#[repr(C)]
1098#[derive(Clone, Copy)]
1099pub struct F32x8(pub [f32; 8]);
1100unsafe impl cudarc::driver::DeviceRepr for F32x8 {}
1101
1102/// Harness timing contract: wall nanos of the LAST generate/generate_spec prompt prime on this
1103/// process. Bench binaries read it right after the call to print gen-only throughput without the
1104/// prime-subtraction hack (which amplifies prime jitter into the gen number at long prompts).
1105pub static PRIME_NANOS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1106
1107impl Engine {
1108    pub fn new(ordinal: usize) -> Result<Self, Box<dyn std::error::Error>> {
1109        let gpu = memra_runtime::Gpu::new(ordinal)?;
1110        // ARCH GUARD (unified dual-arch engine): the fatbins carry single-arch SASS, so a
1111        // binary/device mismatch otherwise dies at first module load with an opaque CUDA
1112        // error. Fail early with the rebuild hint instead. MEMRA_ARCH_CHECK=0 skips.
1113        if std::env::var("MEMRA_ARCH_CHECK").as_deref() != Ok("0") {
1114            use cudarc::driver::sys::CUdevice_attribute_enum as A;
1115            let (maj, min) = cudarc::driver::result::device::get(ordinal as i32)
1116                .and_then(|d| unsafe {
1117                    Ok((
1118                        cudarc::driver::result::device::get_attribute(
1119                            d,
1120                            A::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR,
1121                        )?,
1122                        cudarc::driver::result::device::get_attribute(
1123                            d,
1124                            A::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR,
1125                        )?,
1126                    ))
1127                })
1128                .unwrap_or((0, 0));
1129            let built = env!("MEMRA_BUILT_CUDA_ARCH");
1130            let ok = matches!(
1131                (built, maj, min),
1132                ("120a", 12, 0) | ("120a", 12, 1) | ("100a", 10, 0) | ("90a", 9, 0) | ("89", 8, 9)
1133            );
1134            if !ok {
1135                return Err(format!(
1136                    "memra was built for sm_{built} but device {ordinal} reports compute \
1137                     capability {maj}.{min}. Rebuild on this machine (MEMRA_CUDA_ARCH \
1138                     auto-detects the GPU) or set MEMRA_ARCH_CHECK=0 to bypass."
1139                )
1140                .into());
1141            }
1142        }
1143        // Default async-pool RELEASE_THRESHOLD is 0: freed blocks return to the OS at every
1144        // sync, so cuMemAllocAsync NODES inside captured graphs re-map memory on EVERY
1145        // cuGraphLaunch (measured 226us/launch on the gemma graph door, 2026-07-23 osrt).
1146        // Pinning the threshold keeps the pool cached -> alloc nodes become pointer bumps.
1147        unsafe {
1148            use cudarc::driver::sys;
1149            let dev: sys::CUdevice = ordinal as sys::CUdevice;
1150            let mut pool: sys::CUmemoryPool = std::ptr::null_mut();
1151            if sys::cuDeviceGetDefaultMemPool(&mut pool, dev) == sys::CUresult::CUDA_SUCCESS {
1152                let mut thresh: u64 = u64::MAX;
1153                let _ = sys::cuMemPoolSetAttribute(
1154                    pool,
1155                    sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_RELEASE_THRESHOLD,
1156                    &mut thresh as *mut u64 as *mut core::ffi::c_void,
1157                );
1158            }
1159        }
1160        let module = gpu.ctx.load_module(Ptx::from_binary(FATBIN.to_vec()))?;
1161        let hybrid = gpu
1162            .ctx
1163            .load_module(Ptx::from_binary(HYBRID_FATBIN.to_vec()))?;
1164        let qmatvec = gpu
1165            .ctx
1166            .load_module(Ptx::from_binary(QMATVEC_FATBIN.to_vec()))?;
1167        let flash = gpu
1168            .ctx
1169            .load_module(Ptx::from_binary(flash_fatbin_bytes().to_vec()))?;
1170        let gemm = gpu
1171            .ctx
1172            .load_module(Ptx::from_binary(gemm_fatbin_bytes().into_owned()))?;
1173        let router = gpu
1174            .ctx
1175            .load_module(Ptx::from_binary(ROUTER_FATBIN.to_vec()))?;
1176        let sample = gpu
1177            .ctx
1178            .load_module(Ptx::from_binary(SAMPLE_FATBIN.to_vec()))?;
1179        let copy_stream = gpu.ctx.new_stream()?;
1180        // DECODE EVENT-TRACKING ELISION — DEFAULT ON (2026-07-05; MEMRA_EVT=1 = escape hatch).
1181        // cudarc is in multi-stream mode (main stream +
1182        // copy_stream are both created streams), so with tracking on EVERY launch arg records a
1183        // read/write CudaEvent and inserts cuStreamWaitEvent on prior events. On the 35B MoE decode
1184        // that is ~19k cuStreamWaitEvent + ~9k cuEventRecord + ~6k event create/destroy per token
1185        // (~7 ms/tok host time, measured nsys 2026-07-04 g7e), and +4.6% measured on 27B decode —
1186        // protecting NOTHING: every hot-path kernel/memcpy runs on the ONE gpu.stream.
1187        // CROSS-STREAM HAZARD AUDIT: MoeSlotCache in-memory prefetch uses copy_stream. Every
1188        // overwrite explicitly records the prior compute point and makes copy_stream wait; every
1189        // consumer explicitly waits for the copy completion event. The opt-in positioned-read
1190        // proof stays on gpu.stream and retains an explicit event solely to guard pinned-source
1191        // reuse. Graph-capture sites use only gpu.stream, so these handoffs never rely on cudarc's
1192        // implicit event tracking.
1193        // SAFETY: single-stream ordering is total; the runtime mem-pool is configured with
1194        // internal-dependency reuse (memra-runtime), so alloc reuse is stream-ordered too.
1195        if std::env::var("MEMRA_EVT")
1196            .map(|v| v == "1")
1197            .unwrap_or(false)
1198        {
1199            // escape hatch: keep cudarc's implicit cross-stream event tracking.
1200        } else {
1201            unsafe {
1202                gpu.ctx.disable_event_tracking();
1203            }
1204        }
1205        Ok(Self {
1206            gpu,
1207            module,
1208            hybrid,
1209            qmatvec,
1210            flash,
1211            flash_g: std::sync::OnceLock::new(),
1212            gemm,
1213            router,
1214            sample,
1215            moe_cache: Mutex::new(None),
1216            moe_cache_layout: Mutex::new(None),
1217            copy_stream,
1218            capture_keep_on: std::sync::atomic::AtomicBool::new(false),
1219            verify_exact: std::sync::atomic::AtomicBool::new(false),
1220            capture_keep: Mutex::new(Vec::new()),
1221            argmax_partials: Mutex::new(None),
1222            prime_deqw_ws: Mutex::new(None),
1223            router_stage: Mutex::new(None),
1224            fp8_scratch: Mutex::new(None),
1225            fa_vf16_scratch: Mutex::new(None),
1226            fa_part_pool: Mutex::new(None),
1227            fa_part_retired: Mutex::new(Vec::new()),
1228            fn_cache: Mutex::new(Default::default()),
1229            f16_scratch: Mutex::new(None),
1230            #[cfg(memra_cutlass)]
1231            cutlass_scratch: Mutex::new(None),
1232        })
1233    }
1234
1235    pub fn ctx(&self) -> &Arc<CudaContext> {
1236        &self.gpu.ctx
1237    }
1238
1239    /// Bytes the async pool holds MAPPED but NOT LIVE (reserved - used), i.e. freed blocks
1240    /// parked in the pool because `Engine::new` pins RELEASE_THRESHOLD to u64::MAX above.
1241    ///
1242    /// Why this is a public engine surface: `mem_get_info`'s `free` DOES NOT SEE these bytes —
1243    /// they are mapped to this process, so `free` counts them as gone, yet the very next
1244    /// `alloc_u8` is satisfied from them without touching `free` at all. Any admission or
1245    /// budget decision that reads `free` alone therefore under-counts real headroom by exactly
1246    /// this amount. Effective allocatable headroom is `free + pool_cached_bytes()`.
1247    ///
1248    /// MEASURED SIZE (c=64 serve burst, 9B NVFP4 + draft, 24GB card, 2026-08-06): 34-89 MB
1249    /// during the burst — SMALL. The admission gate adds it because a term that can only ever
1250    /// under-count headroom does not belong in a gate that queues real work, but the honest
1251    /// reading of this number is that pool caching is NOT where a long-running server's VRAM
1252    /// hides on this path: reserved ~= used throughout, so the memory the driver reports as
1253    /// gone is genuinely LIVE (see `pool_reserved_used` for the diagnostic pair).
1254    ///
1255    /// Returns 0 if the pool cannot be queried (never a false-positive headroom claim).
1256    pub fn pool_cached_bytes(&self) -> usize {
1257        let (reserved, used) = self.pool_reserved_used();
1258        reserved.saturating_sub(used)
1259    }
1260
1261    /// Raw async-pool occupancy: (RESERVED_MEM_CURRENT, USED_MEM_CURRENT) in bytes. Reserved is
1262    /// what the pool has mapped from the driver; used is what is live inside it. Exposed for
1263    /// admission/VRAM diagnostics — the pair distinguishes "memory is parked in the pool and
1264    /// `free` cannot see it" (reserved >> used) from "memory is genuinely held live by some
1265    /// owner" (reserved ~= used), which are opposite bugs with opposite fixes.
1266    /// (0, 0) if the pool cannot be queried.
1267    pub fn pool_reserved_used(&self) -> (usize, usize) {
1268        use cudarc::driver::sys;
1269        unsafe {
1270            let mut pool: sys::CUmemoryPool = std::ptr::null_mut();
1271            if sys::cuDeviceGetDefaultMemPool(&mut pool, self.gpu.ctx.ordinal() as sys::CUdevice)
1272                != sys::CUresult::CUDA_SUCCESS
1273            {
1274                return (0, 0);
1275            }
1276            let (mut reserved, mut used) = (0u64, 0u64);
1277            if sys::cuMemPoolGetAttribute(
1278                pool,
1279                sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_RESERVED_MEM_CURRENT,
1280                &mut reserved as *mut u64 as *mut core::ffi::c_void,
1281            ) != sys::CUresult::CUDA_SUCCESS
1282            {
1283                return (0, 0);
1284            }
1285            if sys::cuMemPoolGetAttribute(
1286                pool,
1287                sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_USED_MEM_CURRENT,
1288                &mut used as *mut u64 as *mut core::ffi::c_void,
1289            ) != sys::CUresult::CUDA_SUCCESS
1290            {
1291                return (0, 0);
1292            }
1293            (reserved as usize, used as usize)
1294        }
1295    }
1296
1297    /// Ambient stream (by value since M1-PP2 increment 2): the thread's pp2 stage stream
1298    /// when a stage scope is active, else the main compute stream — see `Gpu::stream`.
1299    pub fn stream(&self) -> Arc<CudaStream> {
1300        self.gpu.stream()
1301    }
1302    /// FP8-GLOBALS switch (MEMRA_GEMMA_GKV, default ON): gemma global (hd512) layers keep
1303    /// their KV in e4m3 — the dequant-latency arc (HANDOVER). Windowed layers stay q8_0/q5_1.
1304    pub fn gkv_on() -> bool {
1305        memra_kv::gkv_on()
1306    }
1307
1308    /// FP8-WINDOWED switch (MEMRA_GEMMA_WKV — measured 2026-07-12 in a validity-gated
1309    /// window: 1.7k 174.1-174.4 vs 168.6-169.4 default (+3%), 4.9k 158.7-160.4; vs llama
1310    /// same-window 159.5-160.2 / 140.6 = 1.09x / 1.13x): gemma windowed (hd256 SWA)
1311    /// layers hold e4m3 KV and ride the format-aware v4 lane from the kf8vf8 module.
1312    /// SERVING-MODE DEFAULT (2026-07-12, the 31B spec unlock): fp8-windowed KV GUTS the
1313    /// MTP drafter's acceptance — its single swa attention reads the windowed cache and
1314    /// e4m3 noise flips its argmaxes (31B short accept .758 -> 1.000 with q8/q5, spec 88
1315    /// -> 122.7 vs llama-mtp 112; depth .59 -> .78; 26B depth .57 -> .89). So the default
1316    /// keys on serving intent: SPEC serving (MEMRA_DRAFT set) -> OFF, plain -> ON (its
1317    /// depth-plain +3% stands). Explicit MEMRA_GEMMA_WKV always wins. GKV (globals) stays
1318    /// ON for both — no acceptance cost measured.
1319    pub fn wkv_on() -> bool {
1320        memra_kv::wkv_on()
1321    }
1322
1323    /// QWEN FP8-KV switch (MEMRA_KV_FP8 explicit; else the per-model KV_FP8_FORCE door set
1324    /// at model load; else OFF). Non-gemma full-attn layers hold e4m3 K/V via the kf8vf8
1325    /// module. Per-model verdict 2026-07-12: 9B +0.7-4% scaling with depth, 27B flat,
1326    /// 35B −2% (fp8 format-gates its v3 dp4a lane) — so the 9B class defaults ON
1327    /// (adopted 2026-07-28 with the deferred acceptance battery), others stay OFF.
1328    pub fn kv_fp8_on() -> bool {
1329        memra_kv::kv_fp8_on()
1330    }
1331
1332    /// fa kernel routed by head_dim: hd512 (gemma globals) resolves from the kf8vf8 module
1333    /// when the fp8-globals arm is on; everything else from the default flash module.
1334    fn fa_func(&self, name: &str, head_dim: usize) -> CudaFunction {
1335        if head_dim == 512 && Self::gkv_on() {
1336            self.func_g(name)
1337        } else {
1338            self.func(name)
1339        }
1340    }
1341
1342    /// Kernel from the FP8-GLOBALS (kf8vf8) flash module — gemma global-layer arm only.
1343    /// Format-AGNOSTIC kernels (e.g. fa_decode_combine_f32) are not compiled into the
1344    /// per-format fatbins; fall back to the base modules for those.
1345    fn func_g(&self, name: &str) -> CudaFunction {
1346        let m = self.flash_g.get_or_init(|| {
1347            self.gpu
1348                .ctx
1349                .load_module(cudarc::nvrtc::Ptx::from_binary(
1350                    FLASH_FATBIN_KF8VF8.to_vec(),
1351                ))
1352                .expect("load kf8vf8 flash fatbin (fp8-globals arm)")
1353        });
1354        let key = format!("g:{name}");
1355        if let Some(f) = self.fn_cache.lock().unwrap().get(&key) {
1356            return f.clone();
1357        }
1358        let f = match m.load_function(name) {
1359            Ok(f) => f,
1360            Err(_) => self.func(name),
1361        };
1362        self.fn_cache.lock().unwrap().insert(key, f.clone());
1363        f
1364    }
1365
1366    fn func(&self, name: &str) -> CudaFunction {
1367        // Resolution cache: cuModuleGetFunction fails inside a CUDA-graph capture region,
1368        // so capture-time lookups MUST be host-memory hits (warmups populate the cache).
1369        if let Some(f) = self.fn_cache.lock().unwrap().get(name) {
1370            return f.clone();
1371        }
1372        let f = self
1373            .module
1374            .load_function(name)
1375            .or_else(|_| self.hybrid.load_function(name))
1376            .or_else(|_| self.qmatvec.load_function(name))
1377            .or_else(|_| self.flash.load_function(name))
1378            .or_else(|_| self.gemm.load_function(name))
1379            .or_else(|_| self.router.load_function(name))
1380            .or_else(|_| self.sample.load_function(name))
1381            .unwrap_or_else(|_| panic!("kernel {name} not in any fatbin"));
1382        self.fn_cache
1383            .lock()
1384            .unwrap()
1385            .insert(name.to_string(), f.clone());
1386        f
1387    }
1388
1389    /// Scatter trimmed draft logits into full-vocab space: dst = -inf everywhere, then
1390    /// dst[d2t[i]] = src[i]. Two launches (fill, scatter) — no grid-wide sync needed.
1391    pub fn scatter_trim_logits(
1392        &self,
1393        src: &CudaSlice<f32>,
1394        d2t: &CudaSlice<u32>,
1395        dst: &mut CudaSlice<f32>,
1396        d_vocab: usize,
1397        n_vocab: usize,
1398    ) -> Result<(), Box<dyn std::error::Error>> {
1399        let f1 = self.func("scatter_trim_logits_f32");
1400        let f2 = self.func("scatter_trim_logits_pass2_f32");
1401        let (dv, nv) = (d_vocab as i32, n_vocab as i32);
1402        let cfg1 = LaunchConfig {
1403            grid_dim: (256, 1, 1),
1404            block_dim: (256, 1, 1),
1405            shared_mem_bytes: 0,
1406        };
1407        let __s_b1 = self.gpu.stream();
1408        let mut b1 = __s_b1.launch_builder(&f1);
1409        b1.arg(src).arg(d2t).arg(&mut *dst).arg(&dv).arg(&nv);
1410        unsafe {
1411            b1.launch(cfg1)?;
1412        }
1413        let cfg2 = LaunchConfig {
1414            grid_dim: (d_vocab.div_ceil(256) as u32, 1, 1),
1415            block_dim: (256, 1, 1),
1416            shared_mem_bytes: 0,
1417        };
1418        let __s_b2 = self.gpu.stream();
1419        let mut b2 = __s_b2.launch_builder(&f2);
1420        b2.arg(src).arg(d2t).arg(&mut *dst).arg(&dv);
1421        unsafe {
1422            b2.launch(cfg2)?;
1423        }
1424        Ok(())
1425    }
1426
1427    // ---- FILTERED-SPEC (feat/filtered-spec): top-k/p/min-p transforms applied symmetrically
1428    // to p and q — rejection sampling stays distribution-exact for the filtered target. ----
1429
1430    /// Per-row filtered-softmax stats: out[r] = (threshold_e, renorm_mass_e, row_max) for the
1431    /// filter (top_k, top_p, min_p) at `temp`. Rows index into x with row_stride f32s.
1432    #[allow(clippy::too_many_arguments)]
1433    pub fn filter_stats(
1434        &self,
1435        x: &CudaSlice<f32>,
1436        row_stride: usize,
1437        rows: &CudaSlice<i32>,
1438        out_th: &mut CudaSlice<f32>,
1439        out_z: &mut CudaSlice<f32>,
1440        out_max: &mut CudaSlice<f32>,
1441        n: usize,
1442        nrow: usize,
1443        temp: f32,
1444        top_k: i32,
1445        top_p: f32,
1446        min_p: f32,
1447    ) -> Result<(), Box<dyn std::error::Error>> {
1448        // A top-K-selection form of this kernel (3 vocab passes vs the search's ~51) was
1449        // implemented and REFUTED on 2026-08-21 (lane/moebatch-q35moe): the 248k-vocab row is
1450        // L2-resident, so the extra passes are near-free while the per-thread selection list
1451        // spills to local memory — B=8 tick 12.8/11.2 ms (cap 64/32) vs 10.4 ms for this
1452        // kernel, and serve c8 agg ~648 vs ~666. The receipts row is the record; the real
1453        // filtered-sampling win was batching the per-row launches (decode_batch.rs).
1454        //
1455        // COOPERATIVE MULTI-BLOCK FORM (lane/samplat, 2026-08-21): the surviving inefficiency
1456        // was WIDTH, not passes — one block per row left ~94% of the device idle for ~620us
1457        // per B=8 serve tick (5.9% of the tick, box4 nsys receipt). filter_stats_coop_f32
1458        // splits each row across 16 blocks with grid-synced bisection totals — same algorithm,
1459        // slice-partial f32 sums (accepted device-sampling class; sample-check arbitrates).
1460        // Admission: cooperative grid must co-reside (16*nrow blocks vs SM count).
1461        // MEMRA_FILTER_COOP=0 is the rollback seam to the single-block form.
1462        static COOP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1463        let coop_on =
1464            *COOP_ON.get_or_init(|| std::env::var("MEMRA_FILTER_COOP").as_deref() != Ok("0"));
1465        let (ni, nr, rs) = (n as i32, nrow as i32, row_stride as i64);
1466        if coop_on && 16 * nrow <= self.sm_count() as usize {
1467            let f = self.func("filter_stats_coop_f32");
1468            let mut ws = self.alloc_uninit::<f32>(nrow * (2 * 16 + 2))?;
1469            let cfg = LaunchConfig {
1470                grid_dim: (16, nrow as u32, 1),
1471                block_dim: (512, 1, 1),
1472                shared_mem_bytes: 0,
1473            };
1474            let __s_b = self.gpu.stream();
1475            let mut b = __s_b.launch_builder(&f);
1476            b.arg(x)
1477                .arg(&rs)
1478                .arg(rows)
1479                .arg(&mut *out_th)
1480                .arg(&mut *out_z)
1481                .arg(&mut *out_max)
1482                .arg(&mut ws)
1483                .arg(&ni)
1484                .arg(&nr)
1485                .arg(&temp)
1486                .arg(&top_k)
1487                .arg(&top_p)
1488                .arg(&min_p);
1489            unsafe {
1490                b.launch_cooperative(cfg)?;
1491            }
1492            return Ok(());
1493        }
1494        let f = self.func("filter_stats_f32");
1495        let cfg = LaunchConfig {
1496            grid_dim: (nrow as u32, 1, 1),
1497            block_dim: (1024, 1, 1),
1498            shared_mem_bytes: 0,
1499        };
1500        let __s_b = self.gpu.stream();
1501        let mut b = __s_b.launch_builder(&f);
1502        b.arg(x)
1503            .arg(&rs)
1504            .arg(rows)
1505            .arg(&mut *out_th)
1506            .arg(&mut *out_z)
1507            .arg(&mut *out_max)
1508            .arg(&ni)
1509            .arg(&nr)
1510            .arg(&temp)
1511            .arg(&top_k)
1512            .arg(&top_p)
1513            .arg(&min_p);
1514        unsafe {
1515            b.launch(cfg)?;
1516        }
1517        Ok(())
1518    }
1519
1520    /// out[pair] = filtered-softmax prob of ids[pair] in row rows[pair] (th/z per PAIR).
1521    #[allow(clippy::too_many_arguments)]
1522    pub fn softmax_gather_filtered(
1523        &self,
1524        x: &CudaSlice<f32>,
1525        row_stride: usize,
1526        ids: &CudaSlice<u32>,
1527        rows: &CudaSlice<i32>,
1528        th: &CudaSlice<f32>,
1529        z: &CudaSlice<f32>,
1530        out: &mut CudaSlice<f32>,
1531        n: usize,
1532        npair: usize,
1533        temp: f32,
1534    ) -> Result<(), Box<dyn std::error::Error>> {
1535        let f = self.func("softmax_gather_filtered_f32");
1536        let (ni, np, rs) = (n as i32, npair as i32, row_stride as i64);
1537        let cfg = LaunchConfig {
1538            grid_dim: (npair as u32, 1, 1),
1539            block_dim: (256, 1, 1),
1540            shared_mem_bytes: 0,
1541        };
1542        let __s_b = self.gpu.stream();
1543        let mut b = __s_b.launch_builder(&f);
1544        b.arg(x)
1545            .arg(&rs)
1546            .arg(ids)
1547            .arg(rows)
1548            .arg(th)
1549            .arg(z)
1550            .arg(&mut *out)
1551            .arg(&ni)
1552            .arg(&np)
1553            .arg(&temp);
1554        unsafe {
1555            b.launch(cfg)?;
1556        }
1557        Ok(())
1558    }
1559
1560    /// Filtered residual sample: token ~ norm(max(0, fp - fq)) with fp/fq the filtered softmaxes.
1561    #[allow(clippy::too_many_arguments)]
1562    pub fn residual_sample_filtered(
1563        &self,
1564        p: &CudaSlice<f32>,
1565        q: Option<&CudaSlice<f32>>,
1566        n: usize,
1567        temp: f32,
1568        seed: u64,
1569        stream_pos: u32,
1570        p_stats: (f32, f32, f32),
1571        q_stats: (f32, f32, f32),
1572        out_tok: &mut CudaSlice<u32>,
1573    ) -> Result<(), Box<dyn std::error::Error>> {
1574        let f = self.func("residual_sample_filtered_f32");
1575        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
1576        let has_q: i32 = q.is_some() as i32;
1577        let qbuf = q.unwrap_or(p);
1578        let (pm, pth, pz) = p_stats;
1579        let (qm, qth, qz) = q_stats;
1580        let cfg = LaunchConfig {
1581            grid_dim: (1, 1, 1),
1582            block_dim: (1024, 1, 1),
1583            shared_mem_bytes: 0,
1584        };
1585        let __s_b = self.gpu.stream();
1586        let mut b = __s_b.launch_builder(&f);
1587        b.arg(p)
1588            .arg(qbuf)
1589            .arg(&has_q)
1590            .arg(&ni)
1591            .arg(&temp)
1592            .arg(&slo)
1593            .arg(&shi)
1594            .arg(&stream_pos)
1595            .arg(&pm)
1596            .arg(&pth)
1597            .arg(&pz)
1598            .arg(&qm)
1599            .arg(&qth)
1600            .arg(&qz)
1601            .arg(&mut *out_tok);
1602        unsafe {
1603            b.launch(cfg)?;
1604        }
1605        Ok(())
1606    }
1607
1608    /// Gumbel-max draw from the FILTERED distribution (masked perturb; argmax after).
1609    #[allow(clippy::too_many_arguments)]
1610    pub fn gumbel_perturb_filtered(
1611        &self,
1612        x: &CudaSlice<f32>,
1613        y: &mut CudaSlice<f32>,
1614        n: usize,
1615        seed: u64,
1616        stream_pos: u32,
1617        temp: f32,
1618        row_max: f32,
1619        th: f32,
1620    ) -> Result<(), Box<dyn std::error::Error>> {
1621        let f = self.func("gumbel_perturb_filtered_f32");
1622        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
1623        let cfg = LaunchConfig {
1624            grid_dim: (n.div_ceil(256) as u32, 1, 1),
1625            block_dim: (256, 1, 1),
1626            shared_mem_bytes: 0,
1627        };
1628        let __s_b = self.gpu.stream();
1629        let mut b = __s_b.launch_builder(&f);
1630        b.arg(x)
1631            .arg(&mut *y)
1632            .arg(&ni)
1633            .arg(&slo)
1634            .arg(&shi)
1635            .arg(&stream_pos)
1636            .arg(&temp)
1637            .arg(&row_max)
1638            .arg(&th);
1639        unsafe {
1640            b.launch(cfg)?;
1641        }
1642        Ok(())
1643    }
1644
1645    /// Keskar penalties applied IN PLACE to a logits buffer: history token ids get
1646    /// rep-divided/multiplied + freq*count + presence subtracted. Symmetric p/q usage keeps
1647    /// filtered rejection sampling exact for the penalized target.
1648    #[allow(clippy::too_many_arguments)]
1649    pub fn penalize_logits(
1650        &self,
1651        x: &mut CudaSlice<f32>,
1652        hist: &CudaSlice<u32>,
1653        n_hist: usize,
1654        rep: f32,
1655        freq: f32,
1656        present: f32,
1657        n: usize,
1658    ) -> Result<(), Box<dyn std::error::Error>> {
1659        if n_hist == 0 {
1660            return Ok(());
1661        }
1662        let f = self.func("penalize_logits_f32");
1663        let (nh, ni) = (n_hist as i32, n as i32);
1664        let cfg = LaunchConfig {
1665            grid_dim: (n_hist.div_ceil(128) as u32, 1, 1),
1666            block_dim: (128, 1, 1),
1667            shared_mem_bytes: 0,
1668        };
1669        let __s_b = self.gpu.stream();
1670        let mut b = __s_b.launch_builder(&f);
1671        b.arg(&mut *x)
1672            .arg(hist)
1673            .arg(&nh)
1674            .arg(&rep)
1675            .arg(&freq)
1676            .arg(&present)
1677            .arg(&ni);
1678        unsafe {
1679            b.launch(cfg)?;
1680        }
1681        Ok(())
1682    }
1683
1684    /// Rows variant: penalize `nrow` contiguous rows of length n in one launch.
1685    #[allow(clippy::too_many_arguments)]
1686    pub fn penalize_logits_rows(
1687        &self,
1688        x: &mut CudaSlice<f32>,
1689        hist: &CudaSlice<u32>,
1690        n_hist: usize,
1691        rep: f32,
1692        freq: f32,
1693        present: f32,
1694        n: usize,
1695        nrow: usize,
1696    ) -> Result<(), Box<dyn std::error::Error>> {
1697        if n_hist == 0 || nrow == 0 {
1698            return Ok(());
1699        }
1700        let f = self.func("penalize_logits_rows_f32");
1701        let (nh, ni, nr) = (n_hist as i32, n as i32, nrow as i32);
1702        let cfg = LaunchConfig {
1703            grid_dim: (n_hist.div_ceil(128) as u32, nrow as u32, 1),
1704            block_dim: (128, 1, 1),
1705            shared_mem_bytes: 0,
1706        };
1707        let __s_b = self.gpu.stream();
1708        let mut b = __s_b.launch_builder(&f);
1709        b.arg(&mut *x)
1710            .arg(hist)
1711            .arg(&nh)
1712            .arg(&rep)
1713            .arg(&freq)
1714            .arg(&present)
1715            .arg(&ni)
1716            .arg(&nr);
1717        unsafe {
1718            b.launch(cfg)?;
1719        }
1720        Ok(())
1721    }
1722
1723    /// WEIGHT PREFETCH (SOTA item 3, 2026-07-13, DEFAULT ON): during a bandwidth-idle
1724    /// window (the fa launch reads KV, not weights) prefetch the NEXT matvec's
1725    /// decode-plane bytes into L2 so it reads L2-warm. Value-free scheduling op — same
1726    /// class as prefetch_l2 (numerics untouched by construction). Wired only where it
1727    /// measured positive: the E4B dc attn arm (+0.65%). 26B (flat — MoE ffn dominates),
1728    /// 31B (−0.2% — decode at the DRAM wall) and the ffn gate/up cascade (−1% — 29MB/layer
1729    /// floods the fill path) all probed and NOT wired. MEMRA_WPF=0 rollback seam.
1730    pub fn wpf_level() -> u32 {
1731        static ON: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
1732        *ON.get_or_init(|| {
1733            std::env::var("MEMRA_WPF")
1734                .ok()
1735                .and_then(|v| v.parse().ok())
1736                .unwrap_or(1)
1737        })
1738    }
1739
1740    /// PDL launch arm (SOTA item 2, 2026-07-13, DEFAULT ON): the six MEMRA_PDL_ENTRY glue
1741    /// kernels launch through cuLaunchKernelEx with PROGRAMMATIC_STREAM_SERIALIZATION — the
1742    /// grid launches while the predecessor drains (~120ns/kernel back, pdl_probe), the
1743    /// kernels' entry grid-dep sync restores read order (SASS-audited: ACQBULK precedes
1744    /// every LDG in all six). Valid windows: E4B +1.0-1.2% (128 AND 384-tok gens);
1745    /// 26B/31B/qwen flat no-harm. Battery: kernel-check GREEN, run-gen tokens IDENTICAL x3
1746    /// gemma, spec 64/64 E4B K=1/4/8 + 26B/31B K=4 + qwen PASS. Works eager AND under
1747    /// capture (capture encodes native programmatic edges — the post-capture edge-REWRITE
1748    /// arm died: engine graphs hold cuMemAllocAsync alloc nodes, edge edits on those return
1749    /// CUDA_ERROR_NOT_SUPPORTED). MEMRA_PDL=0 rollback seam.
1750    /// See the `verify_exact` field. Scoped by the dflash round around its t=16 verify.
1751    pub fn set_verify_exact(&self, on: bool) {
1752        self.verify_exact
1753            .store(on, std::sync::atomic::Ordering::Relaxed);
1754    }
1755    pub(crate) fn verify_exact_on(&self) -> bool {
1756        self.verify_exact.load(std::sync::atomic::Ordering::Relaxed)
1757    }
1758
1759    /// m=1 norm+rope+append fold seam (2026-07-23): MEMRA_QKV_APPEND=0 reverts to the
1760    /// fused-norm-rope + standalone-append pair (the exact-oracle bisect arm).
1761    pub fn qkv_append_on() -> bool {
1762        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1763        *ON.get_or_init(|| {
1764            std::env::var("MEMRA_QKV_APPEND")
1765                .map(|v| v != "0")
1766                .unwrap_or(true)
1767        })
1768    }
1769
1770    /// PDL wave-B1a seam: the four dense-glue kernels (rms_norm_f32, add_rms_norm_f32,
1771    /// add_scale_rms_norm_q8_1, quantize_q8_1). MEMRA_PDL_WB=0 reverts alone.
1772    pub fn pdl_wb_on() -> bool {
1773        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1774        *ON.get_or_init(|| {
1775            std::env::var("MEMRA_PDL_WB")
1776                .map(|v| v != "0")
1777                .unwrap_or(true)
1778        })
1779    }
1780
1781    /// PDL wave-A seam: the mmvq matvec PDL launches only (the six glue kernels keep
1782    /// their own MEMRA_PDL master seam). MEMRA_PDL_MMVQ=0 reverts wave-A alone — the
1783    /// per-model no-harm bisect knob.
1784    pub fn pdl_mmvq_on() -> bool {
1785        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1786        *ON.get_or_init(|| {
1787            std::env::var("MEMRA_PDL_MMVQ")
1788                .map(|v| v != "0")
1789                .unwrap_or(true)
1790        })
1791    }
1792
1793    pub fn pdl_on() -> bool {
1794        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1795        *ON.get_or_init(|| std::env::var("MEMRA_PDL").map(|v| v != "0").unwrap_or(true))
1796    }
1797
1798    /// PDL wave-B seam (gap-diagnosis arc, GAP-DIAGNOSIS.md verdict 8): the gemma
1799    /// NVFP4mix decode chain's hot kernels — nvfp4 fused2/mr2 and the q8_0 `_rp`
1800    /// singles — join the wave-A launch class. Scheduling-only (the entry macro waits
1801    /// on the producer before any read), bit-identical by construction.
1802    /// MEMRA_PDL_NVFP4=0 reverts wave-B alone.
1803    pub fn pdl_nvfp4q8_on() -> bool {
1804        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1805        *ON.get_or_init(|| {
1806            std::env::var("MEMRA_PDL_NVFP4")
1807                .map(|v| v != "0")
1808                .unwrap_or(true)
1809        })
1810    }
1811
1812    /// Raw CUfunction for a PDL-attributed launch: the SAME kernels.fatbin loaded once more
1813    /// through the raw driver API (cudarc hides its CUfunction handles; a duplicate module
1814    /// of tiny glue kernels is free). Resolved lazily per name, cached process-wide.
1815    /// Fused t=1 q4_0 mr policy: env MEMRA_Q40_MR wins (1/2); else the per-model
1816    /// FUSED_MR1_DEFAULT (dense gemma = mr1, MoE = mr2 — see the static's doc).
1817    fn q40_mr1_on() -> bool {
1818        static Q40MR: std::sync::OnceLock<Option<u32>> = std::sync::OnceLock::new();
1819        match *Q40MR.get_or_init(|| {
1820            std::env::var("MEMRA_Q40_MR")
1821                .ok()
1822                .and_then(|v| v.parse().ok())
1823        }) {
1824            Some(v) => v == 1,
1825            None => crate::FUSED_MR1_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
1826        }
1827    }
1828
1829    /// PDL wave-B2: flash-module PDL functions. `g` selects the kf8vf8 flavor — the
1830    /// caller MUST pass the SAME flavor its builder launch would resolve (fa_func/func_g
1831    /// mirror); the flavors differ semantically (KV byte formats), a wrong-module launch
1832    /// writes wrong bytes silently.
1833    fn pdl_func_flash(
1834        &self,
1835        g: bool,
1836        name: &'static str,
1837    ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
1838        use cudarc::driver::sys as cu;
1839        // PER-CONTEXT caches (M1-PP2 cross-device fix, 8x box 2026-08-02): CUmodule and
1840        // CUfunction handles are CONTEXT-scoped, and a remote-stage Engine
1841        // (MEMRA_PP_DEVICES=a,b) lives in the other device's primary context. The old
1842        // process-wide OnceLock cache handed stage 1 the dev-a handles, so every stage-1
1843        // launch_pdl* died CUDA_ERROR_INVALID_HANDLE. Key module + function caches by
1844        // this engine's CUcontext; single-context runs behave exactly as before.
1845        static MODS: std::sync::Mutex<Option<std::collections::HashMap<(usize, bool), usize>>> =
1846            std::sync::Mutex::new(None);
1847        static FNS: std::sync::Mutex<
1848            Option<std::collections::HashMap<(usize, bool, &'static str), usize>>,
1849        > = std::sync::Mutex::new(None);
1850        let ctx_key = self.ctx().cu_ctx() as usize;
1851        if let Some(&f) = FNS
1852            .lock()
1853            .unwrap()
1854            .get_or_insert_with(Default::default)
1855            .get(&(ctx_key, g, name))
1856        {
1857            return Ok(f as cu::CUfunction);
1858        }
1859        let module = {
1860            let mut mods = MODS.lock().unwrap();
1861            let map = mods.get_or_insert_with(Default::default);
1862            match map.get(&(ctx_key, g)) {
1863                Some(&m) => m,
1864                None => {
1865                    let m = self.pdl_load_module_in_ctx(if g {
1866                        FLASH_FATBIN_KF8VF8
1867                    } else {
1868                        FLASH_FATBIN
1869                    })?;
1870                    map.insert((ctx_key, g), m);
1871                    m
1872                }
1873            }
1874        };
1875        let cname = std::ffi::CString::new(name)?;
1876        let mut f: cu::CUfunction = std::ptr::null_mut();
1877        let r = unsafe { cu::cuModuleGetFunction(&mut f, module as cu::CUmodule, cname.as_ptr()) };
1878        if r != cu::CUresult::CUDA_SUCCESS {
1879            return Err(format!("pdl_func_flash {name} (g={g}): {r:?}").into());
1880        }
1881        FNS.lock()
1882            .unwrap()
1883            .get_or_insert_with(Default::default)
1884            .insert((ctx_key, g, name), f as usize);
1885        Ok(f)
1886    }
1887
1888    /// Load a fatbin as a raw CUmodule IN THIS ENGINE'S CONTEXT. `cuModuleLoadData` binds
1889    /// the module to the thread's CURRENT context — a remote-stage engine must not
1890    /// inherit the primary's (the INVALID_HANDLE class above). Restores the caller's
1891    /// current context before returning.
1892    fn pdl_load_module_in_ctx(&self, bytes: &[u8]) -> Result<usize, Box<dyn std::error::Error>> {
1893        use cudarc::driver::sys as cu;
1894        let mut prev: cu::CUcontext = std::ptr::null_mut();
1895        unsafe {
1896            cu::cuCtxGetCurrent(&mut prev).result()?;
1897        }
1898        self.ctx().bind_to_thread()?;
1899        let mut m: cu::CUmodule = std::ptr::null_mut();
1900        let r = unsafe { cu::cuModuleLoadData(&mut m, bytes.as_ptr() as *const std::ffi::c_void) };
1901        let restore = if prev.is_null() {
1902            cu::CUresult::CUDA_SUCCESS
1903        } else {
1904            unsafe { cu::cuCtxSetCurrent(prev) }
1905        };
1906        if r != cu::CUresult::CUDA_SUCCESS {
1907            return Err(format!("pdl module load: {r:?}").into());
1908        }
1909        if restore != cu::CUresult::CUDA_SUCCESS {
1910            return Err(format!("pdl module load: ctx restore {restore:?}").into());
1911        }
1912        Ok(m as usize)
1913    }
1914
1915    fn pdl_func(
1916        &self,
1917        name: &'static str,
1918    ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
1919        use cudarc::driver::sys as cu;
1920        // PER-CONTEXT caches — same M1-PP2 cross-device fix as pdl_func_flash (handles
1921        // are context-scoped; key everything by this engine's CUcontext).
1922        static MODULES: std::sync::Mutex<Option<std::collections::HashMap<usize, usize>>> =
1923            std::sync::Mutex::new(None);
1924        // PDL wave-A: the mmvq kernels live in the qmatvec fatbin, not kernels.cu — second
1925        // duplicate module, loaded lazily on the first kernels-module miss.
1926        static QMODULES: std::sync::Mutex<Option<std::collections::HashMap<usize, usize>>> =
1927            std::sync::Mutex::new(None);
1928        static FNS: std::sync::Mutex<
1929            Option<std::collections::HashMap<(usize, &'static str), usize>>,
1930        > = std::sync::Mutex::new(None);
1931        let ctx_key = self.ctx().cu_ctx() as usize;
1932        if let Some(&f) = FNS
1933            .lock()
1934            .unwrap()
1935            .get_or_insert_with(Default::default)
1936            .get(&(ctx_key, name))
1937        {
1938            return Ok(f as cu::CUfunction);
1939        }
1940        let module = {
1941            let mut mods = MODULES.lock().unwrap();
1942            let map = mods.get_or_insert_with(Default::default);
1943            match map.get(&ctx_key) {
1944                Some(&m) => m,
1945                None => {
1946                    let m = self.pdl_load_module_in_ctx(FATBIN)?;
1947                    map.insert(ctx_key, m);
1948                    m
1949                }
1950            }
1951        };
1952        let cname = std::ffi::CString::new(name)?;
1953        let mut f: cu::CUfunction = std::ptr::null_mut();
1954        let mut r =
1955            unsafe { cu::cuModuleGetFunction(&mut f, module as cu::CUmodule, cname.as_ptr()) };
1956        if r == cu::CUresult::CUDA_ERROR_NOT_FOUND {
1957            let qmodule = {
1958                let mut mods = QMODULES.lock().unwrap();
1959                let map = mods.get_or_insert_with(Default::default);
1960                match map.get(&ctx_key) {
1961                    Some(&m) => m,
1962                    None => {
1963                        let m = self.pdl_load_module_in_ctx(QMATVEC_FATBIN)?;
1964                        map.insert(ctx_key, m);
1965                        m
1966                    }
1967                }
1968            };
1969            r = unsafe { cu::cuModuleGetFunction(&mut f, qmodule as cu::CUmodule, cname.as_ptr()) };
1970        }
1971        if r != cu::CUresult::CUDA_SUCCESS {
1972            return Err(format!("pdl_func {name}: {r:?}").into());
1973        }
1974        FNS.lock()
1975            .unwrap()
1976            .get_or_insert_with(Default::default)
1977            .insert((ctx_key, name), f as usize);
1978        Ok(f)
1979    }
1980
1981    /// cuLaunchKernelEx with CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION on the
1982    /// compute stream. ONLY legal for kernels whose entry carries MEMRA_PDL_ENTRY.
1983    ///
1984    /// # Safety
1985    /// `params` must match the kernel's exact parameter list (order, types, count) —
1986    /// a mismatch corrupts the launch silently.
1987    /// Flash-module twin of `launch_pdl` — `g` picks the kf8vf8 flavor (must mirror the
1988    /// builder path's fa_func/func_g choice exactly).
1989    ///
1990    /// # Safety
1991    /// Same contract as `launch_pdl`.
1992    unsafe fn launch_pdl_flash(
1993        &self,
1994        g: bool,
1995        name: &'static str,
1996        grid: (u32, u32, u32),
1997        block: (u32, u32, u32),
1998        smem: u32,
1999        params: &mut [*mut std::ffi::c_void],
2000    ) -> Result<(), Box<dyn std::error::Error>> {
2001        use cudarc::driver::sys as cu;
2002        let f = self.pdl_func_flash(g, name)?;
2003        if smem > 0 {
2004            // mirror the builder path's opt-in ceiling (idempotent host-side set).
2005            let r =
2006                unsafe {
2007                    cu::cuFuncSetAttribute(f,
2008                cu::CUfunction_attribute_enum::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
2009                smem as i32)
2010                };
2011            if r != cu::CUresult::CUDA_SUCCESS {
2012                return Err(format!("pdl smem attr {name}: {r:?}").into());
2013            }
2014        }
2015        let mut attr = cu::CUlaunchAttribute {
2016            id: cu::CUlaunchAttributeID::CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION,
2017            pad: [0; 4],
2018            value: cu::CUlaunchAttributeValue {
2019                programmaticStreamSerializationAllowed: 1,
2020            },
2021        };
2022        let cfg = cu::CUlaunchConfig {
2023            gridDimX: grid.0,
2024            gridDimY: grid.1,
2025            gridDimZ: grid.2,
2026            blockDimX: block.0,
2027            blockDimY: block.1,
2028            blockDimZ: block.2,
2029            sharedMemBytes: smem,
2030            hStream: self.gpu.stream().cu_stream(),
2031            attrs: &mut attr,
2032            numAttrs: 1,
2033        };
2034        let r = unsafe { cu::cuLaunchKernelEx(&cfg, f, params.as_mut_ptr(), std::ptr::null_mut()) };
2035        if r != cu::CUresult::CUDA_SUCCESS {
2036            return Err(format!("launch_pdl_flash {name}: {r:?}").into());
2037        }
2038        Ok(())
2039    }
2040
2041    unsafe fn launch_pdl(
2042        &self,
2043        name: &'static str,
2044        grid: (u32, u32, u32),
2045        block: (u32, u32, u32),
2046        params: &mut [*mut std::ffi::c_void],
2047    ) -> Result<(), Box<dyn std::error::Error>> {
2048        use cudarc::driver::sys as cu;
2049        let f = self.pdl_func(name)?;
2050        let mut attr = cu::CUlaunchAttribute {
2051            id: cu::CUlaunchAttributeID::CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION,
2052            pad: [0; 4],
2053            value: cu::CUlaunchAttributeValue {
2054                programmaticStreamSerializationAllowed: 1,
2055            },
2056        };
2057        let cfg = cu::CUlaunchConfig {
2058            gridDimX: grid.0,
2059            gridDimY: grid.1,
2060            gridDimZ: grid.2,
2061            blockDimX: block.0,
2062            blockDimY: block.1,
2063            blockDimZ: block.2,
2064            sharedMemBytes: 0,
2065            hStream: self.gpu.stream().cu_stream(),
2066            attrs: &mut attr,
2067            numAttrs: 1,
2068        };
2069        let r = unsafe { cu::cuLaunchKernelEx(&cfg, f, params.as_mut_ptr(), std::ptr::null_mut()) };
2070        if r != cu::CUresult::CUDA_SUCCESS {
2071            return Err(format!("launch_pdl {name}: {r:?}").into());
2072        }
2073        Ok(())
2074    }
2075
2076    /// L2-prefetch a quant weight's DECODE plane (the rp4 split-plane mirror when present —
2077    /// that is what the m<=8 dispatch reads — else the raw block bytes). No-op on float arms.
2078    pub fn prefetch_weight_l2(
2079        &self,
2080        w: &crate::model::GpuTensor,
2081    ) -> Result<(), Box<dyn std::error::Error>> {
2082        if let crate::model::GpuTensor::Quant { bytes, rp4, .. } = w {
2083            let p = rp4.as_ref().unwrap_or(bytes);
2084            self.prefetch_l2(p, p.len())?;
2085        }
2086        Ok(())
2087    }
2088
2089    /// DSpark markov chain ops (dflash lane): gather one bf16 row of a [V, rank] table
2090    /// by the DEVICE token id at tok[idx] into f32.
2091    pub fn gather_row_bf16(
2092        &self,
2093        table: &CudaSlice<u8>,
2094        tok: &CudaSlice<u32>,
2095        idx: usize,
2096        dst: &mut CudaSlice<f32>,
2097        ncols: usize,
2098    ) -> Result<(), Box<dyn std::error::Error>> {
2099        let f = self.func("gather_row_bf16_f32");
2100        let cfg = LaunchConfig {
2101            grid_dim: (ncols.div_ceil(256) as u32, 1, 1),
2102            block_dim: (256, 1, 1),
2103            shared_mem_bytes: 0,
2104        };
2105        let (nc, ix) = (ncols as i32, idx as i32);
2106        let __s_b = self.gpu.stream();
2107        let mut b = __s_b.launch_builder(&f);
2108        b.arg(table).arg(tok).arg(&ix).arg(dst).arg(&nc);
2109        unsafe {
2110            b.launch(cfg)?;
2111        }
2112        Ok(())
2113    }
2114
2115    /// logits[row_off .. row_off+n] += bias[0..n] (in place, one row).
2116    pub fn add_row_inplace(
2117        &self,
2118        logits: &mut CudaSlice<f32>,
2119        bias: &CudaSlice<f32>,
2120        n: usize,
2121        row_off: usize,
2122    ) -> Result<(), Box<dyn std::error::Error>> {
2123        let f = self.func("add_row_inplace_f32");
2124        let cfg = LaunchConfig {
2125            grid_dim: (n.div_ceil(256) as u32, 1, 1),
2126            block_dim: (256, 1, 1),
2127            shared_mem_bytes: 0,
2128        };
2129        let (ni, off) = (n as i32, row_off as i64);
2130        let __s_b = self.gpu.stream();
2131        let mut b = __s_b.launch_builder(&f);
2132        b.arg(logits).arg(bias).arg(&ni).arg(&off);
2133        unsafe {
2134            b.launch(cfg)?;
2135        }
2136        Ok(())
2137    }
2138
2139    /// L2 prefetch of a device byte range (latency-hiding arc; value-free scheduling op).
2140    pub fn prefetch_l2(
2141        &self,
2142        p: &CudaSlice<u8>,
2143        n: usize,
2144    ) -> Result<(), Box<dyn std::error::Error>> {
2145        let f = self.func("prefetch_l2_bytes");
2146        let lines = n.div_ceil(128);
2147        let ni = n as i64;
2148        let cfg = LaunchConfig {
2149            grid_dim: (lines.div_ceil(256) as u32, 1, 1),
2150            block_dim: (256, 1, 1),
2151            shared_mem_bytes: 0,
2152        };
2153        let __s_b = self.gpu.stream();
2154        let mut b = __s_b.launch_builder(&f);
2155        b.arg(p).arg(&ni);
2156        unsafe {
2157            b.launch(cfg)?;
2158        }
2159        Ok(())
2160    }
2161
2162    /// MoE router GEMV (MEMRA_ROUTER_KERNEL): deterministic warp-per-(expert,token) f32 dot.
2163    /// Different FP order than the cuBLAS path it replaces — battery-gated numeric config.
2164    pub fn router_gemv(
2165        &self,
2166        w: &CudaSlice<f32>,
2167        x: &CudaSlice<f32>,
2168        n_embd: usize,
2169        n_experts: usize,
2170        t: usize,
2171    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2172        // float4 v2 probed 2026-07-14: +0.25% but flips near-tie routing (new FP order,
2173        // stream differs) — too small to justify a numeric config change; deleted.
2174        // w8 twin (2026-07-31): on the 132-SM H100 the lone-warp form is 14.8% of the q35
2175        // decode step (latency-bound) — the calculus flipped. MEMRA_ROUTER_V2=0 reverts to
2176        // the warp form (rollback seam; new FP order, battery-arbitrated per model).
2177        let w8 = match std::env::var("MEMRA_ROUTER_V2").as_deref() {
2178            Ok("0") => false,
2179            Ok(_) => true,
2180            Err(_) => ROUTER_W8_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
2181        };
2182        // FAST-ROUTER batch twin (lane/fast-router, 2026-08-02): at prefill m the per-(e,tok)
2183        // w8 form re-streams both operand rows per output (GEMV program at GEMM shape — the
2184        // concat-prime exactness fix paid -10% q35 board-2048 prefill through it). The batch
2185        // twin (8x8 expert-x-token register tile) is BIT-IDENTICAL per row (same k order,
2186        // same tree, same fold — kernel-check sweeps m=1..2048 on real router weights), so
2187        // the crossover is pure perf, not a numeric config. MIN_T from the on-box sweep
2188        // (research/fast-router-20260802/crossover-router*.jsonl); decode t=1 and small-t
2189        // spec verify keep the plain w8 form. MEMRA_ROUTER_BATCH=0: rollback seam
2190        // (perf-only, bits equal).
2191        let batch = w8 && t >= ROUTER_BATCH_MIN_T && router_batch_on();
2192        self.router_gemv_form(w, x, n_embd, n_experts, t, w8, batch)
2193    }
2194
2195    /// Form-explicit router GEMV launch (kernel-check bit-identity gate + crossover bench
2196    /// force both forms; `batch` requires `w8`).
2197    pub fn router_gemv_form(
2198        &self,
2199        w: &CudaSlice<f32>,
2200        x: &CudaSlice<f32>,
2201        n_embd: usize,
2202        n_experts: usize,
2203        t: usize,
2204        w8: bool,
2205        batch: bool,
2206    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2207        debug_assert!(!batch || w8, "batch twin exists for the w8 form only");
2208        let mut y = self.alloc_uninit::<f32>(t * n_experts)?;
2209        let f = if batch {
2210            self.func("router_gemv_f32_w8_batch")
2211        } else if w8 {
2212            self.func("router_gemv_f32_w8")
2213        } else {
2214            self.func("router_gemv_f32")
2215        };
2216        let (ne, nx, ti) = (n_embd as i32, n_experts as i32, t as i32);
2217        let cfg = if batch {
2218            LaunchConfig {
2219                grid_dim: (n_experts.div_ceil(8) as u32, t.div_ceil(8) as u32, 1),
2220                block_dim: (32, 8, 1),
2221                shared_mem_bytes: 0,
2222            }
2223        } else {
2224            LaunchConfig {
2225                grid_dim: (n_experts as u32, t as u32, 1),
2226                block_dim: (32, if w8 { 8 } else { 1 }, 1),
2227                shared_mem_bytes: 0,
2228            }
2229        };
2230        let __s_b = self.gpu.stream();
2231        let mut b = __s_b.launch_builder(&f);
2232        b.arg(w).arg(x).arg(&mut y).arg(&ne).arg(&nx).arg(&ti);
2233        unsafe {
2234            b.launch(cfg)?;
2235        }
2236        Ok(y)
2237    }
2238
2239    /// f32 row permute: dst[idx[i], :] = src[i, :] (grouped-GEMM CSR -> pair-id reorder).
2240    pub fn rows_permute(
2241        &self,
2242        src: &CudaSlice<f32>,
2243        idx: &CudaSlice<i32>,
2244        nrows: usize,
2245        ncols: usize,
2246    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2247        let mut dst = self.alloc_uninit::<f32>(nrows * ncols)?;
2248        let f = self.func("rows_permute_f32");
2249        let (nc, nr) = (ncols as i32, nrows as i32);
2250        let cfg = LaunchConfig {
2251            grid_dim: (nrows as u32, 1, 1),
2252            block_dim: (256, 1, 1),
2253            shared_mem_bytes: 0,
2254        };
2255        let __s_b = self.gpu.stream();
2256        let mut b = __s_b.launch_builder(&f);
2257        b.arg(src).arg(idx).arg(&mut dst).arg(&nc).arg(&nr);
2258        unsafe {
2259            b.launch(cfg)?;
2260        }
2261        Ok(dst)
2262    }
2263
2264    /// shexp gate fused dot: g[tok] = sigmoid(dot(x[tok,:], w)) — replaces the per-layer
2265    /// cuBLASLt m=1 GEMM + separate sigmoid launch on the qwen35moe decode path (the
2266    /// splitKreduce x40/step dig, 2026-07-31). One fold order for every t, so the t=1
2267    /// decode chain and the small-t spec-verify chain match per row by construction.
2268    pub fn sigmoid_dot_rows(
2269        &self,
2270        x: &CudaSlice<f32>,
2271        w: &CudaSlice<f32>,
2272        n_embd: usize,
2273        t: usize,
2274    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2275        // MEMRA_SHEXP_DOT=0: rollback seam to the cuBLASLt linear + sigmoid pair (numeric
2276        // config; same class as MEMRA_ROUTER_V2).
2277        static OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2278        if *OFF.get_or_init(|| std::env::var("MEMRA_SHEXP_DOT").as_deref() == Ok("0")) {
2279            let gs = self.linear(x, w, t, n_embd, 1)?;
2280            let mut g = self.uninit(t)?;
2281            self.sigmoid(&gs, &mut g, t)?;
2282            return Ok(g);
2283        }
2284        // FAST-ROUTER lane note (2026-08-02): a register-tiled 8-token batch twin of this
2285        // kernel was built, proven bit-identical, and measured SLOWER at every prefill t on
2286        // the 5090 (0.62-0.89x — launch-latency-bound op, ~7us/layer at m=2048;
2287        // research/fast-router-20260802/crossover-router.jsonl). Dispatch arm killed per
2288        // flags doctrine; this per-token form serves every t.
2289        let mut g = self.alloc_uninit::<f32>(t)?;
2290        let f = self.func("sigmoid_dot_rows_f32");
2291        let (ne, ti) = (n_embd as i32, t as i32);
2292        let cfg = LaunchConfig {
2293            grid_dim: (t as u32, 1, 1),
2294            block_dim: (32, 8, 1),
2295            shared_mem_bytes: 0,
2296        };
2297        let __s_b = self.gpu.stream();
2298        let mut b = __s_b.launch_builder(&f);
2299        b.arg(x).arg(w).arg(&mut g).arg(&ne).arg(&ti);
2300        unsafe {
2301            b.launch(cfg)?;
2302        }
2303        Ok(g)
2304    }
2305
2306    /// ROUND-STREAM stream rollback: all counters <- pos_start + base + n_acc.
2307    pub fn spec_rollback_stream(
2308        &self,
2309        len_ptrs: &CudaSlice<u64>,
2310        pos_start: &CudaSlice<i32>,
2311        acc: &CudaSlice<u32>,
2312        base: usize,
2313        n_rows: usize,
2314    ) -> Result<(), Box<dyn std::error::Error>> {
2315        let f = self.func("spec_rollback_stream");
2316        let (b, nr) = (base as i32, n_rows as i32);
2317        let cfg = LaunchConfig {
2318            grid_dim: (n_rows.div_ceil(64) as u32, 1, 1),
2319            block_dim: (64, 1, 1),
2320            shared_mem_bytes: 0,
2321        };
2322        let __s_bl = self.gpu.stream();
2323        let mut bl = __s_bl.launch_builder(&f);
2324        bl.arg(len_ptrs).arg(pos_start).arg(acc).arg(&b).arg(&nr);
2325        unsafe {
2326            bl.launch(cfg)?;
2327        }
2328        Ok(())
2329    }
2330
2331    /// PLAIN-DECODE GRAPH ring store: ring[(pos_start - base) % cap] = vam[0].
2332    pub fn plain_tok_ring(
2333        &self,
2334        vam: &CudaSlice<u32>,
2335        pos_start: &CudaSlice<i32>,
2336        base: usize,
2337        ring: &mut CudaSlice<u32>,
2338    ) -> Result<(), Box<dyn std::error::Error>> {
2339        let f = self.func("plain_tok_ring");
2340        let (b, cap) = (base as i32, ring.len() as i32);
2341        let cfg = LaunchConfig {
2342            grid_dim: (1, 1, 1),
2343            block_dim: (32, 1, 1),
2344            shared_mem_bytes: 0,
2345        };
2346        let __s_bl = self.gpu.stream();
2347        let mut bl = __s_bl.launch_builder(&f);
2348        bl.arg(vam).arg(pos_start).arg(&b).arg(&mut *ring).arg(&cap);
2349        unsafe {
2350            bl.launch(cfg)?;
2351        }
2352        Ok(())
2353    }
2354
2355    /// ROUND-STREAM stage (c) 4 epilogue: ring commit + tiny counter copies.
2356    pub fn spec_ring_commit(
2357        &self,
2358        vtok: &CudaSlice<u32>,
2359        acc: &CudaSlice<u32>,
2360        brk: &CudaSlice<u32>,
2361        ring: &mut CudaSlice<u32>,
2362        pend: &mut CudaSlice<u32>,
2363    ) -> Result<(), Box<dyn std::error::Error>> {
2364        let f = self.func("spec_ring_commit");
2365        let cfg = LaunchConfig {
2366            grid_dim: (1, 1, 1),
2367            block_dim: (32, 1, 1),
2368            shared_mem_bytes: 0,
2369        };
2370        let __s_b = self.gpu.stream();
2371        let mut b = __s_b.launch_builder(&f);
2372        b.arg(vtok).arg(acc).arg(brk).arg(ring).arg(pend);
2373        unsafe {
2374            b.launch(cfg)?;
2375        }
2376        Ok(())
2377    }
2378    pub fn i32_copy_add(
2379        &self,
2380        src: &CudaSlice<i32>,
2381        dst: &mut CudaSlice<i32>,
2382        delta: i32,
2383    ) -> Result<(), Box<dyn std::error::Error>> {
2384        let f = self.func("i32_copy_add");
2385        let cfg = LaunchConfig {
2386            grid_dim: (1, 1, 1),
2387            block_dim: (32, 1, 1),
2388            shared_mem_bytes: 0,
2389        };
2390        let __s_b = self.gpu.stream();
2391        let mut b = __s_b.launch_builder(&f);
2392        b.arg(src).arg(dst).arg(&delta);
2393        unsafe {
2394            b.launch(cfg)?;
2395        }
2396        Ok(())
2397    }
2398    pub fn u32_copy(
2399        &self,
2400        src: &CudaSlice<u32>,
2401        dst: &mut CudaSlice<u32>,
2402    ) -> Result<(), Box<dyn std::error::Error>> {
2403        let f = self.func("u32_copy");
2404        let cfg = LaunchConfig {
2405            grid_dim: (1, 1, 1),
2406            block_dim: (32, 1, 1),
2407            shared_mem_bytes: 0,
2408        };
2409        let __s_b = self.gpu.stream();
2410        let mut b = __s_b.launch_builder(&f);
2411        b.arg(src).arg(dst);
2412        unsafe {
2413            b.launch(cfg)?;
2414        }
2415        Ok(())
2416    }
2417
2418    /// ROUND-GRAPH adaptive depth: brk[0] <- clamp(acc[0] + 1, floor, cap) — the host
2419    /// adaptive policy as a captured device op (policy-identical: the accept walk depth
2420    /// caps acceptance exactly like drafting fewer tokens).
2421    pub fn spec_adapt_k(
2422        &self,
2423        acc: &CudaSlice<u32>,
2424        brk: &mut CudaSlice<u32>,
2425        floor: usize,
2426        cap: usize,
2427    ) -> Result<(), Box<dyn std::error::Error>> {
2428        let f = self.func("spec_adapt_k");
2429        let (fl, cp) = (floor as i32, cap as i32);
2430        let cfg = LaunchConfig {
2431            grid_dim: (1, 1, 1),
2432            block_dim: (32, 1, 1),
2433            shared_mem_bytes: 0,
2434        };
2435        let __s_b = self.gpu.stream();
2436        let mut b = __s_b.launch_builder(&f);
2437        b.arg(acc).arg(brk).arg(&fl).arg(&cp);
2438        unsafe {
2439            b.launch(cfg)?;
2440        }
2441        Ok(())
2442    }
2443
2444    /// ROUND-STREAM stage (c) 3: accept walk fully device-driven (brk + assembled vtok).
2445    pub fn spec_accept_greedy_dc(
2446        &self,
2447        preds: &CudaSlice<u32>,
2448        vtok: &CudaSlice<u32>,
2449        last_pred: &CudaSlice<u32>,
2450        brk: &CudaSlice<u32>,
2451        out: &mut CudaSlice<u32>,
2452    ) -> Result<(), Box<dyn std::error::Error>> {
2453        let f = self.func("spec_accept_greedy_dc");
2454        let cfg = LaunchConfig {
2455            grid_dim: (1, 1, 1),
2456            block_dim: (32, 1, 1),
2457            shared_mem_bytes: 0,
2458        };
2459        let __s_b = self.gpu.stream();
2460        let mut b = __s_b.launch_builder(&f);
2461        b.arg(preds).arg(vtok).arg(last_pred).arg(brk).arg(out);
2462        unsafe {
2463            b.launch(cfg)?;
2464        }
2465        Ok(())
2466    }
2467
2468    /// ROUND-STREAM stage (c) 2: verify-chain device-pos entries.
2469    pub fn pos_iota(
2470        &self,
2471        pos0: &CudaSlice<i32>,
2472        out: &mut CudaSlice<i32>,
2473        t: usize,
2474    ) -> Result<(), Box<dyn std::error::Error>> {
2475        let f = self.func("pos_iota_i32");
2476        let ti = t as i32;
2477        let cfg = LaunchConfig {
2478            grid_dim: (1, 1, 1),
2479            block_dim: (t.max(1) as u32, 1, 1),
2480            shared_mem_bytes: 0,
2481        };
2482        let __s_b = self.gpu.stream();
2483        let mut b = __s_b.launch_builder(&f);
2484        b.arg(pos0).arg(out).arg(&ti);
2485        unsafe {
2486            b.launch(cfg)?;
2487        }
2488        Ok(())
2489    }
2490    #[allow(clippy::too_many_arguments)]
2491    pub fn append_kv_quantized_rows_dc(
2492        &self,
2493        k_rows: &CudaSlice<f32>,
2494        v_rows: &CudaSlice<f32>,
2495        kc: &mut CudaSlice<u8>,
2496        vc: &mut CudaSlice<u8>,
2497        t0_dev: &CudaSlice<i32>,
2498        t: usize,
2499        kv_dim_k: usize,
2500        kv_dim_v: usize,
2501        k_tok_bytes: usize,
2502        v_tok_bytes: usize,
2503        g: bool,
2504    ) -> Result<(), Box<dyn std::error::Error>> {
2505        let f = if g {
2506            self.func_g("append_quantize_kv_q8_0_q5_1_rows_dc")
2507        } else {
2508            self.func("append_quantize_kv_q8_0_q5_1_rows_dc")
2509        };
2510        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
2511        let cfg = LaunchConfig {
2512            grid_dim: (nblk, t as u32, 1),
2513            block_dim: (32, 1, 1),
2514            shared_mem_bytes: 0,
2515        };
2516        let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
2517        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
2518        let __s_b = self.gpu.stream();
2519        let mut b = __s_b.launch_builder(&f);
2520        b.arg(k_rows)
2521            .arg(v_rows)
2522            .arg(kc)
2523            .arg(vc)
2524            .arg(t0_dev)
2525            .arg(&kdk)
2526            .arg(&kdv)
2527            .arg(&ktb)
2528            .arg(&vtb);
2529        unsafe {
2530            b.launch(cfg)?;
2531        }
2532        Ok(())
2533    }
2534
2535    /// t=1 dc append with a FUSED len_d increment (wave 5c) — one launch replaces
2536    /// append_rows_dc + inc_seqlen. Single block (read-before-inc ordering).
2537    #[allow(clippy::too_many_arguments)]
2538    pub fn append_kv_quantized_row_dc_inc(
2539        &self,
2540        k_row: &CudaSlice<f32>,
2541        v_row: &CudaSlice<f32>,
2542        kc: &mut CudaSlice<u8>,
2543        vc: &mut CudaSlice<u8>,
2544        t0_dev: &mut CudaSlice<i32>,
2545        kv_dim_k: usize,
2546        kv_dim_v: usize,
2547        k_tok_bytes: usize,
2548        v_tok_bytes: usize,
2549        g: bool,
2550    ) -> Result<(), Box<dyn std::error::Error>> {
2551        let f = if g {
2552            self.func_g("append_quantize_kv_q8_0_q5_1_dc_inc")
2553        } else {
2554            self.func("append_quantize_kv_q8_0_q5_1_dc_inc")
2555        };
2556        let nthreads = ((kv_dim_k.max(kv_dim_v) / 32) * 32).min(1024) as u32;
2557        let cfg = LaunchConfig {
2558            grid_dim: (1, 1, 1),
2559            block_dim: (nthreads, 1, 1),
2560            shared_mem_bytes: 0,
2561        };
2562        let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
2563        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
2564        let __s_b = self.gpu.stream();
2565        let mut b = __s_b.launch_builder(&f);
2566        b.arg(k_row)
2567            .arg(v_row)
2568            .arg(kc)
2569            .arg(vc)
2570            .arg(t0_dev)
2571            .arg(&kdk)
2572            .arg(&kdv)
2573            .arg(&ktb)
2574            .arg(&vtb);
2575        unsafe {
2576            b.launch(cfg)?;
2577        }
2578        Ok(())
2579    }
2580
2581    /// ROUND-STREAM: draft-chain pack + in-graph d2t remap (see kernels.cu headers).
2582    pub fn pack_tok_p(
2583        &self,
2584        tok: &CudaSlice<u32>,
2585        p: &CudaSlice<f32>,
2586        out: &mut CudaSlice<u32>,
2587        slot: usize,
2588    ) -> Result<(), Box<dyn std::error::Error>> {
2589        let f = self.func("pack_tok_p");
2590        let sl = slot as i32;
2591        let cfg = LaunchConfig {
2592            grid_dim: (1, 1, 1),
2593            block_dim: (32, 1, 1),
2594            shared_mem_bytes: 0,
2595        };
2596        let __s_b = self.gpu.stream();
2597        let mut b = __s_b.launch_builder(&f);
2598        b.arg(tok).arg(p).arg(out).arg(&sl);
2599        unsafe {
2600            b.launch(cfg)?;
2601        }
2602        Ok(())
2603    }
2604    pub fn tok_map_u32(
2605        &self,
2606        tok: &mut CudaSlice<u32>,
2607        map: &CudaSlice<u32>,
2608    ) -> Result<(), Box<dyn std::error::Error>> {
2609        let f = self.func("tok_map_u32");
2610        let cfg = LaunchConfig {
2611            grid_dim: (1, 1, 1),
2612            block_dim: (32, 1, 1),
2613            shared_mem_bytes: 0,
2614        };
2615        let __s_b = self.gpu.stream();
2616        let mut b = __s_b.launch_builder(&f);
2617        b.arg(tok).arg(map);
2618        unsafe {
2619            b.launch(cfg)?;
2620        }
2621        Ok(())
2622    }
2623
2624    /// ROUND-STREAM stage (c) 1: device verify-token assembly + p-min break derivation.
2625    #[allow(clippy::too_many_arguments)]
2626    pub fn spec_assemble_verify(
2627        &self,
2628        tokp: &CudaSlice<u32>,
2629        pend: &CudaSlice<u32>,
2630        d2t: Option<&CudaSlice<u32>>,
2631        vtok: &mut CudaSlice<u32>,
2632        brk: &mut CudaSlice<u32>,
2633        p_min: f32,
2634        k: usize,
2635        pmin0: bool,
2636    ) -> Result<(), Box<dyn std::error::Error>> {
2637        let f = self.func("spec_assemble_verify");
2638        let (ki, pm) = (k as i32, if pmin0 { 1i32 } else { 0i32 });
2639        let cfg = LaunchConfig {
2640            grid_dim: (1, 1, 1),
2641            block_dim: (32, 1, 1),
2642            shared_mem_bytes: 0,
2643        };
2644        let __s_b = self.gpu.stream();
2645        let mut b = __s_b.launch_builder(&f);
2646        match d2t {
2647            Some(m) => {
2648                b.arg(tokp)
2649                    .arg(pend)
2650                    .arg(m)
2651                    .arg(vtok)
2652                    .arg(brk)
2653                    .arg(&p_min)
2654                    .arg(&ki)
2655                    .arg(&pm);
2656                unsafe {
2657                    b.launch(cfg)?;
2658                }
2659            }
2660            None => {
2661                let null: u64 = 0;
2662                b.arg(tokp)
2663                    .arg(pend)
2664                    .arg(&null)
2665                    .arg(vtok)
2666                    .arg(brk)
2667                    .arg(&p_min)
2668                    .arg(&ki)
2669                    .arg(&pm);
2670                unsafe {
2671                    b.launch(cfg)?;
2672                }
2673            }
2674        }
2675        Ok(())
2676    }
2677
2678    /// ROUND-STREAM stage (b) 3b: recur-restore twins with device-j (see hybrid.cu headers).
2679    #[allow(clippy::too_many_arguments)]
2680    pub fn ssm_conv_ring_rebuild_dc(
2681        &self,
2682        qkv_tm: &CudaSlice<f32>,
2683        ring_old: &CudaSlice<f32>,
2684        conv_state: &mut CudaSlice<f32>,
2685        conv_dim: usize,
2686        acc: &CudaSlice<u32>,
2687        base: usize,
2688        t_v: usize,
2689        d_conv: usize,
2690    ) -> Result<(), Box<dyn std::error::Error>> {
2691        let f = self.func("ssm_conv_ring_rebuild_f32_dc");
2692        let n = conv_dim * (d_conv - 1);
2693        let cfg = LaunchConfig::for_num_elems(n as u32);
2694        let (cd, b0, tv, dc) = (conv_dim as i32, base as i32, t_v as i32, d_conv as i32);
2695        let __s_b = self.gpu.stream();
2696        let mut b = __s_b.launch_builder(&f);
2697        b.arg(qkv_tm)
2698            .arg(ring_old)
2699            .arg(conv_state)
2700            .arg(&cd)
2701            .arg(acc)
2702            .arg(&b0)
2703            .arg(&tv)
2704            .arg(&dc);
2705        unsafe {
2706            b.launch(cfg)?;
2707        }
2708        Ok(())
2709    }
2710    #[allow(clippy::too_many_arguments)]
2711    pub fn gdn_scan_s128_dc(
2712        &self,
2713        q: &CudaSlice<f32>,
2714        k: &CudaSlice<f32>,
2715        v: &CudaSlice<f32>,
2716        g: &CudaSlice<f32>,
2717        beta: &CudaSlice<f32>,
2718        state_in: &CudaSlice<f32>,
2719        state_out: &mut CudaSlice<f32>,
2720        o: &mut CudaSlice<f32>,
2721        n_head: usize,
2722        acc: &CudaSlice<u32>,
2723        base: usize,
2724        t_v: usize,
2725        scale: f32,
2726    ) -> Result<(), Box<dyn std::error::Error>> {
2727        let f = self.func("gdn_scan_s128_dc");
2728        const S_V: u32 = 128;
2729        const WARP: u32 = 32;
2730        const COLS_PER_BLOCK: u32 = 4;
2731        let cfg = LaunchConfig {
2732            grid_dim: (n_head as u32, 1, S_V / COLS_PER_BLOCK),
2733            block_dim: (WARP, COLS_PER_BLOCK, 1),
2734            shared_mem_bytes: 0,
2735        };
2736        let (h, b0, tv) = (n_head as i32, base as i32, t_v as i32);
2737        let __s_b = self.gpu.stream();
2738        let mut b = __s_b.launch_builder(&f);
2739        b.arg(q)
2740            .arg(k)
2741            .arg(v)
2742            .arg(g)
2743            .arg(beta)
2744            .arg(state_in)
2745            .arg(state_out)
2746            .arg(o)
2747            .arg(&h)
2748            .arg(acc)
2749            .arg(&b0)
2750            .arg(&tv)
2751            .arg(&scale);
2752        unsafe {
2753            b.launch(cfg)?;
2754        }
2755        Ok(())
2756    }
2757
2758    /// ROUND-STREAM stage (b) 3a: device per-layer KV-len rollback (see spec_rollback_kv).
2759    pub fn spec_rollback_kv(
2760        &self,
2761        len_ptrs: &CudaSlice<u64>,
2762        saved: &CudaSlice<i32>,
2763        acc: &CudaSlice<u32>,
2764        base: usize,
2765        n_layer: usize,
2766    ) -> Result<(), Box<dyn std::error::Error>> {
2767        let f = self.func("spec_rollback_kv");
2768        let (b, nl) = (base as i32, n_layer as i32);
2769        let cfg = LaunchConfig {
2770            grid_dim: (n_layer.div_ceil(64) as u32, 1, 1),
2771            block_dim: (64, 1, 1),
2772            shared_mem_bytes: 0,
2773        };
2774        let __s_bl = self.gpu.stream();
2775        let mut bl = __s_bl.launch_builder(&f);
2776        bl.arg(len_ptrs).arg(saved).arg(acc).arg(&b).arg(&nl);
2777        unsafe {
2778            bl.launch(cfg)?;
2779        }
2780        Ok(())
2781    }
2782
2783    /// OPTIPIPE increment 1: derive the K=1 successor-valid bit on device.
2784    pub fn spec_fork_valid(
2785        &self,
2786        acc: &CudaSlice<u32>,
2787        optimistic_pending: u32,
2788        valid: &mut CudaSlice<u32>,
2789    ) -> Result<(), Box<dyn std::error::Error>> {
2790        let f = self.func("spec_fork_valid");
2791        let cfg = LaunchConfig {
2792            grid_dim: (1, 1, 1),
2793            block_dim: (1, 1, 1),
2794            shared_mem_bytes: 0,
2795        };
2796        let __s_bl = self.gpu.stream();
2797        let mut bl = __s_bl.launch_builder(&f);
2798        bl.arg(acc).arg(&optimistic_pending).arg(valid);
2799        unsafe {
2800            bl.launch(cfg)?;
2801        }
2802        Ok(())
2803    }
2804
2805    /// OPTIPIPE increment 1: leave stage-local KV lengths on hit, restore them on miss.
2806    pub fn spec_fork_reconcile_kv(
2807        &self,
2808        len_ptrs: &CudaSlice<u64>,
2809        saved: &CudaSlice<i32>,
2810        acc: &CudaSlice<u32>,
2811        valid: &CudaSlice<u32>,
2812        base: usize,
2813        n_layer: usize,
2814    ) -> Result<(), Box<dyn std::error::Error>> {
2815        let f = self.func("spec_fork_reconcile_kv");
2816        let (b, nl) = (base as i32, n_layer as i32);
2817        let cfg = LaunchConfig {
2818            grid_dim: (n_layer.div_ceil(64) as u32, 1, 1),
2819            block_dim: (64, 1, 1),
2820            shared_mem_bytes: 0,
2821        };
2822        let __s_bl = self.gpu.stream();
2823        let mut bl = __s_bl.launch_builder(&f);
2824        bl.arg(len_ptrs)
2825            .arg(saved)
2826            .arg(acc)
2827            .arg(valid)
2828            .arg(&b)
2829            .arg(&nl);
2830        unsafe {
2831            bl.launch(cfg)?;
2832        }
2833        Ok(())
2834    }
2835
2836    /// OPTIPIPE increment 1: conditionally restore one stage-owned recurrent-state buffer.
2837    pub fn spec_fork_restore_f32(
2838        &self,
2839        snapshot: &CudaSlice<f32>,
2840        state: &mut CudaSlice<f32>,
2841        valid: &CudaSlice<u32>,
2842    ) -> Result<(), Box<dyn std::error::Error>> {
2843        assert_eq!(
2844            snapshot.len(),
2845            state.len(),
2846            "fork recurrent snapshot shape mismatch"
2847        );
2848        let f = self.func("spec_fork_restore_f32");
2849        let n = state.len() as i32;
2850        let blocks = state.len().div_ceil(256).min(65535).max(1) as u32;
2851        let cfg = LaunchConfig {
2852            grid_dim: (blocks, 1, 1),
2853            block_dim: (256, 1, 1),
2854            shared_mem_bytes: 0,
2855        };
2856        let __s_bl = self.gpu.stream();
2857        let mut bl = __s_bl.launch_builder(&f);
2858        bl.arg(snapshot).arg(state).arg(valid).arg(&n);
2859        unsafe {
2860            bl.launch(cfg)?;
2861        }
2862        Ok(())
2863    }
2864
2865    /// ROUND-STREAM stage (b): device next-round seed gather (see spec_seed_gather header).
2866    /// Caller D2Ds h_seed into fill_prev after (both slots carry the same value in every arm).
2867    pub fn spec_seed_gather(
2868        &self,
2869        vx: &CudaSlice<f32>,
2870        fill_prev: &CudaSlice<f32>,
2871        acc: &CudaSlice<u32>,
2872        h_seed: &mut CudaSlice<f32>,
2873        base: usize,
2874        n_embd: usize,
2875    ) -> Result<(), Box<dyn std::error::Error>> {
2876        let f = self.func("spec_seed_gather");
2877        let (b, ne) = (base as i32, n_embd as i32);
2878        let cfg = LaunchConfig {
2879            grid_dim: (n_embd.div_ceil(256) as u32, 1, 1),
2880            block_dim: (256, 1, 1),
2881            shared_mem_bytes: 0,
2882        };
2883        let __s_bl = self.gpu.stream();
2884        let mut bl = __s_bl.launch_builder(&f);
2885        bl.arg(vx)
2886            .arg(fill_prev)
2887            .arg(acc)
2888            .arg(h_seed)
2889            .arg(&b)
2890            .arg(&ne);
2891        unsafe {
2892            bl.launch(cfg)?;
2893        }
2894        Ok(())
2895    }
2896
2897    /// ROUND-STREAM stage (a): device greedy accept walk (see spec_accept_greedy header).
2898    pub fn spec_accept_greedy(
2899        &self,
2900        preds: &CudaSlice<u32>,
2901        draft: &CudaSlice<u32>,
2902        last_pred: u32,
2903        base: usize,
2904        k_round: usize,
2905        out: &mut CudaSlice<u32>,
2906    ) -> Result<(), Box<dyn std::error::Error>> {
2907        let f = self.func("spec_accept_greedy");
2908        let (b, k) = (base as i32, k_round as i32);
2909        let cfg = LaunchConfig {
2910            grid_dim: (1, 1, 1),
2911            block_dim: (32, 1, 1),
2912            shared_mem_bytes: 0,
2913        };
2914        let __s_bl = self.gpu.stream();
2915        let mut bl = __s_bl.launch_builder(&f);
2916        bl.arg(preds)
2917            .arg(draft)
2918            .arg(&last_pred)
2919            .arg(&b)
2920            .arg(&k)
2921            .arg(out);
2922        unsafe {
2923            bl.launch(cfg)?;
2924        }
2925        Ok(())
2926    }
2927
2928    // ================= SAMPLED-SPEC PRIMITIVES (spec_sample.cu, piece A) =================
2929    // Counter-based randomness: every call takes (seed, stream_pos) — the caller owns the
2930    // event counter (one per sampled token). temp <= 0 arms are exact greedy limits.
2931
2932    /// y = x/temp + Gumbel(Philox(seed, stream_pos)) over n logits (then run device argmax on y
2933    /// = one categorical sample at temperature `temp`). temp<=0: y = x (pure copy).
2934    pub fn gumbel_perturb(
2935        &self,
2936        x: &CudaSlice<f32>,
2937        y: &mut CudaSlice<f32>,
2938        n: usize,
2939        seed: u64,
2940        stream_pos: u32,
2941        temp: f32,
2942    ) -> Result<(), Box<dyn std::error::Error>> {
2943        let f = self.func("gumbel_perturb_f32");
2944        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
2945        let cfg = LaunchConfig {
2946            grid_dim: (n.div_ceil(256) as u32, 1, 1),
2947            block_dim: (256, 1, 1),
2948            shared_mem_bytes: 0,
2949        };
2950        let __s_b = self.gpu.stream();
2951        let mut b = __s_b.launch_builder(&f);
2952        b.arg(x)
2953            .arg(&mut *y)
2954            .arg(&ni)
2955            .arg(&slo)
2956            .arg(&shi)
2957            .arg(&stream_pos)
2958            .arg(&temp);
2959        unsafe {
2960            b.launch(cfg)?;
2961        }
2962        Ok(())
2963    }
2964
2965    /// GRAMMAR TOKEN MASK (constrained decoding, lane/constrained-full): ban every vocab id
2966    /// whose bit is unset in the packed llguidance bitset, IN PLACE on row `col` of a stacked
2967    /// [B, n_vocab] logits buffer. `mask` = the SimpleVob u32 words H2D'd verbatim
2968    /// (~n_vocab/8 bytes/step — trivial on PCIe); ids >= 32*mask_words (padded lm_head tail)
2969    /// are banned too, the device twin of constrained::apply_mask. Banned value -FLT_MAX ==
2970    /// the argmax/gumbel kernels' init sentinel, so a fully-banned tail can never win and
2971    /// ordering matches the host -inf mask bit-for-bit for every finite logit.
2972    pub fn mask_logits_col(
2973        &self,
2974        logits: &mut CudaSlice<f32>,
2975        mask: &CudaSlice<u32>,
2976        col: usize,
2977        n: usize,
2978        mask_words: usize,
2979    ) -> Result<(), Box<dyn std::error::Error>> {
2980        let f = self.func("mask_logits_f32");
2981        let (ci, ni, mw) = (col as i32, n as i32, mask_words as i32);
2982        let cfg = LaunchConfig {
2983            grid_dim: (n.div_ceil(256).min(1024) as u32, 1, 1),
2984            block_dim: (256, 1, 1),
2985            shared_mem_bytes: 0,
2986        };
2987        let __s_b = self.gpu.stream();
2988        let mut b = __s_b.launch_builder(&f);
2989        b.arg(&mut *logits).arg(mask).arg(&ci).arg(&ni).arg(&mw);
2990        unsafe {
2991            b.launch(cfg)?;
2992        }
2993        Ok(())
2994    }
2995
2996    /// Column-`col` twin of `gumbel_perturb` over stacked logits [B, n_vocab] (the batched
2997    /// serving tick's device sampler): y = x[col]/temp + gumbel(seed, stream_pos, lane).
2998    /// SAME kernel/Philox mapping as `gumbel_perturb` — bit-identical perturbation for the
2999    /// same (seed, stream_pos, temp) regardless of which batch column the row sits in
3000    /// (the lane index is the in-row position; `col` only moves the input pointer). That
3001    /// pointer-invariance IS the serving isolation contract for sampled rows.
3002    pub fn gumbel_perturb_col(
3003        &self,
3004        x: &CudaSlice<f32>,
3005        col: usize,
3006        y: &mut CudaSlice<f32>,
3007        n: usize,
3008        seed: u64,
3009        stream_pos: u32,
3010        temp: f32,
3011    ) -> Result<(), Box<dyn std::error::Error>> {
3012        let f = self.func("gumbel_perturb_f32");
3013        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3014        let col_view = x.slice(col * n..(col + 1) * n);
3015        let cfg = LaunchConfig {
3016            grid_dim: (n.div_ceil(256) as u32, 1, 1),
3017            block_dim: (256, 1, 1),
3018            shared_mem_bytes: 0,
3019        };
3020        let __s_b = self.gpu.stream();
3021        let mut b = __s_b.launch_builder(&f);
3022        b.arg(&col_view)
3023            .arg(&mut *y)
3024            .arg(&ni)
3025            .arg(&slo)
3026            .arg(&shi)
3027            .arg(&stream_pos)
3028            .arg(&temp);
3029        unsafe {
3030            b.launch(cfg)?;
3031        }
3032        Ok(())
3033    }
3034
3035    /// Filtered twin of `gumbel_perturb_col`: the per-row (row_max, th) floor comes from
3036    /// DEVICE buffers (`filter_stats` output slots at `stat_idx`) — one filtered draw from
3037    /// the top-k/top-p/min-p-truncated softmax with no stat D2H and no row copy. Same
3038    /// Philox mapping as every gumbel kernel (pointer-invariant across batch columns —
3039    /// the serving isolation contract for sampled rows).
3040    #[allow(clippy::too_many_arguments)]
3041    pub fn gumbel_perturb_filtered_col(
3042        &self,
3043        x: &CudaSlice<f32>,
3044        col: usize,
3045        y: &mut CudaSlice<f32>,
3046        n: usize,
3047        seed: u64,
3048        stream_pos: u32,
3049        temp: f32,
3050        stat_max: &CudaSlice<f32>,
3051        stat_th: &CudaSlice<f32>,
3052        stat_idx: usize,
3053    ) -> Result<(), Box<dyn std::error::Error>> {
3054        let f = self.func("gumbel_perturb_filtered_col_f32");
3055        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3056        let (ci, si) = (col as i32, stat_idx as i32);
3057        let cfg = LaunchConfig {
3058            grid_dim: (n.div_ceil(256) as u32, 1, 1),
3059            block_dim: (256, 1, 1),
3060            shared_mem_bytes: 0,
3061        };
3062        let __s_b = self.gpu.stream();
3063        let mut b = __s_b.launch_builder(&f);
3064        b.arg(x)
3065            .arg(&ci)
3066            .arg(&mut *y)
3067            .arg(&ni)
3068            .arg(&slo)
3069            .arg(&shi)
3070            .arg(&stream_pos)
3071            .arg(&temp)
3072            .arg(stat_max)
3073            .arg(stat_th)
3074            .arg(&si);
3075        unsafe {
3076            b.launch(cfg)?;
3077        }
3078        Ok(())
3079    }
3080
3081    /// In-graph sampling-event counter bump (spec_sample.cu kernel 5): ctr[0] += 1. The sampled
3082    /// graph-draft chain replays with FIXED kernel args, so the Philox event counter must be
3083    /// DEVICE data — the host seeds it once per round; every replay bumps it before the perturb
3084    /// reads it (counter is data, not state — graph-replay-safe).
3085    pub fn sctr_inc(&self, ctr: &mut CudaSlice<u32>) -> Result<(), Box<dyn std::error::Error>> {
3086        let f = self.func("memra_sctr_inc");
3087        let cfg = LaunchConfig {
3088            grid_dim: (1, 1, 1),
3089            block_dim: (1, 1, 1),
3090            shared_mem_bytes: 0,
3091        };
3092        let __s_b = self.gpu.stream();
3093        let mut b = __s_b.launch_builder(&f);
3094        b.arg(&mut *ctr);
3095        unsafe {
3096            b.launch(cfg)?;
3097        }
3098        Ok(())
3099    }
3100
3101    /// Graph-capturable `gumbel_perturb`: the sampling-event counter comes from DEVICE memory
3102    /// (`ctr[0]`) instead of a host scalar. Identical math to `gumbel_perturb` at
3103    /// stream_pos == ctr[0] (same Philox call, same lane mapping) — the eager and graph sampled
3104    /// chains produce bit-identical perturbations for the same (seed, counter, temp).
3105    pub fn gumbel_perturb_ctr(
3106        &self,
3107        x: &CudaSlice<f32>,
3108        y: &mut CudaSlice<f32>,
3109        n: usize,
3110        seed: u64,
3111        ctr: &CudaSlice<u32>,
3112        temp: f32,
3113    ) -> Result<(), Box<dyn std::error::Error>> {
3114        let f = self.func("gumbel_perturb_ctr_f32");
3115        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3116        let cfg = LaunchConfig {
3117            grid_dim: (n.div_ceil(256) as u32, 1, 1),
3118            block_dim: (256, 1, 1),
3119            shared_mem_bytes: 0,
3120        };
3121        let __s_b = self.gpu.stream();
3122        let mut b = __s_b.launch_builder(&f);
3123        b.arg(x)
3124            .arg(&mut *y)
3125            .arg(&ni)
3126            .arg(&slo)
3127            .arg(&shi)
3128            .arg(ctr)
3129            .arg(&temp);
3130        unsafe {
3131            b.launch(cfg)?;
3132        }
3133        Ok(())
3134    }
3135
3136    /// out[pair] = softmax_temp(x[rows[pair]])[ids[pair]] for npair (row, id) pairs; rows index
3137    /// into x with `row_stride` f32s per row. temp<=0: out = 1.0 iff id is the row argmax
3138    /// (smallest-index tie-break — matches the argmax-gate contract).
3139    pub fn softmax_gather(
3140        &self,
3141        x: &CudaSlice<f32>,
3142        row_stride: usize,
3143        ids: &CudaSlice<u32>,
3144        rows: &CudaSlice<i32>,
3145        out: &mut CudaSlice<f32>,
3146        n: usize,
3147        npair: usize,
3148        temp: f32,
3149    ) -> Result<(), Box<dyn std::error::Error>> {
3150        let f = self.func("softmax_gather_f32");
3151        let (ni, rs) = (n as i32, row_stride as i64);
3152        let np = npair as i32;
3153        let cfg = LaunchConfig {
3154            grid_dim: (npair as u32, 1, 1),
3155            block_dim: (256, 1, 1),
3156            shared_mem_bytes: 0,
3157        };
3158        let __s_b = self.gpu.stream();
3159        let mut b = __s_b.launch_builder(&f);
3160        b.arg(x)
3161            .arg(&rs)
3162            .arg(ids)
3163            .arg(rows)
3164            .arg(&mut *out)
3165            .arg(&ni)
3166            .arg(&np)
3167            .arg(&temp);
3168        unsafe {
3169            b.launch(cfg)?;
3170        }
3171        Ok(())
3172    }
3173
3174    /// Sample token from norm(max(0, softmax_temp(p) - softmax_temp(q))) (q = None -> plain
3175    /// categorical from softmax_temp(p)). Row stats (max, sumexp at temp) must be precomputed
3176    /// (softmax_gather's pass-1 values; see spec.rs caller). Deterministic fixed-order CDF walk.
3177    pub fn residual_sample(
3178        &self,
3179        p: &CudaSlice<f32>,
3180        q: Option<&CudaSlice<f32>>,
3181        n: usize,
3182        temp: f32,
3183        seed: u64,
3184        stream_pos: u32,
3185        out_tok: &mut CudaSlice<u32>,
3186    ) -> Result<(), Box<dyn std::error::Error>> {
3187        let f = self.func("residual_sample_f32");
3188        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3189        let nth = 1024u32;
3190        let cfg = LaunchConfig {
3191            grid_dim: (1, 1, 1),
3192            block_dim: (nth, 1, 1),
3193            shared_mem_bytes: 0,
3194        };
3195        let has_q: i32 = q.is_some() as i32;
3196        let qbuf = q.unwrap_or(p); // dummy when absent; kernel gates on has_q
3197        let __s_b = self.gpu.stream();
3198        let mut b = __s_b.launch_builder(&f);
3199        b.arg(p)
3200            .arg(qbuf)
3201            .arg(&has_q)
3202            .arg(&ni)
3203            .arg(&temp)
3204            .arg(&slo)
3205            .arg(&shi)
3206            .arg(&stream_pos)
3207            .arg(&mut *out_tok);
3208        unsafe {
3209            b.launch(cfg)?;
3210        }
3211        Ok(())
3212    }
3213
3214    /// Access the shared MoE residency cache (EDGE-1 §B), building it on first use under
3215    /// MEMRA_MOE_CACHE. The closure runs while the lock is held — keep it to lookup/admit/issue, not
3216    /// the GEMM. `max_block_bytes` sizes the slots (largest of gate/up/down). Returns the closure's
3217    /// result. If MEMRA_MOE_CACHE is unset this is never called (the caller checks the env first).
3218    pub fn with_moe_cache<R>(
3219        &self,
3220        max_block_bytes: usize,
3221        f: impl FnOnce(
3222            &mut crate::moe_cache::MoeSlotCache,
3223            &Engine,
3224        ) -> Result<R, Box<dyn std::error::Error>>,
3225    ) -> Result<R, Box<dyn std::error::Error>> {
3226        let mut guard = self.moe_cache.lock().unwrap();
3227        if guard.is_none() {
3228            *guard = Some(crate::moe_cache::MoeSlotCache::new(self, max_block_bytes)?);
3229        }
3230        let cache = guard.as_mut().unwrap();
3231        f(cache, self)
3232    }
3233
3234    /// Freeze the already-built MoE residency set. This never constructs a cache: callers use it
3235    /// only after a real prefill has populated the machine-specific CPU/GPU working set.
3236    pub fn freeze_moe_cache(&self) {
3237        if let Some(cache) = self.moe_cache.lock().unwrap().as_mut() {
3238            cache.freeze();
3239        }
3240    }
3241
3242    /// The current residency set as (layer, proj, ex) triples, or None if no cache was built.
3243    /// Never constructs a cache.
3244    pub fn export_moe_residency(&self) -> Option<Vec<(u16, u8, u16)>> {
3245        self.moe_cache
3246            .lock()
3247            .unwrap()
3248            .as_ref()
3249            .map(crate::moe_cache::MoeSlotCache::export_residency)
3250    }
3251
3252    pub(crate) fn moe_cache_frozen(&self) -> bool {
3253        self.moe_cache
3254            .lock()
3255            .unwrap()
3256            .as_ref()
3257            .is_some_and(crate::moe_cache::MoeSlotCache::is_frozen)
3258    }
3259
3260    /// A frozen heterogeneous CPU/GPU expert split cannot use Hy3's ordinary batched prefill
3261    /// efficiently: T>=PRIME_MIN_T bypasses the CPU backend and transiently rereads every missing
3262    /// expert through the GPU spill path. Replay the short prompt through decode after freezing,
3263    /// while leaving the profiling warmup's established batched behavior untouched.
3264    /// (`pub`: run-gen's #46 batched-prime gate skips itself when generation will take the
3265    /// tokenwise arm anyway.)
3266    pub fn frozen_cpu_experts_prefer_tokenwise_prime(&self) -> bool {
3267        crate::cpu_experts::configured()
3268            && self.moe_cache_frozen()
3269            && std::env::var("MEMRA_CPU_EXPERT_BATCHED_PRIME").as_deref() != Ok("1")
3270    }
3271
3272    /// Install the loaded model's exact retained expert-block inventory before lazy cache build.
3273    pub(crate) fn configure_moe_cache_layout(&self, block_bytes: Vec<usize>) {
3274        assert!(
3275            self.moe_cache.lock().unwrap().is_none(),
3276            "MoE cache layout configured after cache construction"
3277        );
3278        *self.moe_cache_layout.lock().unwrap() = Some(block_bytes);
3279    }
3280
3281    pub(crate) fn moe_cache_layout(&self) -> Option<Vec<usize>> {
3282        self.moe_cache_layout.lock().unwrap().clone()
3283    }
3284
3285    /// True if the MoE residency cache is enabled (MEMRA_MOE_CACHE set).
3286    pub fn moe_cache_enabled() -> bool {
3287        std::env::var("MEMRA_MOE_CACHE").as_deref() != Ok("0")
3288    }
3289
3290    /// Snapshot the MoE cache counters (hits, misses, staged_bytes, n_slots) for the §D.4 PCIe gate.
3291    /// Returns None if the cache was never built (disabled or no MoE forward ran).
3292    pub fn moe_cache_stats(&self) -> Option<(u64, u64, u64, usize)> {
3293        let guard = self.moe_cache.lock().unwrap();
3294        guard
3295            .as_ref()
3296            .map(|c| (c.hits, c.misses, c.staged_bytes, c.n_slots()))
3297    }
3298
3299    /// Experimental CPU expert backend counters: completed layer calls, experts served, and the
3300    /// sum of backend wall nanoseconds. The timer includes explicit disk->RAM fills on cache misses;
3301    /// callers compare a before/after snapshot around a decode window.
3302    pub fn cpu_expert_stats(
3303        &self,
3304    ) -> Option<(u64, u64, u64, u64, u64, u64, u64, u64, u64, u64, u64)> {
3305        crate::cpu_experts::configured().then(crate::cpu_experts::stats)
3306    }
3307
3308    /// Caller-blocked nanoseconds at CPU expert joins. Compare before/after snapshots to measure
3309    /// the backend tail that resident-GPU expert work did not hide.
3310    pub fn cpu_expert_predictor_stats(&self) -> (u64, u64) {
3311        crate::cpu_experts::predictor_stats()
3312    }
3313
3314    pub fn cpu_expert_exposed_wait_ns(&self) -> Option<u64> {
3315        crate::cpu_experts::configured().then(crate::cpu_experts::exposed_wait_ns)
3316    }
3317
3318    /// CPU-routed expert selections grouped by how many of their three projections were already
3319    /// resident in HBM. This makes otherwise-stranded partial residency visible to tuning runs.
3320    pub fn cpu_expert_gpu_residency_stats(&self) -> Option<(u64, u64, u64)> {
3321        crate::cpu_experts::configured().then(crate::cpu_experts::incomplete_gpu_residency_stats)
3322    }
3323
3324    /// Positioned-read proof-backend counters:
3325    /// `(reads, bytes, read_errors, short_reads, mmap_fallbacks, buffer_waits, ring_full)`.
3326    pub fn moe_pread_stats(&self) -> Option<(u64, u64, u64, u64, u64, u64, u64)> {
3327        let guard = self.moe_cache.lock().unwrap();
3328        guard
3329            .as_ref()
3330            .and_then(|cache| cache.pread_stats())
3331            .map(|stats| {
3332                (
3333                    stats.reads,
3334                    stats.bytes,
3335                    stats.read_errors,
3336                    stats.short_reads,
3337                    stats.fallbacks,
3338                    stats.buffer_waits,
3339                    stats.ring_full,
3340                )
3341            })
3342    }
3343
3344    /// Spill configuration values that warned and substituted their documented defaults.
3345    pub fn spill_config_fallbacks(&self) -> u64 {
3346        crate::spill_pread::config_fallbacks()
3347    }
3348
3349    /// Reset the MoE cache perf counters (to separate warmup from steady-state windows).
3350    pub fn moe_cache_reset_counters(&self) {
3351        if let Some(c) = self.moe_cache.lock().unwrap().as_mut() {
3352            c.reset_counters();
3353        }
3354    }
3355
3356    pub fn htod_bytes(&self, v: &[u8]) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
3357        Ok(self.gpu.stream().clone_htod(v)?)
3358    }
3359
3360    /// `htod_bytes` with a mapped (uninit) tail pad: the wide-load expert dots read up to 6B
3361    /// past the final q4_0 block through their aligned window — the bytes never reach a
3362    /// result (funnelshift discards them) but must be mapped memory.
3363    pub fn htod_bytes_padded(
3364        &self,
3365        v: &[u8],
3366        pad: usize,
3367    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
3368        let mut d = self.alloc_u8_uninit(v.len() + pad)?;
3369        {
3370            let mut view = d.slice_mut(0..v.len());
3371            self.gpu.stream().memcpy_htod(v, &mut view)?;
3372        }
3373        Ok(d)
3374    }
3375
3376    /// Device-to-device copy of `src` into `dst[off..off+len]` (f32). For in-place KV append.
3377    pub fn copy_into(
3378        &self,
3379        dst: &mut CudaSlice<f32>,
3380        off: usize,
3381        src: &CudaSlice<f32>,
3382        len: usize,
3383    ) -> Result<(), Box<dyn std::error::Error>> {
3384        let mut view = dst.slice_mut(off..off + len);
3385        self.gpu
3386            .stream()
3387            .memcpy_dtod(&src.slice(0..len), &mut view)?;
3388        Ok(())
3389    }
3390
3391    /// View a sub-range of a device buffer (for attending over [0..len) of a KV cache).
3392    /// u8 twin of copy_into (D2D byte-range copy at an offset).
3393    pub fn copy_u8_into(
3394        &self,
3395        dst: &mut CudaSlice<u8>,
3396        off: usize,
3397        src: &CudaSlice<u8>,
3398        len: usize,
3399    ) -> Result<(), Box<dyn std::error::Error>> {
3400        let mut view = dst.slice_mut(off..off + len);
3401        self.gpu
3402            .stream()
3403            .memcpy_dtod(&src.slice(0..len), &mut view)?;
3404        Ok(())
3405    }
3406
3407    /// D2D byte-range copy with explicit source and destination offsets.
3408    pub fn copy_u8_range_into(
3409        &self,
3410        dst: &mut CudaSlice<u8>,
3411        dst_off: usize,
3412        src: &CudaSlice<u8>,
3413        src_off: usize,
3414        len: usize,
3415    ) -> Result<(), Box<dyn std::error::Error>> {
3416        let mut dst_view = dst.slice_mut(dst_off..dst_off + len);
3417        self.gpu
3418            .stream()
3419            .memcpy_dtod(&src.slice(src_off..src_off + len), &mut dst_view)?;
3420        Ok(())
3421    }
3422
3423    /// Resolve an absolute append slot to the Step35 SWA layer's physical rows. At wrap, copy
3424    /// only the aligned live prefix through temporary device storage and rebase it at row zero,
3425    /// keeping the audited attention range contiguous without changing its absolute start.
3426    pub fn prepare_kv_append(
3427        &self,
3428        kv: &mut crate::cache::KvLayer,
3429        retain_from: usize,
3430        append_rows: usize,
3431    ) -> Result<usize, Box<dyn std::error::Error>> {
3432        let Some(plan) = kv
3433            .ring
3434            .as_ref()
3435            .map(|ring| ring.append_plan(kv.len, retain_from, append_rows))
3436            .transpose()?
3437        else {
3438            return Ok(kv.len);
3439        };
3440        match plan {
3441            crate::cache::KvRingAppend::Contiguous { write_row } => Ok(write_row),
3442            crate::cache::KvRingAppend::Rebase {
3443                src_row,
3444                keep_rows,
3445                new_base,
3446                write_row,
3447            } => {
3448                if keep_rows > 0 {
3449                    let k_len = keep_rows * kv.k_tok_bytes;
3450                    let v_len = keep_rows * kv.v_tok_bytes;
3451                    let mut k_tmp = self.alloc_u8_uninit(k_len)?;
3452                    let mut v_tmp = self.alloc_u8_uninit(v_len)?;
3453                    self.copy_u8_range_into(&mut k_tmp, 0, &kv.k, src_row * kv.k_tok_bytes, k_len)?;
3454                    self.copy_u8_range_into(&mut v_tmp, 0, &kv.v, src_row * kv.v_tok_bytes, v_len)?;
3455                    self.copy_u8_into(&mut kv.k, 0, &k_tmp, k_len)?;
3456                    self.copy_u8_into(&mut kv.v, 0, &v_tmp, v_len)?;
3457                }
3458                kv.ring.as_mut().unwrap().apply_rebase(new_base);
3459                Ok(write_row)
3460            }
3461        }
3462    }
3463
3464    /// H2D write of `src` into `dst[off..off+src.len()]` (u8). In-place row updates for the
3465    /// adaptive trim head: no realloc, so captured graphs keep their baked addresses.
3466    pub fn htod_u8_into(
3467        &self,
3468        dst: &mut CudaSlice<u8>,
3469        off: usize,
3470        src: &[u8],
3471    ) -> Result<(), Box<dyn std::error::Error>> {
3472        let mut view = dst.slice_mut(off..off + src.len());
3473        self.gpu.stream().memcpy_htod(src, &mut view)?;
3474        Ok(())
3475    }
3476
3477    pub fn view<'a>(&self, b: &'a CudaSlice<f32>, len: usize) -> cudarc::driver::CudaView<'a, f32> {
3478        b.slice(0..len)
3479    }
3480
3481    /// View the first `len` BYTES of a u8 device buffer (quantized KV cache: [0..t_kv*tok_bytes)).
3482    /// Byte-range view (gemma4 R6 window offset into the quantized KV stream).
3483    pub fn view_u8_range<'a>(
3484        &self,
3485        b: &'a CudaSlice<u8>,
3486        start: usize,
3487        end: usize,
3488    ) -> cudarc::driver::CudaView<'a, u8> {
3489        b.slice(start..end)
3490    }
3491    pub fn view_u8<'a>(
3492        &self,
3493        b: &'a CudaSlice<u8>,
3494        len: usize,
3495    ) -> cudarc::driver::CudaView<'a, u8> {
3496        b.slice(0..len)
3497    }
3498
3499    /// Append-quantize ONE token's post-RoPE K (q8_0) and V (q5_1) into the resident byte caches at
3500    /// token index `t` (KVQUANT-PLAN §C). One CTA (one warp) per 32-element block; the kernel writes
3501    /// the f16 scale(s) + packed quants for K and V. k_row/v_row are f32 [kv_dim_k]/[kv_dim_v].
3502    pub fn append_kv_quantized(
3503        &self,
3504        k_row: &CudaSlice<f32>,
3505        v_row: &CudaSlice<f32>,
3506        kc: &mut CudaSlice<u8>,
3507        vc: &mut CudaSlice<u8>,
3508        t: usize,
3509        kv_dim_k: usize,
3510        kv_dim_v: usize,
3511        k_tok_bytes: usize,
3512        v_tok_bytes: usize,
3513        g: bool,
3514    ) -> Result<(), Box<dyn std::error::Error>> {
3515        let f = if g {
3516            self.func_g("append_quantize_kv_q8_0_q5_1")
3517        } else {
3518            self.func("append_quantize_kv_q8_0_q5_1")
3519        };
3520        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
3521        let cfg = LaunchConfig {
3522            grid_dim: (nblk, 1, 1),
3523            block_dim: (32, 1, 1),
3524            shared_mem_bytes: 0,
3525        };
3526        let (ti, kdk, kdv) = (t as i32, kv_dim_k as i32, kv_dim_v as i32);
3527        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3528        let __s_b = self.gpu.stream();
3529        let mut b = __s_b.launch_builder(&f);
3530        b.arg(k_row)
3531            .arg(v_row)
3532            .arg(kc)
3533            .arg(vc)
3534            .arg(&ti)
3535            .arg(&kdk)
3536            .arg(&kdv)
3537            .arg(&ktb)
3538            .arg(&vtb);
3539        unsafe {
3540            b.launch(cfg)?;
3541        }
3542        Ok(())
3543    }
3544
3545    /// Device-counter variant of `append_kv_quantized` (CUDA-GRAPH-PLAN Phase 2): the write slot
3546    /// `t` is read from `t_dev[0]` (a resident device i32[1]) instead of a host int arg, so the
3547    /// launch args are FIXED across decode steps (graph-capturable). Identical quant math.
3548    pub fn append_kv_quantized_dc(
3549        &self,
3550        k_row: &CudaSlice<f32>,
3551        v_row: &CudaSlice<f32>,
3552        kc: &mut CudaSlice<u8>,
3553        vc: &mut CudaSlice<u8>,
3554        t_dev: &CudaSlice<i32>,
3555        kv_dim_k: usize,
3556        kv_dim_v: usize,
3557        k_tok_bytes: usize,
3558        v_tok_bytes: usize,
3559        g: bool,
3560    ) -> Result<(), Box<dyn std::error::Error>> {
3561        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
3562        let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
3563        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3564        // PDL wave-B2: flash-module flavor mirrors the builder path's g flag exactly.
3565        if Self::pdl_on() && Self::pdl_wb_on() {
3566            use cudarc::driver::{DevicePtr, DevicePtrMut};
3567            let s = &self.gpu.stream();
3568            let (pk, _g0) = k_row.device_ptr(s);
3569            let (pv, _g1) = v_row.device_ptr(s);
3570            let (pkc, _g2) = kc.device_ptr_mut(s);
3571            let (pvc, _g3) = vc.device_ptr_mut(s);
3572            let (pt, _g4) = t_dev.device_ptr(s);
3573            let mut ps = [
3574                &pk as *const _ as *mut std::ffi::c_void,
3575                &pv as *const _ as *mut _,
3576                &pkc as *const _ as *mut _,
3577                &pvc as *const _ as *mut _,
3578                &pt as *const _ as *mut _,
3579                &kdk as *const _ as *mut _,
3580                &kdv as *const _ as *mut _,
3581                &ktb as *const _ as *mut _,
3582                &vtb as *const _ as *mut _,
3583            ];
3584            unsafe {
3585                self.launch_pdl_flash(
3586                    g,
3587                    "append_quantize_kv_q8_0_q5_1_dc",
3588                    (nblk, 1, 1),
3589                    (32, 1, 1),
3590                    0,
3591                    &mut ps,
3592                )?;
3593            }
3594            return Ok(());
3595        }
3596        let f = if g {
3597            self.func_g("append_quantize_kv_q8_0_q5_1_dc")
3598        } else {
3599            self.func("append_quantize_kv_q8_0_q5_1_dc")
3600        };
3601        let cfg = LaunchConfig {
3602            grid_dim: (nblk, 1, 1),
3603            block_dim: (32, 1, 1),
3604            shared_mem_bytes: 0,
3605        };
3606        let __s_b = self.gpu.stream();
3607        let mut b = __s_b.launch_builder(&f);
3608        b.arg(k_row)
3609            .arg(v_row)
3610            .arg(kc)
3611            .arg(vc)
3612            .arg(t_dev)
3613            .arg(&kdk)
3614            .arg(&kdv)
3615            .arg(&ktb)
3616            .arg(&vtb);
3617        unsafe {
3618            b.launch(cfg)?;
3619        }
3620        Ok(())
3621    }
3622
3623    /// Append-quantize T token rows in one shot (BATCHED PROMPT PRIME). k_rows/v_rows are
3624    /// token-major [T, kv_dim] post-RoPE f32; rows land at cache slots t0..t0+T. Default = the
3625    /// batched `_rows` kernel: one (nblk, T) launch whose per-(block,token) warp program is the
3626    /// per-token append kernel verbatim -> every written row is BIT-IDENTICAL to T sequential
3627    /// `append_kv_quantized_view` calls (kernel_check pins the bytes). MEMRA_PRIME_APPEND_LOOP=1
3628    /// forces the T-launch per-row loop (the A/B seam that measured the launch overhead).
3629    #[allow(clippy::too_many_arguments)]
3630    pub fn append_kv_quantized_rows(
3631        &self,
3632        k_rows: &CudaSlice<f32>,
3633        v_rows: &CudaSlice<f32>,
3634        kc: &mut CudaSlice<u8>,
3635        vc: &mut CudaSlice<u8>,
3636        t0: usize,
3637        t: usize,
3638        kv_dim_k: usize,
3639        kv_dim_v: usize,
3640        k_tok_bytes: usize,
3641        v_tok_bytes: usize,
3642        g: bool,
3643    ) -> Result<(), Box<dyn std::error::Error>> {
3644        if std::env::var("MEMRA_PRIME_APPEND_LOOP").is_ok() {
3645            for i in 0..t {
3646                let k_row = k_rows.slice(i * kv_dim_k..(i + 1) * kv_dim_k);
3647                let v_row = v_rows.slice(i * kv_dim_v..(i + 1) * kv_dim_v);
3648                self.append_kv_quantized_view(
3649                    &k_row,
3650                    &v_row,
3651                    kc,
3652                    vc,
3653                    t0 + i,
3654                    kv_dim_k,
3655                    kv_dim_v,
3656                    k_tok_bytes,
3657                    v_tok_bytes,
3658                    g,
3659                )?;
3660            }
3661            return Ok(());
3662        }
3663        let f = if g {
3664            self.func_g("append_quantize_kv_q8_0_q5_1_rows")
3665        } else {
3666            self.func("append_quantize_kv_q8_0_q5_1_rows")
3667        };
3668        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
3669        let cfg = LaunchConfig {
3670            grid_dim: (nblk, t as u32, 1),
3671            block_dim: (32, 1, 1),
3672            shared_mem_bytes: 0,
3673        };
3674        let (t0i, kdk, kdv) = (t0 as i32, kv_dim_k as i32, kv_dim_v as i32);
3675        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3676        let __s_b = self.gpu.stream();
3677        let mut b = __s_b.launch_builder(&f);
3678        b.arg(k_rows)
3679            .arg(v_rows)
3680            .arg(kc)
3681            .arg(vc)
3682            .arg(&t0i)
3683            .arg(&kdk)
3684            .arg(&kdv)
3685            .arg(&ktb)
3686            .arg(&vtb);
3687        unsafe {
3688            b.launch(cfg)?;
3689        }
3690        Ok(())
3691    }
3692
3693    /// Increment a device i32[1] counter in place (p[0] += 1) via the resident `inc_i32` kernel.
3694    /// Used to advance the device-resident seqlen/pos counters inside the decode-dc path (and,
3695    /// later, inside a captured graph) without a host round-trip.
3696    pub fn inc_seqlen(&self, p: &mut CudaSlice<i32>) -> Result<(), Box<dyn std::error::Error>> {
3697        let f = self.func("inc_i32");
3698        let cfg = LaunchConfig {
3699            grid_dim: (1, 1, 1),
3700            block_dim: (1, 1, 1),
3701            shared_mem_bytes: 0,
3702        };
3703        let __s_b = self.gpu.stream();
3704        let mut b = __s_b.launch_builder(&f);
3705        b.arg(p);
3706        unsafe {
3707            b.launch(cfg)?;
3708        }
3709        Ok(())
3710    }
3711
3712    /// Like `append_kv_quantized` but k_row/v_row are CudaViews (one token's row sliced out of a
3713    /// token-major [T, kv_dim] activation buffer — the MTP verify path appends T tokens).
3714    pub fn append_kv_quantized_view(
3715        &self,
3716        k_row: &cudarc::driver::CudaView<f32>,
3717        v_row: &cudarc::driver::CudaView<f32>,
3718        kc: &mut CudaSlice<u8>,
3719        vc: &mut CudaSlice<u8>,
3720        t: usize,
3721        kv_dim_k: usize,
3722        kv_dim_v: usize,
3723        k_tok_bytes: usize,
3724        v_tok_bytes: usize,
3725        g: bool,
3726    ) -> Result<(), Box<dyn std::error::Error>> {
3727        let f = if g {
3728            self.func_g("append_quantize_kv_q8_0_q5_1")
3729        } else {
3730            self.func("append_quantize_kv_q8_0_q5_1")
3731        };
3732        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
3733        let cfg = LaunchConfig {
3734            grid_dim: (nblk, 1, 1),
3735            block_dim: (32, 1, 1),
3736            shared_mem_bytes: 0,
3737        };
3738        let (ti, kdk, kdv) = (t as i32, kv_dim_k as i32, kv_dim_v as i32);
3739        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3740        let __s_b = self.gpu.stream();
3741        let mut b = __s_b.launch_builder(&f);
3742        b.arg(k_row)
3743            .arg(v_row)
3744            .arg(kc)
3745            .arg(vc)
3746            .arg(&ti)
3747            .arg(&kdk)
3748            .arg(&kdv)
3749            .arg(&ktb)
3750            .arg(&vtb);
3751        unsafe {
3752            b.launch(cfg)?;
3753        }
3754        Ok(())
3755    }
3756
3757    /// Device-to-device copy of a CudaView `src` into `dst[off..off+len]` (f32). Like `copy_into`
3758    /// but the source is a sub-view (e.g. one column of a token-major activation buffer).
3759    pub fn copy_view_into(
3760        &self,
3761        dst: &mut CudaSlice<f32>,
3762        off: usize,
3763        src: &cudarc::driver::CudaView<f32>,
3764        len: usize,
3765    ) -> Result<(), Box<dyn std::error::Error>> {
3766        let mut view = dst.slice_mut(off..off + len);
3767        self.gpu
3768            .stream()
3769            .memcpy_dtod(&src.slice(0..len), &mut view)?;
3770        Ok(())
3771    }
3772
3773    /// Real device-to-device COPY of `src` into a freshly allocated buffer (NOT an Arc clone).
3774    /// Used for cache snapshots (MTP-PLAN §D.4): `CudaSlice::clone()` only bumps a refcount and
3775    /// would alias the live buffer; this allocs new device memory and memcpy_dtod's the contents.
3776    pub fn clone_dtod(
3777        &self,
3778        src: &CudaSlice<f32>,
3779    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3780        let mut dst = self.gpu.stream().alloc_zeros::<f32>(src.len())?;
3781        self.gpu.stream().memcpy_dtod(src, &mut dst)?;
3782        Ok(dst)
3783    }
3784
3785    /// D2D row extraction: copy a view (e.g. one row of a [B, n] batch buffer) into `dst`.
3786    /// Stream-ordered, async — decode_batch's per-sequence row plumbing.
3787    pub fn dtod_copy_view(
3788        &self,
3789        src: &cudarc::driver::CudaView<f32>,
3790        dst: &mut CudaSlice<f32>,
3791    ) -> Result<(), Box<dyn std::error::Error>> {
3792        self.gpu.stream().memcpy_dtod(src, dst)?;
3793        Ok(())
3794    }
3795
3796    /// D2D i8 twin of `dtod_copy_view` (q8_1 activation rows).
3797    pub fn dtod_copy_view_i8(
3798        &self,
3799        src: &cudarc::driver::CudaView<i8>,
3800        dst: &mut CudaSlice<i8>,
3801    ) -> Result<(), Box<dyn std::error::Error>> {
3802        self.gpu.stream().memcpy_dtod(src, dst)?;
3803        Ok(())
3804    }
3805
3806    /// D2D row placement: copy `src` into `dst[offset .. offset+src.len()]`.
3807    pub fn dtod_copy_into(
3808        &self,
3809        src: &CudaSlice<f32>,
3810        dst: &mut CudaSlice<f32>,
3811        offset: usize,
3812    ) -> Result<(), Box<dyn std::error::Error>> {
3813        let n = src.len();
3814        let mut dv = dst.slice_mut(offset..offset + n);
3815        self.gpu.stream().memcpy_dtod(src, &mut dv)?;
3816        Ok(())
3817    }
3818
3819    /// Batched uniform-size D2D copy (engine-bundle slice 1, DSF-ROUNDCOST-20260820 §1.1):
3820    /// `n` disjoint regions of `words` f32 each; `table` = [src_0..src_{n-1}, dst_0..dst_{n-1}]
3821    /// raw device pointers. ONE kernel launch replaces `n` memcpy_dtod dispatches — the dspark
3822    /// round's snap/commit copy dribble (~0.9 ms/round of dispatch serialization measured).
3823    /// Bytes and stream order are identical to the memcpy sequence it replaces.
3824    pub fn copy_batch_uniform_f32(
3825        &self,
3826        table: &CudaSlice<u64>,
3827        n: usize,
3828        words: usize,
3829    ) -> Result<(), Box<dyn std::error::Error>> {
3830        if n == 0 || words == 0 {
3831            return Ok(());
3832        }
3833        debug_assert!(
3834            table.len() >= 2 * n,
3835            "pointer table must hold n srcs + n dsts"
3836        );
3837        let f = self.func("copy_batch_uniform_f32");
3838        // Enough blocks to stream a multi-MB region, few enough that (chunks x n) stays a
3839        // sane grid: 512K-word ssm regions get 48 grid-striding blocks each.
3840        let chunks = (words / 4).max(1).div_ceil(256).min(48) as u32;
3841        let (ni, wi) = (n as i32, words as i32);
3842        let cfg = LaunchConfig {
3843            grid_dim: (chunks, n as u32, 1),
3844            block_dim: (256, 1, 1),
3845            shared_mem_bytes: 0,
3846        };
3847        let __s = self.gpu.stream();
3848        let mut b = __s.launch_builder(&f);
3849        b.arg(table).arg(&ni).arg(&wi);
3850        unsafe {
3851            b.launch(cfg)?;
3852        }
3853        Ok(())
3854    }
3855
3856    /// H2D refresh of an EXISTING u64 pointer table IN PLACE (stable pointer — the batched
3857    /// state-copy tables are refreshed per round because the GDN ssm handles ping-pong).
3858    pub fn htod_u64_into(
3859        &self,
3860        v: &[u64],
3861        dst: &mut CudaSlice<u64>,
3862    ) -> Result<(), Box<dyn std::error::Error>> {
3863        let mut view = dst.slice_mut(0..v.len());
3864        self.gpu.stream().memcpy_htod(v, &mut view)?;
3865        Ok(())
3866    }
3867
3868    /// Indirect-source copy (engine-bundle slice 3): the src ADDRESS is loaded from a
3869    /// device pointer-table entry at run time, so a captured graph follows the gdn
3870    /// ping-pong through the same table its scan kernels read — a baked memcpy node
3871    /// would keep the capture-time physical buffer. `dst_off`/`words` in f32 elements.
3872    pub fn copy_indirect_src_f32(
3873        &self,
3874        src_entry: &cudarc::driver::CudaView<u64>,
3875        dst: &mut CudaSlice<f32>,
3876        dst_off: usize,
3877        words: usize,
3878    ) -> Result<(), Box<dyn std::error::Error>> {
3879        let f = self.func("copy_indirect_src_f32");
3880        let chunks = (words / 4).max(1).div_ceil(256).min(48) as u32;
3881        let wi = words as i32;
3882        let cfg = LaunchConfig {
3883            grid_dim: (chunks, 1, 1),
3884            block_dim: (256, 1, 1),
3885            shared_mem_bytes: 0,
3886        };
3887        let mut dv = dst.slice_mut(dst_off..dst_off + words);
3888        let __s = self.gpu.stream();
3889        let mut b = __s.launch_builder(&f);
3890        b.arg(src_entry).arg(&mut dv).arg(&wi);
3891        unsafe {
3892            b.launch(cfg)?;
3893        }
3894        Ok(())
3895    }
3896
3897    /// Uninitialized i8 device buffer (decode_batch q8_1 row scratch).
3898    pub fn uninit_i8(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
3899        self.alloc_uninit::<i8>(n)
3900    }
3901
3902    /// Resident-quantized linear (Stage-A: f32 dequant-in-kernel). y[m,out]=x[m,in]@W[out,in]^T.
3903    pub fn qmatvec(
3904        &self,
3905        w: &CudaSlice<u8>,
3906        x: &CudaSlice<f32>,
3907        m: usize,
3908        in_f: usize,
3909        out_f: usize,
3910        qtype: i32,
3911        row_bytes: usize,
3912    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3913        let f = self.func("qmatvec_f32");
3914        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
3915        let cfg = LaunchConfig {
3916            grid_dim: (out_f as u32, m as u32, 1),
3917            block_dim: (256, 1, 1),
3918            shared_mem_bytes: 0,
3919        };
3920        let (inf, outf, mi, qt, rb) =
3921            (in_f as i32, out_f as i32, m as i32, qtype, row_bytes as i64);
3922        let __s_b = self.gpu.stream();
3923        let mut b = __s_b.launch_builder(&f);
3924        b.arg(w)
3925            .arg(x)
3926            .arg(&mut y)
3927            .arg(&inf)
3928            .arg(&outf)
3929            .arg(&mi)
3930            .arg(&qt)
3931            .arg(&rb);
3932        unsafe {
3933            b.launch(cfg)?;
3934        }
3935        Ok(y)
3936    }
3937
3938    /// Allocate a reusable u8 GPU scratch buffer (for staged expert weights).
3939    pub fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
3940        let s = self.gpu.stream().alloc_zeros::<u8>(n)?;
3941        self.keep_if_capturing(&s);
3942        Ok(s)
3943    }
3944
3945    /// Uninitialized u8 scratch — skips alloc_zeros' memset. ONLY for staging buffers whose read
3946    /// range is fully overwritten by a stage_expert H2D before any kernel reads it (LAUNCH-STRUCTURE
3947    /// STAGE 2: the per-layer MoE scratch trio was 3 dead ~1MB memsets per layer per decode token).
3948    pub fn alloc_u8_uninit(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
3949        let s = unsafe { self.gpu.stream().alloc::<u8>(n)? };
3950        self.keep_if_capturing(&s);
3951        Ok(s)
3952    }
3953
3954    /// Zero a SUB-RANGE of an f32 buffer (CudaViewMut) — the row-sized memset the moe_out
3955    /// memset-elision uses for tokens that fall off the gdec fast path (LAUNCH-STRUCTURE STAGE 2).
3956    pub fn memset_zeros_view(
3957        &self,
3958        dst: &mut cudarc::driver::CudaViewMut<f32>,
3959    ) -> Result<(), Box<dyn std::error::Error>> {
3960        self.gpu.stream().memset_zeros(dst)?;
3961        Ok(())
3962    }
3963
3964    /// EDGE-1 staging: copy `host_bytes` (a sub-slice of a HostExps buffer) into `scratch`
3965    /// at byte offset `off` (async H2D on the default stream). Length is host_bytes.len().
3966    /// The qmatvec_view that reads `scratch[off..]` is enqueued on the SAME stream after this,
3967    /// so ordering is guaranteed without an explicit sync (Stage-1; Stage-2 prefetch on a 2nd
3968    /// stream would require an event).
3969    pub fn stage_expert(
3970        &self,
3971        host_bytes: &[u8],
3972        scratch: &mut CudaSlice<u8>,
3973        off: usize,
3974    ) -> Result<(), Box<dyn std::error::Error>> {
3975        let mut dst = scratch.slice_mut(off..off + host_bytes.len()); // CudaViewMut<u8>
3976        self.gpu.stream().memcpy_htod(host_bytes, &mut dst)?; // accepts &[u8] HostSlice src
3977        Ok(())
3978    }
3979
3980    /// EDGE-1 §A: fused MoE router. `logits` is the router output [t, n_expert] (device, f32, the
3981    /// `gate_inp @ z` result). Returns (sel_idx [t, n_used] i32, sel_w [t, n_used] f32): the top-k
3982    /// expert ids (DESC by prob, ascending-index tiebreak) and renormalized weights. Replaces the
3983    /// host dtoh + softmax-256 + stable DESC top-8 sort + renorm (hybrid_forward.rs ~281-298).
3984    /// One CTA per token row, 256 threads (one per expert).
3985    pub fn moe_router_topk(
3986        &self,
3987        logits: &CudaSlice<f32>,
3988        t: usize,
3989        n_expert: usize,
3990        n_used: usize,
3991    ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3992        let f = self.func("moe_router_topk_f32");
3993        let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?; // kernel fully overwrites
3994        let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?; // kernel fully overwrites
3995        let cfg = LaunchConfig {
3996            grid_dim: (t as u32, 1, 1),
3997            block_dim: (n_expert as u32, 1, 1),
3998            shared_mem_bytes: 0,
3999        };
4000        let (ne, nu) = (n_expert as i32, n_used as i32);
4001        let __s_b = self.gpu.stream();
4002        let mut b = __s_b.launch_builder(&f);
4003        b.arg(logits)
4004            .arg(&mut sel_idx)
4005            .arg(&mut sel_w)
4006            .arg(&ne)
4007            .arg(&nu);
4008        unsafe {
4009            b.launch(cfg)?;
4010        }
4011        Ok((sel_idx, sel_w))
4012    }
4013
4014    /// gemma4 twin: per-expert output scale folded into the topk renorm write (replaces the
4015    /// separate moe_w_exscale launch; value chain identical: (w/ws) * s[sel]).
4016    pub fn moe_router_topk_scaled(
4017        &self,
4018        logits: &CudaSlice<f32>,
4019        t: usize,
4020        n_expert: usize,
4021        n_used: usize,
4022        ex_scale: &CudaSlice<f32>,
4023    ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4024        // barrier-lean v2 twin (per-warp top-k + one-warp merge) FALSIFIED 2026-07-14:
4025        // bit-identical streams but −1.4% (26B plain N=3 interleaved) — at t=1 the grid is
4026        // ONE block, so the 6.6us is launch/dependency overhead, not the barrier chain;
4027        // fewer barriers bought nothing and the merge structure cost. jsonl is the record.
4028        let f = self.func("moe_router_topk_scaled_f32");
4029        let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?;
4030        let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?;
4031        let cfg = LaunchConfig {
4032            grid_dim: (t as u32, 1, 1),
4033            block_dim: (n_expert as u32, 1, 1),
4034            shared_mem_bytes: 0,
4035        };
4036        let (ne, nu) = (n_expert as i32, n_used as i32);
4037        let __s_b = self.gpu.stream();
4038        let mut b = __s_b.launch_builder(&f);
4039        b.arg(logits)
4040            .arg(&mut sel_idx)
4041            .arg(&mut sel_w)
4042            .arg(&ne)
4043            .arg(&nu)
4044            .arg(ex_scale);
4045        unsafe {
4046            b.launch(cfg)?;
4047        }
4048        Ok((sel_idx, sel_w))
4049    }
4050
4051    /// LAUNCH-STRUCTURE STAGE 1 (2026-07-05): fused router + SINGLE-SYNC host readback. The old
4052    /// MEMRA_FUSED_ROUTER path lost 2% at t=1 because it paid TWO full stream syncs (dtoh_i32 then
4053    /// dtoh, each = clone_dtoh + synchronize) + two alloc_zeros memsets per MoE layer, where the
4054    /// host route pays ONE sync on the 1KB logits dtoh. This variant: uninit outputs (kernel fully
4055    /// overwrites), both DtoH copies issued ASYNC into a persistent PINNED host staging buffer
4056    /// (flags=0 — cacheable, NOT cudarc's WRITECOMBINED default, so the host-side reads of sel/w
4057    /// stay cached), then ONE synchronize. Numerics identical to `moe_router_topk` (same kernel).
4058    pub fn moe_router_topk_host(
4059        &self,
4060        logits: &CudaSlice<f32>,
4061        t: usize,
4062        n_expert: usize,
4063        n_used: usize,
4064    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
4065        let f = self.func("moe_router_topk_f32");
4066        let n = t * n_used;
4067        let mut sel_idx = self.alloc_uninit::<i32>(n)?;
4068        let mut sel_w = self.alloc_uninit::<f32>(n)?;
4069        let cfg = LaunchConfig {
4070            grid_dim: (t as u32, 1, 1),
4071            block_dim: (n_expert as u32, 1, 1),
4072            shared_mem_bytes: 0,
4073        };
4074        let (ne, nu) = (n_expert as i32, n_used as i32);
4075        let __s_b = self.gpu.stream();
4076        let mut b = __s_b.launch_builder(&f);
4077        b.arg(logits)
4078            .arg(&mut sel_idx)
4079            .arg(&mut sel_w)
4080            .arg(&ne)
4081            .arg(&nu);
4082        unsafe {
4083            b.launch(cfg)?;
4084        }
4085        // single-sync readback: sel (i32) at offset 0, w (f32) at offset n*4 of the pinned stage.
4086        let bytes = n * 8;
4087        let mut guard = self.router_stage.lock().unwrap();
4088        if guard.as_ref().map(|p| p.cap < bytes).unwrap_or(true) {
4089            *guard = Some(PinnedStage::new(bytes.max(4096))?);
4090        }
4091        let stage = guard.as_mut().unwrap();
4092        let (si, sw) = unsafe {
4093            (
4094                std::slice::from_raw_parts_mut(stage.ptr as *mut i32, n),
4095                std::slice::from_raw_parts_mut(stage.ptr.add(n * 4) as *mut f32, n),
4096            )
4097        };
4098        self.gpu.stream().memcpy_dtoh(&sel_idx, si)?; // async (pinned dst)
4099        self.gpu.stream().memcpy_dtoh(&sel_w, sw)?; // async (pinned dst)
4100        self.gpu.stream().synchronize()?; // ONE sync for both
4101        Ok((si.iter().map(|&i| i as u32).collect(), sw.to_vec()))
4102    }
4103
4104    /// Device sigmoid router for Step-3.7 / DeepSeek-V3-class MoEs. `correction_bias` is added
4105    /// only to the top-k key; returned weights use the un-biased sigmoid score. `active` masks
4106    /// original expert ids before top-k. Exact key ties choose the smaller original id.
4107    #[allow(clippy::too_many_arguments)]
4108    pub fn moe_router_sigmoid_topk(
4109        &self,
4110        logits: &CudaSlice<f32>,
4111        t: usize,
4112        n_expert: usize,
4113        n_used: usize,
4114        active_count: usize,
4115        correction_bias: &CudaSlice<f32>,
4116        active: &CudaSlice<u8>,
4117        scaling_factor: f32,
4118        route_norm: bool,
4119    ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4120        crate::sigrouter_contract::validate_active_count(n_used, active_count)?;
4121        if n_expert == 0 || n_expert > 1024 || n_used == 0 || n_used > n_expert {
4122            return Err(format!(
4123                "sigmoid router shape unsupported: n_expert={n_expert}, n_used={n_used}",
4124            )
4125            .into());
4126        }
4127        if logits.len() < t * n_expert
4128            || correction_bias.len() != n_expert
4129            || active.len() != n_expert
4130        {
4131            return Err(format!(
4132                "sigmoid router buffer mismatch: logits={} bias={} active={} expected logits>={} row={}",
4133                logits.len(), correction_bias.len(), active.len(), t * n_expert, n_expert,
4134            ).into());
4135        }
4136        let f = self.func("moe_router_sigmoid_topk_f32");
4137        let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?;
4138        let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?;
4139        let threads = n_expert.div_ceil(32) * 32;
4140        let cfg = LaunchConfig {
4141            grid_dim: (t as u32, 1, 1),
4142            block_dim: (threads as u32, 1, 1),
4143            shared_mem_bytes: 0,
4144        };
4145        let (ne, nu, rn) = (n_expert as i32, n_used as i32, i32::from(route_norm));
4146        let __s_b = self.gpu.stream();
4147        let mut b = __s_b.launch_builder(&f);
4148        b.arg(logits)
4149            .arg(correction_bias)
4150            .arg(active)
4151            .arg(&mut sel_idx)
4152            .arg(&mut sel_w)
4153            .arg(&ne)
4154            .arg(&nu)
4155            .arg(&scaling_factor)
4156            .arg(&rn);
4157        unsafe {
4158            b.launch(cfg)?;
4159        }
4160        Ok((sel_idx, sel_w))
4161    }
4162
4163    /// Single-sync pinned readback twin of `moe_router_sigmoid_topk`. This preserves the existing
4164    /// grouped/staged dispatch contract while replacing the full-logit DtoH plus host sigmoid/sort.
4165    #[allow(clippy::too_many_arguments)]
4166    pub fn moe_router_sigmoid_topk_host(
4167        &self,
4168        logits: &CudaSlice<f32>,
4169        t: usize,
4170        n_expert: usize,
4171        n_used: usize,
4172        active_count: usize,
4173        correction_bias: &CudaSlice<f32>,
4174        active: &CudaSlice<u8>,
4175        scaling_factor: f32,
4176        route_norm: bool,
4177    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
4178        let (sel_idx, sel_w) = self.moe_router_sigmoid_topk(
4179            logits,
4180            t,
4181            n_expert,
4182            n_used,
4183            active_count,
4184            correction_bias,
4185            active,
4186            scaling_factor,
4187            route_norm,
4188        )?;
4189        let n = t * n_used;
4190        let bytes = n * 8;
4191        let mut guard = self.router_stage.lock().unwrap();
4192        if guard.as_ref().map(|p| p.cap < bytes).unwrap_or(true) {
4193            *guard = Some(PinnedStage::new(bytes.max(4096))?);
4194        }
4195        let stage = guard.as_mut().unwrap();
4196        let (si, sw) = unsafe {
4197            (
4198                std::slice::from_raw_parts_mut(stage.ptr as *mut i32, n),
4199                std::slice::from_raw_parts_mut(stage.ptr.add(n * 4) as *mut f32, n),
4200            )
4201        };
4202        self.gpu.stream().memcpy_dtoh(&sel_idx, si)?;
4203        self.gpu.stream().memcpy_dtoh(&sel_w, sw)?;
4204        self.gpu.stream().synchronize()?;
4205        Ok((si.iter().map(|&i| i as u32).collect(), sw.to_vec()))
4206    }
4207
4208    /// EDGE-1 §C.2: async H2D of `host_bytes` into `scratch[off..]` on the COPY stream, returning a
4209    /// recorded event the compute stream can `wait` on before the dependent GEMM. Used for in-token
4210    /// expert prefetch (pipeline by one). `host_bytes` should be pinned for a true DMA (§C.1).
4211    pub fn stage_expert_async(
4212        &self,
4213        host_bytes: &[u8],
4214        scratch: &mut CudaSlice<u8>,
4215        off: usize,
4216    ) -> Result<cudarc::driver::CudaEvent, Box<dyn std::error::Error>> {
4217        let mut dst = scratch.slice_mut(off..off + host_bytes.len());
4218        self.copy_stream.memcpy_htod(host_bytes, &mut dst)?;
4219        Ok(self.copy_stream.record_event(None)?)
4220    }
4221
4222    /// Make the compute stream wait for an async copy event (the consumer side of `stage_expert_async`).
4223    pub fn compute_wait(
4224        &self,
4225        ev: &cudarc::driver::CudaEvent,
4226    ) -> Result<(), Box<dyn std::error::Error>> {
4227        self.gpu.stream().wait(ev)?;
4228        Ok(())
4229    }
4230
4231    /// qmatvec over a byte sub-range of a (resident/scratch) CudaSlice<u8> holding ONE expert
4232    /// matrix. x is a CudaView<f32> (a sliced row of z, or a sliced activation). Reuses the
4233    /// validated qmatvec_f32 dequant path (NOT a fast path — the correctness gate). The
4234    /// CudaView base+offset pointer is honored by the launch arg.
4235    pub fn qmatvec_view(
4236        &self,
4237        w: &CudaSlice<u8>,
4238        range: std::ops::Range<usize>,
4239        x: &cudarc::driver::CudaView<f32>,
4240        m: usize,
4241        in_f: usize,
4242        out_f: usize,
4243        qtype: i32,
4244        row_bytes: usize,
4245    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4246        let f = self.func("qmatvec_f32");
4247        let wv = w.slice(range); // CudaView<u8>, offset honored
4248        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
4249        let cfg = LaunchConfig {
4250            grid_dim: (out_f as u32, m as u32, 1),
4251            block_dim: (256, 1, 1),
4252            shared_mem_bytes: 0,
4253        };
4254        let (inf, outf, mi, qt, rb) =
4255            (in_f as i32, out_f as i32, m as i32, qtype, row_bytes as i64);
4256        let __s_b = self.gpu.stream();
4257        let mut b = __s_b.launch_builder(&f);
4258        b.arg(&wv)
4259            .arg(x)
4260            .arg(&mut y)
4261            .arg(&inf)
4262            .arg(&outf)
4263            .arg(&mi)
4264            .arg(&qt)
4265            .arg(&rb);
4266        unsafe {
4267            b.launch(cfg)?;
4268        }
4269        Ok(y)
4270    }
4271
4272    /// STAGE-2 GROUPED DECODE (2026-07-04): one MoE layer's gate+up+SiLU for all `n_used` routed
4273    /// experts of ONE token in ONE launch (replaces 8x qmatvec(gate) + 8x qmatvec(up) + 8x
4274    /// silu_mul = 24 launches). `gp`/`up` are the 8 expert weight-block device pointers (SLRU
4275    /// cache slots — fixed-address, stable for the launch). Returns act [n_used, n_ff].
4276    /// BIT-IDENTICAL to the sequential chain: each dot reproduces qmatvec_f32's exact 256-thread
4277    /// reduction; the SiLU epilogue is silu_mul_f32's exact expression (see kernel header).
4278    #[allow(clippy::too_many_arguments)]
4279    /// dp4a q8 twins (MoE expert dp4a arc, 2026-07-06): same contract as the _f32 versions but
4280    /// consume a PRE-QUANTIZED q8_1 activation. FP-order differs from _f32 (int dot + warp tree)
4281    /// — the argmax/stream-identity battery arbitrates; MEMRA_MOE_Q8=0 restores f32.
4282    pub fn moe_gate_up_silu8_q8(
4283        &self,
4284        gp: WPtr8,
4285        up: WPtr8,
4286        aq: &CudaSlice<i8>,
4287        ad: &CudaSlice<f32>,
4288        in_f: usize,
4289        n_ff: usize,
4290        n_used: usize,
4291        qt_g: i32,
4292        qt_u: i32,
4293        rb_g: usize,
4294        rb_u: usize,
4295    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4296        let f = self.func("moe_gate_up_silu8_q8");
4297        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
4298        let cfg = LaunchConfig {
4299            grid_dim: (n_ff as u32, n_used as u32, 1),
4300            block_dim: (32, 1, 1),
4301            shared_mem_bytes: 0,
4302        };
4303        let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
4304        let __s_b = self.gpu.stream();
4305        let mut b = __s_b.launch_builder(&f);
4306        b.arg(&gp)
4307            .arg(&up)
4308            .arg(aq)
4309            .arg(ad)
4310            .arg(&mut act)
4311            .arg(&inf)
4312            .arg(&nff)
4313            .arg(&qt_g)
4314            .arg(&qt_u)
4315            .arg(&rbg)
4316            .arg(&rbu);
4317        unsafe {
4318            b.launch(cfg)?;
4319        }
4320        Ok(act)
4321    }
4322
4323    #[allow(clippy::too_many_arguments)]
4324    pub fn moe_down8_fma_q8(
4325        &self,
4326        dp: WPtr8,
4327        w: F32x8,
4328        aq2: &CudaSlice<i8>,
4329        ad2: &CudaSlice<f32>,
4330        dst: &mut cudarc::driver::CudaViewMut<f32>,
4331        in_f: usize,
4332        out_f: usize,
4333        n_used: usize,
4334        qt: i32,
4335        rb: usize,
4336    ) -> Result<(), Box<dyn std::error::Error>> {
4337        let f = self.func("moe_down8_fma_q8");
4338        let cfg = LaunchConfig {
4339            grid_dim: (out_f as u32, 1, 1),
4340            block_dim: (32, 1, 1),
4341            shared_mem_bytes: 0,
4342        };
4343        let (inf, outf, nu, rbi) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
4344        let __s_b = self.gpu.stream();
4345        let mut b = __s_b.launch_builder(&f);
4346        b.arg(&dp)
4347            .arg(&w)
4348            .arg(aq2)
4349            .arg(ad2)
4350            .arg(dst)
4351            .arg(&inf)
4352            .arg(&outf)
4353            .arg(&nu)
4354            .arg(&qt)
4355            .arg(&rbi);
4356        unsafe {
4357            b.launch(cfg)?;
4358        }
4359        Ok(())
4360    }
4361
4362    /// q8 sequential expert matvec (staged path twin of qmatvec_view for IQ3_S/IQ4_XS).
4363    pub fn qmatvec_expert_q8(
4364        &self,
4365        w: &CudaSlice<u8>,
4366        range: std::ops::Range<usize>,
4367        aq: &CudaSlice<i8>,
4368        ad: &CudaSlice<f32>,
4369        m: usize,
4370        in_f: usize,
4371        out_f: usize,
4372        qtype: i32,
4373        row_bytes: usize,
4374    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4375        let f = self.func("qmatvec_expert_q8");
4376        let wv = w.slice(range);
4377        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
4378        const ROWS: u32 = 4; // MEMRA_MMVQ_ROWS
4379        let cfg = LaunchConfig {
4380            grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, m as u32, 1),
4381            block_dim: (32, ROWS, 1),
4382            shared_mem_bytes: 0,
4383        };
4384        let (inf, outf, mi, rbi) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
4385        let __s_b = self.gpu.stream();
4386        let mut b = __s_b.launch_builder(&f);
4387        b.arg(&wv)
4388            .arg(aq)
4389            .arg(ad)
4390            .arg(&mut y)
4391            .arg(&inf)
4392            .arg(&outf)
4393            .arg(&mi)
4394            .arg(&qtype)
4395            .arg(&rbi);
4396        unsafe {
4397            b.launch(cfg)?;
4398        }
4399        Ok(y)
4400    }
4401
4402    pub fn moe_gate_up_silu8(
4403        &self,
4404        gp: WPtr8,
4405        up: WPtr8,
4406        x: &cudarc::driver::CudaView<f32>,
4407        in_f: usize,
4408        n_ff: usize,
4409        n_used: usize,
4410        qt_g: i32,
4411        qt_u: i32,
4412        rb_g: usize,
4413        rb_u: usize,
4414    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4415        let f = self.func("moe_gate_up_silu8_f32");
4416        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?; // fully overwritten
4417        let cfg = LaunchConfig {
4418            grid_dim: (n_ff as u32, n_used as u32, 1),
4419            block_dim: (256, 1, 1),
4420            shared_mem_bytes: 0,
4421        };
4422        let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
4423        let __s_b = self.gpu.stream();
4424        let mut b = __s_b.launch_builder(&f);
4425        b.arg(&gp)
4426            .arg(&up)
4427            .arg(x)
4428            .arg(&mut act)
4429            .arg(&inf)
4430            .arg(&nff)
4431            .arg(&qt_g)
4432            .arg(&qt_u)
4433            .arg(&rbg)
4434            .arg(&rbu);
4435        unsafe {
4436            b.launch(cfg)?;
4437        }
4438        Ok(act)
4439    }
4440
4441    /// STAGE-2 GROUPED DECODE: one MoE layer's down-proj + weighted accumulation for all `n_used`
4442    /// routed experts in ONE launch (replaces 8x qmatvec(down) + 8x axpy = 16 launches), writing
4443    /// the token's moe_out row DIRECTLY (`dst` is the zeroed row; the in-kernel slot-ordered
4444    /// __fmaf_rn chain starting at 0.0f reproduces the sequential axpy_f32 accumulation into the
4445    /// zeroed row bit-for-bit — the A2 byte-identity scheme at m=1).
4446    #[allow(clippy::too_many_arguments)]
4447    pub fn moe_down8_fma_into(
4448        &self,
4449        dp: WPtr8,
4450        w: F32x8,
4451        act: &CudaSlice<f32>,
4452        dst: &mut cudarc::driver::CudaViewMut<f32>,
4453        in_f: usize,
4454        out_f: usize,
4455        n_used: usize,
4456        qt: i32,
4457        rb: usize,
4458    ) -> Result<(), Box<dyn std::error::Error>> {
4459        let f = self.func("moe_down8_fma_f32");
4460        let cfg = LaunchConfig {
4461            grid_dim: (out_f as u32, 1, 1),
4462            block_dim: (256, 1, 1),
4463            shared_mem_bytes: 0,
4464        };
4465        let (inf, outf, nu, rbv) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
4466        let __s_b = self.gpu.stream();
4467        let mut b = __s_b.launch_builder(&f);
4468        b.arg(&dp)
4469            .arg(&w)
4470            .arg(act)
4471            .arg(dst)
4472            .arg(&inf)
4473            .arg(&outf)
4474            .arg(&nu)
4475            .arg(&qt)
4476            .arg(&rbv);
4477        unsafe {
4478            b.launch(cfg)?;
4479        }
4480        Ok(())
4481    }
4482
4483    /// LAUNCH-STRUCTURE STAGE 3: device-dispatch twin of `moe_gate_up_silu8` for FULLY-RESIDENT
4484    /// layers. The expert ids come from the router kernel's DEVICE `sel` output (no DtoH) and the
4485    /// weight pointers from the per-layer device table `[3, n_expert]` of slot base addresses.
4486    /// BIT-IDENTICAL math (same grid/block/reduction; only the pointer/id source differs).
4487    #[allow(clippy::too_many_arguments)]
4488    /// dp4a q8 twin of the _dev pair (resident-experts arc).
4489    ///
4490    /// GEOMETRY VARIANTS (multirow/occupancy arc 2026-07-05): all outputs are BIT-IDENTICAL to
4491    /// the base one-warp-per-(row,slot) kernel (same expert_dot_g g-order + warp tree per row;
4492    /// down's FMA chain stays slot-ordered serial). Seams:
4493    ///   MEMRA_MOE_DEVQ8_GU   = 0(base) | 1 | 2 | 4 -> _r{1,2,4} multirow twin (RPW rows/warp)
4494    ///                       | s2 (gate/up warp split) | s2z (s2 + WPB rows packed per block)
4495    ///                       | gs4 (gate/up x low/high-group 4-warp split, nsb==64 only)
4496    ///                       | u64 (nsb==64 unrolled ILP twin, geometry unchanged)
4497    ///   MEMRA_MOE_DEVQ8_WPB  = warps per block for _r twins / z-rows for s2z (default 4)
4498    ///   MEMRA_MOE_DEVQ8_DOWN = auto(default: w8h2 when in_f==512 & n_used<=8 — measured +3.8%
4499    ///                       decode on 35B/G7e) | 0 (base one-warp serial-slot) | 1 | 2 | 4 ->
4500    ///                       _w8r{1,2,4} slot-parallel twin | h2 (half-warp dual-row, nsb==16
4501    ///                       only) | w8h2 (h2 x slot-parallel)
4502    #[allow(clippy::too_many_arguments)]
4503    /// MoE PREFILL pair-batch matvec: one launch covers all (token,expert) pairs for one proj.
4504    #[allow(clippy::too_many_arguments)]
4505    pub fn moe_pairs_matvec_q8(
4506        &self,
4507        table: &CudaSlice<u64>,
4508        proj: i32,
4509        pair_tok: &CudaSlice<i32>,
4510        pair_ex: &CudaSlice<i32>,
4511        aq: &CudaSlice<i8>,
4512        ad: &CudaSlice<f32>,
4513        in_f: usize,
4514        out_f: usize,
4515        n_expert: usize,
4516        n_pairs: usize,
4517        qtype: i32,
4518        row_bytes: usize,
4519    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4520        let f = self.func("moe_pairs_matvec_q8");
4521        let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
4522        const ROWS: u32 = 4;
4523        let cfg = LaunchConfig {
4524            grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_pairs as u32, 1),
4525            block_dim: (32, ROWS, 1),
4526            shared_mem_bytes: 0,
4527        };
4528        let (inf, outf, ne, np, rbi) = (
4529            in_f as i32,
4530            out_f as i32,
4531            n_expert as i32,
4532            n_pairs as i32,
4533            row_bytes as i64,
4534        );
4535        let __s_b = self.gpu.stream();
4536        let mut b = __s_b.launch_builder(&f);
4537        b.arg(table)
4538            .arg(&proj)
4539            .arg(pair_tok)
4540            .arg(pair_ex)
4541            .arg(aq)
4542            .arg(ad)
4543            .arg(&mut y)
4544            .arg(&inf)
4545            .arg(&outf)
4546            .arg(&ne)
4547            .arg(&np)
4548            .arg(&qtype)
4549            .arg(&rbi);
4550        unsafe {
4551            b.launch(cfg)?;
4552        }
4553        Ok(y)
4554    }
4555
4556    /// Expert-major pair matvec (weight-reuse across each expert's token group).
4557    #[allow(clippy::too_many_arguments)]
4558    pub fn moe_pairs_matvec_q8_em(
4559        &self,
4560        table: &CudaSlice<u64>,
4561        proj: i32,
4562        ex_ids: &CudaSlice<i32>,
4563        ex_off: &CudaSlice<i32>,
4564        ex_pairs: &CudaSlice<i32>,
4565        pair_tok: &CudaSlice<i32>,
4566        aq: &CudaSlice<i8>,
4567        ad: &CudaSlice<f32>,
4568        in_f: usize,
4569        out_f: usize,
4570        n_expert: usize,
4571        n_active: usize,
4572        n_pairs: usize,
4573        qtype: i32,
4574        row_bytes: usize,
4575    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4576        let f = self.func("moe_pairs_matvec_q8_em");
4577        let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
4578        const ROWS: u32 = 4;
4579        let cfg = LaunchConfig {
4580            grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_active as u32, 1),
4581            block_dim: (32, ROWS, 1),
4582            shared_mem_bytes: 0,
4583        };
4584        let (inf, outf, ne, na, rbi) = (
4585            in_f as i32,
4586            out_f as i32,
4587            n_expert as i32,
4588            n_active as i32,
4589            row_bytes as i64,
4590        );
4591        let __s_b = self.gpu.stream();
4592        let mut b = __s_b.launch_builder(&f);
4593        b.arg(table)
4594            .arg(&proj)
4595            .arg(ex_ids)
4596            .arg(ex_off)
4597            .arg(ex_pairs)
4598            .arg(pair_tok)
4599            .arg(aq)
4600            .arg(ad)
4601            .arg(&mut y)
4602            .arg(&inf)
4603            .arg(&outf)
4604            .arg(&ne)
4605            .arg(&na)
4606            .arg(&qtype)
4607            .arg(&rbi);
4608        unsafe {
4609            b.launch(cfg)?;
4610        }
4611        Ok(y)
4612    }
4613
4614    // Decode-once expert-major MMQ (rung 3). Same CSR inputs/geometry as _em; kernel dequants each
4615    // weight group once per (row,group) then dp4a's across the expert's token group.
4616    #[allow(clippy::too_many_arguments)]
4617    pub fn moe_pairs_matvec_q8_dec(
4618        &self,
4619        table: &CudaSlice<u64>,
4620        proj: i32,
4621        ex_ids: &CudaSlice<i32>,
4622        ex_off: &CudaSlice<i32>,
4623        ex_pairs: &CudaSlice<i32>,
4624        pair_tok: &CudaSlice<i32>,
4625        aq: &CudaSlice<i8>,
4626        ad: &CudaSlice<f32>,
4627        in_f: usize,
4628        out_f: usize,
4629        n_expert: usize,
4630        n_active: usize,
4631        n_pairs: usize,
4632        qtype: i32,
4633        row_bytes: usize,
4634    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4635        let f = self.func("moe_pairs_matvec_q8_dec");
4636        let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
4637        const ROWS: u32 = 4;
4638        let cfg = LaunchConfig {
4639            grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_active as u32, 1),
4640            block_dim: (32, ROWS, 1),
4641            shared_mem_bytes: 0,
4642        };
4643        let (inf, outf, ne, na, rbi) = (
4644            in_f as i32,
4645            out_f as i32,
4646            n_expert as i32,
4647            n_active as i32,
4648            row_bytes as i64,
4649        );
4650        let __s_b = self.gpu.stream();
4651        let mut b = __s_b.launch_builder(&f);
4652        b.arg(table)
4653            .arg(&proj)
4654            .arg(ex_ids)
4655            .arg(ex_off)
4656            .arg(ex_pairs)
4657            .arg(pair_tok)
4658            .arg(aq)
4659            .arg(ad)
4660            .arg(&mut y)
4661            .arg(&inf)
4662            .arg(&outf)
4663            .arg(&ne)
4664            .arg(&na)
4665            .arg(&qtype)
4666            .arg(&rbi);
4667        unsafe {
4668            b.launch(cfg)?;
4669        }
4670        Ok(y)
4671    }
4672
4673    pub fn moe_pairs_gelu_mul(
4674        &self,
4675        gate: &CudaSlice<f32>,
4676        up: &CudaSlice<f32>,
4677        n: usize,
4678    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4679        let f = self.func("moe_pairs_gelu_mul");
4680        let mut act = self.alloc_uninit::<f32>(n)?;
4681        let cfg = LaunchConfig::for_num_elems(n as u32);
4682        let nl = n as i64;
4683        let __s_b = self.gpu.stream();
4684        let mut b = __s_b.launch_builder(&f);
4685        b.arg(gate).arg(up).arg(&mut act).arg(&nl);
4686        unsafe {
4687            b.launch(cfg)?;
4688        }
4689        Ok(act)
4690    }
4691
4692    pub fn moe_pairs_silu_mul(
4693        &self,
4694        gate: &CudaSlice<f32>,
4695        up: &CudaSlice<f32>,
4696        n: usize,
4697    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4698        let f = self.func("moe_pairs_silu_mul");
4699        let mut act = self.alloc_uninit::<f32>(n)?;
4700        let cfg = LaunchConfig::for_num_elems(n as u32);
4701        let nl = n as i64;
4702        let __s_b = self.gpu.stream();
4703        let mut b = __s_b.launch_builder(&f);
4704        b.arg(gate).arg(up).arg(&mut act).arg(&nl);
4705        unsafe {
4706            b.launch(cfg)?;
4707        }
4708        Ok(act)
4709    }
4710
4711    #[allow(clippy::too_many_arguments)]
4712    pub fn moe_pairs_scatter(
4713        &self,
4714        y_down: &CudaSlice<f32>,
4715        pair_w: &CudaSlice<f32>,
4716        tok_pair_off: &CudaSlice<i32>,
4717        tok_pair_ids: &CudaSlice<i32>,
4718        moe_out: &mut CudaSlice<f32>,
4719        t: usize,
4720        n_embd: usize,
4721    ) -> Result<(), Box<dyn std::error::Error>> {
4722        let f = self.func("moe_pairs_scatter");
4723        let cfg = LaunchConfig {
4724            grid_dim: (((n_embd + 255) / 256) as u32, t as u32, 1),
4725            block_dim: (256, 1, 1),
4726            shared_mem_bytes: 0,
4727        };
4728        let ne = n_embd as i32;
4729        let __s_b = self.gpu.stream();
4730        let mut b = __s_b.launch_builder(&f);
4731        b.arg(y_down)
4732            .arg(pair_w)
4733            .arg(tok_pair_off)
4734            .arg(tok_pair_ids)
4735            .arg(moe_out)
4736            .arg(&ne);
4737        unsafe {
4738            b.launch(cfg)?;
4739        }
4740        Ok(())
4741    }
4742
4743    /// gemma4 GELU twin of moe_gate_up_silu8_dev_q8 (base geometry — slot-packed j8/j8r2
4744    /// twins probed 2026-08-01 g26 decode dig: bit-identical rows, -2.5%/-2.9% whole-model
4745    /// decode x3 interleaved -> refuted and killed; research/g26-decode-20260801/receipts.md).
4746    #[allow(clippy::too_many_arguments)]
4747    pub fn moe_gate_up_gelu8_dev_q8(
4748        &self,
4749        table: &CudaSlice<u64>,
4750        sel: &cudarc::driver::CudaView<i32>,
4751        aq: &CudaSlice<i8>,
4752        ad: &CudaSlice<f32>,
4753        in_f: usize,
4754        n_ff: usize,
4755        n_used: usize,
4756        n_expert: usize,
4757        qt_g: i32,
4758        qt_u: i32,
4759        rb_g: usize,
4760        rb_u: usize,
4761    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4762        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
4763        let (inf, nff, ne, rbg, rbu) = (
4764            in_f as i32,
4765            n_ff as i32,
4766            n_expert as i32,
4767            rb_g as i64,
4768            rb_u as i64,
4769        );
4770        let f = self.func("moe_gate_up_gelu8_dev_q8");
4771        let cfg = LaunchConfig {
4772            grid_dim: (n_ff as u32, n_used as u32, 1),
4773            block_dim: (32, 1, 1),
4774            shared_mem_bytes: 0,
4775        };
4776        let __s_b = self.gpu.stream();
4777        let mut b = __s_b.launch_builder(&f);
4778        b.arg(table)
4779            .arg(sel)
4780            .arg(aq)
4781            .arg(ad)
4782            .arg(&mut act)
4783            .arg(&inf)
4784            .arg(&nff)
4785            .arg(&ne)
4786            .arg(&qt_g)
4787            .arg(&qt_u)
4788            .arg(&rbg)
4789            .arg(&rbu);
4790        unsafe {
4791            b.launch(cfg)?;
4792        }
4793        Ok(act)
4794    }
4795
4796    /// gemma4 GELU rows twin (verify): one launch over (n_ff, n_used, t).
4797    #[allow(clippy::too_many_arguments)]
4798    pub fn moe_gate_up_gelu8_dev_q8_rows(
4799        &self,
4800        table: &CudaSlice<u64>,
4801        sel: &CudaSlice<i32>,
4802        aq: &CudaSlice<i8>,
4803        ad: &CudaSlice<f32>,
4804        t: usize,
4805        in_f: usize,
4806        n_ff: usize,
4807        n_used: usize,
4808        n_expert: usize,
4809        qt_g: i32,
4810        qt_u: i32,
4811        rb_g: usize,
4812        rb_u: usize,
4813    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4814        let mut act = self.alloc_uninit::<f32>(t * n_used * n_ff)?;
4815        let (inf, nff, ne, rbg, rbu, nu) = (
4816            in_f as i32,
4817            n_ff as i32,
4818            n_expert as i32,
4819            rb_g as i64,
4820            rb_u as i64,
4821            n_used as i32,
4822        );
4823        let f = self.func("moe_gate_up_gelu8_dev_q8_rows");
4824        let cfg = LaunchConfig {
4825            grid_dim: (n_ff as u32, n_used as u32, t as u32),
4826            block_dim: (32, 1, 1),
4827            shared_mem_bytes: 0,
4828        };
4829        let __s_b = self.gpu.stream();
4830        let mut b = __s_b.launch_builder(&f);
4831        b.arg(table)
4832            .arg(sel)
4833            .arg(aq)
4834            .arg(ad)
4835            .arg(&mut act)
4836            .arg(&inf)
4837            .arg(&nff)
4838            .arg(&ne)
4839            .arg(&qt_g)
4840            .arg(&qt_u)
4841            .arg(&rbg)
4842            .arg(&rbu)
4843            .arg(&nu);
4844        unsafe {
4845            b.launch(cfg)?;
4846        }
4847        Ok(act)
4848    }
4849
4850    /// gemma4 GELU CSR twin (verify dedup: owner block serves every pair of its expert).
4851    #[allow(clippy::too_many_arguments)]
4852    pub fn moe_gate_up_gelu8_dev_q8_csr(
4853        &self,
4854        table: &CudaSlice<u64>,
4855        sel: &CudaSlice<i32>,
4856        aq: &CudaSlice<i8>,
4857        ad: &CudaSlice<f32>,
4858        n_pairs: usize,
4859        in_f: usize,
4860        n_ff: usize,
4861        n_used: usize,
4862        n_expert: usize,
4863        qt_g: i32,
4864        qt_u: i32,
4865        rb_g: usize,
4866        rb_u: usize,
4867    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4868        let mut act = self.alloc_uninit::<f32>(n_pairs * n_ff)?;
4869        let (inf, nff, ne, rbg, rbu, nu, npi) = (
4870            in_f as i32,
4871            n_ff as i32,
4872            n_expert as i32,
4873            rb_g as i64,
4874            rb_u as i64,
4875            n_used as i32,
4876            n_pairs as i32,
4877        );
4878        let f = self.func("moe_gate_up_gelu8_dev_q8_csr");
4879        let cfg = LaunchConfig {
4880            grid_dim: (n_ff as u32, n_pairs as u32, 1),
4881            block_dim: (32, 1, 1),
4882            shared_mem_bytes: 0,
4883        };
4884        let __s_b = self.gpu.stream();
4885        let mut b = __s_b.launch_builder(&f);
4886        b.arg(table)
4887            .arg(sel)
4888            .arg(aq)
4889            .arg(ad)
4890            .arg(&mut act)
4891            .arg(&inf)
4892            .arg(&nff)
4893            .arg(&ne)
4894            .arg(&qt_g)
4895            .arg(&qt_u)
4896            .arg(&rbg)
4897            .arg(&rbu)
4898            .arg(&nu)
4899            .arg(&npi);
4900        unsafe {
4901            b.launch(cfg)?;
4902        }
4903        Ok(act)
4904    }
4905
4906    /// gemma4 generic down rows twin (verify): one launch over (out_f, 1, t).
4907    #[allow(clippy::too_many_arguments)]
4908    pub fn moe_down8_fma_dev_q8_rows_g(
4909        &self,
4910        table: &CudaSlice<u64>,
4911        sel: &CudaSlice<i32>,
4912        w: &CudaSlice<f32>,
4913        aq2: &CudaSlice<i8>,
4914        ad2: &CudaSlice<f32>,
4915        dst: &mut CudaSlice<f32>,
4916        t: usize,
4917        in_f: usize,
4918        out_f: usize,
4919        n_used: usize,
4920        n_expert: usize,
4921        qt: i32,
4922        rb: usize,
4923    ) -> Result<(), Box<dyn std::error::Error>> {
4924        let (inf, outf, nu, ne, rbi) = (
4925            in_f as i32,
4926            out_f as i32,
4927            n_used as i32,
4928            n_expert as i32,
4929            rb as i64,
4930        );
4931        // Exact Step-3.7 B=1 shape: expose the eight independent slot dots as
4932        // eight warps, then replay the original slot-ordered FMA chain. Every
4933        // other shape retains the generic one-warp rows kernel.
4934        let step_b1_w8 = t == 1 && in_f == 1280 && out_f == 4096 && n_used == 8 && qt == QT_IQ4_XS;
4935        let f = self.func(if step_b1_w8 {
4936            "moe_down8_fma_dev_q8_rows_w8"
4937        } else {
4938            "moe_down8_fma_dev_q8_rows_g"
4939        });
4940        let cfg = LaunchConfig {
4941            grid_dim: (out_f as u32, 1, t as u32),
4942            block_dim: (32, if step_b1_w8 { 8 } else { 1 }, 1),
4943            shared_mem_bytes: 0,
4944        };
4945        let __s_b = self.gpu.stream();
4946        let mut b = __s_b.launch_builder(&f);
4947        b.arg(table)
4948            .arg(sel)
4949            .arg(w)
4950            .arg(aq2)
4951            .arg(ad2)
4952            .arg(dst)
4953            .arg(&inf)
4954            .arg(&outf)
4955            .arg(&nu)
4956            .arg(&ne)
4957            .arg(&qt)
4958            .arg(&rbi);
4959        unsafe {
4960            b.launch(cfg)?;
4961        }
4962        Ok(())
4963    }
4964
4965    /// rp_q4 microprobe (2026-07-10 verify-trunk lever): b4 GGUF-block layout vs the Q4_0
4966    /// split-plane twin on the wq-class shape. Returns (blk_us, rp_us) after asserting bitwise
4967    /// identity. Bench-only surface (rp_q4_probe bin); no production dispatch reads this.
4968    pub fn rp_probe_q4(&self, m: usize) -> Result<(f64, f64), Box<dyn std::error::Error>> {
4969        let (out_f, in_f) = (2048usize, 2816usize);
4970        let nblk = in_f / 32;
4971        let mut seed = 0x9E3779B97F4A7C15u64;
4972        let mut rng = move || {
4973            seed = seed
4974                .wrapping_mul(6364136223846793005)
4975                .wrapping_add(1442695040888963407);
4976            (seed >> 33) as u8
4977        };
4978        let mut w = vec![0u8; out_f * nblk * 18];
4979        for b in w.iter_mut() {
4980            *b = rng();
4981        }
4982        for r in 0..out_f {
4983            for g in 0..nblk {
4984                let off = (r * nblk + g) * 18;
4985                w[off] = 0x00;
4986                w[off + 1] = 0x2C; // sane half d
4987            }
4988        }
4989        let qplane = out_f * nblk * 16;
4990        let mut wrp = vec![0u8; w.len()];
4991        for r in 0..out_f {
4992            for g in 0..nblk {
4993                let src = &w[(r * nblk + g) * 18..(r * nblk + g) * 18 + 18];
4994                wrp[qplane + (r * nblk + g) * 2..qplane + (r * nblk + g) * 2 + 2]
4995                    .copy_from_slice(&src[0..2]);
4996                wrp[(r * nblk + g) * 16..(r * nblk + g) * 16 + 16].copy_from_slice(&src[2..18]);
4997            }
4998        }
4999        let w_d = self.htod_bytes(&w)?;
5000        let wrp_d = self.htod_bytes(&wrp)?;
5001        let mut aq = vec![0i8; m * in_f];
5002        for v in aq.iter_mut() {
5003            *v = rng() as i8;
5004        }
5005        let aq_d = self.htod_i8(&aq)?;
5006        let ad_d = self.htod(&vec![0.03125f32; m * nblk])?;
5007        let mut y0 = self.alloc_uninit::<f32>(m * out_f)?;
5008        let mut y1 = self.alloc_uninit::<f32>(m * out_f)?;
5009        const RPB: u32 = 4;
5010        let cfg = LaunchConfig {
5011            grid_dim: ((out_f as u32).div_ceil(RPB), 1, 1),
5012            block_dim: (32, RPB, 1),
5013            shared_mem_bytes: 0,
5014        };
5015        let (inf, outf, mi) = (in_f as i32, out_f as i32, m as i32);
5016        let (rb, qp) = ((nblk * 18) as i64, qplane as i64);
5017        let fb = self.func("qmatvec_q4_0_mmvq_b4");
5018        let fr = self.func("qmatvec_q4_0_mmvq_b4_rp");
5019        {
5020            let __s_b = self.gpu.stream();
5021            let mut b = __s_b.launch_builder(&fb);
5022            b.arg(&w_d)
5023                .arg(&aq_d)
5024                .arg(&ad_d)
5025                .arg(&mut y0)
5026                .arg(&inf)
5027                .arg(&outf)
5028                .arg(&mi)
5029                .arg(&rb);
5030            unsafe {
5031                b.launch(cfg)?;
5032            }
5033            let __s_b = self.gpu.stream();
5034            let mut b = __s_b.launch_builder(&fr);
5035            b.arg(&wrp_d)
5036                .arg(&aq_d)
5037                .arg(&ad_d)
5038                .arg(&mut y1)
5039                .arg(&inf)
5040                .arg(&outf)
5041                .arg(&mi)
5042                .arg(&qp);
5043            unsafe {
5044                b.launch(cfg)?;
5045            }
5046        }
5047        self.gpu.stream().synchronize()?;
5048        let (h0, h1) = (self.dtoh(&y0)?, self.dtoh(&y1)?);
5049        let nd = h0
5050            .iter()
5051            .zip(&h1)
5052            .filter(|(a, b)| a.to_bits() != b.to_bits())
5053            .count();
5054        if nd != 0 {
5055            return Err(format!("rp twin not bitwise: {nd}/{} diffs", h0.len()).into());
5056        }
5057        let mut time = |rp: bool| -> Result<f64, Box<dyn std::error::Error>> {
5058            self.gpu.stream().synchronize()?;
5059            let t0 = std::time::Instant::now();
5060            for _ in 0..500 {
5061                if rp {
5062                    let __s_b = self.gpu.stream();
5063                    let mut b = __s_b.launch_builder(&fr);
5064                    b.arg(&wrp_d)
5065                        .arg(&aq_d)
5066                        .arg(&ad_d)
5067                        .arg(&mut y1)
5068                        .arg(&inf)
5069                        .arg(&outf)
5070                        .arg(&mi)
5071                        .arg(&qp);
5072                    unsafe {
5073                        b.launch(cfg)?;
5074                    }
5075                } else {
5076                    let __s_b = self.gpu.stream();
5077                    let mut b = __s_b.launch_builder(&fb);
5078                    b.arg(&w_d)
5079                        .arg(&aq_d)
5080                        .arg(&ad_d)
5081                        .arg(&mut y0)
5082                        .arg(&inf)
5083                        .arg(&outf)
5084                        .arg(&mi)
5085                        .arg(&rb);
5086                    unsafe {
5087                        b.launch(cfg)?;
5088                    }
5089                }
5090            }
5091            self.gpu.stream().synchronize()?;
5092            Ok(t0.elapsed().as_secs_f64() * 1e6 / 500.0)
5093        };
5094        let _ = time(false)?;
5095        let _ = time(true)?; // warm
5096        Ok((time(false)?, time(true)?))
5097    }
5098
5099    /// Build the Q4_0 split-plane decode mirror for a 2D Quant tensor (device-side permutation,
5100    /// q4_0_split_rp_build). Raw bytes stay resident (prefill/gemm/Stage-A); the m<=8 decode
5101    /// dispatch prefers the mirror (_rp twins). No-op unless (Q4_0, 2D, mirror absent).
5102    /// VRAM cost == the tensor's weight size. MEMRA_Q4RP=0 disables at the call sites.
5103    pub fn build_q4_rp4(
5104        &self,
5105        t: &mut crate::model::GpuTensor,
5106    ) -> Result<(), Box<dyn std::error::Error>> {
5107        use crate::model::GpuTensor;
5108        let GpuTensor::Quant {
5109            bytes,
5110            qtype,
5111            row_bytes,
5112            ne,
5113            rp4,
5114            ..
5115        } = t
5116        else {
5117            return Ok(());
5118        };
5119        if *qtype != QT_Q4_0 || rp4.is_some() || ne.len() != 2 {
5120            return Ok(());
5121        }
5122        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5123        if in_f % 32 != 0 || *row_bytes != (in_f / 32) * 18 {
5124            return Ok(());
5125        }
5126        let nblk = in_f / 32;
5127        let mut dst = self.alloc_uninit::<u8>(out_f * nblk * 18)?;
5128        let f = self.func("q4_0_split_rp_build");
5129        let n = (out_f * nblk) as i32;
5130        let cfg = LaunchConfig {
5131            grid_dim: (((out_f * nblk) as u32).div_ceil(256), 1, 1),
5132            block_dim: (256, 1, 1),
5133            shared_mem_bytes: 0,
5134        };
5135        let (of, nb) = (out_f as i32, nblk as i32);
5136        let _ = n;
5137        let __s_b = self.gpu.stream();
5138        let mut b = __s_b.launch_builder(&f);
5139        b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
5140        unsafe {
5141            b.launch(cfg)?;
5142        }
5143        *rp4 = Some(dst);
5144        Ok(())
5145    }
5146
5147    /// Q8_0 twin of `build_q4_rp4` (H100 coalescing fix, 2026-07-26 ncu: GGUF 34B-stride
5148    /// weight loads hold Max Bandwidth at 41-46%; the split mirror makes them aligned 16B
5149    /// ldcs). Raw bytes stay resident (prefill GEMM/MMQ/fused m=1 launches read GGUF layout);
5150    /// the mmvq/batched decode arms prefer the mirror via `rp4`. Bit-identical outputs.
5151    pub fn build_q8_rp4(
5152        &self,
5153        t: &mut crate::model::GpuTensor,
5154    ) -> Result<(), Box<dyn std::error::Error>> {
5155        use crate::model::GpuTensor;
5156        let GpuTensor::Quant {
5157            bytes,
5158            qtype,
5159            row_bytes,
5160            ne,
5161            rp4,
5162            ..
5163        } = t
5164        else {
5165            return Ok(());
5166        };
5167        if *qtype != QT_Q8_0 || rp4.is_some() || ne.len() != 2 {
5168            return Ok(());
5169        }
5170        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5171        if in_f % 32 != 0 || *row_bytes != (in_f / 32) * 34 {
5172            return Ok(());
5173        }
5174        *rp4 = Some(self.build_q8_rp4_raw(bytes, in_f, out_f)?);
5175        Ok(())
5176    }
5177
5178    /// Raw rp-mirror build for gates/benches: split GGUF Q8_0 bytes into the qplane+dplane
5179    /// mirror without a GpuTensor (same kernel the loader path above uses).
5180    pub fn build_q8_rp4_raw(
5181        &self,
5182        bytes: &CudaSlice<u8>,
5183        in_f: usize,
5184        out_f: usize,
5185    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
5186        assert!(in_f % 32 == 0);
5187        let nblk = in_f / 32;
5188        let mut dst = self.alloc_uninit::<u8>(out_f * nblk * 34)?;
5189        let f = self.func("q8_0_split_rp_build");
5190        let cfg = LaunchConfig {
5191            grid_dim: (((out_f * nblk) as u32).div_ceil(256), 1, 1),
5192            block_dim: (256, 1, 1),
5193            shared_mem_bytes: 0,
5194        };
5195        let (of, nb) = (out_f as i32, nblk as i32);
5196        let __s_b = self.gpu.stream();
5197        let mut b = __s_b.launch_builder(&f);
5198        b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
5199        unsafe {
5200            b.launch(cfg)?;
5201        }
5202        Ok(dst)
5203    }
5204
5205    /// K-quant twins of `build_q8_rp4` (H100 K-quant coalescing fix, 2026-08-01 ncu on the
5206    /// q27 Q4_K_M decode: q4_K mmvq DRAM 41-54% with 65% excessive sectors, q6_K 40% with
5207    /// 78% — the 144B/210B superblock strides land every 4B weight load off-sector). The
5208    /// mirror re-packs each tensor into planes (q4_K: qs ++ 16B meta; q6_K: ql ++ qh ++
5209    /// scales ++ d — same total bytes) so every quant fetch is an aligned 16B ldcs. Raw
5210    /// bytes stay resident (prefill GEMM/dequant/Stage-A read GGUF layout); the mmvq/batched
5211    /// decode arms prefer the mirror via `rp4`. Bit-identical outputs.
5212    pub fn build_q4k_rp4(
5213        &self,
5214        t: &mut crate::model::GpuTensor,
5215    ) -> Result<(), Box<dyn std::error::Error>> {
5216        use crate::model::GpuTensor;
5217        let GpuTensor::Quant {
5218            bytes,
5219            qtype,
5220            row_bytes,
5221            ne,
5222            rp4,
5223            ..
5224        } = t
5225        else {
5226            return Ok(());
5227        };
5228        if *qtype != QT_Q4_K || rp4.is_some() || ne.len() != 2 {
5229            return Ok(());
5230        }
5231        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5232        if in_f % 256 != 0 || *row_bytes != (in_f / 256) * 144 {
5233            return Ok(());
5234        }
5235        *rp4 = Some(self.build_kq_rp4_raw(bytes, in_f, out_f, QT_Q4_K)?);
5236        Ok(())
5237    }
5238
5239    pub fn build_q6k_rp4(
5240        &self,
5241        t: &mut crate::model::GpuTensor,
5242    ) -> Result<(), Box<dyn std::error::Error>> {
5243        use crate::model::GpuTensor;
5244        let GpuTensor::Quant {
5245            bytes,
5246            qtype,
5247            row_bytes,
5248            ne,
5249            rp4,
5250            ..
5251        } = t
5252        else {
5253            return Ok(());
5254        };
5255        if *qtype != QT_Q6_K || rp4.is_some() || ne.len() != 2 {
5256            return Ok(());
5257        }
5258        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5259        if in_f % 256 != 0 || *row_bytes != (in_f / 256) * 210 {
5260            return Ok(());
5261        }
5262        *rp4 = Some(self.build_kq_rp4_raw(bytes, in_f, out_f, QT_Q6_K)?);
5263        Ok(())
5264    }
5265
5266    /// Raw K-quant rp-mirror build for gates/benches (same kernels the loader path uses).
5267    pub fn build_kq_rp4_raw(
5268        &self,
5269        bytes: &CudaSlice<u8>,
5270        in_f: usize,
5271        out_f: usize,
5272        qtype: i32,
5273    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
5274        assert!(in_f % 256 == 0);
5275        let nsbk = in_f / 256;
5276        let (sb_bytes, kname) = match qtype {
5277            QT_Q4_K => (144usize, "q4_K_split_rp_build"),
5278            QT_Q6_K => (210usize, "q6_K_split_rp_build"),
5279            _ => return Err(format!("build_kq_rp4_raw: qtype {qtype} has no rp mirror").into()),
5280        };
5281        let mut dst = self.alloc_uninit::<u8>(out_f * nsbk * sb_bytes)?;
5282        let f = self.func(kname);
5283        let cfg = LaunchConfig {
5284            grid_dim: (((out_f * nsbk) as u32).div_ceil(256), 1, 1),
5285            block_dim: (256, 1, 1),
5286            shared_mem_bytes: 0,
5287        };
5288        let (of, nb) = (out_f as i32, nsbk as i32);
5289        let __s_b = self.gpu.stream();
5290        let mut b = __s_b.launch_builder(&f);
5291        b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
5292        unsafe {
5293            b.launch(cfg)?;
5294        }
5295        Ok(dst)
5296    }
5297
5298    /// MEMRA_KQRP seam: the K-quant (q4_K/q6_K) split-plane decode mirrors at model load.
5299    /// Default follows the Q8RP convention — ON on the Hopper lane (80GB pays the mirror
5300    /// VRAM), OFF elsewhere (a 24GB card cannot hold model + mirror + KV for the big trunks).
5301    pub fn kqrp_enabled() -> bool {
5302        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5303        *ON.get_or_init(|| match std::env::var("MEMRA_KQRP").as_deref() {
5304            Ok("0") => false,
5305            Ok(_) => true,
5306            Err(_) => cfg!(memra_hopper_mma),
5307        })
5308    }
5309
5310    /// IN-PLACE split-plane swap (the 31B dense arc): build the split layout and REPLACE the
5311    /// GGUF bytes (zero extra steady-state VRAM — the transient peak is one tensor's size).
5312    /// The tensor's `rp` flag then routes every consumer (mmvq/batched `_rp` twins, the
5313    /// `qmatvec_gemm_q4_0_rp` prefill kernel). Callers gate on the fast path being active —
5314    /// the Stage-A f32 oracle (`MEMRA_FAST=0`) reads GGUF layout and must never see a swap.
5315    pub fn build_q4_rp_swap(
5316        &self,
5317        t: &mut crate::model::GpuTensor,
5318    ) -> Result<bool, Box<dyn std::error::Error>> {
5319        use crate::model::GpuTensor;
5320        // QTYPE GUARD IN THE SWAP ITSELF (gemma4 NVFP4mix prefill-NaN, 2026-08-17):
5321        // `rp4` is a SHARED field — the Q8RP walk parks Q8_0 MIRRORS there, and this
5322        // fn used to `take()` whatever it found. On a Q8_0-carrying gemma4-dense trunk
5323        // the swap hijacked those mirrors: `bytes` became split-plane in place, the
5324        // m<=16 `_rp` dispatch read them correctly (masking the corruption from every
5325        // decode pin), and every GGUF-layout prefill consumer (MMQ + GEMM) read the
5326        // fp16 d-plane as weights -> layer-0 NaN, <pad>-spam serving. Only a tensor
5327        // this fn's OWN builder serves may ever be swapped; everything else refuses
5328        // here, regardless of walk ordering.
5329        if !matches!(t, GpuTensor::Quant { qtype, .. } if *qtype == QT_Q4_0) {
5330            return Ok(false);
5331        }
5332        self.build_q4_rp4(t)?;
5333        self.gpu.stream().synchronize()?; // build kernel reads the GGUF bytes — drain BEFORE dropping them
5334        let GpuTensor::Quant { bytes, rp4, rp, .. } = t else {
5335            return Ok(false);
5336        };
5337        match rp4.take() {
5338            Some(split) => {
5339                *bytes = split; // the GGUF-layout buffer drops here
5340                *rp = true;
5341                Ok(true)
5342            }
5343            None => Ok(false),
5344        }
5345    }
5346
5347    /// MEMRA_Q4RP seam (default ON): the Q4_0 split-plane decode mirror at model load.
5348    pub fn q4rp_enabled() -> bool {
5349        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5350        *ON.get_or_init(|| {
5351            std::env::var("MEMRA_Q4RP")
5352                .map(|v| v != "0")
5353                .unwrap_or(true)
5354        })
5355    }
5356
5357    /// gemma4-E4B: dense [t][row_elems] gather of layer il's rows from the strided prologue
5358    /// buffer ([t][n_layer][n_epl]; off = il*n_epl, stride = n_layer*n_epl).
5359    pub fn copy_rows_strided(
5360        &self,
5361        src: &CudaSlice<f32>,
5362        dst: &mut CudaSlice<f32>,
5363        row_elems: usize,
5364        n_rows: usize,
5365        src_stride: usize,
5366        src_off: usize,
5367    ) -> Result<(), Box<dyn std::error::Error>> {
5368        let f = self.func("copy_rows_strided_f32");
5369        let cfg = LaunchConfig {
5370            grid_dim: (((row_elems as u32 + 255) / 256).max(1), n_rows as u32, 1),
5371            block_dim: (256, 1, 1),
5372            shared_mem_bytes: 0,
5373        };
5374        let (re, nr) = (row_elems as i32, n_rows as i32);
5375        let (st, off) = (src_stride as i64, src_off as i64);
5376        let __s_b = self.gpu.stream();
5377        let mut b = __s_b.launch_builder(&f);
5378        b.arg(src)
5379            .arg(&mut *dst)
5380            .arg(&re)
5381            .arg(&nr)
5382            .arg(&st)
5383            .arg(&off);
5384        unsafe {
5385            b.launch(cfg)?;
5386        }
5387        Ok(())
5388    }
5389
5390    /// Async device u32 store (value rides the kernel ARG — no host-memory transfer/sync).
5391    pub fn u32_set_k(
5392        &self,
5393        dst: &mut CudaSlice<u32>,
5394        v: u32,
5395        idx: usize,
5396    ) -> Result<(), Box<dyn std::error::Error>> {
5397        let f = self.func("u32_set_k");
5398        let cfg = LaunchConfig {
5399            grid_dim: (1, 1, 1),
5400            block_dim: (1, 1, 1),
5401            shared_mem_bytes: 0,
5402        };
5403        let ii = idx as i32;
5404        let __s_b = self.gpu.stream();
5405        let mut b = __s_b.launch_builder(&f);
5406        b.arg(dst).arg(&v).arg(&ii);
5407        unsafe {
5408            b.launch(cfg)?;
5409        }
5410        Ok(())
5411    }
5412
5413    /// counter += v (device-slot append advance; the +1 twin is `inc_seqlen`).
5414    pub fn i32_add_k(
5415        &self,
5416        d: &mut CudaSlice<i32>,
5417        v: i32,
5418    ) -> Result<(), Box<dyn std::error::Error>> {
5419        let f = self.func("i32_add_k");
5420        let cfg = LaunchConfig {
5421            grid_dim: (1, 1, 1),
5422            block_dim: (32, 1, 1),
5423            shared_mem_bytes: 0,
5424        };
5425        let __s_b = self.gpu.stream();
5426        let mut b = __s_b.launch_builder(&f);
5427        b.arg(d).arg(&v);
5428        unsafe {
5429            b.launch(cfg)?;
5430        }
5431        Ok(())
5432    }
5433
5434    /// pos rows from a device counter: dst[i] = ctr[0] + i (verify-stream rope positions).
5435    pub fn i32_iota_from(
5436        &self,
5437        ctr: &CudaSlice<i32>,
5438        dst: &mut CudaSlice<i32>,
5439        n: usize,
5440    ) -> Result<(), Box<dyn std::error::Error>> {
5441        let f = self.func("i32_iota_from");
5442        let cfg = LaunchConfig::for_num_elems(n as u32);
5443        let ni = n as i32;
5444        let __s_b = self.gpu.stream();
5445        let mut b = __s_b.launch_builder(&f);
5446        b.arg(ctr).arg(dst).arg(&ni);
5447        unsafe {
5448            b.launch(cfg)?;
5449        }
5450        Ok(())
5451    }
5452
5453    /// In-place trim-id translate: buf[idx] = map[buf[idx]] (FR-Spec d2t, async single-slot).
5454    pub fn u32_map_k(
5455        &self,
5456        buf: &mut CudaSlice<u32>,
5457        map: &CudaSlice<u32>,
5458        idx: usize,
5459    ) -> Result<(), Box<dyn std::error::Error>> {
5460        let f = self.func("u32_map_k");
5461        let cfg = LaunchConfig {
5462            grid_dim: (1, 1, 1),
5463            block_dim: (1, 1, 1),
5464            shared_mem_bytes: 0,
5465        };
5466        let ii = idx as i32;
5467        let __s_b = self.gpu.stream();
5468        let mut b = __s_b.launch_builder(&f);
5469        b.arg(buf).arg(map).arg(&ii);
5470        unsafe {
5471            b.launch(cfg)?;
5472        }
5473        Ok(())
5474    }
5475
5476    /// Pack a[off..off+n1] ++ b[0..n2] into one buffer (single dtoh follows).
5477    #[allow(clippy::too_many_arguments)]
5478    pub fn u32_pack2(
5479        &self,
5480        a: &CudaSlice<u32>,
5481        off_a: usize,
5482        n1: usize,
5483        b_in: &CudaSlice<u32>,
5484        n2: usize,
5485        out: &mut CudaSlice<u32>,
5486    ) -> Result<(), Box<dyn std::error::Error>> {
5487        let f = self.func("u32_pack2");
5488        let cfg = LaunchConfig::for_num_elems((n1 + n2) as u32);
5489        let (oa, i1, i2) = (off_a as i32, n1 as i32, n2 as i32);
5490        let __s_b = self.gpu.stream();
5491        let mut b = __s_b.launch_builder(&f);
5492        b.arg(a).arg(&oa).arg(&i1).arg(b_in).arg(&i2).arg(out);
5493        unsafe {
5494            b.launch(cfg)?;
5495        }
5496        Ok(())
5497    }
5498
5499    /// gemma4 R3 device fold: w[i] *= s[sel[i]] over the router's [n] (sel, w) pair.
5500    pub fn moe_w_exscale(
5501        &self,
5502        w: &mut CudaSlice<f32>,
5503        sel: &CudaSlice<i32>,
5504        s: &CudaSlice<f32>,
5505        n: usize,
5506    ) -> Result<(), Box<dyn std::error::Error>> {
5507        let f = self.func("moe_w_exscale");
5508        let cfg = LaunchConfig::for_num_elems(n as u32);
5509        let ni = n as i32;
5510        let __s_b = self.gpu.stream();
5511        let mut b = __s_b.launch_builder(&f);
5512        b.arg(w).arg(sel).arg(s).arg(&ni);
5513        unsafe {
5514            b.launch(cfg)?;
5515        }
5516        Ok(())
5517    }
5518
5519    /// Down-projection macro fold: w[i] *= macros[2*n_expert + sel[i]] on the device router
5520    /// weights (one launch per MoE layer, only for macro-carrying artifacts — see MoeWeights).
5521    pub fn moe_w_scale_by_expert(
5522        &self,
5523        w: &mut CudaSlice<f32>,
5524        sel: &CudaSlice<i32>,
5525        macros: &CudaSlice<f32>,
5526        n_expert: usize,
5527        n: usize,
5528    ) -> Result<(), Box<dyn std::error::Error>> {
5529        let f = self.func("moe_w_scale_by_expert");
5530        let cfg = LaunchConfig {
5531            grid_dim: (n.div_ceil(64) as u32, 1, 1),
5532            block_dim: (64, 1, 1),
5533            shared_mem_bytes: 0,
5534        };
5535        let (ne, nn) = (n_expert as i32, n as i32);
5536        let __s_b = self.gpu.stream();
5537        let mut b = __s_b.launch_builder(&f);
5538        b.arg(w).arg(sel).arg(macros).arg(&ne).arg(&nn);
5539        unsafe {
5540            b.launch(cfg)?;
5541        }
5542        Ok(())
5543    }
5544
5545    pub fn moe_gate_up_silu8_dev_q8(
5546        &self,
5547        table: &CudaSlice<u64>,
5548        sel: &cudarc::driver::CudaView<i32>,
5549        aq: &CudaSlice<i8>,
5550        ad: &CudaSlice<f32>,
5551        in_f: usize,
5552        n_ff: usize,
5553        n_used: usize,
5554        n_expert: usize,
5555        qt_g: i32,
5556        qt_u: i32,
5557        rb_g: usize,
5558        rb_u: usize,
5559        macros: &CudaSlice<f32>,
5560    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5561        static GU: std::sync::OnceLock<(String, u32)> = std::sync::OnceLock::new();
5562        let (mode, wpb) = GU.get_or_init(|| {
5563            let mode = std::env::var("MEMRA_MOE_DEVQ8_GU").unwrap_or_default();
5564            let wpb = std::env::var("MEMRA_MOE_DEVQ8_WPB")
5565                .ok()
5566                .and_then(|v| v.parse().ok())
5567                .unwrap_or(4u32)
5568                .clamp(1, 16);
5569            (mode, wpb)
5570        });
5571        let (mode, wpb) = (mode.as_str(), *wpb);
5572        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
5573        let (inf, nff, ne, rbg, rbu) = (
5574            in_f as i32,
5575            n_ff as i32,
5576            n_expert as i32,
5577            rb_g as i64,
5578            rb_u as i64,
5579        );
5580        let (f, cfg) = match mode {
5581            "1" | "2" | "4" => {
5582                let rpw: u32 = mode.parse().unwrap();
5583                let f = self.func(match rpw {
5584                    1 => "moe_gate_up_silu8_dev_q8_r1",
5585                    2 => "moe_gate_up_silu8_dev_q8_r2",
5586                    _ => "moe_gate_up_silu8_dev_q8_r4",
5587                });
5588                let rows_per_block = (rpw * wpb) as usize;
5589                let gx = n_ff.div_ceil(rows_per_block) as u32;
5590                (
5591                    f,
5592                    LaunchConfig {
5593                        grid_dim: (gx, n_used as u32, 1),
5594                        block_dim: (32, wpb, 1),
5595                        shared_mem_bytes: 0,
5596                    },
5597                )
5598            }
5599            "j8" if n_used <= 32 => (
5600                self.func("moe_gate_up_silu8_dev_q8_j8"),
5601                LaunchConfig {
5602                    grid_dim: (n_ff as u32, 1, 1),
5603                    block_dim: (32, n_used as u32, 1),
5604                    shared_mem_bytes: 0,
5605                },
5606            ),
5607            // SMEM-GRID twins (IQ3_S 2KB grid copied to shared, static smem — bit-identical dots)
5608            "vsm2" => {
5609                let f = self.func("moe_gate_up_silu8_dev_q8_vsm2");
5610                let sh = (rb_g + rb_u) as u32;
5611                use cudarc::driver::sys::CUfunction_attribute_enum as A;
5612                f.set_attribute(
5613                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
5614                    sh as i32,
5615                )?;
5616                (
5617                    f,
5618                    LaunchConfig {
5619                        grid_dim: (n_ff as u32, n_used as u32, 1),
5620                        block_dim: (32, 1, 1),
5621                        shared_mem_bytes: sh,
5622                    },
5623                )
5624            }
5625            "vsm" => {
5626                let f = self.func("moe_gate_up_silu8_dev_q8_vsm");
5627                let sh = (rb_g + rb_u) as u32;
5628                use cudarc::driver::sys::CUfunction_attribute_enum as A;
5629                f.set_attribute(
5630                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
5631                    sh as i32,
5632                )?;
5633                (
5634                    f,
5635                    LaunchConfig {
5636                        grid_dim: (n_ff as u32, n_used as u32, 1),
5637                        block_dim: (32, 1, 1),
5638                        shared_mem_bytes: sh,
5639                    },
5640                )
5641            }
5642            "sg" => (
5643                self.func("moe_gate_up_silu8_dev_q8_sg"),
5644                LaunchConfig {
5645                    grid_dim: (n_ff as u32, n_used as u32, 1),
5646                    block_dim: (32, 1, 1),
5647                    shared_mem_bytes: 0,
5648                },
5649            ),
5650            "j8sg" if n_used <= 32 => (
5651                self.func("moe_gate_up_silu8_dev_q8_j8sg"),
5652                LaunchConfig {
5653                    grid_dim: (n_ff as u32, 1, 1),
5654                    block_dim: (32, n_used as u32, 1),
5655                    shared_mem_bytes: 0,
5656                },
5657            ),
5658            "u64" if in_f == 2048 => (
5659                self.func("moe_gate_up_silu8_dev_q8_u64"),
5660                LaunchConfig {
5661                    grid_dim: (n_ff as u32, n_used as u32, 1),
5662                    block_dim: (32, 1, 1),
5663                    shared_mem_bytes: 0,
5664                },
5665            ),
5666            "gs4" if in_f == 2048 => (
5667                self.func("moe_gate_up_silu8_dev_q8_gs4"),
5668                LaunchConfig {
5669                    grid_dim: (n_ff as u32, n_used as u32, 1),
5670                    block_dim: (32, 4, 1),
5671                    shared_mem_bytes: 0,
5672                },
5673            ),
5674            // _v twin (down8 lane 2026-07-08): wide-load IQ4_XS dot, base geometry, bit-identical.
5675            "v" | "" => (
5676                self.func("moe_gate_up_silu8_dev_q8_v"),
5677                LaunchConfig {
5678                    grid_dim: (n_ff as u32, n_used as u32, 1),
5679                    block_dim: (32, 1, 1),
5680                    shared_mem_bytes: 0,
5681                },
5682            ),
5683            "s2" => (
5684                self.func("moe_gate_up_silu8_dev_q8_s2"),
5685                LaunchConfig {
5686                    grid_dim: (n_ff as u32, n_used as u32, 1),
5687                    block_dim: (32, 2, 1),
5688                    shared_mem_bytes: 0,
5689                },
5690            ),
5691            "s2z" => {
5692                let rz = wpb.min(16); // s2z smem tile is [16][2]
5693                (
5694                    self.func("moe_gate_up_silu8_dev_q8_s2z"),
5695                    LaunchConfig {
5696                        grid_dim: (n_ff.div_ceil(rz as usize) as u32, n_used as u32, 1),
5697                        block_dim: (32, 2, rz),
5698                        shared_mem_bytes: 0,
5699                    },
5700                )
5701            }
5702            _ => (
5703                self.func("moe_gate_up_silu8_dev_q8"),
5704                LaunchConfig {
5705                    grid_dim: (n_ff as u32, n_used as u32, 1),
5706                    block_dim: (32, 1, 1),
5707                    shared_mem_bytes: 0,
5708                },
5709            ),
5710        };
5711        let __s_b = self.gpu.stream();
5712        let mut b = __s_b.launch_builder(&f);
5713        b.arg(table)
5714            .arg(sel)
5715            .arg(aq)
5716            .arg(ad)
5717            .arg(&mut act)
5718            .arg(&inf)
5719            .arg(&nff)
5720            .arg(&ne)
5721            .arg(&qt_g)
5722            .arg(&qt_u)
5723            .arg(&rbg)
5724            .arg(&rbu)
5725            .arg(macros);
5726        unsafe {
5727            b.launch(cfg)?;
5728        }
5729        Ok(act)
5730    }
5731
5732    #[allow(clippy::too_many_arguments)]
5733    pub fn moe_down8_fma_dev_q8(
5734        &self,
5735        table: &CudaSlice<u64>,
5736        sel: &cudarc::driver::CudaView<i32>,
5737        w: &cudarc::driver::CudaView<f32>,
5738        aq2: &CudaSlice<i8>,
5739        ad2: &CudaSlice<f32>,
5740        dst: &mut cudarc::driver::CudaViewMut<f32>,
5741        in_f: usize,
5742        out_f: usize,
5743        n_used: usize,
5744        n_expert: usize,
5745        qt: i32,
5746        rb: usize,
5747    ) -> Result<(), Box<dyn std::error::Error>> {
5748        static DOWN: std::sync::OnceLock<String> = std::sync::OnceLock::new();
5749        let mode = DOWN.get_or_init(|| std::env::var("MEMRA_MOE_DEVQ8_DOWN").unwrap_or_default());
5750        let (inf, outf, nu, ne, rbi) = (
5751            in_f as i32,
5752            out_f as i32,
5753            n_used as i32,
5754            n_expert as i32,
5755            rb as i64,
5756        );
5757        // the w8 twins' smem tile is [RPW][8] — n_used must fit the 8-slot tile;
5758        // the h2 twins are nsb==16 (in_f==512) shape-gated.
5759        let (f, cfg) = match mode.as_str() {
5760            m @ ("1" | "2" | "4") if n_used <= 8 => {
5761                let rpw: usize = m.parse().unwrap();
5762                let f = self.func(match rpw {
5763                    1 => "moe_down8_fma_dev_q8_w8r1",
5764                    2 => "moe_down8_fma_dev_q8_w8r2",
5765                    _ => "moe_down8_fma_dev_q8_w8r4",
5766                });
5767                (
5768                    f,
5769                    LaunchConfig {
5770                        grid_dim: (out_f.div_ceil(rpw) as u32, 1, 1),
5771                        block_dim: (32, n_used as u32, 1),
5772                        shared_mem_bytes: 0,
5773                    },
5774                )
5775            }
5776            "h2" if in_f == 512 => (
5777                self.func("moe_down8_fma_dev_q8_h2"),
5778                LaunchConfig {
5779                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
5780                    block_dim: (32, 1, 1),
5781                    shared_mem_bytes: 0,
5782                },
5783            ),
5784            // "" = AUTO gemma shape (in_f==704): w8r2 measured +1 tok/s vs base (sweep
5785            // 1/2/4 -> 133.6/134.2/133.6, 2026-07-10); slot-ordered chain preserved.
5786            "" if in_f == 704 && n_used <= 8 => (
5787                self.func("moe_down8_fma_dev_q8_w8r2"),
5788                LaunchConfig {
5789                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
5790                    block_dim: (32, n_used as u32, 1),
5791                    shared_mem_bytes: 0,
5792                },
5793            ),
5794            // "" = AUTO: the measured winner for the 35B expert shape (arc 2026-07-05, +3.8%);
5795            // any shape the h2 kernels can't take (nsb!=16 / n_used>8) falls to base via `_`.
5796            // _v twins (down8 lane 2026-07-08): wide-load IQ4_XS dot, bit-identical outputs.
5797            "w8h2v" | "" if in_f == 512 && n_used <= 8 => (
5798                self.func("moe_down8_fma_dev_q8_w8h2v"),
5799                LaunchConfig {
5800                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
5801                    block_dim: (32, n_used as u32, 1),
5802                    shared_mem_bytes: 0,
5803                },
5804            ),
5805            "w8h2r2v" if in_f == 512 && n_used <= 8 => (
5806                self.func("moe_down8_fma_dev_q8_w8h2r2v"),
5807                LaunchConfig {
5808                    grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
5809                    block_dim: (32, n_used as u32, 1),
5810                    shared_mem_bytes: 0,
5811                },
5812            ),
5813            "w8h2r2" if in_f == 512 && n_used <= 8 => (
5814                self.func("moe_down8_fma_dev_q8_w8h2r2"),
5815                LaunchConfig {
5816                    grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
5817                    block_dim: (32, n_used as u32, 1),
5818                    shared_mem_bytes: 0,
5819                },
5820            ),
5821            "w8h2" if in_f == 512 && n_used <= 8 => (
5822                self.func("moe_down8_fma_dev_q8_w8h2"),
5823                LaunchConfig {
5824                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
5825                    block_dim: (32, n_used as u32, 1),
5826                    shared_mem_bytes: 0,
5827                },
5828            ),
5829            _ => (
5830                self.func("moe_down8_fma_dev_q8"),
5831                LaunchConfig {
5832                    grid_dim: (out_f as u32, 1, 1),
5833                    block_dim: (32, 1, 1),
5834                    shared_mem_bytes: 0,
5835                },
5836            ),
5837        };
5838        let __s_b = self.gpu.stream();
5839        let mut b = __s_b.launch_builder(&f);
5840        b.arg(table)
5841            .arg(sel)
5842            .arg(w)
5843            .arg(aq2)
5844            .arg(ad2)
5845            .arg(dst)
5846            .arg(&inf)
5847            .arg(&outf)
5848            .arg(&nu)
5849            .arg(&ne)
5850            .arg(&qt)
5851            .arg(&rbi);
5852        unsafe {
5853            b.launch(cfg)?;
5854        }
5855        Ok(())
5856    }
5857
5858    /// SMALL-M VERIFY rows twin (MEMRA_SPEC_M2, lane/spec-m2): ONE launch covers all `t` tokens
5859    /// of the spec verify's MoE dev gate/up (grid.z = token) — the _v geometry per token, with
5860    /// tok-offset sel/aq/ad/act pointers matching the serial loop's slices. BIT-IDENTICAL per
5861    /// token (see the kernel header). aq/ad are the BATCHED z-quantize ([t, in_f] rows —
5862    /// quantize_q8_1's per-32-block program is row-independent, so batched rows == the serial
5863    /// loop's per-token quantize_q8_1_view bytes). Returns act [t, n_used, n_ff].
5864    #[allow(clippy::too_many_arguments)]
5865    pub fn moe_gate_up_silu8_dev_q8_rows(
5866        &self,
5867        table: &CudaSlice<u64>,
5868        sel: &CudaSlice<i32>,
5869        aq: &CudaSlice<i8>,
5870        ad: &CudaSlice<f32>,
5871        t: usize,
5872        in_f: usize,
5873        n_ff: usize,
5874        n_used: usize,
5875        n_expert: usize,
5876        qt_g: i32,
5877        qt_u: i32,
5878        rb_g: usize,
5879        rb_u: usize,
5880        macros: &CudaSlice<f32>,
5881    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5882        let f = self.func("moe_gate_up_silu8_dev_q8_v_rows");
5883        let mut act = self.alloc_uninit::<f32>(t * n_used * n_ff)?;
5884        let cfg = LaunchConfig {
5885            grid_dim: (n_ff as u32, n_used as u32, t as u32),
5886            block_dim: (32, 1, 1),
5887            shared_mem_bytes: 0,
5888        };
5889        let (inf, nff, ne, nu, rbg, rbu) = (
5890            in_f as i32,
5891            n_ff as i32,
5892            n_expert as i32,
5893            n_used as i32,
5894            rb_g as i64,
5895            rb_u as i64,
5896        );
5897        let __s_b = self.gpu.stream();
5898        let mut b = __s_b.launch_builder(&f);
5899        b.arg(table)
5900            .arg(sel)
5901            .arg(aq)
5902            .arg(ad)
5903            .arg(&mut act)
5904            .arg(&inf)
5905            .arg(&nff)
5906            .arg(&ne)
5907            .arg(&qt_g)
5908            .arg(&qt_u)
5909            .arg(&rbg)
5910            .arg(&rbu)
5911            .arg(&nu)
5912            .arg(macros);
5913        unsafe {
5914            b.launch(cfg)?;
5915        }
5916        Ok(act)
5917    }
5918
5919    /// SMALL-M VERIFY rows twin of the down proj: w8h2v geometry per token on a grid.z token
5920    /// axis. Caller gates the w8h2v shape contract (in_f == 512, n_used <= 8) — same gate as
5921    /// the AUTO dispatch in `moe_down8_fma_dev_q8`. aq2/ad2 = batched act quantize
5922    /// ([t*n_used, in_f] rows). dst rows are FULLY overwritten per token.
5923    #[allow(clippy::too_many_arguments)]
5924    pub fn moe_down8_fma_dev_q8_rows(
5925        &self,
5926        table: &CudaSlice<u64>,
5927        sel: &CudaSlice<i32>,
5928        w: &CudaSlice<f32>,
5929        aq2: &CudaSlice<i8>,
5930        ad2: &CudaSlice<f32>,
5931        dst: &mut CudaSlice<f32>,
5932        t: usize,
5933        in_f: usize,
5934        out_f: usize,
5935        n_used: usize,
5936        n_expert: usize,
5937        qt: i32,
5938        rb: usize,
5939    ) -> Result<(), Box<dyn std::error::Error>> {
5940        assert!(
5941            in_f == 512 && n_used <= 8,
5942            "down rows twin is w8h2v shape-gated"
5943        );
5944        let f = self.func("moe_down8_fma_dev_q8_w8h2v_rows");
5945        let cfg = LaunchConfig {
5946            grid_dim: (out_f.div_ceil(2) as u32, 1, t as u32),
5947            block_dim: (32, n_used as u32, 1),
5948            shared_mem_bytes: 0,
5949        };
5950        let (inf, outf, nu, ne, rbi) = (
5951            in_f as i32,
5952            out_f as i32,
5953            n_used as i32,
5954            n_expert as i32,
5955            rb as i64,
5956        );
5957        let __s_b = self.gpu.stream();
5958        let mut b = __s_b.launch_builder(&f);
5959        b.arg(table)
5960            .arg(sel)
5961            .arg(w)
5962            .arg(aq2)
5963            .arg(ad2)
5964            .arg(dst)
5965            .arg(&inf)
5966            .arg(&outf)
5967            .arg(&nu)
5968            .arg(&ne)
5969            .arg(&qt)
5970            .arg(&rbi);
5971        unsafe {
5972            b.launch(cfg)?;
5973        }
5974        Ok(())
5975    }
5976
5977    /// CSR gate/up v3 (owner-scan dedup, no build kernel): qtypes {IQ4_XS, IQ3_S} (caller
5978    /// gates), grid.y = pair index; the first pair of each expert serves all its pairs.
5979    /// Bit-identical to moe_gate_up_silu8_dev_q8_v_rows (explicit-intrinsic accumulate).
5980    #[allow(clippy::too_many_arguments)]
5981    pub fn moe_gate_up_silu8_dev_q8_csr(
5982        &self,
5983        table: &CudaSlice<u64>,
5984        sel: &CudaSlice<i32>,
5985        aq: &CudaSlice<i8>,
5986        ad: &CudaSlice<f32>,
5987        n_pairs: usize,
5988        in_f: usize,
5989        n_ff: usize,
5990        n_used: usize,
5991        n_expert: usize,
5992        qt_g: i32,
5993        qt_u: i32,
5994        rb_g: usize,
5995        rb_u: usize,
5996    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5997        // NVFP4 experts take the NVFP4-specialized owner-scan twin (lane/moebatch-q35moe);
5998        // host gate guarantees qt_g == qt_u within a supported class.
5999        let f = if qt_g == crate::QT_NVFP4 {
6000            self.func("moe_gate_up_silu8_dev_q8_csr_nvfp4")
6001        } else {
6002            self.func("moe_gate_up_silu8_dev_q8_csr_iq4")
6003        };
6004        let mut act = self.alloc_uninit::<f32>(n_pairs * n_ff)?;
6005        let cfg = LaunchConfig {
6006            grid_dim: (n_ff as u32, n_pairs as u32, 1),
6007            block_dim: (32, 1, 1),
6008            shared_mem_bytes: 0,
6009        };
6010        let (inf, nff, ne, nu, npi, rbg, rbu) = (
6011            in_f as i32,
6012            n_ff as i32,
6013            n_expert as i32,
6014            n_used as i32,
6015            n_pairs as i32,
6016            rb_g as i64,
6017            rb_u as i64,
6018        );
6019        let __s_b = self.gpu.stream();
6020        let mut b = __s_b.launch_builder(&f);
6021        b.arg(table)
6022            .arg(sel)
6023            .arg(aq)
6024            .arg(ad)
6025            .arg(&mut act)
6026            .arg(&inf)
6027            .arg(&nff)
6028            .arg(&ne)
6029            .arg(&qt_g)
6030            .arg(&qt_u)
6031            .arg(&rbg)
6032            .arg(&rbu)
6033            .arg(&nu)
6034            .arg(&npi);
6035        unsafe {
6036            b.launch(cfg)?;
6037        }
6038        Ok(act)
6039    }
6040
6041    /// TEST SEAM (down8 lane 2026-07-08): launch a down dev_q8 variant BY NAME with its
6042    /// canonical geometry, bypassing the env-cached dispatch so moe-devq8-check can byte-
6043    /// compare variants in one process. Variants: "base", "w8h2", "w8h2r2", "w8h2v", "w8h2r2v".
6044    #[allow(clippy::too_many_arguments)]
6045    pub fn moe_down8_fma_dev_q8_variant(
6046        &self,
6047        variant: &str,
6048        table: &CudaSlice<u64>,
6049        sel: &cudarc::driver::CudaView<i32>,
6050        w: &cudarc::driver::CudaView<f32>,
6051        aq2: &CudaSlice<i8>,
6052        ad2: &CudaSlice<f32>,
6053        dst: &mut cudarc::driver::CudaViewMut<f32>,
6054        in_f: usize,
6055        out_f: usize,
6056        n_used: usize,
6057        n_expert: usize,
6058        qt: i32,
6059        rb: usize,
6060    ) -> Result<(), Box<dyn std::error::Error>> {
6061        let (inf, outf, nu, ne, rbi) = (
6062            in_f as i32,
6063            out_f as i32,
6064            n_used as i32,
6065            n_expert as i32,
6066            rb as i64,
6067        );
6068        let (f, cfg) = match variant {
6069            "w8h2" | "w8h2v" => (
6070                self.func(if variant == "w8h2" {
6071                    "moe_down8_fma_dev_q8_w8h2"
6072                } else {
6073                    "moe_down8_fma_dev_q8_w8h2v"
6074                }),
6075                LaunchConfig {
6076                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6077                    block_dim: (32, n_used as u32, 1),
6078                    shared_mem_bytes: 0,
6079                },
6080            ),
6081            "w8h2r2" | "w8h2r2v" => (
6082                self.func(if variant == "w8h2r2" {
6083                    "moe_down8_fma_dev_q8_w8h2r2"
6084                } else {
6085                    "moe_down8_fma_dev_q8_w8h2r2v"
6086                }),
6087                LaunchConfig {
6088                    grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
6089                    block_dim: (32, n_used as u32, 1),
6090                    shared_mem_bytes: 0,
6091                },
6092            ),
6093            _ => (
6094                self.func("moe_down8_fma_dev_q8"),
6095                LaunchConfig {
6096                    grid_dim: (out_f as u32, 1, 1),
6097                    block_dim: (32, 1, 1),
6098                    shared_mem_bytes: 0,
6099                },
6100            ),
6101        };
6102        let __s_b = self.gpu.stream();
6103        let mut b = __s_b.launch_builder(&f);
6104        b.arg(table)
6105            .arg(sel)
6106            .arg(w)
6107            .arg(aq2)
6108            .arg(ad2)
6109            .arg(dst)
6110            .arg(&inf)
6111            .arg(&outf)
6112            .arg(&nu)
6113            .arg(&ne)
6114            .arg(&qt)
6115            .arg(&rbi);
6116        unsafe {
6117            b.launch(cfg)?;
6118        }
6119        Ok(())
6120    }
6121
6122    /// TEST SEAM (down8 lane): gate_up twin of the above. Variants: "base", "v".
6123    #[allow(clippy::too_many_arguments)]
6124    pub fn moe_gate_up_silu8_dev_q8_variant(
6125        &self,
6126        variant: &str,
6127        table: &CudaSlice<u64>,
6128        sel: &cudarc::driver::CudaView<i32>,
6129        aq: &CudaSlice<i8>,
6130        ad: &CudaSlice<f32>,
6131        in_f: usize,
6132        n_ff: usize,
6133        n_used: usize,
6134        n_expert: usize,
6135        qt_g: i32,
6136        qt_u: i32,
6137        rb_g: usize,
6138        rb_u: usize,
6139    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6140        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
6141        let (inf, nff, ne, rbg, rbu) = (
6142            in_f as i32,
6143            n_ff as i32,
6144            n_expert as i32,
6145            rb_g as i64,
6146            rb_u as i64,
6147        );
6148        let f = self.func(if variant == "v" {
6149            "moe_gate_up_silu8_dev_q8_v"
6150        } else {
6151            "moe_gate_up_silu8_dev_q8"
6152        });
6153        let cfg = LaunchConfig {
6154            grid_dim: (n_ff as u32, n_used as u32, 1),
6155            block_dim: (32, 1, 1),
6156            shared_mem_bytes: 0,
6157        };
6158        let __s_b = self.gpu.stream();
6159        let mut b = __s_b.launch_builder(&f);
6160        b.arg(table)
6161            .arg(sel)
6162            .arg(aq)
6163            .arg(ad)
6164            .arg(&mut act)
6165            .arg(&inf)
6166            .arg(&nff)
6167            .arg(&ne)
6168            .arg(&qt_g)
6169            .arg(&qt_u)
6170            .arg(&rbg)
6171            .arg(&rbu);
6172        unsafe {
6173            b.launch(cfg)?;
6174        }
6175        Ok(act)
6176    }
6177
6178    pub fn moe_gate_up_silu8_dev(
6179        &self,
6180        table: &CudaSlice<u64>,
6181        sel: &cudarc::driver::CudaView<i32>,
6182        x: &cudarc::driver::CudaView<f32>,
6183        in_f: usize,
6184        n_ff: usize,
6185        n_used: usize,
6186        n_expert: usize,
6187        qt_g: i32,
6188        qt_u: i32,
6189        rb_g: usize,
6190        rb_u: usize,
6191        macros: &CudaSlice<f32>,
6192    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6193        let f = self.func("moe_gate_up_silu8_dev");
6194        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?; // fully overwritten
6195        let cfg = LaunchConfig {
6196            grid_dim: (n_ff as u32, n_used as u32, 1),
6197            block_dim: (256, 1, 1),
6198            shared_mem_bytes: 0,
6199        };
6200        let (inf, nff, ne, rbg, rbu) = (
6201            in_f as i32,
6202            n_ff as i32,
6203            n_expert as i32,
6204            rb_g as i64,
6205            rb_u as i64,
6206        );
6207        let __s_b = self.gpu.stream();
6208        let mut b = __s_b.launch_builder(&f);
6209        b.arg(table)
6210            .arg(sel)
6211            .arg(x)
6212            .arg(&mut act)
6213            .arg(&inf)
6214            .arg(&nff)
6215            .arg(&ne)
6216            .arg(&qt_g)
6217            .arg(&qt_u)
6218            .arg(&rbg)
6219            .arg(&rbu)
6220            .arg(macros);
6221        unsafe {
6222            b.launch(cfg)?;
6223        }
6224        Ok(act)
6225    }
6226
6227    /// LAUNCH-STRUCTURE STAGE 3: device-dispatch twin of `moe_down8_fma_into` — expert ids AND
6228    /// renormalized weights read from the router kernel's device output. BIT-IDENTICAL chain.
6229    #[allow(clippy::too_many_arguments)]
6230    pub fn moe_down8_fma_dev(
6231        &self,
6232        table: &CudaSlice<u64>,
6233        sel: &cudarc::driver::CudaView<i32>,
6234        w: &cudarc::driver::CudaView<f32>,
6235        act: &CudaSlice<f32>,
6236        dst: &mut cudarc::driver::CudaViewMut<f32>,
6237        in_f: usize,
6238        out_f: usize,
6239        n_used: usize,
6240        n_expert: usize,
6241        qt: i32,
6242        rb: usize,
6243    ) -> Result<(), Box<dyn std::error::Error>> {
6244        let f = self.func("moe_down8_fma_dev");
6245        let cfg = LaunchConfig {
6246            grid_dim: (out_f as u32, 1, 1),
6247            block_dim: (256, 1, 1),
6248            shared_mem_bytes: 0,
6249        };
6250        let (inf, outf, nu, ne, rbv) = (
6251            in_f as i32,
6252            out_f as i32,
6253            n_used as i32,
6254            n_expert as i32,
6255            rb as i64,
6256        );
6257        let __s_b = self.gpu.stream();
6258        let mut b = __s_b.launch_builder(&f);
6259        b.arg(table)
6260            .arg(sel)
6261            .arg(w)
6262            .arg(act)
6263            .arg(dst)
6264            .arg(&inf)
6265            .arg(&outf)
6266            .arg(&nu)
6267            .arg(&ne)
6268            .arg(&qt)
6269            .arg(&rbv);
6270        unsafe {
6271            b.launch(cfg)?;
6272        }
6273        Ok(())
6274    }
6275
6276    /// dst[i] += alpha * src[i], i in 0..n. dst is a CudaViewMut (a row of moe_out).
6277    pub fn axpy_into(
6278        &self,
6279        src: &CudaSlice<f32>,
6280        alpha: f32,
6281        dst: &mut cudarc::driver::CudaViewMut<f32>,
6282        n: usize,
6283    ) -> Result<(), Box<dyn std::error::Error>> {
6284        let f = self.func("axpy_f32");
6285        let cfg = LaunchConfig::for_num_elems(n as u32);
6286        let (a, ni) = (alpha, n as i32);
6287        let __s_b = self.gpu.stream();
6288        let mut b = __s_b.launch_builder(&f);
6289        b.arg(src).arg(dst).arg(&a).arg(&ni);
6290        unsafe {
6291            b.launch(cfg)?;
6292        }
6293        Ok(())
6294    }
6295
6296    /// dst[r*ncols + c] += src[r*ncols + c] * scale[r]. Per-row scalar accumulate (shared expert).
6297    pub fn add_scaled_rows(
6298        &self,
6299        src: &CudaSlice<f32>,
6300        scale: &CudaSlice<f32>,
6301        dst: &mut CudaSlice<f32>,
6302        ncols: usize,
6303        nrows: usize,
6304    ) -> Result<(), Box<dyn std::error::Error>> {
6305        let f = self.func("add_scaled_rows_f32");
6306        let cfg = LaunchConfig::for_num_elems((ncols * nrows) as u32);
6307        let (nc, nr) = (ncols as i32, nrows as i32);
6308        let __s_b = self.gpu.stream();
6309        let mut b = __s_b.launch_builder(&f);
6310        b.arg(src).arg(scale).arg(dst).arg(&nc).arg(&nr);
6311        unsafe {
6312            b.launch(cfg)?;
6313        }
6314        Ok(())
6315    }
6316
6317    // ======== A2 GROUPED MoE PREFILL KERNELS ========
6318
6319    /// Gather m_e rows from src[T, ncols] into dst[m_e, ncols] using index array idx[m_e].
6320    pub fn gather_rows(
6321        &self,
6322        src: &CudaSlice<f32>,
6323        idx: &CudaSlice<i32>,
6324        dst: &mut CudaSlice<f32>,
6325        ncols: usize,
6326        m_e: usize,
6327    ) -> Result<(), Box<dyn std::error::Error>> {
6328        let f = self.func("gather_rows_f32");
6329        let cfg = LaunchConfig::for_num_elems((m_e * ncols) as u32);
6330        let (nc, me) = (ncols as i32, m_e as i32);
6331        let __s_b = self.gpu.stream();
6332        let mut b = __s_b.launch_builder(&f);
6333        b.arg(src).arg(idx).arg(dst).arg(&nc).arg(&me);
6334        unsafe {
6335            b.launch(cfg)?;
6336        }
6337        Ok(())
6338    }
6339
6340    /// Scatter expert outputs into per-token slots: dst[tok_idx[r], slot_idx[r], :] = src[r, :] * weight[r].
6341    /// dst is [T, n_used, ncols], zero-initialized. Each (expert, token) pair maps to a unique slot.
6342    /// Scatter expert outputs into per-token slots (raw copy, no weight multiply).
6343    /// Weight stored into wbuf[tok*n_used + slot] for FMA in reduce step.
6344    pub fn scatter_slot(
6345        &self,
6346        src: &CudaSlice<f32>,
6347        tok_idx: &CudaSlice<i32>,
6348        slot_idx: &CudaSlice<i32>,
6349        weight: &CudaSlice<f32>,
6350        dst: &mut CudaSlice<f32>,
6351        wbuf: &mut CudaSlice<f32>,
6352        ncols: usize,
6353        n_used: usize,
6354        m_e: usize,
6355    ) -> Result<(), Box<dyn std::error::Error>> {
6356        let f = self.func("scatter_add_slot_f32");
6357        let cfg = LaunchConfig::for_num_elems((m_e * ncols) as u32);
6358        let (nc, nu, me) = (ncols as i32, n_used as i32, m_e as i32);
6359        let __s_b = self.gpu.stream();
6360        let mut b = __s_b.launch_builder(&f);
6361        b.arg(src)
6362            .arg(tok_idx)
6363            .arg(slot_idx)
6364            .arg(weight)
6365            .arg(dst)
6366            .arg(wbuf)
6367            .arg(&nc)
6368            .arg(&nu)
6369            .arg(&me);
6370        unsafe {
6371            b.launch(cfg)?;
6372        }
6373        Ok(())
6374    }
6375
6376    /// Reduce n_used slots per token: dst[t, col] = sum_s slots[t, s, col].
6377    /// Reduce n_used slots per token: dst[t, col] = sum_s FMA(wbuf[t,s], slots[t,s,col], acc).
6378    /// Uses FMA for bit-identity with the sequential axpy path.
6379    pub fn reduce_slots(
6380        &self,
6381        slots: &CudaSlice<f32>,
6382        wbuf: &CudaSlice<f32>,
6383        dst: &mut CudaSlice<f32>,
6384        ncols: usize,
6385        n_used: usize,
6386        t: usize,
6387    ) -> Result<(), Box<dyn std::error::Error>> {
6388        let f = self.func("reduce_slots_f32");
6389        let cfg = LaunchConfig::for_num_elems((t * ncols) as u32);
6390        let (nc, nu, ti) = (ncols as i32, n_used as i32, t as i32);
6391        let __s_b = self.gpu.stream();
6392        let mut b = __s_b.launch_builder(&f);
6393        b.arg(slots).arg(wbuf).arg(dst).arg(&nc).arg(&nu).arg(&ti);
6394        unsafe {
6395            b.launch(cfg)?;
6396        }
6397        Ok(())
6398    }
6399
6400    /// Stage-B: quantize activation [m,in] f32 -> q8_1 (int8 qs + per-block f32 scale).
6401    /// Quantize an activation [m, in_f] to q8_1 (int8 qs + per-32 f32 scale). Public so the
6402    /// forward can quantize a SHARED activation ONCE and feed it to several matmuls (gate+up
6403    /// share `z`; q/k/v and wqkv/gate/beta/alpha share `h`) — quantize_q8_1 was 13.5% of decode
6404    /// GPU time, ~half of it redundant re-quantization of the same row.
6405    /// quantize_q8_1 over a CudaView (a sliced z-row) — same kernel, offset-honoring arg.
6406    pub fn quantize_q8_1_view(
6407        &self,
6408        x: &cudarc::driver::CudaView<f32>,
6409        m: usize,
6410        in_f: usize,
6411    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6412        let f = self.func("quantize_q8_1");
6413        let nblk = in_f / 32;
6414        let mut q = self.alloc_uninit::<i8>(m * in_f)?;
6415        let mut d = self.alloc_uninit::<f32>(m * nblk)?;
6416        let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
6417        let (inf, mi) = (in_f as i32, m as i32);
6418        let __s_b = self.gpu.stream();
6419        let mut b = __s_b.launch_builder(&f);
6420        b.arg(x).arg(&mut q).arg(&mut d).arg(&inf).arg(&mi);
6421        unsafe {
6422            b.launch(cfg)?;
6423        }
6424        Ok((q, d))
6425    }
6426
6427    pub fn quantize_q8_1(
6428        &self,
6429        x: &CudaSlice<f32>,
6430        m: usize,
6431        in_f: usize,
6432    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6433        let nblk = in_f / 32;
6434        let mut q = self.alloc_uninit::<i8>(m * in_f)?; // full-overwrite output: skip memset
6435        let mut d = self.alloc_uninit::<f32>(m * nblk)?; // full-overwrite output: skip memset
6436        // WARP-PER-BLOCK kernel: one warp per 32-block -> m*in_f threads total.
6437        let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
6438        let (inf, mi) = (in_f as i32, m as i32);
6439        if Self::pdl_on() && Self::pdl_wb_on() {
6440            {
6441                use cudarc::driver::{DevicePtr, DevicePtrMut};
6442                let s = &self.gpu.stream();
6443                let (px, _g0) = x.device_ptr(s);
6444                let (pq, _g1) = q.device_ptr_mut(s);
6445                let (pd, _g2) = d.device_ptr_mut(s);
6446                let mut ps = [
6447                    &px as *const _ as *mut std::ffi::c_void,
6448                    &pq as *const _ as *mut _,
6449                    &pd as *const _ as *mut _,
6450                    &inf as *const _ as *mut _,
6451                    &mi as *const _ as *mut _,
6452                ];
6453                unsafe {
6454                    self.launch_pdl("quantize_q8_1", cfg.grid_dim, cfg.block_dim, &mut ps)?;
6455                }
6456            }
6457            return Ok((q, d));
6458        }
6459        let f = self.func("quantize_q8_1");
6460        let __s_b = self.gpu.stream();
6461        let mut b = __s_b.launch_builder(&f);
6462        b.arg(x).arg(&mut q).arg(&mut d).arg(&inf).arg(&mi);
6463        unsafe {
6464            b.launch(cfg)?;
6465        }
6466        Ok((q, d))
6467    }
6468
6469    /// Stage-C FP4: quantize activation [m,in] f32 -> e2m1 nibbles (aq4: u32 [m, in/8]) + per-16
6470    /// UE4M3 scale (ad4: u8 [m, in/16]), the layout the mxf4nvf4 block-scale GEMM B-operand wants.
6471    /// in_f must be a multiple of 64 (one NVFP4 K-block). One thread per (token, 16-block).
6472    pub fn quantize_fp4_act(
6473        &self,
6474        x: &CudaSlice<f32>,
6475        m: usize,
6476        in_f: usize,
6477    ) -> Result<(CudaSlice<u32>, CudaSlice<u8>), Box<dyn std::error::Error>> {
6478        let f = self.func("quantize_fp4_act");
6479        let nb16 = in_f / 16;
6480        let mut aq4 = self.alloc_uninit::<u32>(m * (in_f / 8))?; // full-overwrite output: skip memset
6481        let mut ad4 = self.alloc_uninit::<u8>(m * nb16)?; // full-overwrite output: skip memset
6482        let cfg = LaunchConfig::for_num_elems((m * nb16) as u32);
6483        let (inf, mi) = (in_f as i32, m as i32);
6484        let __s_b = self.gpu.stream();
6485        let mut b = __s_b.launch_builder(&f);
6486        b.arg(x).arg(&mut aq4).arg(&mut ad4).arg(&inf).arg(&mi);
6487        unsafe {
6488            b.launch(cfg)?;
6489        }
6490        Ok((aq4, ad4))
6491    }
6492
6493    /// Stage-C FP4 GEMM (NVFP4 weights): native mxf4nvf4 block-scale tensor-core matmul. Feeds raw
6494    /// e2m1 weight nibbles + raw UE4M3 micro-scales directly to mma.sync.m16n8k64 (762 TFLOP/s peak,
6495    /// 3.5x int8). Activation `x` is quantized to FP4 e2m1 here. NVFP4 per-tensor macro-scale applied
6496    /// post (scale==1.0 -> no-op). `bytes` = raw NVFP4 weight rows. Used by the MEMRA_FP4 prefill path.
6497    pub fn qmatvec_gemm_nvfp4_fp4(
6498        &self,
6499        bytes: &CudaSlice<u8>,
6500        x: &CudaSlice<f32>,
6501        m: usize,
6502        in_f: usize,
6503        out_f: usize,
6504        row_bytes: usize,
6505        scale: f32,
6506    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6507        assert!(
6508            in_f % 64 == 0,
6509            "FP4 GEMM requires in_f % 64 == 0, got {in_f}"
6510        );
6511        let (aq4, ad4) = self.quantize_fp4_act(x, m, in_f)?;
6512        let mut y = self.fp4_gemm_launch(bytes, &aq4, &ad4, m, in_f, out_f, row_bytes)?;
6513        if scale != 1.0 {
6514            self.scale_inplace(&mut y, scale, m * out_f)?;
6515        }
6516        Ok(y)
6517    }
6518
6519    /// Shared mxf4 GEMM launch (pre-quantized FP4 activation aq4/ad4). Same CTA tile as the int8 GEMM
6520    /// (BM=64 rows x BN=128 tokens, 4 warps). No macro-scale applied here.
6521    fn fp4_gemm_launch(
6522        &self,
6523        bytes: &CudaSlice<u8>,
6524        aq4: &CudaSlice<u32>,
6525        ad4: &CudaSlice<u8>,
6526        m: usize,
6527        in_f: usize,
6528        out_f: usize,
6529        row_bytes: usize,
6530    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6531        let f = self.func("qmatvec_gemm_nvfp4_fp4");
6532        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
6533        const BM: u32 = 64;
6534        const BN: u32 = 256;
6535        let cfg = LaunchConfig {
6536            grid_dim: ((out_f as u32 + BM - 1) / BM, (m as u32 + BN - 1) / BN, 1),
6537            block_dim: (32, 4, 1),
6538            shared_mem_bytes: 0,
6539        };
6540        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
6541        let __s_b = self.gpu.stream();
6542        let mut b = __s_b.launch_builder(&f);
6543        b.arg(bytes)
6544            .arg(aq4)
6545            .arg(ad4)
6546            .arg(&mut y)
6547            .arg(&inf)
6548            .arg(&outf)
6549            .arg(&mi)
6550            .arg(&rb);
6551        unsafe {
6552            b.launch(cfg)?;
6553        }
6554        Ok(y)
6555    }
6556
6557    /// Test entry (kernel_check): run the FP4 GEMM from raw bytes; NO macro-scale (caller compares bare).
6558    pub fn qmatvec_gemm_nvfp4_fp4_raw(
6559        &self,
6560        bytes: &CudaSlice<u8>,
6561        x: &CudaSlice<f32>,
6562        m: usize,
6563        in_f: usize,
6564        out_f: usize,
6565        row_bytes: usize,
6566    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6567        assert!(
6568            in_f % 64 == 0,
6569            "FP4 GEMM requires in_f % 64 == 0, got {in_f}"
6570        );
6571        let (aq4, ad4) = self.quantize_fp4_act(x, m, in_f)?;
6572        self.fp4_gemm_launch(bytes, &aq4, &ad4, m, in_f, out_f, row_bytes)
6573    }
6574
6575    /// Stage-B: Q8_0 weight x q8_1 activation int8 dp4a matmul. y[m,out]=x@W^T.
6576    pub fn qmatvec_q8_0_fast(
6577        &self,
6578        w: &CudaSlice<u8>,
6579        x: &CudaSlice<f32>,
6580        m: usize,
6581        in_f: usize,
6582        out_f: usize,
6583        row_bytes: usize,
6584    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6585        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
6586        let f = self.func("qmatvec_q8_0_dp4a");
6587        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
6588        let cfg = LaunchConfig {
6589            grid_dim: (out_f as u32, m as u32, 1),
6590            block_dim: (128, 1, 1),
6591            shared_mem_bytes: 0,
6592        };
6593        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
6594        let __s_b = self.gpu.stream();
6595        let mut b = __s_b.launch_builder(&f);
6596        b.arg(w)
6597            .arg(&aq)
6598            .arg(&ad)
6599            .arg(&mut y)
6600            .arg(&inf)
6601            .arg(&outf)
6602            .arg(&mi)
6603            .arg(&rb);
6604        unsafe {
6605            b.launch(cfg)?;
6606        }
6607        Ok(y)
6608    }
6609
6610    /// Stage-B: Q4_K weight x q8_1 activation int8 dp4a (decode). Min-offset via q8_1 sum term.
6611    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
6612    pub fn qmatvec_q4_K_fast(
6613        &self,
6614        w: &CudaSlice<u8>,
6615        x: &CudaSlice<f32>,
6616        m: usize,
6617        in_f: usize,
6618        out_f: usize,
6619        row_bytes: usize,
6620    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6621        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
6622        let f = self.func("qmatvec_q4_K_dp4a");
6623        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
6624        let cfg = LaunchConfig {
6625            grid_dim: (out_f as u32, m as u32, 1),
6626            block_dim: (128, 1, 1),
6627            shared_mem_bytes: 0,
6628        };
6629        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
6630        let __s_b = self.gpu.stream();
6631        let mut b = __s_b.launch_builder(&f);
6632        b.arg(w)
6633            .arg(&aq)
6634            .arg(&ad)
6635            .arg(&mut y)
6636            .arg(&inf)
6637            .arg(&outf)
6638            .arg(&mi)
6639            .arg(&rb);
6640        unsafe {
6641            b.launch(cfg)?;
6642        }
6643        Ok(y)
6644    }
6645
6646    /// Stage-B: Q6_K weight x q8_1 activation int8 dp4a (decode, symmetric).
6647    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
6648    pub fn qmatvec_q6_K_fast(
6649        &self,
6650        w: &CudaSlice<u8>,
6651        x: &CudaSlice<f32>,
6652        m: usize,
6653        in_f: usize,
6654        out_f: usize,
6655        row_bytes: usize,
6656    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6657        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
6658        let f = self.func("qmatvec_q6_K_dp4a");
6659        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
6660        let cfg = LaunchConfig {
6661            grid_dim: (out_f as u32, m as u32, 1),
6662            block_dim: (128, 1, 1),
6663            shared_mem_bytes: 0,
6664        };
6665        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
6666        let __s_b = self.gpu.stream();
6667        let mut b = __s_b.launch_builder(&f);
6668        b.arg(w)
6669            .arg(&aq)
6670            .arg(&ad)
6671            .arg(&mut y)
6672            .arg(&inf)
6673            .arg(&outf)
6674            .arg(&mi)
6675            .arg(&rb);
6676        unsafe {
6677            b.launch(cfg)?;
6678        }
6679        Ok(y)
6680    }
6681
6682    /// Stage-B: Q5_K weight x q8_1 activation int8 dp4a (decode). Min-offset via q8_1 sum term.
6683    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
6684    pub fn qmatvec_q5_K_fast(
6685        &self,
6686        w: &CudaSlice<u8>,
6687        x: &CudaSlice<f32>,
6688        m: usize,
6689        in_f: usize,
6690        out_f: usize,
6691        row_bytes: usize,
6692    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6693        self.qmatvec_dp4a_named("qmatvec_q5_K_dp4a", w, x, m, in_f, out_f, row_bytes)
6694    }
6695    /// Stage-B: Q3_K weight x q8_1 activation int8 dp4a (decode, symmetric).
6696    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
6697    pub fn qmatvec_q3_K_fast(
6698        &self,
6699        w: &CudaSlice<u8>,
6700        x: &CudaSlice<f32>,
6701        m: usize,
6702        in_f: usize,
6703        out_f: usize,
6704        row_bytes: usize,
6705    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6706        self.qmatvec_dp4a_named("qmatvec_q3_K_dp4a", w, x, m, in_f, out_f, row_bytes)
6707    }
6708    /// A6 split-plane twin of `qmatvec_nvfp4_fast` (weights repacked; used by the rp gates).
6709    pub fn qmatvec_nvfp4_fast_rp(
6710        &self,
6711        w: &CudaSlice<u8>,
6712        x: &CudaSlice<f32>,
6713        m: usize,
6714        in_f: usize,
6715        out_f: usize,
6716        row_bytes: usize,
6717    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6718        assert!(
6719            in_f % 64 == 0,
6720            "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
6721        );
6722        self.qmatvec_dp4a_named("qmatvec_nvfp4_dp4a_rp", w, x, m, in_f, out_f, row_bytes)
6723    }
6724    /// Stage-B: NVFP4 weight x q8_1 activation int8 dp4a (decode, symmetric, codebook lookup).
6725    pub fn qmatvec_nvfp4_fast(
6726        &self,
6727        w: &CudaSlice<u8>,
6728        x: &CudaSlice<f32>,
6729        m: usize,
6730        in_f: usize,
6731        out_f: usize,
6732        row_bytes: usize,
6733    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6734        // B1: the NVFP4 dp4a kernel maps two 32-elem q8_1 blocks onto one 64-elem block_nvfp4
6735        // (sblk = g >> 1). in_f must be a multiple of 64 or the last block reads a partial superblock.
6736        assert!(
6737            in_f % 64 == 0,
6738            "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
6739        );
6740        self.qmatvec_dp4a_named("qmatvec_nvfp4_dp4a", w, x, m, in_f, out_f, row_bytes)
6741    }
6742    /// Stage-B (optional perf): IQ4_XS codebook int8 dp4a.
6743    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
6744    pub fn qmatvec_iq4_XS_fast(
6745        &self,
6746        w: &CudaSlice<u8>,
6747        x: &CudaSlice<f32>,
6748        m: usize,
6749        in_f: usize,
6750        out_f: usize,
6751        row_bytes: usize,
6752    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6753        self.qmatvec_dp4a_named("qmatvec_iq4_XS_dp4a", w, x, m, in_f, out_f, row_bytes)
6754    }
6755
6756    /// Shared dp4a launcher: quantize_q8_1 then call the named kernel (grid (out,m), block 64).
6757    fn qmatvec_dp4a_named(
6758        &self,
6759        name: &str,
6760        w: &CudaSlice<u8>,
6761        x: &CudaSlice<f32>,
6762        m: usize,
6763        in_f: usize,
6764        out_f: usize,
6765        row_bytes: usize,
6766    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6767        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
6768        let f = self.func(name);
6769        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
6770        let cfg = LaunchConfig {
6771            grid_dim: (out_f as u32, m as u32, 1),
6772            block_dim: (128, 1, 1),
6773            shared_mem_bytes: 0,
6774        };
6775        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
6776        let __s_b = self.gpu.stream();
6777        let mut b = __s_b.launch_builder(&f);
6778        b.arg(w)
6779            .arg(&aq)
6780            .arg(&ad)
6781            .arg(&mut y)
6782            .arg(&inf)
6783            .arg(&outf)
6784            .arg(&mi)
6785            .arg(&rb);
6786        unsafe {
6787            b.launch(cfg)?;
6788        }
6789        Ok(y)
6790    }
6791
6792    pub fn htod(&self, v: &[f32]) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6793        Ok(self.gpu.stream().clone_htod(v)?)
6794    }
6795    pub fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
6796        Ok(self.gpu.stream().clone_htod(v)?)
6797    }
6798    /// i8 upload (moe-devq8-check: synthetic q8_1 activation bytes).
6799    pub fn htod_i8(&self, v: &[i8]) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
6800        Ok(self.gpu.stream().clone_htod(v)?)
6801    }
6802    pub fn htod_u64(&self, v: &[u64]) -> Result<CudaSlice<u64>, Box<dyn std::error::Error>> {
6803        Ok(self.gpu.stream().clone_htod(v)?)
6804    }
6805    /// View twin of `dtoh` (lean-logits component 3: D2H one row of a [B, n_vocab] stack).
6806    pub fn dtoh_view(
6807        &self,
6808        d: &cudarc::driver::CudaView<f32>,
6809    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
6810        let v = self.gpu.stream().clone_dtoh(d)?;
6811        self.gpu.stream().synchronize()?;
6812        Ok(v)
6813    }
6814    pub fn dtoh(&self, d: &CudaSlice<f32>) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
6815        let v = self.gpu.stream().clone_dtoh(d)?;
6816        self.gpu.stream().synchronize()?;
6817        Ok(v)
6818    }
6819    /// Queue two f32 device-to-host copies on the compute stream, then establish one host
6820    /// boundary for both. Hy3's CPU/GPU expert split needs the router logits and the MoE input;
6821    /// issuing them together avoids a second stream synchronization in every trunk layer.
6822    pub fn dtoh_pair(
6823        &self,
6824        a: &CudaSlice<f32>,
6825        b: &CudaSlice<f32>,
6826    ) -> Result<(Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
6827        let av = self.gpu.stream().clone_dtoh(a)?;
6828        let bv = self.gpu.stream().clone_dtoh(b)?;
6829        self.gpu.stream().synchronize()?;
6830        Ok((av, bv))
6831    }
6832    /// Device-to-host copy of an i32 buffer (fused-router sel_idx readback).
6833    pub fn dtoh_i32(&self, d: &CudaSlice<i32>) -> Result<Vec<i32>, Box<dyn std::error::Error>> {
6834        let v = self.gpu.stream().clone_dtoh(d)?;
6835        self.gpu.stream().synchronize()?;
6836        Ok(v)
6837    }
6838    /// Device-to-host copy of a u8 buffer (used to read back the quantized KV cache for validation).
6839    pub fn dtoh_u8(&self, d: &CudaSlice<u8>) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
6840        let v = self.gpu.stream().clone_dtoh(d)?;
6841        self.gpu.stream().synchronize()?;
6842        Ok(v)
6843    }
6844    pub fn dtoh_u8_view(
6845        &self,
6846        d: &cudarc::driver::CudaView<u8>,
6847    ) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
6848        let v = self.gpu.stream().clone_dtoh(d)?;
6849        self.gpu.stream().synchronize()?;
6850        Ok(v)
6851    }
6852    pub fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6853        let s = self.gpu.stream().alloc_zeros::<f32>(n)?;
6854        self.keep_if_capturing(&s);
6855        Ok(s)
6856    }
6857
6858    /// GPU-resident greedy argmax (CUDA-GRAPH-PLAN Phase 1): logits[n_vocab] -> token id in a
6859    /// resident device u32 [1]. PARALLEL 2-pass (RANK1 LEVER): the old single-CTA scan (one 256-thread
6860    /// block on one SM over 248K logits) was memory-starved at ~426us/token. Now pass 1 fans NB=256
6861    /// blocks across the SMs to saturate HBM, pass 2 reduces the NB partials. Bit-identical to host
6862    /// `argmax` (smallest index on tie). The whole point is NOT to dtoh logits — only a [1] u32 is read
6863    /// back (or kept resident for graph replay). Returns the device token buffer.
6864    /// Softmax probability of the (already-argmaxed) token `tok` under `logits` — the spec-decode
6865    /// p-min confidence signal. 2-pass like the parallel argmax; returns a device [1] f32.
6866    pub fn prob_of_token_device(
6867        &self,
6868        logits: &CudaSlice<f32>,
6869        tok: &CudaSlice<u32>,
6870        n_vocab: usize,
6871    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6872        let nb = ARGMAX_NB;
6873        let mut part = self.alloc_uninit::<f32>(nb)?;
6874        let mut p = self.alloc_uninit::<f32>(1)?;
6875        let f1 = self.func("prob_of_token_partial_f32");
6876        let cfg1 = LaunchConfig {
6877            grid_dim: (nb as u32, 1, 1),
6878            block_dim: (256, 1, 1),
6879            shared_mem_bytes: 0,
6880        };
6881        let nv = n_vocab as i32;
6882        let __s_b1 = self.gpu.stream();
6883        let mut b1 = __s_b1.launch_builder(&f1);
6884        b1.arg(logits).arg(tok).arg(&mut part).arg(&nv);
6885        unsafe {
6886            b1.launch(cfg1)?;
6887        }
6888        let f2 = self.func("prob_of_token_final_f32");
6889        let cfg2 = LaunchConfig {
6890            grid_dim: (1, 1, 1),
6891            block_dim: (256, 1, 1),
6892            shared_mem_bytes: 0,
6893        };
6894        let nbi = nb as i32;
6895        let __s_b2 = self.gpu.stream();
6896        let mut b2 = __s_b2.launch_builder(&f2);
6897        b2.arg(&part).arg(&mut p).arg(&nbi);
6898        unsafe {
6899            b2.launch(cfg2)?;
6900        }
6901        Ok(p)
6902    }
6903
6904    /// Like `prob_of_token_device` but writes into a PERSISTENT `p_out` buffer (stable pointer).
6905    /// Required for CUDA-graph capture of the draft chain: the captured prob kernels must write
6906    /// where the host reads the p-min confidence between replays. Same kernels, same math.
6907    /// Slot-addressed twin of `prob_of_token_device_into`: token read from `tok_all[tok_idx]`
6908    /// (a view at the slot), probability written to `p_out[p_idx]` — same two kernels, the
6909    /// pointers just land mid-buffer. Zero-sync (gemma confidence-adaptive draft depth).
6910    pub fn prob_of_token_device_col(
6911        &self,
6912        logits: &CudaSlice<f32>,
6913        tok_all: &CudaSlice<u32>,
6914        tok_idx: usize,
6915        p_out: &mut CudaSlice<f32>,
6916        p_idx: usize,
6917        n_vocab: usize,
6918    ) -> Result<(), Box<dyn std::error::Error>> {
6919        let tok_v = tok_all.slice(tok_idx..tok_idx + 1);
6920        let mut p_v = p_out.slice_mut(p_idx..p_idx + 1);
6921        let nb = ARGMAX_NB;
6922        let mut part = self.alloc_uninit::<f32>(nb)?;
6923        let f1 = self.func("prob_of_token_partial_f32");
6924        let cfg1 = LaunchConfig {
6925            grid_dim: (nb as u32, 1, 1),
6926            block_dim: (256, 1, 1),
6927            shared_mem_bytes: 0,
6928        };
6929        let nv = n_vocab as i32;
6930        let __s_b1 = self.gpu.stream();
6931        let mut b1 = __s_b1.launch_builder(&f1);
6932        b1.arg(logits).arg(&tok_v).arg(&mut part).arg(&nv);
6933        unsafe {
6934            b1.launch(cfg1)?;
6935        }
6936        let f2 = self.func("prob_of_token_final_f32");
6937        let cfg2 = LaunchConfig {
6938            grid_dim: (1, 1, 1),
6939            block_dim: (256, 1, 1),
6940            shared_mem_bytes: 0,
6941        };
6942        let nbi = nb as i32;
6943        let __s_b2 = self.gpu.stream();
6944        let mut b2 = __s_b2.launch_builder(&f2);
6945        b2.arg(&part).arg(&mut p_v).arg(&nbi);
6946        unsafe {
6947            b2.launch(cfg2)?;
6948        }
6949        Ok(())
6950    }
6951
6952    pub fn prob_of_token_device_into(
6953        &self,
6954        logits: &CudaSlice<f32>,
6955        tok: &CudaSlice<u32>,
6956        p_out: &mut CudaSlice<f32>,
6957        n_vocab: usize,
6958    ) -> Result<(), Box<dyn std::error::Error>> {
6959        let nb = ARGMAX_NB;
6960        let mut part = self.alloc_uninit::<f32>(nb)?;
6961        let f1 = self.func("prob_of_token_partial_f32");
6962        let cfg1 = LaunchConfig {
6963            grid_dim: (nb as u32, 1, 1),
6964            block_dim: (256, 1, 1),
6965            shared_mem_bytes: 0,
6966        };
6967        let nv = n_vocab as i32;
6968        let __s_b1 = self.gpu.stream();
6969        let mut b1 = __s_b1.launch_builder(&f1);
6970        b1.arg(logits).arg(tok).arg(&mut part).arg(&nv);
6971        unsafe {
6972            b1.launch(cfg1)?;
6973        }
6974        let f2 = self.func("prob_of_token_final_f32");
6975        let cfg2 = LaunchConfig {
6976            grid_dim: (1, 1, 1),
6977            block_dim: (256, 1, 1),
6978            shared_mem_bytes: 0,
6979        };
6980        let nbi = nb as i32;
6981        let __s_b2 = self.gpu.stream();
6982        let mut b2 = __s_b2.launch_builder(&f2);
6983        b2.arg(&part).arg(p_out).arg(&nbi);
6984        unsafe {
6985            b2.launch(cfg2)?;
6986        }
6987        Ok(())
6988    }
6989
6990    pub fn argmax_token_device(
6991        &self,
6992        logits: &CudaSlice<f32>,
6993        n_vocab: usize,
6994    ) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
6995        let mut tok = unsafe { self.gpu.stream().alloc::<u32>(1)? };
6996        self.argmax_token_device_into(logits, &mut tok, n_vocab)?;
6997        Ok(tok)
6998    }
6999    /// Like `argmax_token_device` but writes into a PERSISTENT `tok` buffer (stable pointer) instead
7000    /// of allocating a fresh one. Required for CUDA-graph capture: the captured argmax must write the
7001    /// next token into the SAME device buffer the next replay's embed_gather reads, so the buffer
7002    /// pointer is baked once and the token id never round-trips to host inside steady state. The
7003    /// pass-1 partials scratch (`argmax_partials`) is also a resident stable-pointer buffer so both
7004    /// captured passes bake fixed addresses.
7005    pub fn argmax_token_device_into(
7006        &self,
7007        logits: &CudaSlice<f32>,
7008        tok: &mut CudaSlice<u32>,
7009        n_vocab: usize,
7010    ) -> Result<(), Box<dyn std::error::Error>> {
7011        let nb = ARGMAX_NB;
7012        let f1 = self.func("argmax_partial_f32");
7013        let f2 = self.func("argmax_final_f32");
7014        let mut guard = self.argmax_partials.lock().unwrap();
7015        if guard.is_none() {
7016            // allocate ONCE; under generate_graph this runs in the tracking-off prime window so the
7017            // buffers carry no cudarc events (illegal inside capture).
7018            let pv = self.gpu.stream().alloc_zeros::<f32>(nb)?;
7019            let pi = self.gpu.stream().alloc_zeros::<i32>(nb)?;
7020            *guard = Some((pv, pi));
7021        }
7022        let (part_v, part_i) = guard.as_mut().unwrap();
7023        let nv = n_vocab as i32;
7024        let nbi = nb as i32;
7025        // pass 1: NB blocks x 256 threads grid-stride scan -> per-block (val, idx) partials.
7026        let cfg1 = LaunchConfig {
7027            grid_dim: (nb as u32, 1, 1),
7028            block_dim: (256, 1, 1),
7029            shared_mem_bytes: 0,
7030        };
7031        let __s_b1 = self.gpu.stream();
7032        let mut b1 = __s_b1.launch_builder(&f1);
7033        b1.arg(logits).arg(&mut *part_v).arg(&mut *part_i).arg(&nv);
7034        unsafe {
7035            b1.launch(cfg1)?;
7036        }
7037        // pass 2: one block reduces NB partials -> token_out[0].
7038        let cfg2 = LaunchConfig {
7039            grid_dim: (1, 1, 1),
7040            block_dim: (256, 1, 1),
7041            shared_mem_bytes: 0,
7042        };
7043        let __s_b2 = self.gpu.stream();
7044        let mut b2 = __s_b2.launch_builder(&f2);
7045        b2.arg(&*part_v).arg(&*part_i).arg(tok).arg(&nbi);
7046        unsafe {
7047            b2.launch(cfg2)?;
7048        }
7049        Ok(())
7050    }
7051    /// Column-`col` device argmax over a stacked verify-logits buffer [t, n_vocab] (spec accept
7052    /// walk): toks[out_idx] = argmax(logits[col*n_vocab .. (col+1)*n_vocab]). SAME 2-pass kernels
7053    /// and tie-break contract as `argmax_token_device_into` (bit-identical to host argmax,
7054    /// argmax_gate-validated) — only the input pointer (a column view) and the output slot differ.
7055    /// Lets the accept walk read ONE [t] u32 instead of dtoh'ing the full [t, n_vocab] logits.
7056    pub fn argmax_token_device_col(
7057        &self,
7058        logits: &CudaSlice<f32>,
7059        col: usize,
7060        n_vocab: usize,
7061        toks: &mut CudaSlice<u32>,
7062        out_idx: usize,
7063    ) -> Result<(), Box<dyn std::error::Error>> {
7064        let nb = ARGMAX_NB;
7065        let f1 = self.func("argmax_partial_f32");
7066        let f2 = self.func("argmax_final_f32");
7067        let mut guard = self.argmax_partials.lock().unwrap();
7068        if guard.is_none() {
7069            let pv = self.gpu.stream().alloc_zeros::<f32>(nb)?;
7070            let pi = self.gpu.stream().alloc_zeros::<i32>(nb)?;
7071            *guard = Some((pv, pi));
7072        }
7073        let (part_v, part_i) = guard.as_mut().unwrap();
7074        let col_view = logits.slice(col * n_vocab..(col + 1) * n_vocab);
7075        let nv = n_vocab as i32;
7076        let nbi = nb as i32;
7077        let cfg1 = LaunchConfig {
7078            grid_dim: (nb as u32, 1, 1),
7079            block_dim: (256, 1, 1),
7080            shared_mem_bytes: 0,
7081        };
7082        let __s_b1 = self.gpu.stream();
7083        let mut b1 = __s_b1.launch_builder(&f1);
7084        b1.arg(&col_view)
7085            .arg(&mut *part_v)
7086            .arg(&mut *part_i)
7087            .arg(&nv);
7088        unsafe {
7089            b1.launch(cfg1)?;
7090        }
7091        let mut tok_view = toks.slice_mut(out_idx..out_idx + 1);
7092        let cfg2 = LaunchConfig {
7093            grid_dim: (1, 1, 1),
7094            block_dim: (256, 1, 1),
7095            shared_mem_bytes: 0,
7096        };
7097        let __s_b2 = self.gpu.stream();
7098        let mut b2 = __s_b2.launch_builder(&f2);
7099        b2.arg(&*part_v).arg(&*part_i).arg(&mut tok_view).arg(&nbi);
7100        unsafe {
7101            b2.launch(cfg2)?;
7102        }
7103        Ok(())
7104    }
7105    /// Read back a device u32 buffer (the spec accept walk's [t] per-column argmax tokens).
7106    pub fn htod_u32_v(&self, v: &[u32]) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
7107        Ok(self.gpu.stream().clone_htod(v)?)
7108    }
7109    pub fn dtoh_u32(&self, d: &CudaSlice<u32>) -> Result<Vec<u32>, Box<dyn std::error::Error>> {
7110        let v = self.gpu.stream().clone_dtoh(d)?;
7111        self.gpu.stream().synchronize()?;
7112        Ok(v)
7113    }
7114    /// Allocate a zeroed device u32 buffer (persistent spec-loop prediction slots).
7115    /// H2D into an EXISTING u32 buffer (stable pointer — the per-step grammar-mask upload:
7116    /// contents change every step, the address must not, so a captured graph can read it).
7117    pub fn htod_u32_into(
7118        &self,
7119        dst: &mut CudaSlice<u32>,
7120        src: &[u32],
7121    ) -> Result<(), Box<dyn std::error::Error>> {
7122        let mut view = dst.slice_mut(0..src.len());
7123        self.gpu.stream().memcpy_htod(src, &mut view)?;
7124        Ok(())
7125    }
7126
7127    /// H2D into an existing i32 buffer. OPTIPIPE uses this to refresh a stage-local saved-len
7128    /// table without changing the device address its reconcile kernel consumes.
7129    pub fn htod_i32_into(
7130        &self,
7131        dst: &mut CudaSlice<i32>,
7132        src: &[i32],
7133    ) -> Result<(), Box<dyn std::error::Error>> {
7134        let mut view = dst.slice_mut(0..src.len());
7135        self.gpu.stream().memcpy_htod(src, &mut view)?;
7136        Ok(())
7137    }
7138
7139    pub fn alloc_u32_zeroed(&self, n: usize) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
7140        let s = self.gpu.stream().alloc_zeros::<u32>(n)?;
7141        self.keep_if_capturing(&s);
7142        Ok(s)
7143    }
7144    /// embed_gather into a PERSISTENT `x_out` buffer (stable pointer) for CUDA-graph capture (the
7145    /// embed output starts the per-step kernel chain and must be at a fixed address across replays).
7146    pub fn embed_gather_device_into(
7147        &self,
7148        embd: &CudaSlice<u8>,
7149        token_d: &CudaSlice<u32>,
7150        x_out: &mut CudaSlice<f32>,
7151        n_embd: usize,
7152        qtype: i32,
7153        row_bytes: usize,
7154    ) -> Result<(), Box<dyn std::error::Error>> {
7155        let f = self.func("embed_gather_u32");
7156        let cfg = LaunchConfig {
7157            grid_dim: (((n_embd as u32 + 255) / 256).max(1), 1, 1),
7158            block_dim: (256, 1, 1),
7159            shared_mem_bytes: 0,
7160        };
7161        let (ne, qt, rb) = (n_embd as i32, qtype, row_bytes as i64);
7162        let __s_b = self.gpu.stream();
7163        let mut b = __s_b.launch_builder(&f);
7164        b.arg(embd)
7165            .arg(token_d)
7166            .arg(x_out)
7167            .arg(&ne)
7168            .arg(&qt)
7169            .arg(&rb);
7170        unsafe {
7171            b.launch(cfg)?;
7172        }
7173        Ok(())
7174    }
7175    /// Read a [1] i32 device counter (pos / seqlen) back to host. Tiny D2H + sync.
7176    pub fn dtoh_i32_one(&self, d: &CudaSlice<i32>) -> Result<i32, Box<dyn std::error::Error>> {
7177        let v = self.gpu.stream().clone_dtoh(d)?;
7178        self.gpu.stream().synchronize()?;
7179        Ok(v[0])
7180    }
7181    /// Set a [1] i32 device counter IN PLACE (keeps the buffer pointer stable — required for the
7182    /// graph-resident pos/seqlen counters whose addresses are baked into captured graphs). Restores
7183    /// the counter value after the throwaway capture warmups corrupt it.
7184    /// ASYNC i32 single-slot store (value rides the kernel arg — no host-memory transfer/sync).
7185    /// The graph-arc device-len counters use this; set_i32_one below is the SYNCING pageable
7186    /// copy (fine at stream-idle boundaries, poison mid-round).
7187    pub fn i32_set_k(
7188        &self,
7189        dst: &mut CudaSlice<i32>,
7190        v: i32,
7191    ) -> Result<(), Box<dyn std::error::Error>> {
7192        let f = self.func("i32_set_k");
7193        let cfg = LaunchConfig {
7194            grid_dim: (1, 1, 1),
7195            block_dim: (1, 1, 1),
7196            shared_mem_bytes: 0,
7197        };
7198        let idx = 0i32;
7199        let __s_b = self.gpu.stream();
7200        let mut b = __s_b.launch_builder(&f);
7201        b.arg(dst).arg(&v).arg(&idx);
7202        unsafe {
7203            b.launch(cfg)?;
7204        }
7205        Ok(())
7206    }
7207
7208    pub fn set_i32_one(
7209        &self,
7210        d: &mut CudaSlice<i32>,
7211        v: i32,
7212    ) -> Result<(), Box<dyn std::error::Error>> {
7213        self.gpu.stream().memcpy_htod(&[v], d)?;
7214        Ok(())
7215    }
7216    /// Set a [1] u32 device buffer IN PLACE (stable pointer) — for the resident `token_d` counter
7217    /// during priming / capture-state restore.
7218    pub fn set_u32_one(
7219        &self,
7220        d: &mut CudaSlice<u32>,
7221        v: u32,
7222    ) -> Result<(), Box<dyn std::error::Error>> {
7223        self.gpu.stream().memcpy_htod(&[v], d)?;
7224        Ok(())
7225    }
7226    /// Read back a [1] u32 device buffer (the argmax token). One tiny D2H + sync.
7227    pub fn dtoh_u32_one(&self, d: &CudaSlice<u32>) -> Result<u32, Box<dyn std::error::Error>> {
7228        let v = self.gpu.stream().clone_dtoh(d)?;
7229        self.gpu.stream().synchronize()?;
7230        Ok(v[0])
7231    }
7232    /// Upload raw bytes to a resident device u8 buffer (e.g. the embed table for device gather).
7233    pub fn upload_u8(&self, bytes: &[u8]) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
7234        Ok(self.gpu.stream().clone_htod(bytes)?)
7235    }
7236    /// Embed-from-device (CUDA-GRAPH-PLAN Phase 1): gather+dequant the row for the token id in
7237    /// `token_d[0]` from the resident embed table `embd` -> x_out[n_embd]. Bit-identical to host
7238    /// EmbedHost::gather (same per-dtype `deq`). No host round-trip of the token id.
7239    pub fn embed_gather_device(
7240        &self,
7241        embd: &CudaSlice<u8>,
7242        token_d: &CudaSlice<u32>,
7243        n_embd: usize,
7244        qtype: i32,
7245        row_bytes: usize,
7246    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7247        let f = self.func("embed_gather_u32");
7248        let mut x = self.alloc_uninit::<f32>(n_embd)?;
7249        let cfg = LaunchConfig {
7250            grid_dim: (((n_embd as u32 + 255) / 256).max(1), 1, 1),
7251            block_dim: (256, 1, 1),
7252            shared_mem_bytes: 0,
7253        };
7254        let (ne, qt, rb) = (n_embd as i32, qtype, row_bytes as i64);
7255        let __s_b = self.gpu.stream();
7256        let mut b = __s_b.launch_builder(&f);
7257        b.arg(embd)
7258            .arg(token_d)
7259            .arg(&mut x)
7260            .arg(&ne)
7261            .arg(&qt)
7262            .arg(&rb);
7263        unsafe {
7264            b.launch(cfg)?;
7265        }
7266        Ok(x)
7267    }
7268
7269    /// T-token device embed gather (spec verify/replay): tokens uploaded as a tiny [T] u32 htod,
7270    /// rows dequanted on-device -> x[T, n_embd]. Replaces host per-row dequant + T*n_embd*4B htod
7271    /// (nsys: 84% of spec API time was HtoD). Bit-identical rows (same per-dtype deq).
7272    pub fn embed_gather_device_t(
7273        &self,
7274        embd: &CudaSlice<u8>,
7275        tokens: &[u32],
7276        n_embd: usize,
7277        qtype: i32,
7278        row_bytes: usize,
7279    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7280        let t = tokens.len();
7281        let tok_d = self.gpu.stream().clone_htod(tokens)?;
7282        let f = self.func("embed_gather_u32_t");
7283        let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
7284        let cfg = LaunchConfig {
7285            grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
7286            block_dim: (256, 1, 1),
7287            shared_mem_bytes: 0,
7288        };
7289        let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
7290        let __s_b = self.gpu.stream();
7291        let mut b = __s_b.launch_builder(&f);
7292        b.arg(embd)
7293            .arg(&tok_d)
7294            .arg(&mut x)
7295            .arg(&ne)
7296            .arg(&qt)
7297            .arg(&rb)
7298            .arg(&ti);
7299        unsafe {
7300            b.launch(cfg)?;
7301        }
7302        Ok(x)
7303    }
7304
7305    /// T-token embed gather from a DEVICE token buffer (round-stream stage c: the verify tokens
7306    /// are assembled on-device from the draft-chain pack slots; no host round trip). Same kernel
7307    /// as embed_gather_device_t — bit-identical rows.
7308    /// embed_gather over a token VIEW (spec round: tokens live in the round's batch buffer).
7309    pub fn embed_gather_device_tv(
7310        &self,
7311        embd: &CudaSlice<u8>,
7312        tok_v: &cudarc::driver::CudaView<u32>,
7313        t: usize,
7314        n_embd: usize,
7315        qtype: i32,
7316        row_bytes: usize,
7317    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7318        let f = self.func("embed_gather_u32_t");
7319        let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
7320        let cfg = LaunchConfig {
7321            grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
7322            block_dim: (256, 1, 1),
7323            shared_mem_bytes: 0,
7324        };
7325        let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
7326        let __s_b = self.gpu.stream();
7327        let mut b = __s_b.launch_builder(&f);
7328        b.arg(embd)
7329            .arg(tok_v)
7330            .arg(&mut x)
7331            .arg(&ne)
7332            .arg(&qt)
7333            .arg(&rb)
7334            .arg(&ti);
7335        unsafe {
7336            b.launch(cfg)?;
7337        }
7338        Ok(x)
7339    }
7340
7341    pub fn embed_gather_device_td(
7342        &self,
7343        embd: &CudaSlice<u8>,
7344        tok_d: &CudaSlice<u32>,
7345        t: usize,
7346        n_embd: usize,
7347        qtype: i32,
7348        row_bytes: usize,
7349    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7350        let f = self.func("embed_gather_u32_t");
7351        let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
7352        let cfg = LaunchConfig {
7353            grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
7354            block_dim: (256, 1, 1),
7355            shared_mem_bytes: 0,
7356        };
7357        let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
7358        let __s_b = self.gpu.stream();
7359        let mut b = __s_b.launch_builder(&f);
7360        b.arg(embd)
7361            .arg(tok_d)
7362            .arg(&mut x)
7363            .arg(&ne)
7364            .arg(&qt)
7365            .arg(&rb)
7366            .arg(&ti);
7367        unsafe {
7368            b.launch(cfg)?;
7369        }
7370        Ok(x)
7371    }
7372
7373    /// Uninitialized device buffer — SKIPS the memset that `alloc_zeros` always issues. Decode
7374    /// profile (nsys): ~1050 memsets/token = 6.5% of decode GPU time + ~half the launch count, the
7375    /// dominant contributor to the 19% inter-kernel idle gap and a blocker for clean CUDA-graph
7376    /// capture. Use ONLY for buffers a kernel FULLY overwrites (every element written, no `+=`).
7377    /// SAFETY: caller guarantees the producing kernel writes every element before any read.
7378    #[inline]
7379    /// Keep an allocation alive for the current capture (no-op when retain mode is off).
7380    fn keep_if_capturing<T: cudarc::driver::DeviceRepr + Send + 'static>(&self, s: &CudaSlice<T>) {
7381        if self
7382            .capture_keep_on
7383            .load(std::sync::atomic::Ordering::Relaxed)
7384        {
7385            self.capture_keep.lock().unwrap().push(Box::new(s.clone()));
7386        }
7387    }
7388
7389    fn alloc_uninit<T: cudarc::driver::DeviceRepr + Send + 'static>(
7390        &self,
7391        n: usize,
7392    ) -> Result<CudaSlice<T>, Box<dyn std::error::Error>> {
7393        let mut s = unsafe { self.gpu.stream().alloc::<T>(n)? };
7394        // MEMRA_DEBUG_ZERO_ALLOCS=1 (task #14 defect hunt): memset EVERY engine allocation —
7395        // the global uninit-read discriminator (the prime-fn-scoped zeroing experiment could
7396        // not cover engine-internal buffers). Debug-only: massive launch overhead.
7397        {
7398            static Z: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7399            if *Z.get_or_init(|| std::env::var("MEMRA_DEBUG_ZERO_ALLOCS").as_deref() == Ok("1")) {
7400                // raw D8 memset (T lacks ValidAsZeroBits in the generic bound)
7401                use cudarc::driver::DevicePtrMut;
7402                let n_bytes = s.len() * std::mem::size_of::<T>();
7403                let stream = self.gpu.stream();
7404                let (p_, _g) = s.device_ptr_mut(&stream);
7405                unsafe {
7406                    cudarc::driver::sys::cuMemsetD8Async(p_, 0, n_bytes, stream.cu_stream())
7407                        .result()?;
7408                }
7409            }
7410        }
7411        self.keep_if_capturing(&s);
7412        Ok(s)
7413    }
7414
7415    /// Public f32 uninitialized scratch (see `alloc_uninit`). For decode/forward scratch a kernel
7416    /// fully overwrites. SAFETY: producing kernel must write every element before any read.
7417    /// Uninitialized q8_1 activation pair (int8 + per-32 scales) — the fa combine q8-emit
7418    /// consumers alloc through this (m=1 decode arms).
7419    pub fn uninit_q8_pair(
7420        &self,
7421        n: usize,
7422    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7423        Ok((
7424            self.alloc_uninit::<i8>(n)?,
7425            self.alloc_uninit::<f32>(n / 32)?,
7426        ))
7427    }
7428
7429    pub fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7430        self.alloc_uninit::<f32>(n)
7431    }
7432
7433    /// i8 uninitialized scratch (same contract as `uninit`).
7434    pub fn alloc_i8_uninit(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
7435        self.alloc_uninit::<i8>(n)
7436    }
7437
7438    /// RMSNorm: x[ncols,nrows] row-major, weight[ncols] -> dst. One block/row, 256 threads.
7439    /// gemma4: 3 rms_norms of the SAME input in one launch (one reduction, three weights).
7440    /// Per-output bit-identical to three rms_norm calls (verbatim reduction/scale chain).
7441    #[allow(clippy::too_many_arguments)]
7442    pub fn rms_norm3(
7443        &self,
7444        x: &CudaSlice<f32>,
7445        w0: &CudaSlice<f32>,
7446        w1: &CudaSlice<f32>,
7447        w2: &CudaSlice<f32>,
7448        d0: &mut CudaSlice<f32>,
7449        d1: &mut CudaSlice<f32>,
7450        d2: &mut CudaSlice<f32>,
7451        ncols: usize,
7452        nrows: usize,
7453        eps: f32,
7454    ) -> Result<(), Box<dyn std::error::Error>> {
7455        let f = self.func("rms_norm3_f32");
7456        let cfg = LaunchConfig {
7457            grid_dim: (nrows as u32, 1, 1),
7458            block_dim: (rms_block(), 1, 1),
7459            shared_mem_bytes: 0,
7460        };
7461        let (nc, e) = (ncols as i32, eps);
7462        let __s_b = self.gpu.stream();
7463        let mut b = __s_b.launch_builder(&f);
7464        b.arg(x)
7465            .arg(w0)
7466            .arg(w1)
7467            .arg(w2)
7468            .arg(d0)
7469            .arg(d1)
7470            .arg(d2)
7471            .arg(&nc)
7472            .arg(&e);
7473        unsafe {
7474            b.launch(cfg)?;
7475        }
7476        Ok(())
7477    }
7478
7479    /// gemma4 fused q/k/v head norms (one launch, per-row rms_norm_f32-verbatim).
7480    #[allow(clippy::too_many_arguments)]
7481    /// True when the warp-per-row qkv norm would engage for (rows, ncols) — the emit lane
7482    /// piggybacks on the same conditions.
7483    pub fn qkvnorm_w_on_prefill(rows: usize, ncols: usize) -> bool {
7484        static WARP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7485        *WARP_ON.get_or_init(|| {
7486            std::env::var("MEMRA_QKVNORM_W")
7487                .map(|v| v != "0")
7488                .unwrap_or(true)
7489        }) && ncols % 4 == 0
7490            && rows >= 64
7491    }
7492
7493    /// w4 norm with bf16 V EMIT (31B glue lane): the v segment also writes its normed rows as
7494    /// bf16 (the FA V operand — bit-identical to a post-hoc f32_to_bf16). Prefill-depth only.
7495    #[allow(clippy::too_many_arguments)]
7496    pub fn rms_norm_qkv_w4b(
7497        &self,
7498        q: &CudaSlice<f32>,
7499        k: &CudaSlice<f32>,
7500        v: &CudaSlice<f32>,
7501        wq: &CudaSlice<f32>,
7502        wk: &CudaSlice<f32>,
7503        wv: &CudaSlice<f32>,
7504        dq: &mut CudaSlice<f32>,
7505        dk: &mut CudaSlice<f32>,
7506        dv: &mut CudaSlice<f32>,
7507        dvb: &mut CudaSlice<u8>,
7508        ncols: usize,
7509        rq: usize,
7510        rk: usize,
7511        eps: f32,
7512        vf16: bool,
7513    ) -> Result<(), Box<dyn std::error::Error>> {
7514        assert!(ncols % 4 == 0 && rq + 2 * rk >= 64);
7515        let f = self.func("rms_norm_qkv_w4b_f32");
7516        let rows = (rq + 2 * rk) as u32;
7517        let cfg = LaunchConfig {
7518            grid_dim: (rows.div_ceil(8), 1, 1),
7519            block_dim: (256, 1, 1),
7520            shared_mem_bytes: 0,
7521        };
7522        let (nc, rqi, rki, rvi, e) = (ncols as i32, rq as i32, rk as i32, rk as i32, eps);
7523        let vf = vf16 as i32;
7524        let __s_b = self.gpu.stream();
7525        let mut b = __s_b.launch_builder(&f);
7526        b.arg(q)
7527            .arg(k)
7528            .arg(v)
7529            .arg(wq)
7530            .arg(wk)
7531            .arg(wv)
7532            .arg(dq)
7533            .arg(dk)
7534            .arg(dv)
7535            .arg(&mut *dvb)
7536            .arg(&nc)
7537            .arg(&rqi)
7538            .arg(&rki)
7539            .arg(&rvi)
7540            .arg(&e)
7541            .arg(&vf);
7542        unsafe {
7543            b.launch(cfg)?;
7544        }
7545        Ok(())
7546    }
7547
7548    pub fn rms_norm_qkv(
7549        &self,
7550        q: &CudaSlice<f32>,
7551        k: &CudaSlice<f32>,
7552        v: &CudaSlice<f32>,
7553        wq: &CudaSlice<f32>,
7554        wk: &CudaSlice<f32>,
7555        wv: &CudaSlice<f32>,
7556        dq: &mut CudaSlice<f32>,
7557        dk: &mut CudaSlice<f32>,
7558        dv: &mut CudaSlice<f32>,
7559        ncols: usize,
7560        rq: usize,
7561        rk: usize,
7562        eps: f32,
7563    ) -> Result<(), Box<dyn std::error::Error>> {
7564        // Warp-per-row float4 twin (default; MEMRA_QKVNORM_W=0 reverts): the block-per-row form
7565        // spends 767us/launch on 17k+ 2KB rows at prefill depth (launch/reduce latency-bound,
7566        // ~92GB/s). Own numeric config (reduce order differs) — battery-gated.
7567        static WARP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7568        let warp_on = *WARP_ON.get_or_init(|| {
7569            std::env::var("MEMRA_QKVNORM_W")
7570                .map(|v| v != "0")
7571                .unwrap_or(true)
7572        });
7573        // rows >= 64 keeps decode (nh + 2*nkv rows) on the block-tree kernel — decode/verify/
7574        // replay numerics are untouched on every model; only prefill depth takes the new config.
7575        if warp_on && ncols % 4 == 0 && rq + 2 * rk >= 64 {
7576            let f = self.func("rms_norm_qkv_w4_f32");
7577            let rows = (rq + 2 * rk) as u32;
7578            let cfg = LaunchConfig {
7579                grid_dim: (rows.div_ceil(8), 1, 1),
7580                block_dim: (256, 1, 1),
7581                shared_mem_bytes: 0,
7582            };
7583            let (nc, rqi, rki, rvi, e) = (ncols as i32, rq as i32, rk as i32, rk as i32, eps);
7584            let __s_b = self.gpu.stream();
7585            let mut b = __s_b.launch_builder(&f);
7586            b.arg(q)
7587                .arg(k)
7588                .arg(v)
7589                .arg(wq)
7590                .arg(wk)
7591                .arg(wv)
7592                .arg(dq)
7593                .arg(dk)
7594                .arg(dv)
7595                .arg(&nc)
7596                .arg(&rqi)
7597                .arg(&rki)
7598                .arg(&rvi)
7599                .arg(&e);
7600            unsafe {
7601                b.launch(cfg)?;
7602            }
7603            return Ok(());
7604        }
7605        let f = self.func("rms_norm_qkv_f32");
7606        let grid = (rq + 2 * rk) as u32;
7607        let cfg = LaunchConfig {
7608            grid_dim: (grid, 1, 1),
7609            block_dim: (rms_block(), 1, 1),
7610            shared_mem_bytes: 0,
7611        };
7612        let (nc, rqi, rki, e) = (ncols as i32, rq as i32, rk as i32, eps);
7613        let __s_b = self.gpu.stream();
7614        let mut b = __s_b.launch_builder(&f);
7615        b.arg(q)
7616            .arg(k)
7617            .arg(v)
7618            .arg(wq)
7619            .arg(wk)
7620            .arg(wv)
7621            .arg(dq)
7622            .arg(dk)
7623            .arg(dv)
7624            .arg(&nc)
7625            .arg(&rqi)
7626            .arg(&rki)
7627            .arg(&e);
7628        unsafe {
7629            b.launch(cfg)?;
7630        }
7631        Ok(())
7632    }
7633
7634    /// gemma4 fused pair of rms_norms over two different inputs (same width).
7635    #[allow(clippy::too_many_arguments)]
7636    pub fn rms_norm2x(
7637        &self,
7638        a: &CudaSlice<f32>,
7639        bb: &CudaSlice<f32>,
7640        wa: &CudaSlice<f32>,
7641        wb: &CudaSlice<f32>,
7642        da: &mut CudaSlice<f32>,
7643        db: &mut CudaSlice<f32>,
7644        ncols: usize,
7645        nrows: usize,
7646        eps: f32,
7647    ) -> Result<(), Box<dyn std::error::Error>> {
7648        let f = self.func("rms_norm2x_f32");
7649        let cfg = LaunchConfig {
7650            grid_dim: (2 * nrows as u32, 1, 1),
7651            block_dim: (rms_block(), 1, 1),
7652            shared_mem_bytes: 0,
7653        };
7654        let (nc, nr, e) = (ncols as i32, nrows as i32, eps);
7655        let __s_b = self.gpu.stream();
7656        let mut b = __s_b.launch_builder(&f);
7657        b.arg(a)
7658            .arg(bb)
7659            .arg(wa)
7660            .arg(wb)
7661            .arg(da)
7662            .arg(db)
7663            .arg(&nc)
7664            .arg(&nr)
7665            .arg(&e);
7666        unsafe {
7667            b.launch(cfg)?;
7668        }
7669        Ok(())
7670    }
7671
7672    /// gemma4 R4: in-place final-logit softcap y = cap*tanh(y/cap).
7673    pub fn softcap(
7674        &self,
7675        y: &mut CudaSlice<f32>,
7676        cap: f32,
7677        n: usize,
7678    ) -> Result<(), Box<dyn std::error::Error>> {
7679        let f = self.func("softcap_f32");
7680        let cfg = LaunchConfig::for_num_elems(n as u32);
7681        let ni = n as i32;
7682        let __s_b = self.gpu.stream();
7683        let mut b = __s_b.launch_builder(&f);
7684        b.arg(y).arg(&cap).arg(&ni);
7685        unsafe {
7686            b.launch(cfg)?;
7687        }
7688        Ok(())
7689    }
7690
7691    /// gemma4 suppress-token mask: y[row][ids[j]] = -inf over t logits rows (fixed-arg launch —
7692    /// graph-capture safe; NOT monotonic like softcap, so it must run before any argmax).
7693    pub fn mask_ids_rows(
7694        &self,
7695        y: &mut CudaSlice<f32>,
7696        ids: &CudaSlice<i32>,
7697        n_ids: usize,
7698        n_vocab: usize,
7699        t: usize,
7700    ) -> Result<(), Box<dyn std::error::Error>> {
7701        let f = self.func("mask_ids_rows_f32");
7702        let cfg = LaunchConfig::for_num_elems((n_ids * t) as u32);
7703        let (ni, nv, ti) = (n_ids as i32, n_vocab as i32, t as i32);
7704        let __s_b = self.gpu.stream();
7705        let mut b = __s_b.launch_builder(&f);
7706        b.arg(y).arg(ids).arg(&ni).arg(&nv).arg(&ti);
7707        unsafe {
7708            b.launch(cfg)?;
7709        }
7710        Ok(())
7711    }
7712
7713    /// gemma4: res = (a+b)*c AND dst = rms_norm(res, w) in one launch.
7714    #[allow(clippy::too_many_arguments)]
7715    pub fn add_scale_rms_norm(
7716        &self,
7717        a: &CudaSlice<f32>,
7718        b_in: &CudaSlice<f32>,
7719        c: f32,
7720        w: &CudaSlice<f32>,
7721        res: &mut CudaSlice<f32>,
7722        dst: &mut CudaSlice<f32>,
7723        ncols: usize,
7724        nrows: usize,
7725        eps: f32,
7726    ) -> Result<(), Box<dyn std::error::Error>> {
7727        let f = self.func("add_scale_rms_norm_f32");
7728        let cfg = LaunchConfig {
7729            grid_dim: (nrows as u32, 1, 1),
7730            block_dim: (rms_block(), 1, 1),
7731            shared_mem_bytes: 0,
7732        };
7733        let (nc, e2) = (ncols as i32, eps);
7734        let __s_b = self.gpu.stream();
7735        let mut b = __s_b.launch_builder(&f);
7736        b.arg(a)
7737            .arg(b_in)
7738            .arg(&c)
7739            .arg(w)
7740            .arg(res)
7741            .arg(dst)
7742            .arg(&nc)
7743            .arg(&e2);
7744        unsafe {
7745            b.launch(cfg)?;
7746        }
7747        Ok(())
7748    }
7749
7750    /// gemma4: res = (a+b)*c AND the next layer's attn_norm EMITTED q8_1 in one launch.
7751    /// Quantize epilogue bit-identical to quantize_q8_1 (the rms_norm_q8_1 form).
7752    #[allow(clippy::too_many_arguments)]
7753    pub fn add_scale_rms_norm_q8_1(
7754        &self,
7755        a: &CudaSlice<f32>,
7756        b_in: &CudaSlice<f32>,
7757        c: f32,
7758        w: &CudaSlice<f32>,
7759        res: &mut CudaSlice<f32>,
7760        ncols: usize,
7761        nrows: usize,
7762        eps: f32,
7763    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7764        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
7765        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
7766        let (nc, e2) = (ncols as i32, eps);
7767        if Self::pdl_on() && Self::pdl_wb_on() {
7768            {
7769                use cudarc::driver::{DevicePtr, DevicePtrMut};
7770                let s = &self.gpu.stream();
7771                let (pa, _g0) = a.device_ptr(s);
7772                let (pb, _g1) = b_in.device_ptr(s);
7773                let (pw, _g2) = w.device_ptr(s);
7774                let (pr, _g3) = res.device_ptr_mut(s);
7775                let (pq, _g4) = out_q.device_ptr_mut(s);
7776                let (pd, _g5) = out_d.device_ptr_mut(s);
7777                let mut ps = [
7778                    &pa as *const _ as *mut std::ffi::c_void,
7779                    &pb as *const _ as *mut _,
7780                    &c as *const _ as *mut _,
7781                    &pw as *const _ as *mut _,
7782                    &pr as *const _ as *mut _,
7783                    &pq as *const _ as *mut _,
7784                    &pd as *const _ as *mut _,
7785                    &nc as *const _ as *mut _,
7786                    &e2 as *const _ as *mut _,
7787                ];
7788                unsafe {
7789                    self.launch_pdl(
7790                        "add_scale_rms_norm_q8_1",
7791                        (nrows as u32, 1, 1),
7792                        (rms_block(), 1, 1),
7793                        &mut ps,
7794                    )?;
7795                }
7796            }
7797            return Ok((out_q, out_d));
7798        }
7799        let f = self.func("add_scale_rms_norm_q8_1");
7800        let cfg = LaunchConfig {
7801            grid_dim: (nrows as u32, 1, 1),
7802            block_dim: (rms_block(), 1, 1),
7803            shared_mem_bytes: 0,
7804        };
7805        let __s_b = self.gpu.stream();
7806        let mut b = __s_b.launch_builder(&f);
7807        b.arg(a)
7808            .arg(b_in)
7809            .arg(&c)
7810            .arg(w)
7811            .arg(res)
7812            .arg(&mut out_q)
7813            .arg(&mut out_d)
7814            .arg(&nc)
7815            .arg(&e2);
7816        unsafe {
7817            b.launch(cfg)?;
7818        }
7819        Ok((out_q, out_d))
7820    }
7821
7822    /// Slot-fed add_scale_rms_norm_q8_1 twin (alloc-free capture lane).
7823    #[allow(clippy::too_many_arguments)]
7824    pub fn add_scale_rms_norm_q8_1_into(
7825        &self,
7826        a: &CudaSlice<f32>,
7827        b_in: &CudaSlice<f32>,
7828        c: f32,
7829        w: &CudaSlice<f32>,
7830        res: &mut CudaSlice<f32>,
7831        ncols: usize,
7832        nrows: usize,
7833        eps: f32,
7834        out_q: &mut CudaSlice<i8>,
7835        out_d: &mut CudaSlice<f32>,
7836    ) -> Result<(), Box<dyn std::error::Error>> {
7837        debug_assert!(out_q.len() >= nrows * ncols && out_d.len() >= nrows * (ncols / 32));
7838        let (nc, e2) = (ncols as i32, eps);
7839        if Self::pdl_on() && Self::pdl_wb_on() {
7840            use cudarc::driver::{DevicePtr, DevicePtrMut};
7841            let s = &self.gpu.stream();
7842            let (pa, _g0) = a.device_ptr(s);
7843            let (pb, _g1) = b_in.device_ptr(s);
7844            let (pw, _g2) = w.device_ptr(s);
7845            let (pr, _g3) = res.device_ptr_mut(s);
7846            let (pq, _g4) = out_q.device_ptr_mut(s);
7847            let (pd, _g5) = out_d.device_ptr_mut(s);
7848            let mut ps = [
7849                &pa as *const _ as *mut std::ffi::c_void,
7850                &pb as *const _ as *mut _,
7851                &c as *const _ as *mut _,
7852                &pw as *const _ as *mut _,
7853                &pr as *const _ as *mut _,
7854                &pq as *const _ as *mut _,
7855                &pd as *const _ as *mut _,
7856                &nc as *const _ as *mut _,
7857                &e2 as *const _ as *mut _,
7858            ];
7859            unsafe {
7860                self.launch_pdl(
7861                    "add_scale_rms_norm_q8_1",
7862                    (nrows as u32, 1, 1),
7863                    (rms_block(), 1, 1),
7864                    &mut ps,
7865                )?;
7866            }
7867            return Ok(());
7868        }
7869        let f = self.func("add_scale_rms_norm_q8_1");
7870        let cfg = LaunchConfig {
7871            grid_dim: (nrows as u32, 1, 1),
7872            block_dim: (rms_block(), 1, 1),
7873            shared_mem_bytes: 0,
7874        };
7875        let __s_b = self.gpu.stream();
7876        let mut b = __s_b.launch_builder(&f);
7877        b.arg(a)
7878            .arg(b_in)
7879            .arg(&c)
7880            .arg(w)
7881            .arg(res)
7882            .arg(&mut *out_q)
7883            .arg(&mut *out_d)
7884            .arg(&nc)
7885            .arg(&e2);
7886        unsafe {
7887            b.launch(cfg)?;
7888        }
7889        Ok(())
7890    }
7891
7892    /// E4B glue fusion: rms(a, wa) prologue + the add_scale_rms_norm_q8_1 program — one launch
7893    /// replaces the per-layer rms_norm_f32(y) + emit pair in the PLE tail.
7894    #[allow(clippy::too_many_arguments)]
7895    pub fn rms_pre_add_scale_rms_norm_q8_1(
7896        &self,
7897        a: &CudaSlice<f32>,
7898        wa: &CudaSlice<f32>,
7899        b_in: &CudaSlice<f32>,
7900        c: f32,
7901        w: &CudaSlice<f32>,
7902        res: &mut CudaSlice<f32>,
7903        ncols: usize,
7904        nrows: usize,
7905        eps: f32,
7906    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7907        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
7908        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
7909        let (nc, e2) = (ncols as i32, eps);
7910        if Self::pdl_on() {
7911            {
7912                use cudarc::driver::{DevicePtr, DevicePtrMut};
7913                let s = &self.gpu.stream();
7914                let (pa, _g0) = a.device_ptr(s);
7915                let (pwa, _g1) = wa.device_ptr(s);
7916                let (pb, _g2) = b_in.device_ptr(s);
7917                let (pw, _g3) = w.device_ptr(s);
7918                let (pr, _g4) = res.device_ptr_mut(s);
7919                let (pq, _g5) = out_q.device_ptr_mut(s);
7920                let (pd, _g6) = out_d.device_ptr_mut(s);
7921                let mut ps = [
7922                    &pa as *const _ as *mut std::ffi::c_void,
7923                    &pwa as *const _ as *mut _,
7924                    &pb as *const _ as *mut _,
7925                    &c as *const _ as *mut _,
7926                    &pw as *const _ as *mut _,
7927                    &pr as *const _ as *mut _,
7928                    &pq as *const _ as *mut _,
7929                    &pd as *const _ as *mut _,
7930                    &nc as *const _ as *mut _,
7931                    &e2 as *const _ as *mut _,
7932                ];
7933                unsafe {
7934                    self.launch_pdl(
7935                        "rms_pre_add_scale_rms_norm_q8_1",
7936                        (nrows as u32, 1, 1),
7937                        (rms_block(), 1, 1),
7938                        &mut ps,
7939                    )?;
7940                }
7941            }
7942            return Ok((out_q, out_d));
7943        }
7944        let f = self.func("rms_pre_add_scale_rms_norm_q8_1");
7945        let cfg = LaunchConfig {
7946            grid_dim: (nrows as u32, 1, 1),
7947            block_dim: (rms_block(), 1, 1),
7948            shared_mem_bytes: 0,
7949        };
7950        let __s_b = self.gpu.stream();
7951        let mut b = __s_b.launch_builder(&f);
7952        b.arg(a)
7953            .arg(wa)
7954            .arg(b_in)
7955            .arg(&c)
7956            .arg(w)
7957            .arg(res)
7958            .arg(&mut out_q)
7959            .arg(&mut out_d)
7960            .arg(&nc)
7961            .arg(&e2);
7962        unsafe {
7963            b.launch(cfg)?;
7964        }
7965        Ok((out_q, out_d))
7966    }
7967
7968    /// GELU(tanh)*up with the activation emitted q8_1 alongside f32 (glue-fusion lane): the
7969    /// consumer matmul rides matmul_pre, killing its standalone quantize_q8_1 launch.
7970    pub fn gelu_tanh_mul_q8_1(
7971        &self,
7972        gate: &CudaSlice<f32>,
7973        up: &cudarc::driver::CudaView<f32>,
7974        act: &mut CudaSlice<f32>,
7975        ncols: usize,
7976        nrows: usize,
7977    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7978        debug_assert!(ncols % 128 == 0);
7979        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
7980        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
7981        let nc = ncols as i32;
7982        if Self::pdl_on() {
7983            {
7984                use cudarc::driver::{DevicePtr, DevicePtrMut};
7985                let s = &self.gpu.stream();
7986                let (pg, _g0) = gate.device_ptr(s);
7987                let (pu, _g1) = up.device_ptr(s);
7988                let (pact, _g2) = act.device_ptr_mut(s);
7989                let (pq, _g3) = out_q.device_ptr_mut(s);
7990                let (pd, _g4) = out_d.device_ptr_mut(s);
7991                let mut ps = [
7992                    &pg as *const _ as *mut std::ffi::c_void,
7993                    &pu as *const _ as *mut _,
7994                    &pact as *const _ as *mut _,
7995                    &pq as *const _ as *mut _,
7996                    &pd as *const _ as *mut _,
7997                    &nc as *const _ as *mut _,
7998                ];
7999                unsafe {
8000                    self.launch_pdl(
8001                        "gelu_tanh_mul_q8_1",
8002                        (nrows as u32, 1, 1),
8003                        (rms_block(), 1, 1),
8004                        &mut ps,
8005                    )?;
8006                }
8007            }
8008            return Ok((out_q, out_d));
8009        }
8010        let f = self.func("gelu_tanh_mul_q8_1");
8011        let cfg = LaunchConfig {
8012            grid_dim: (nrows as u32, 1, 1),
8013            block_dim: (rms_block(), 1, 1),
8014            shared_mem_bytes: 0,
8015        };
8016        let __s_b = self.gpu.stream();
8017        let mut b = __s_b.launch_builder(&f);
8018        b.arg(gate)
8019            .arg(up)
8020            .arg(act)
8021            .arg(&mut out_q)
8022            .arg(&mut out_d)
8023            .arg(&nc);
8024        unsafe {
8025            b.launch(cfg)?;
8026        }
8027        Ok((out_q, out_d))
8028    }
8029
8030    /// Slot-fed gelu_tanh_mul_q8_1 twin (alloc-free capture lane; incl. the PDL arm).
8031    #[allow(clippy::too_many_arguments)]
8032    pub fn gelu_tanh_mul_q8_1_into(
8033        &self,
8034        gate: &CudaSlice<f32>,
8035        up: &cudarc::driver::CudaView<f32>,
8036        act: &mut CudaSlice<f32>,
8037        ncols: usize,
8038        nrows: usize,
8039        out_q: &mut CudaSlice<i8>,
8040        out_d: &mut CudaSlice<f32>,
8041    ) -> Result<(), Box<dyn std::error::Error>> {
8042        debug_assert!(ncols % 128 == 0);
8043        debug_assert!(out_q.len() >= nrows * ncols && out_d.len() >= nrows * (ncols / 32));
8044        let nc = ncols as i32;
8045        if Self::pdl_on() {
8046            use cudarc::driver::{DevicePtr, DevicePtrMut};
8047            let s = &self.gpu.stream();
8048            let (pg, _g0) = gate.device_ptr(s);
8049            let (pu, _g1) = up.device_ptr(s);
8050            let (pact, _g2) = act.device_ptr_mut(s);
8051            let (pq, _g3) = out_q.device_ptr_mut(s);
8052            let (pd, _g4) = out_d.device_ptr_mut(s);
8053            let mut ps = [
8054                &pg as *const _ as *mut std::ffi::c_void,
8055                &pu as *const _ as *mut _,
8056                &pact as *const _ as *mut _,
8057                &pq as *const _ as *mut _,
8058                &pd as *const _ as *mut _,
8059                &nc as *const _ as *mut _,
8060            ];
8061            unsafe {
8062                self.launch_pdl(
8063                    "gelu_tanh_mul_q8_1",
8064                    (nrows as u32, 1, 1),
8065                    (rms_block(), 1, 1),
8066                    &mut ps,
8067                )?;
8068            }
8069            return Ok(());
8070        }
8071        let f = self.func("gelu_tanh_mul_q8_1");
8072        let cfg = LaunchConfig {
8073            grid_dim: (nrows as u32, 1, 1),
8074            block_dim: (rms_block(), 1, 1),
8075            shared_mem_bytes: 0,
8076        };
8077        let __s_b = self.gpu.stream();
8078        let mut b = __s_b.launch_builder(&f);
8079        b.arg(gate)
8080            .arg(up)
8081            .arg(&mut *act)
8082            .arg(&mut *out_q)
8083            .arg(&mut *out_d)
8084            .arg(&nc);
8085        unsafe {
8086            b.launch(cfg)?;
8087        }
8088        Ok(())
8089    }
8090
8091    /// gemma4: add + rms_norm3 with outputs 0/2 emitted q8_1 (zsh + moe_in) and 1 f32 (router).
8092    #[allow(clippy::too_many_arguments)]
8093    pub fn add_rms_norm3_q8z(
8094        &self,
8095        a: &CudaSlice<f32>,
8096        b_in: &CudaSlice<f32>,
8097        w0: &CudaSlice<f32>,
8098        w1: &CudaSlice<f32>,
8099        w2: &CudaSlice<f32>,
8100        res: &mut CudaSlice<f32>,
8101        out1: &mut CudaSlice<f32>,
8102        ncols: usize,
8103        nrows: usize,
8104        eps: f32,
8105    ) -> Result<
8106        (
8107            (CudaSlice<i8>, CudaSlice<f32>),
8108            (CudaSlice<i8>, CudaSlice<f32>),
8109        ),
8110        Box<dyn std::error::Error>,
8111    > {
8112        let mut q0 = self.alloc_uninit::<i8>(nrows * ncols)?;
8113        let mut d0 = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
8114        let mut q2 = self.alloc_uninit::<i8>(nrows * ncols)?;
8115        let mut d2 = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
8116        let f = self.func("add_rms_norm3_q8z_f32");
8117        let cfg = LaunchConfig {
8118            grid_dim: (nrows as u32, 1, 1),
8119            block_dim: (rms_block(), 1, 1),
8120            shared_mem_bytes: 0,
8121        };
8122        let (nc, e2) = (ncols as i32, eps);
8123        let __s_b = self.gpu.stream();
8124        let mut b = __s_b.launch_builder(&f);
8125        b.arg(a)
8126            .arg(b_in)
8127            .arg(w0)
8128            .arg(w1)
8129            .arg(w2)
8130            .arg(res)
8131            .arg(&mut q0)
8132            .arg(&mut d0)
8133            .arg(out1)
8134            .arg(&mut q2)
8135            .arg(&mut d2)
8136            .arg(&nc)
8137            .arg(&e2);
8138        unsafe {
8139            b.launch(cfg)?;
8140        }
8141        Ok(((q0, d0), (q2, d2)))
8142    }
8143
8144    /// gemma4: res = a+b AND the three rms_norms of res in one launch.
8145    #[allow(clippy::too_many_arguments)]
8146    pub fn add_rms_norm3(
8147        &self,
8148        a: &CudaSlice<f32>,
8149        b_in: &CudaSlice<f32>,
8150        w0: &CudaSlice<f32>,
8151        w1: &CudaSlice<f32>,
8152        w2: &CudaSlice<f32>,
8153        res: &mut CudaSlice<f32>,
8154        d0: &mut CudaSlice<f32>,
8155        d1: &mut CudaSlice<f32>,
8156        d2: &mut CudaSlice<f32>,
8157        ncols: usize,
8158        nrows: usize,
8159        eps: f32,
8160    ) -> Result<(), Box<dyn std::error::Error>> {
8161        let f = self.func("add_rms_norm3_f32");
8162        let cfg = LaunchConfig {
8163            grid_dim: (nrows as u32, 1, 1),
8164            block_dim: (rms_block(), 1, 1),
8165            shared_mem_bytes: 0,
8166        };
8167        let (nc, e2) = (ncols as i32, eps);
8168        let __s_b = self.gpu.stream();
8169        let mut b = __s_b.launch_builder(&f);
8170        b.arg(a)
8171            .arg(b_in)
8172            .arg(w0)
8173            .arg(w1)
8174            .arg(w2)
8175            .arg(res)
8176            .arg(d0)
8177            .arg(d1)
8178            .arg(d2)
8179            .arg(&nc)
8180            .arg(&e2);
8181        unsafe {
8182            b.launch(cfg)?;
8183        }
8184        Ok(())
8185    }
8186
8187    /// dst = (a + b) * c (residual add + layer scale, one launch).
8188    pub fn add_scale(
8189        &self,
8190        a: &CudaSlice<f32>,
8191        b_in: &CudaSlice<f32>,
8192        c: f32,
8193        dst: &mut CudaSlice<f32>,
8194        n: usize,
8195    ) -> Result<(), Box<dyn std::error::Error>> {
8196        let f = self.func("add_scale_f32");
8197        let cfg = LaunchConfig::for_num_elems(n as u32);
8198        let ni = n as i32;
8199        let __s_b = self.gpu.stream();
8200        let mut b = __s_b.launch_builder(&f);
8201        b.arg(a).arg(b_in).arg(&c).arg(dst).arg(&ni);
8202        unsafe {
8203            b.launch(cfg)?;
8204        }
8205        Ok(())
8206    }
8207
8208    /// Vision-tower LayerNorm (with bias) over [nrows, ncols] — lane/vision.
8209    pub fn layer_norm_bias(
8210        &self,
8211        x: &CudaSlice<f32>,
8212        w: &CudaSlice<f32>,
8213        b: &CudaSlice<f32>,
8214        dst: &mut CudaSlice<f32>,
8215        ncols: usize,
8216        nrows: usize,
8217        eps: f32,
8218    ) -> Result<(), Box<dyn std::error::Error>> {
8219        let f = self.func("layer_norm_bias_f32");
8220        let (nc, e) = (ncols as i32, eps);
8221        let cfg = LaunchConfig {
8222            grid_dim: (nrows as u32, 1, 1),
8223            block_dim: (256, 1, 1),
8224            shared_mem_bytes: 0,
8225        };
8226        let __s_b = self.gpu.stream();
8227        let mut lb = __s_b.launch_builder(&f);
8228        lb.arg(x).arg(w).arg(b).arg(&mut *dst).arg(&nc).arg(&e);
8229        unsafe {
8230            lb.launch(cfg)?;
8231        }
8232        Ok(())
8233    }
8234
8235    /// gelu_pytorch_tanh elementwise (vision tower MLP activation).
8236    pub fn gelu_tanh(
8237        &self,
8238        x: &CudaSlice<f32>,
8239        dst: &mut CudaSlice<f32>,
8240        n: usize,
8241    ) -> Result<(), Box<dyn std::error::Error>> {
8242        let f = self.func("gelu_tanh_f32");
8243        let ni = n as i64;
8244        let cfg = LaunchConfig {
8245            grid_dim: (n.div_ceil(256) as u32, 1, 1),
8246            block_dim: (256, 1, 1),
8247            shared_mem_bytes: 0,
8248        };
8249        let __s_b = self.gpu.stream();
8250        let mut lb = __s_b.launch_builder(&f);
8251        lb.arg(x).arg(&mut *dst).arg(&ni);
8252        unsafe {
8253            lb.launch(cfg)?;
8254        }
8255        Ok(())
8256    }
8257
8258    /// In-place row softmax over [nrows, ncols] (bidirectional vision attention).
8259    pub fn row_softmax(
8260        &self,
8261        x: &mut CudaSlice<f32>,
8262        ncols: usize,
8263        nrows: usize,
8264    ) -> Result<(), Box<dyn std::error::Error>> {
8265        let f = self.func("row_softmax_f32");
8266        let nc = ncols as i32;
8267        let cfg = LaunchConfig {
8268            grid_dim: (nrows as u32, 1, 1),
8269            block_dim: (256, 1, 1),
8270            shared_mem_bytes: 0,
8271        };
8272        let __s_b = self.gpu.stream();
8273        let mut lb = __s_b.launch_builder(&f);
8274        lb.arg(&mut *x).arg(&nc);
8275        unsafe {
8276            lb.launch(cfg)?;
8277        }
8278        Ok(())
8279    }
8280
8281    pub fn rms_norm(
8282        &self,
8283        x: &CudaSlice<f32>,
8284        w: &CudaSlice<f32>,
8285        dst: &mut CudaSlice<f32>,
8286        ncols: usize,
8287        nrows: usize,
8288        eps: f32,
8289    ) -> Result<(), Box<dyn std::error::Error>> {
8290        let (nc, e) = (ncols as i32, eps);
8291        if Self::pdl_on() && Self::pdl_wb_on() {
8292            use cudarc::driver::{DevicePtr, DevicePtrMut};
8293            let s = &self.gpu.stream();
8294            let (px, _g0) = x.device_ptr(s);
8295            let (pw, _g1) = w.device_ptr(s);
8296            let (pd, _g2) = dst.device_ptr_mut(s);
8297            let mut ps = [
8298                &px as *const _ as *mut std::ffi::c_void,
8299                &pw as *const _ as *mut _,
8300                &pd as *const _ as *mut _,
8301                &nc as *const _ as *mut _,
8302                &e as *const _ as *mut _,
8303            ];
8304            unsafe {
8305                self.launch_pdl(
8306                    "rms_norm_f32",
8307                    (nrows as u32, 1, 1),
8308                    (rms_block(), 1, 1),
8309                    &mut ps,
8310                )?;
8311            }
8312            return Ok(());
8313        }
8314        let f = self.func("rms_norm_f32");
8315        let cfg = LaunchConfig {
8316            grid_dim: (nrows as u32, 1, 1),
8317            block_dim: (rms_block(), 1, 1),
8318            shared_mem_bytes: 0,
8319        };
8320        let __s_b = self.gpu.stream();
8321        let mut b = __s_b.launch_builder(&f);
8322        b.arg(x).arg(w).arg(dst).arg(&nc).arg(&e);
8323        unsafe {
8324            b.launch(cfg)?;
8325        }
8326        Ok(())
8327    }
8328
8329    /// RMS-norm with blockDim=1024 — BIT-IDENTICAL to the fused `rms_norm_q8_1` and
8330    /// `add_rms_norm_q8_1` kernels' sum-of-squares reduction. The spec verify path MUST use this
8331    /// to match decode's FP accumulation order: the standard `rms_norm` at blockDim=256 has a
8332    /// different per-thread stride (ncols/256 partials vs ncols/1024 partials) and therefore a
8333    /// different shfl-tree reduction that can shift `scale = rsqrt(sum/n + eps)` by ULPs, causing
8334    /// divergence through the GDN scan and argmax flips on the 9B text prompt. The underlying
8335    /// `rms_norm_f32` kernel supports any blockDim (generic reduce with shared[32]).
8336    pub fn rms_norm_decode(
8337        &self,
8338        x: &CudaSlice<f32>,
8339        w: &CudaSlice<f32>,
8340        dst: &mut CudaSlice<f32>,
8341        ncols: usize,
8342        nrows: usize,
8343        eps: f32,
8344    ) -> Result<(), Box<dyn std::error::Error>> {
8345        let f = self.func("rms_norm_f32");
8346        let cfg = LaunchConfig {
8347            grid_dim: (nrows as u32, 1, 1),
8348            block_dim: (1024, 1, 1),
8349            shared_mem_bytes: 0,
8350        };
8351        let (nc, e) = (ncols as i32, eps);
8352        let __s_b = self.gpu.stream();
8353        let mut b = __s_b.launch_builder(&f);
8354        b.arg(x).arg(w).arg(dst).arg(&nc).arg(&e);
8355        unsafe {
8356            b.launch(cfg)?;
8357        }
8358        Ok(())
8359    }
8360
8361    /// DECODE GLUE-FUSION LEVER: `z = rms_norm(x)*w` emitted DIRECTLY as q8_1 (no f32 `z` materialized,
8362    /// no standalone quantize_q8_1 launch). Returns (out_q [nrows*ncols i8], out_d [nrows*nblk f32])
8363    /// ready to feed matmul_pre. BIT-IDENTICAL to rms_norm + quantize_q8_1. ncols % 32 == 0.
8364    pub fn rms_norm_q8_1(
8365        &self,
8366        x: &CudaSlice<f32>,
8367        w: &CudaSlice<f32>,
8368        ncols: usize,
8369        nrows: usize,
8370        eps: f32,
8371    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8372        let nblk = ncols / 32;
8373        let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
8374        let mut d = self.alloc_uninit::<f32>(nrows * nblk)?;
8375        let (nc, e) = (ncols as i32, eps);
8376        if Self::pdl_on() {
8377            {
8378                use cudarc::driver::{DevicePtr, DevicePtrMut};
8379                let s = &self.gpu.stream();
8380                let (px, _g0) = x.device_ptr(s);
8381                let (pw, _g1) = w.device_ptr(s);
8382                let (pq, _g2) = q.device_ptr_mut(s);
8383                let (pd, _g3) = d.device_ptr_mut(s);
8384                let mut ps = [
8385                    &px as *const _ as *mut std::ffi::c_void,
8386                    &pw as *const _ as *mut _,
8387                    &pq as *const _ as *mut _,
8388                    &pd as *const _ as *mut _,
8389                    &nc as *const _ as *mut _,
8390                    &e as *const _ as *mut _,
8391                ];
8392                unsafe {
8393                    self.launch_pdl("rms_norm_q8_1", (nrows as u32, 1, 1), (1024, 1, 1), &mut ps)?;
8394                }
8395            }
8396            return Ok((q, d));
8397        }
8398        let f = self.func("rms_norm_q8_1");
8399        // 1024 threads: decode is nrows=1 -> ONE CTA; 32 warps hide the pass1->pass2 latency
8400        // (s[32] reduce already sized for 32 warps). Same shape math at any blockDim.
8401        let cfg = LaunchConfig {
8402            grid_dim: (nrows as u32, 1, 1),
8403            block_dim: (1024, 1, 1),
8404            shared_mem_bytes: 0,
8405        };
8406        let __s_b = self.gpu.stream();
8407        let mut b = __s_b.launch_builder(&f);
8408        b.arg(x).arg(w).arg(&mut q).arg(&mut d).arg(&nc).arg(&e);
8409        unsafe {
8410            b.launch(cfg)?;
8411        }
8412        Ok((q, d))
8413    }
8414
8415    /// Slot-fed rms_norm_q8_1 twin (alloc-free capture lane): identical launch (incl. the
8416    /// PDL arm), caller-owned outputs.
8417    pub fn rms_norm_q8_1_into(
8418        &self,
8419        x: &CudaSlice<f32>,
8420        w: &CudaSlice<f32>,
8421        ncols: usize,
8422        nrows: usize,
8423        eps: f32,
8424        q: &mut CudaSlice<i8>,
8425        d: &mut CudaSlice<f32>,
8426    ) -> Result<(), Box<dyn std::error::Error>> {
8427        let nblk = ncols / 32;
8428        debug_assert!(q.len() >= nrows * ncols && d.len() >= nrows * nblk);
8429        let (nc, e) = (ncols as i32, eps);
8430        if Self::pdl_on() {
8431            use cudarc::driver::{DevicePtr, DevicePtrMut};
8432            let s = &self.gpu.stream();
8433            let (px, _g0) = x.device_ptr(s);
8434            let (pw, _g1) = w.device_ptr(s);
8435            let (pq, _g2) = q.device_ptr_mut(s);
8436            let (pd, _g3) = d.device_ptr_mut(s);
8437            let mut ps = [
8438                &px as *const _ as *mut std::ffi::c_void,
8439                &pw as *const _ as *mut _,
8440                &pq as *const _ as *mut _,
8441                &pd as *const _ as *mut _,
8442                &nc as *const _ as *mut _,
8443                &e as *const _ as *mut _,
8444            ];
8445            unsafe {
8446                self.launch_pdl("rms_norm_q8_1", (nrows as u32, 1, 1), (1024, 1, 1), &mut ps)?;
8447            }
8448            return Ok(());
8449        }
8450        let f = self.func("rms_norm_q8_1");
8451        let cfg = LaunchConfig {
8452            grid_dim: (nrows as u32, 1, 1),
8453            block_dim: (1024, 1, 1),
8454            shared_mem_bytes: 0,
8455        };
8456        let __s_b = self.gpu.stream();
8457        let mut b = __s_b.launch_builder(&f);
8458        b.arg(x).arg(w).arg(&mut *q).arg(&mut *d).arg(&nc).arg(&e);
8459        unsafe {
8460            b.launch(cfg)?;
8461        }
8462        Ok(())
8463    }
8464
8465    /// Slot-fed quantize_q8_1 twin (alloc-free capture lane).
8466    pub fn quantize_q8_1_into(
8467        &self,
8468        x: &CudaSlice<f32>,
8469        m: usize,
8470        in_f: usize,
8471        q: &mut CudaSlice<i8>,
8472        d: &mut CudaSlice<f32>,
8473    ) -> Result<(), Box<dyn std::error::Error>> {
8474        let nblk = in_f / 32;
8475        debug_assert!(q.len() >= m * in_f && d.len() >= m * nblk);
8476        let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
8477        let (inf, mi) = (in_f as i32, m as i32);
8478        if Self::pdl_on() && Self::pdl_wb_on() {
8479            use cudarc::driver::{DevicePtr, DevicePtrMut};
8480            let s = &self.gpu.stream();
8481            let (px, _g0) = x.device_ptr(s);
8482            let (pq, _g1) = q.device_ptr_mut(s);
8483            let (pd, _g2) = d.device_ptr_mut(s);
8484            let mut ps = [
8485                &px as *const _ as *mut std::ffi::c_void,
8486                &pq as *const _ as *mut _,
8487                &pd as *const _ as *mut _,
8488                &inf as *const _ as *mut _,
8489                &mi as *const _ as *mut _,
8490            ];
8491            unsafe {
8492                self.launch_pdl("quantize_q8_1", cfg.grid_dim, cfg.block_dim, &mut ps)?;
8493            }
8494            return Ok(());
8495        }
8496        let f = self.func("quantize_q8_1");
8497        let __s_b = self.gpu.stream();
8498        let mut b = __s_b.launch_builder(&f);
8499        b.arg(x).arg(&mut *q).arg(&mut *d).arg(&inf).arg(&mi);
8500        unsafe {
8501            b.launch(cfg)?;
8502        }
8503        Ok(())
8504    }
8505
8506    /// DECODE GLUE-FUSION LEVER: `res = a+b; z = rms_norm(res)*w` with z emitted as q8_1. `res` is
8507    /// still written (the post-ffn residual add reads it). Fuses add_rms_norm + quantize_q8_1.
8508    /// Returns (out_q, out_d) for matmul_pre. BIT-IDENTICAL. ncols % 32 == 0.
8509    pub fn add_rms_norm_q8_1(
8510        &self,
8511        a: &CudaSlice<f32>,
8512        b_in: &CudaSlice<f32>,
8513        w: &CudaSlice<f32>,
8514        res: &mut CudaSlice<f32>,
8515        ncols: usize,
8516        nrows: usize,
8517        eps: f32,
8518    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8519        let nblk = ncols / 32;
8520        let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
8521        let mut d = self.alloc_uninit::<f32>(nrows * nblk)?;
8522        let f = self.func("add_rms_norm_q8_1");
8523        // 1024 threads: same single-CTA-at-decode reasoning as rms_norm_q8_1.
8524        let cfg = LaunchConfig {
8525            grid_dim: (nrows as u32, 1, 1),
8526            block_dim: (1024, 1, 1),
8527            shared_mem_bytes: 0,
8528        };
8529        let (nc, e) = (ncols as i32, eps);
8530        let __s_bld = self.gpu.stream();
8531        let mut bld = __s_bld.launch_builder(&f);
8532        bld.arg(a)
8533            .arg(b_in)
8534            .arg(w)
8535            .arg(res)
8536            .arg(&mut q)
8537            .arg(&mut d)
8538            .arg(&nc)
8539            .arg(&e);
8540        unsafe {
8541            bld.launch(cfg)?;
8542        }
8543        Ok((q, d))
8544    }
8545
8546    /// RANK3 LEVER (add+rmsnorm fuse): `res = a + b; dst = rms_norm(res) * w` in ONE launch. Fuses
8547    /// e.add(a,b,res) + e.rms_norm(res,w,dst), removing one launch + one HBM read of the residual per
8548    /// residual+norm pair. BIT-IDENTICAL to the two-kernel sequence (same IEEE add, same reduction).
8549    pub fn add_rms_norm(
8550        &self,
8551        a: &CudaSlice<f32>,
8552        b: &CudaSlice<f32>,
8553        w: &CudaSlice<f32>,
8554        res: &mut CudaSlice<f32>,
8555        dst: &mut CudaSlice<f32>,
8556        ncols: usize,
8557        nrows: usize,
8558        eps: f32,
8559    ) -> Result<(), Box<dyn std::error::Error>> {
8560        let (nc, e) = (ncols as i32, eps);
8561        if Self::pdl_on() && Self::pdl_wb_on() {
8562            use cudarc::driver::{DevicePtr, DevicePtrMut};
8563            let s = &self.gpu.stream();
8564            let (pa, _g0) = a.device_ptr(s);
8565            let (pb, _g1) = b.device_ptr(s);
8566            let (pw, _g2) = w.device_ptr(s);
8567            let (pr, _g3) = res.device_ptr_mut(s);
8568            let (pd, _g4) = dst.device_ptr_mut(s);
8569            let mut ps = [
8570                &pa as *const _ as *mut std::ffi::c_void,
8571                &pb as *const _ as *mut _,
8572                &pw as *const _ as *mut _,
8573                &pr as *const _ as *mut _,
8574                &pd as *const _ as *mut _,
8575                &nc as *const _ as *mut _,
8576                &e as *const _ as *mut _,
8577            ];
8578            unsafe {
8579                self.launch_pdl(
8580                    "add_rms_norm_f32",
8581                    (nrows as u32, 1, 1),
8582                    (rms_block(), 1, 1),
8583                    &mut ps,
8584                )?;
8585            }
8586            return Ok(());
8587        }
8588        let f = self.func("add_rms_norm_f32");
8589        let cfg = LaunchConfig {
8590            grid_dim: (nrows as u32, 1, 1),
8591            block_dim: (rms_block(), 1, 1),
8592            shared_mem_bytes: 0,
8593        };
8594        let __s_b2 = self.gpu.stream();
8595        let mut b2 = __s_b2.launch_builder(&f);
8596        b2.arg(a)
8597            .arg(b)
8598            .arg(w)
8599            .arg(&mut *res)
8600            .arg(&mut *dst)
8601            .arg(&nc)
8602            .arg(&e);
8603        unsafe {
8604            b2.launch(cfg)?;
8605        }
8606        Ok(())
8607    }
8608
8609    /// E4B glue fusion: rms(a, wa) prologue + add_rms_norm — folds the post-attn norm into
8610    /// the tail entry (res = rms(a)*wa + b; dst = rms(res)*w).
8611    #[allow(clippy::too_many_arguments)]
8612    pub fn rms_pre_add_rms_norm(
8613        &self,
8614        a: &CudaSlice<f32>,
8615        wa: &CudaSlice<f32>,
8616        b: &CudaSlice<f32>,
8617        w: &CudaSlice<f32>,
8618        res: &mut CudaSlice<f32>,
8619        dst: &mut CudaSlice<f32>,
8620        ncols: usize,
8621        nrows: usize,
8622        eps: f32,
8623    ) -> Result<(), Box<dyn std::error::Error>> {
8624        let f = self.func("rms_pre_add_rms_norm_f32");
8625        let cfg = LaunchConfig {
8626            grid_dim: (nrows as u32, 1, 1),
8627            block_dim: (rms_block(), 1, 1),
8628            shared_mem_bytes: 0,
8629        };
8630        let (nc, e) = (ncols as i32, eps);
8631        let __s_b2 = self.gpu.stream();
8632        let mut b2 = __s_b2.launch_builder(&f);
8633        b2.arg(a)
8634            .arg(wa)
8635            .arg(b)
8636            .arg(w)
8637            .arg(&mut *res)
8638            .arg(&mut *dst)
8639            .arg(&nc)
8640            .arg(&e);
8641        unsafe {
8642            b2.launch(cfg)?;
8643        }
8644        Ok(())
8645    }
8646
8647    /// wave-2 fold: rms(a,wa) + add + ffn-norm with zsh EMITTED q8_1 (fused2 consumes it).
8648    #[allow(clippy::too_many_arguments)]
8649    pub fn rms_pre_add_rms_norm_q8z(
8650        &self,
8651        a: &CudaSlice<f32>,
8652        wa: &CudaSlice<f32>,
8653        b: &CudaSlice<f32>,
8654        w: &CudaSlice<f32>,
8655        res: &mut CudaSlice<f32>,
8656        dst: &mut CudaSlice<f32>,
8657        ncols: usize,
8658        nrows: usize,
8659        eps: f32,
8660    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8661        debug_assert!(ncols % 128 == 0);
8662        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
8663        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
8664        let (nc, e) = (ncols as i32, eps);
8665        if Self::pdl_on() {
8666            {
8667                use cudarc::driver::{DevicePtr, DevicePtrMut};
8668                let s = &self.gpu.stream();
8669                let (pa, _g0) = a.device_ptr(s);
8670                let (pwa, _g1) = wa.device_ptr(s);
8671                let (pb, _g2) = b.device_ptr(s);
8672                let (pw, _g3) = w.device_ptr(s);
8673                let (pr, _g4) = res.device_ptr_mut(s);
8674                let (pdst, _g5) = dst.device_ptr_mut(s);
8675                let (pq, _g6) = out_q.device_ptr_mut(s);
8676                let (pd, _g7) = out_d.device_ptr_mut(s);
8677                let mut ps = [
8678                    &pa as *const _ as *mut std::ffi::c_void,
8679                    &pwa as *const _ as *mut _,
8680                    &pb as *const _ as *mut _,
8681                    &pw as *const _ as *mut _,
8682                    &pr as *const _ as *mut _,
8683                    &pdst as *const _ as *mut _,
8684                    &pq as *const _ as *mut _,
8685                    &pd as *const _ as *mut _,
8686                    &nc as *const _ as *mut _,
8687                    &e as *const _ as *mut _,
8688                ];
8689                unsafe {
8690                    self.launch_pdl(
8691                        "rms_pre_add_rms_norm_q8z_f32",
8692                        (nrows as u32, 1, 1),
8693                        (rms_block(), 1, 1),
8694                        &mut ps,
8695                    )?;
8696                }
8697            }
8698            return Ok((out_q, out_d));
8699        }
8700        let f = self.func("rms_pre_add_rms_norm_q8z_f32");
8701        let cfg = LaunchConfig {
8702            grid_dim: (nrows as u32, 1, 1),
8703            block_dim: (rms_block(), 1, 1),
8704            shared_mem_bytes: 0,
8705        };
8706        let __s_b2 = self.gpu.stream();
8707        let mut b2 = __s_b2.launch_builder(&f);
8708        b2.arg(a)
8709            .arg(wa)
8710            .arg(b)
8711            .arg(w)
8712            .arg(&mut *res)
8713            .arg(&mut *dst)
8714            .arg(&mut out_q)
8715            .arg(&mut out_d)
8716            .arg(&nc)
8717            .arg(&e);
8718        unsafe {
8719            b2.launch(cfg)?;
8720        }
8721        Ok((out_q, out_d))
8722    }
8723
8724    /// Slot-fed twin of `rms_pre_add_rms_norm_q8z` (gemma4 pn-fold, slotted/graph arm):
8725    /// identical kernel, caller-owned outputs, PLAIN launch only — the dc_slotted capture
8726    /// body must stay attribute-free (the fused2_into precedent).
8727    #[allow(clippy::too_many_arguments)]
8728    pub fn rms_pre_add_rms_norm_q8z_into(
8729        &self,
8730        a: &CudaSlice<f32>,
8731        wa: &CudaSlice<f32>,
8732        b: &CudaSlice<f32>,
8733        w: &CudaSlice<f32>,
8734        res: &mut CudaSlice<f32>,
8735        dst: &mut CudaSlice<f32>,
8736        ncols: usize,
8737        nrows: usize,
8738        eps: f32,
8739        out_q: &mut CudaSlice<i8>,
8740        out_d: &mut CudaSlice<f32>,
8741    ) -> Result<(), Box<dyn std::error::Error>> {
8742        debug_assert!(ncols % 128 == 0);
8743        let (nc, e) = (ncols as i32, eps);
8744        let f = self.func("rms_pre_add_rms_norm_q8z_f32");
8745        let cfg = LaunchConfig {
8746            grid_dim: (nrows as u32, 1, 1),
8747            block_dim: (rms_block(), 1, 1),
8748            shared_mem_bytes: 0,
8749        };
8750        let __s_b = self.gpu.stream();
8751        let mut b2 = __s_b.launch_builder(&f);
8752        b2.arg(a)
8753            .arg(wa)
8754            .arg(b)
8755            .arg(w)
8756            .arg(&mut *res)
8757            .arg(&mut *dst)
8758            .arg(&mut *out_q)
8759            .arg(&mut *out_d)
8760            .arg(&nc)
8761            .arg(&e);
8762        unsafe {
8763            b2.launch(cfg)?;
8764        }
8765        Ok(())
8766    }
8767
8768    /// Slot-fed twin of `rms_pre_add_scale_rms_norm_q8_1` (gemma4 pn-fold exit, slotted
8769    /// arm): identical kernel, caller-owned outputs, PLAIN launch (capture-safe).
8770    #[allow(clippy::too_many_arguments)]
8771    pub fn rms_pre_add_scale_rms_norm_q8_1_into(
8772        &self,
8773        a: &CudaSlice<f32>,
8774        wa: &CudaSlice<f32>,
8775        b_in: &CudaSlice<f32>,
8776        c: f32,
8777        w: &CudaSlice<f32>,
8778        res: &mut CudaSlice<f32>,
8779        ncols: usize,
8780        nrows: usize,
8781        eps: f32,
8782        out_q: &mut CudaSlice<i8>,
8783        out_d: &mut CudaSlice<f32>,
8784    ) -> Result<(), Box<dyn std::error::Error>> {
8785        debug_assert!(ncols % 128 == 0);
8786        let (nc, e2) = (ncols as i32, eps);
8787        let f = self.func("rms_pre_add_scale_rms_norm_q8_1");
8788        let cfg = LaunchConfig {
8789            grid_dim: (nrows as u32, 1, 1),
8790            block_dim: (rms_block(), 1, 1),
8791            shared_mem_bytes: 0,
8792        };
8793        let __s_b = self.gpu.stream();
8794        let mut b2 = __s_b.launch_builder(&f);
8795        b2.arg(a)
8796            .arg(wa)
8797            .arg(b_in)
8798            .arg(&c)
8799            .arg(w)
8800            .arg(&mut *res)
8801            .arg(&mut *out_q)
8802            .arg(&mut *out_d)
8803            .arg(&nc)
8804            .arg(&e2);
8805        unsafe {
8806            b2.launch(cfg)?;
8807        }
8808        Ok(())
8809    }
8810
8811    /// gemma4 pn-fold seam (GAP-DIAGNOSIS verdict 7, the E4B glue backport): the dense
8812    /// decode/verify/slotted trio folds post_attn_norm into the tail entry
8813    /// (rms_pre_add_rms_norm[_q8z]) and post_ffw_norm into the residual exit
8814    /// (rms_pre_add_scale_rms_norm_q8_1). BITS-CHANGING vs the two-launch chain (the
8815    /// single-phase reduction's expansion rounding — E4B receipts); every arm moves
8816    /// together so decode == verify == graph parity holds BY CONSTRUCTION within either
8817    /// seam value. MEMRA_G4_PNFOLD=0 restores the unfused chain everywhere.
8818    pub fn g4_pnfold_on() -> bool {
8819        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8820        *ON.get_or_init(|| {
8821            std::env::var("MEMRA_G4_PNFOLD")
8822                .map(|v| v != "0")
8823                .unwrap_or(true)
8824        })
8825    }
8826
8827    /// wave-4b: OUT-dim concat of three Q4_0 tensors (same in_features; rows are independent
8828    /// blocks, so the concat is a D2D byte concat of the GGUF-layout planes). Returns None
8829    /// off-class (non-Q4_0, mismatched widths, or any tensor already rp-swapped in place).
8830    pub fn build_q4_out_concat3(
8831        &self,
8832        w0: &crate::model::GpuTensor,
8833        w1: &crate::model::GpuTensor,
8834        w2: &crate::model::GpuTensor,
8835    ) -> Result<Option<crate::model::GpuTensor>, Box<dyn std::error::Error>> {
8836        use crate::model::GpuTensor;
8837        let part = |w: &GpuTensor| -> Option<(usize, usize)> {
8838            match w {
8839                GpuTensor::Quant {
8840                    qtype,
8841                    row_bytes,
8842                    rp,
8843                    ..
8844                } if *qtype == QT_Q4_0 && !*rp => Some((*row_bytes, w.out_features())),
8845                _ => None,
8846            }
8847        };
8848        let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (part(w0), part(w1), part(w2))
8849        else {
8850            return Ok(None);
8851        };
8852        if rb0 != rb1
8853            || rb0 != rb2
8854            || w0.in_features() != w1.in_features()
8855            || w0.in_features() != w2.in_features()
8856        {
8857            return Ok(None);
8858        }
8859        fn bytes_of(w: &crate::model::GpuTensor) -> &CudaSlice<u8> {
8860            match w {
8861                crate::model::GpuTensor::Quant { bytes, .. } => bytes,
8862                _ => unreachable!(),
8863            }
8864        }
8865        let (b0, b1, b2) = (bytes_of(w0), bytes_of(w1), bytes_of(w2));
8866        let total = rb0 * (o0 + o1 + o2);
8867        let mut cat = self.alloc_u8(total)?;
8868        self.copy_u8_into(&mut cat, 0, b0, rb0 * o0)?;
8869        self.copy_u8_into(&mut cat, rb0 * o0, b1, rb1 * o1)?;
8870        self.copy_u8_into(&mut cat, rb0 * (o0 + o1), b2, rb2 * o2)?;
8871        Ok(Some(GpuTensor::Quant {
8872            bytes: cat,
8873            qtype: QT_Q4_0,
8874            row_bytes: rb0,
8875            ne: vec![w0.in_features() as u64, (o0 + o1 + o2) as u64],
8876            scale: 1.0,
8877            rp: false,
8878            #[cfg(memra_cutlass)]
8879            cutlass: None,
8880            fp8: None,
8881            blk: None,
8882            rp4: None,
8883            f16: None,
8884        }))
8885    }
8886
8887    /// FULL-WIDTH-ROPE CONTRACT for the fused rms_norm+qkv+rope kernels
8888    /// (lane/graph-s-key-exactness-20260819, probe O-10 follow-up).
8889    ///
8890    /// `rms_norm_qkv_rope_f32` / `_cat_f32` (`cu/kernels.cu`) and
8891    /// `rms_norm_qkv_rope_append_body` (`cu/flash_attn.cu`) take NO `n_dims`/`n_rot` argument.
8892    /// They compute `int half = ncols / 2` and rotate the FULL head width by construction — the
8893    /// standalone `rope_neox*` kernels take `n_dims` and early-return above it, these do not.
8894    ///
8895    /// Every call site today is a gemma-4 arm, where `n_rot == head_dim`, so nothing is wrong
8896    /// now. But this is the n_rot bug ONE FUSION away: fuse a partial-rotary arch onto these
8897    /// kernels (qwen3.5 = 64 rotary dims of a 256-wide head; step35 full-attn = 64 of 128) and
8898    /// 192 dims that must pass through unrotated get rotated silently — no error, no NaN, just a
8899    /// wrong model. The n_rot lane already paid for that class once, in the config derivation.
8900    ///
8901    /// So the fusions now take the layer's DERIVED rope width and refuse anything but full
8902    /// width. A future partial-rotary caller fails at its first launch with the geometry named
8903    /// instead of serving quietly wrong logits.
8904    fn full_width_rope_only(
8905        kernel: &str,
8906        n_rot: usize,
8907        head_dim: usize,
8908    ) -> Result<(), Box<dyn std::error::Error>> {
8909        if n_rot == head_dim {
8910            return Ok(());
8911        }
8912        Err(format!(
8913            "{kernel}: PARTIAL ROTARY REFUSED — n_rot {n_rot} != head_dim {head_dim}. This fused \
8914             rms_norm+qkv+rope kernel carries no n_dims parameter and rotates the full head \
8915             width (half = ncols/2), so it would rotate dims {n_rot}..{head_dim} that must pass \
8916             through unrotated. Use the split path (rms_norm_qkv + rope_neox/rope_neox2 with \
8917             n_dims={n_rot}), or add an n_dims early-return to the kernel and widen this guard."
8918        )
8919        .into())
8920    }
8921
8922    /// wave-4b: the qkv-cat twin — one contiguous [rq+2*rk, hd] input from the concat matvec.
8923    /// `n_rot` is the layer's derived rotary width and MUST equal `head_dim`
8924    /// ([`Engine::full_width_rope_only`]).
8925    #[allow(clippy::too_many_arguments)]
8926    pub fn rms_norm_qkv_rope_cat(
8927        &self,
8928        qkv: &CudaSlice<f32>,
8929        wq: &CudaSlice<f32>,
8930        wk: &CudaSlice<f32>,
8931        wv: &CudaSlice<f32>,
8932        q: &mut CudaSlice<f32>,
8933        k: &mut CudaSlice<f32>,
8934        v: &mut CudaSlice<f32>,
8935        head_dim: usize,
8936        n_rot: usize,
8937        rq: usize,
8938        rk: usize,
8939        pos: &CudaSlice<i32>,
8940        nh_q: usize,
8941        nh_k: usize,
8942        base: f32,
8943        freq_scale: f32,
8944        ff: Option<&CudaSlice<f32>>,
8945        eps: f32,
8946    ) -> Result<(), Box<dyn std::error::Error>> {
8947        Self::full_width_rope_only("rms_norm_qkv_rope_cat", n_rot, head_dim)?;
8948        let rows = rq + rk + rk;
8949        let theta_scale = base.powf(-2.0 / head_dim as f32);
8950        let (nc, rqi, rki, nhq, nhk) = (
8951            head_dim as i32,
8952            rq as i32,
8953            rk as i32,
8954            nh_q as i32,
8955            nh_k as i32,
8956        );
8957        if Self::pdl_on() {
8958            use cudarc::driver::{DevicePtr, DevicePtrMut};
8959            let s = &self.gpu.stream();
8960            let (pqkv, _g0) = qkv.device_ptr(s);
8961            let (pwq, _g1) = wq.device_ptr(s);
8962            let (pwk, _g2) = wk.device_ptr(s);
8963            let (pwv, _g3) = wv.device_ptr(s);
8964            let (pq, _g4) = q.device_ptr_mut(s);
8965            let (pk, _g5) = k.device_ptr_mut(s);
8966            let (pv, _g6) = v.device_ptr_mut(s);
8967            let (ppos, _g7) = pos.device_ptr(s);
8968            let (pff, _g8) = match ff {
8969                Some(t) => {
8970                    let (p, g) = t.device_ptr(s);
8971                    (p, Some(g))
8972                }
8973                None => (0, None),
8974            };
8975            let mut ps = [
8976                &pqkv as *const _ as *mut std::ffi::c_void,
8977                &pwq as *const _ as *mut _,
8978                &pwk as *const _ as *mut _,
8979                &pwv as *const _ as *mut _,
8980                &pq as *const _ as *mut _,
8981                &pk as *const _ as *mut _,
8982                &pv as *const _ as *mut _,
8983                &nc as *const _ as *mut _,
8984                &rqi as *const _ as *mut _,
8985                &rki as *const _ as *mut _,
8986                &ppos as *const _ as *mut _,
8987                &nhq as *const _ as *mut _,
8988                &nhk as *const _ as *mut _,
8989                &theta_scale as *const _ as *mut _,
8990                &freq_scale as *const _ as *mut _,
8991                &pff as *const _ as *mut _,
8992                &eps as *const _ as *mut _,
8993            ];
8994            unsafe {
8995                self.launch_pdl(
8996                    "rms_norm_qkv_rope_cat_f32",
8997                    (rows as u32, 1, 1),
8998                    (rms_block(), 1, 1),
8999                    &mut ps,
9000                )?;
9001            }
9002            return Ok(());
9003        }
9004        let f = self.func("rms_norm_qkv_rope_cat_f32");
9005        let cfg = LaunchConfig {
9006            grid_dim: (rows as u32, 1, 1),
9007            block_dim: (rms_block(), 1, 1),
9008            shared_mem_bytes: 0,
9009        };
9010        let __s_b = self.gpu.stream();
9011        let mut b = __s_b.launch_builder(&f);
9012        match ff {
9013            Some(t) => {
9014                b.arg(qkv)
9015                    .arg(wq)
9016                    .arg(wk)
9017                    .arg(wv)
9018                    .arg(&mut *q)
9019                    .arg(&mut *k)
9020                    .arg(&mut *v)
9021                    .arg(&nc)
9022                    .arg(&rqi)
9023                    .arg(&rki)
9024                    .arg(pos)
9025                    .arg(&nhq)
9026                    .arg(&nhk)
9027                    .arg(&theta_scale)
9028                    .arg(&freq_scale)
9029                    .arg(t)
9030                    .arg(&eps);
9031                unsafe {
9032                    b.launch(cfg)?;
9033                }
9034            }
9035            None => {
9036                let null: u64 = 0;
9037                b.arg(qkv)
9038                    .arg(wq)
9039                    .arg(wk)
9040                    .arg(wv)
9041                    .arg(&mut *q)
9042                    .arg(&mut *k)
9043                    .arg(&mut *v)
9044                    .arg(&nc)
9045                    .arg(&rqi)
9046                    .arg(&rki)
9047                    .arg(pos)
9048                    .arg(&nhq)
9049                    .arg(&nhk)
9050                    .arg(&theta_scale)
9051                    .arg(&freq_scale)
9052                    .arg(&null)
9053                    .arg(&eps);
9054                unsafe {
9055                    b.launch(cfg)?;
9056                }
9057            }
9058        }
9059        Ok(())
9060    }
9061
9062    /// wave-3 fold: rms_norm_qkv + rope_neox2 in ONE launch (n_dims == head_dim; ff nullable).
9063    /// `n_rot` is the layer's derived rotary width and MUST equal `head_dim`
9064    /// ([`Engine::full_width_rope_only`]).
9065    #[allow(clippy::too_many_arguments)]
9066    pub fn rms_norm_qkv_rope(
9067        &self,
9068        q0: &CudaSlice<f32>,
9069        k0: &CudaSlice<f32>,
9070        v0: &CudaSlice<f32>,
9071        wq: &CudaSlice<f32>,
9072        wk: &CudaSlice<f32>,
9073        wv: &CudaSlice<f32>,
9074        q: &mut CudaSlice<f32>,
9075        k: &mut CudaSlice<f32>,
9076        v: &mut CudaSlice<f32>,
9077        head_dim: usize,
9078        n_rot: usize,
9079        rq: usize,
9080        rk: usize,
9081        pos: &CudaSlice<i32>,
9082        nh_q: usize,
9083        nh_k: usize,
9084        base: f32,
9085        freq_scale: f32,
9086        ff: Option<&CudaSlice<f32>>,
9087        eps: f32,
9088    ) -> Result<(), Box<dyn std::error::Error>> {
9089        Self::full_width_rope_only("rms_norm_qkv_rope", n_rot, head_dim)?;
9090        let f = self.func("rms_norm_qkv_rope_f32");
9091        let rows = rq + rk + rk; // q rows + k rows + v rows (rk == rv)
9092        let cfg = LaunchConfig {
9093            grid_dim: (rows as u32, 1, 1),
9094            block_dim: (rms_block(), 1, 1),
9095            shared_mem_bytes: 0,
9096        };
9097        let theta_scale = base.powf(-2.0 / head_dim as f32);
9098        let (nc, rqi, rki, nhq, nhk) = (
9099            head_dim as i32,
9100            rq as i32,
9101            rk as i32,
9102            nh_q as i32,
9103            nh_k as i32,
9104        );
9105        let __s_b = self.gpu.stream();
9106        let mut b = __s_b.launch_builder(&f);
9107        match ff {
9108            Some(t) => {
9109                b.arg(q0)
9110                    .arg(k0)
9111                    .arg(v0)
9112                    .arg(wq)
9113                    .arg(wk)
9114                    .arg(wv)
9115                    .arg(&mut *q)
9116                    .arg(&mut *k)
9117                    .arg(&mut *v)
9118                    .arg(&nc)
9119                    .arg(&rqi)
9120                    .arg(&rki)
9121                    .arg(pos)
9122                    .arg(&nhq)
9123                    .arg(&nhk)
9124                    .arg(&theta_scale)
9125                    .arg(&freq_scale)
9126                    .arg(t)
9127                    .arg(&eps);
9128                unsafe {
9129                    b.launch(cfg)?;
9130                }
9131            }
9132            None => {
9133                let null: u64 = 0;
9134                b.arg(q0)
9135                    .arg(k0)
9136                    .arg(v0)
9137                    .arg(wq)
9138                    .arg(wk)
9139                    .arg(wv)
9140                    .arg(&mut *q)
9141                    .arg(&mut *k)
9142                    .arg(&mut *v)
9143                    .arg(&nc)
9144                    .arg(&rqi)
9145                    .arg(&rki)
9146                    .arg(pos)
9147                    .arg(&nhq)
9148                    .arg(&nhk)
9149                    .arg(&theta_scale)
9150                    .arg(&freq_scale)
9151                    .arg(&null)
9152                    .arg(&eps);
9153                unsafe {
9154                    b.launch(cfg)?;
9155                }
9156            }
9157        }
9158        Ok(())
9159    }
9160
9161    /// FUSED norm+rope+APPEND (m=1 decode, 2026-07-23): one launch replaces the
9162    /// rms_norm_qkv_rope + append_kv_quantized_dc pair. Kernel lives in the flash fatbins
9163    /// (format-flavored quant tail) — `g` must mirror the append path's flavor exactly.
9164    /// `n_rot` is the layer's derived rotary width and MUST equal `head_dim`
9165    /// ([`Engine::full_width_rope_only`]).
9166    #[allow(clippy::too_many_arguments)]
9167    pub fn rms_norm_qkv_rope_append_dc(
9168        &self,
9169        q0: &CudaSlice<f32>,
9170        k0: &CudaSlice<f32>,
9171        v0: &CudaSlice<f32>,
9172        wq: &CudaSlice<f32>,
9173        wk: &CudaSlice<f32>,
9174        wv: &CudaSlice<f32>,
9175        q: &mut CudaSlice<f32>,
9176        k: &mut CudaSlice<f32>,
9177        v: &mut CudaSlice<f32>,
9178        head_dim: usize,
9179        n_rot: usize,
9180        rq: usize,
9181        rk: usize,
9182        pos: &CudaSlice<i32>,
9183        nh_q: usize,
9184        nh_k: usize,
9185        base: f32,
9186        freq_scale: f32,
9187        ff: Option<&CudaSlice<f32>>,
9188        eps: f32,
9189        kc: &mut CudaSlice<u8>,
9190        vc: &mut CudaSlice<u8>,
9191        t_dev: &CudaSlice<i32>,
9192        k_tok_bytes: usize,
9193        v_tok_bytes: usize,
9194        g: bool,
9195    ) -> Result<(), Box<dyn std::error::Error>> {
9196        Self::full_width_rope_only("rms_norm_qkv_rope_append_dc", n_rot, head_dim)?;
9197        let rows = rq + rk + rk;
9198        let theta_scale = base.powf(-2.0 / head_dim as f32);
9199        let (nc, rqi, rki, nhq, nhk) = (
9200            head_dim as i32,
9201            rq as i32,
9202            rk as i32,
9203            nh_q as i32,
9204            nh_k as i32,
9205        );
9206        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
9207        if Self::pdl_on() && Self::pdl_wb_on() {
9208            use cudarc::driver::{DevicePtr, DevicePtrMut};
9209            let s = &self.gpu.stream();
9210            let (p0, _a0) = q0.device_ptr(s);
9211            let (p1, _a1) = k0.device_ptr(s);
9212            let (p2, _a2) = v0.device_ptr(s);
9213            let (pwq, _a3) = wq.device_ptr(s);
9214            let (pwk, _a4) = wk.device_ptr(s);
9215            let (pwv, _a5) = wv.device_ptr(s);
9216            let (pq, _a6) = q.device_ptr_mut(s);
9217            let (pk, _a7) = k.device_ptr_mut(s);
9218            let (pv, _a8) = v.device_ptr_mut(s);
9219            let (pp, _a9) = pos.device_ptr(s);
9220            let pff: u64 = match ff {
9221                Some(t) => {
9222                    let (p, _gg) = t.device_ptr(s);
9223                    p as u64
9224                }
9225                None => 0,
9226            };
9227            let (pkc, _a10) = kc.device_ptr_mut(s);
9228            let (pvc, _a11) = vc.device_ptr_mut(s);
9229            let (pt, _a12) = t_dev.device_ptr(s);
9230            let mut ps = [
9231                &p0 as *const _ as *mut std::ffi::c_void,
9232                &p1 as *const _ as *mut _,
9233                &p2 as *const _ as *mut _,
9234                &pwq as *const _ as *mut _,
9235                &pwk as *const _ as *mut _,
9236                &pwv as *const _ as *mut _,
9237                &pq as *const _ as *mut _,
9238                &pk as *const _ as *mut _,
9239                &pv as *const _ as *mut _,
9240                &nc as *const _ as *mut _,
9241                &rqi as *const _ as *mut _,
9242                &rki as *const _ as *mut _,
9243                &pp as *const _ as *mut _,
9244                &nhq as *const _ as *mut _,
9245                &nhk as *const _ as *mut _,
9246                &theta_scale as *const _ as *mut _,
9247                &freq_scale as *const _ as *mut _,
9248                &pff as *const _ as *mut _,
9249                &eps as *const _ as *mut _,
9250                &pkc as *const _ as *mut _,
9251                &pvc as *const _ as *mut _,
9252                &pt as *const _ as *mut _,
9253                &ktb as *const _ as *mut _,
9254                &vtb as *const _ as *mut _,
9255            ];
9256            unsafe {
9257                self.launch_pdl_flash(
9258                    g,
9259                    "rms_norm_qkv_rope_append_dc_f32",
9260                    (rows as u32, 1, 1),
9261                    (rms_block(), 1, 1),
9262                    0,
9263                    &mut ps,
9264                )?;
9265            }
9266            return Ok(());
9267        }
9268        let f = if g {
9269            self.func_g("rms_norm_qkv_rope_append_dc_f32")
9270        } else {
9271            self.func("rms_norm_qkv_rope_append_dc_f32")
9272        };
9273        let cfg = LaunchConfig {
9274            grid_dim: (rows as u32, 1, 1),
9275            block_dim: (rms_block(), 1, 1),
9276            shared_mem_bytes: 0,
9277        };
9278        let __s_b = self.gpu.stream();
9279        let mut b = __s_b.launch_builder(&f);
9280        match ff {
9281            Some(t) => {
9282                b.arg(q0)
9283                    .arg(k0)
9284                    .arg(v0)
9285                    .arg(wq)
9286                    .arg(wk)
9287                    .arg(wv)
9288                    .arg(&mut *q)
9289                    .arg(&mut *k)
9290                    .arg(&mut *v)
9291                    .arg(&nc)
9292                    .arg(&rqi)
9293                    .arg(&rki)
9294                    .arg(pos)
9295                    .arg(&nhq)
9296                    .arg(&nhk)
9297                    .arg(&theta_scale)
9298                    .arg(&freq_scale)
9299                    .arg(t)
9300                    .arg(&eps)
9301                    .arg(&mut *kc)
9302                    .arg(&mut *vc)
9303                    .arg(t_dev)
9304                    .arg(&ktb)
9305                    .arg(&vtb);
9306                unsafe {
9307                    b.launch(cfg)?;
9308                }
9309            }
9310            None => {
9311                let null: u64 = 0;
9312                b.arg(q0)
9313                    .arg(k0)
9314                    .arg(v0)
9315                    .arg(wq)
9316                    .arg(wk)
9317                    .arg(wv)
9318                    .arg(&mut *q)
9319                    .arg(&mut *k)
9320                    .arg(&mut *v)
9321                    .arg(&nc)
9322                    .arg(&rqi)
9323                    .arg(&rki)
9324                    .arg(pos)
9325                    .arg(&nhq)
9326                    .arg(&nhk)
9327                    .arg(&theta_scale)
9328                    .arg(&freq_scale)
9329                    .arg(&null)
9330                    .arg(&eps)
9331                    .arg(&mut *kc)
9332                    .arg(&mut *vc)
9333                    .arg(t_dev)
9334                    .arg(&ktb)
9335                    .arg(&vtb);
9336                unsafe {
9337                    b.launch(cfg)?;
9338                }
9339            }
9340        }
9341        Ok(())
9342    }
9343
9344    /// wave-2 fold: a + b with the sum emitted q8_1 alongside f32.
9345    /// Host-len twin of `rms_norm_qkv_rope_append_dc` (zoo-fusion arc): the eager decode
9346    /// arm tracks kv length on host (`kvl.len`), so the append slot rides the launch arg
9347    /// instead of the device counter. Kernel body is SHARED with the _dc entry (one
9348    /// inlined body — bit-identical to the rms_norm_qkv_rope + append pair it replaces,
9349    /// same law as the dc fold). `n_rot` is the layer's derived rotary width and MUST equal
9350    /// `head_dim` ([`Engine::full_width_rope_only`]).
9351    #[allow(clippy::too_many_arguments)]
9352    pub fn rms_norm_qkv_rope_append(
9353        &self,
9354        q0: &CudaSlice<f32>,
9355        k0: &CudaSlice<f32>,
9356        v0: &CudaSlice<f32>,
9357        wq: &CudaSlice<f32>,
9358        wk: &CudaSlice<f32>,
9359        wv: &CudaSlice<f32>,
9360        q: &mut CudaSlice<f32>,
9361        k: &mut CudaSlice<f32>,
9362        v: &mut CudaSlice<f32>,
9363        head_dim: usize,
9364        n_rot: usize,
9365        rq: usize,
9366        rk: usize,
9367        pos: &CudaSlice<i32>,
9368        nh_q: usize,
9369        nh_k: usize,
9370        base: f32,
9371        freq_scale: f32,
9372        ff: Option<&CudaSlice<f32>>,
9373        eps: f32,
9374        kc: &mut CudaSlice<u8>,
9375        vc: &mut CudaSlice<u8>,
9376        t: usize,
9377        k_tok_bytes: usize,
9378        v_tok_bytes: usize,
9379        g: bool,
9380    ) -> Result<(), Box<dyn std::error::Error>> {
9381        Self::full_width_rope_only("rms_norm_qkv_rope_append", n_rot, head_dim)?;
9382        let rows = rq + rk + rk;
9383        let theta_scale = base.powf(-2.0 / head_dim as f32);
9384        let (nc, rqi, rki, nhq, nhk) = (
9385            head_dim as i32,
9386            rq as i32,
9387            rk as i32,
9388            nh_q as i32,
9389            nh_k as i32,
9390        );
9391        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
9392        let ti = t as i32;
9393        if Self::pdl_on() && Self::pdl_wb_on() {
9394            use cudarc::driver::{DevicePtr, DevicePtrMut};
9395            let s = &self.gpu.stream();
9396            let (p0, _a0) = q0.device_ptr(s);
9397            let (p1, _a1) = k0.device_ptr(s);
9398            let (p2, _a2) = v0.device_ptr(s);
9399            let (pwq, _a3) = wq.device_ptr(s);
9400            let (pwk, _a4) = wk.device_ptr(s);
9401            let (pwv, _a5) = wv.device_ptr(s);
9402            let (pq, _a6) = q.device_ptr_mut(s);
9403            let (pk, _a7) = k.device_ptr_mut(s);
9404            let (pv, _a8) = v.device_ptr_mut(s);
9405            let (pp, _a9) = pos.device_ptr(s);
9406            let pff: u64 = match ff {
9407                Some(t) => {
9408                    let (p, _gg) = t.device_ptr(s);
9409                    p as u64
9410                }
9411                None => 0,
9412            };
9413            let (pkc, _a10) = kc.device_ptr_mut(s);
9414            let (pvc, _a11) = vc.device_ptr_mut(s);
9415            let mut ps = [
9416                &p0 as *const _ as *mut std::ffi::c_void,
9417                &p1 as *const _ as *mut _,
9418                &p2 as *const _ as *mut _,
9419                &pwq as *const _ as *mut _,
9420                &pwk as *const _ as *mut _,
9421                &pwv as *const _ as *mut _,
9422                &pq as *const _ as *mut _,
9423                &pk as *const _ as *mut _,
9424                &pv as *const _ as *mut _,
9425                &nc as *const _ as *mut _,
9426                &rqi as *const _ as *mut _,
9427                &rki as *const _ as *mut _,
9428                &pp as *const _ as *mut _,
9429                &nhq as *const _ as *mut _,
9430                &nhk as *const _ as *mut _,
9431                &theta_scale as *const _ as *mut _,
9432                &freq_scale as *const _ as *mut _,
9433                &pff as *const _ as *mut _,
9434                &eps as *const _ as *mut _,
9435                &pkc as *const _ as *mut _,
9436                &pvc as *const _ as *mut _,
9437                &ti as *const _ as *mut _,
9438                &ktb as *const _ as *mut _,
9439                &vtb as *const _ as *mut _,
9440            ];
9441            unsafe {
9442                self.launch_pdl_flash(
9443                    g,
9444                    "rms_norm_qkv_rope_append_f32",
9445                    (rows as u32, 1, 1),
9446                    (rms_block(), 1, 1),
9447                    0,
9448                    &mut ps,
9449                )?;
9450            }
9451            return Ok(());
9452        }
9453        let f = if g {
9454            self.func_g("rms_norm_qkv_rope_append_f32")
9455        } else {
9456            self.func("rms_norm_qkv_rope_append_f32")
9457        };
9458        let cfg = LaunchConfig {
9459            grid_dim: (rows as u32, 1, 1),
9460            block_dim: (rms_block(), 1, 1),
9461            shared_mem_bytes: 0,
9462        };
9463        let __s_b = self.gpu.stream();
9464        let mut b = __s_b.launch_builder(&f);
9465        let null: u64 = 0;
9466        b.arg(q0)
9467            .arg(k0)
9468            .arg(v0)
9469            .arg(wq)
9470            .arg(wk)
9471            .arg(wv)
9472            .arg(&mut *q)
9473            .arg(&mut *k)
9474            .arg(&mut *v)
9475            .arg(&nc)
9476            .arg(&rqi)
9477            .arg(&rki)
9478            .arg(pos)
9479            .arg(&nhq)
9480            .arg(&nhk)
9481            .arg(&theta_scale)
9482            .arg(&freq_scale);
9483        match ff {
9484            Some(t) => {
9485                b.arg(t);
9486            }
9487            None => {
9488                b.arg(&null);
9489            }
9490        }
9491        b.arg(&eps)
9492            .arg(&mut *kc)
9493            .arg(&mut *vc)
9494            .arg(&ti)
9495            .arg(&ktb)
9496            .arg(&vtb);
9497        unsafe {
9498            b.launch(cfg)?;
9499        }
9500        Ok(())
9501    }
9502
9503    pub fn add_q8_1(
9504        &self,
9505        a: &CudaSlice<f32>,
9506        b: &CudaSlice<f32>,
9507        res: &mut CudaSlice<f32>,
9508        ncols: usize,
9509        nrows: usize,
9510    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9511        debug_assert!(ncols % 128 == 0);
9512        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
9513        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9514        let f = self.func("add_q8_1_f32");
9515        let cfg = LaunchConfig {
9516            grid_dim: (nrows as u32, 1, 1),
9517            block_dim: (rms_block(), 1, 1),
9518            shared_mem_bytes: 0,
9519        };
9520        let nc = ncols as i32;
9521        let __s_b2 = self.gpu.stream();
9522        let mut b2 = __s_b2.launch_builder(&f);
9523        b2.arg(a)
9524            .arg(b)
9525            .arg(&mut *res)
9526            .arg(&mut out_q)
9527            .arg(&mut out_d)
9528            .arg(&nc);
9529        unsafe {
9530            b2.launch(cfg)?;
9531        }
9532        Ok((out_q, out_d))
9533    }
9534
9535    /// E4B FFN-tail exit fusion (glue wave 5): resid = b + rms(a, wa) emitted f32 + q8_1 pair
9536    /// in ONE launch — replaces rms_norm(a,wa->sn) + add_q8_1(sn,b). Same rms_block() config
9537    /// as both parents (bit-identity: identical reduction + quad-walk quantize).
9538    pub fn rms_pre_add_q8_1(
9539        &self,
9540        a: &CudaSlice<f32>,
9541        wa: &CudaSlice<f32>,
9542        b: &CudaSlice<f32>,
9543        res: &mut CudaSlice<f32>,
9544        ncols: usize,
9545        nrows: usize,
9546        eps: f32,
9547    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9548        debug_assert!(ncols % 128 == 0);
9549        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
9550        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9551        let f = self.func("rms_pre_add_q8_1_f32");
9552        let cfg = LaunchConfig {
9553            grid_dim: (nrows as u32, 1, 1),
9554            block_dim: (rms_block(), 1, 1),
9555            shared_mem_bytes: 0,
9556        };
9557        let (nc, ep) = (ncols as i32, eps);
9558        let __s_b2 = self.gpu.stream();
9559        let mut b2 = __s_b2.launch_builder(&f);
9560        b2.arg(a)
9561            .arg(wa)
9562            .arg(b)
9563            .arg(&mut *res)
9564            .arg(&mut out_q)
9565            .arg(&mut out_d)
9566            .arg(&nc)
9567            .arg(&ep);
9568        unsafe {
9569            b2.launch(cfg)?;
9570        }
9571        Ok((out_q, out_d))
9572    }
9573
9574    /// L2 norm per row (head_dim), no weight.
9575    /// PREFILL l2 dispatch (round 27): the warp-per-row float4 v2 when the numeric-config
9576    /// seam allows (MEMRA_L2_V2, default ON, d_state==128 only); else the strided kernel.
9577    pub fn l2_v2_on(ncols: usize) -> bool {
9578        ncols == 128 && std::env::var("MEMRA_L2_V2").as_deref() != Ok("0")
9579    }
9580
9581    pub fn l2_norm_pp(
9582        &self,
9583        x: &CudaSlice<f32>,
9584        dst: &mut CudaSlice<f32>,
9585        dst16: Option<&mut CudaSlice<u8>>,
9586        ncols: usize,
9587        nrows: usize,
9588        eps: f32,
9589    ) -> Result<(), Box<dyn std::error::Error>> {
9590        if Self::l2_v2_on(ncols) {
9591            let f = self.func("l2_norm_pp_v2_f32");
9592            let rows_per_block = 8u32; // 256 threads = 8 warps = 8 rows
9593            let cfg = LaunchConfig {
9594                grid_dim: ((nrows as u32).div_ceil(rows_per_block), 1, 1),
9595                block_dim: (256, 1, 1),
9596                shared_mem_bytes: 0,
9597            };
9598            let (nc, nr, e) = (ncols as i32, nrows as i32, eps);
9599            // mirror-fold: bf16 twin address by value (0 = skip; matches the nullable param)
9600            let d16: u64 = match dst16 {
9601                Some(d) => self.addr_u8(d),
9602                None => 0,
9603            };
9604            let __s_b = self.gpu.stream();
9605            let mut b = __s_b.launch_builder(&f);
9606            b.arg(x).arg(dst).arg(&d16).arg(&nc).arg(&nr).arg(&e);
9607            unsafe {
9608                b.launch(cfg)?;
9609            }
9610            return Ok(());
9611        }
9612        self.l2_norm(x, dst, ncols, nrows, eps)
9613    }
9614
9615    pub fn l2_norm(
9616        &self,
9617        x: &CudaSlice<f32>,
9618        dst: &mut CudaSlice<f32>,
9619        ncols: usize,
9620        nrows: usize,
9621        eps: f32,
9622    ) -> Result<(), Box<dyn std::error::Error>> {
9623        let f = self.func("l2_norm_f32");
9624        let cfg = LaunchConfig {
9625            grid_dim: (nrows as u32, 1, 1),
9626            block_dim: (256, 1, 1),
9627            shared_mem_bytes: 0,
9628        };
9629        let (nc, e) = (ncols as i32, eps);
9630        let __s_b = self.gpu.stream();
9631        let mut b = __s_b.launch_builder(&f);
9632        b.arg(x).arg(dst).arg(&nc).arg(&e);
9633        unsafe {
9634            b.launch(cfg)?;
9635        }
9636        Ok(())
9637    }
9638
9639    /// L2-norm with blockDim=32 (warp-tree reduction) — BIT-IDENTICAL to gdn_prep_decode_f32's
9640    /// per-warp L2 norm. The verify path MUST use this to match decode's FP accumulation order:
9641    /// l2_norm at blockDim=256 produces a different shfl-tree reduction of the 128-element
9642    /// squared-sum (pairwise tree vs serial-4-then-warp-tree), causing ULP differences that
9643    /// propagate through gdn_scan and flip argmax on marginal logits.
9644    pub fn l2_norm_decode(
9645        &self,
9646        x: &CudaSlice<f32>,
9647        dst: &mut CudaSlice<f32>,
9648        ncols: usize,
9649        nrows: usize,
9650        eps: f32,
9651    ) -> Result<(), Box<dyn std::error::Error>> {
9652        let f = self.func("l2_norm_f32");
9653        let cfg = LaunchConfig {
9654            grid_dim: (nrows as u32, 1, 1),
9655            block_dim: (32, 1, 1),
9656            shared_mem_bytes: 0,
9657        };
9658        let (nc, e) = (ncols as i32, eps);
9659        let __s_b = self.gpu.stream();
9660        let mut b = __s_b.launch_builder(&f);
9661        b.arg(x).arg(dst).arg(&nc).arg(&e);
9662        unsafe {
9663            b.launch(cfg)?;
9664        }
9665        Ok(())
9666    }
9667
9668    /// RoPE NEOX in-place. x:[head_dim, n_heads, n_tokens], pos:[n_tokens].
9669    pub fn rope_neox(
9670        &self,
9671        x: &mut CudaSlice<f32>,
9672        pos: &CudaSlice<i32>,
9673        head_dim: usize,
9674        n_dims: usize,
9675        n_heads: usize,
9676        n_tokens: usize,
9677        freq_base: f32,
9678        freq_scale: f32,
9679    ) -> Result<(), Box<dyn std::error::Error>> {
9680        let f = self.func("rope_neox_f32");
9681        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
9682        let grid = (n_heads * n_tokens) as u32;
9683        let cfg = LaunchConfig {
9684            grid_dim: (grid, 1, 1),
9685            block_dim: ((head_dim / 2) as u32, 1, 1),
9686            shared_mem_bytes: 0,
9687        };
9688        let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
9689        let __s_b = self.gpu.stream();
9690        let mut b = __s_b.launch_builder(&f);
9691        b.arg(x)
9692            .arg(pos)
9693            .arg(&hd)
9694            .arg(&nd)
9695            .arg(&nh)
9696            .arg(&theta_scale)
9697            .arg(&freq_scale);
9698        unsafe {
9699            b.launch(cfg)?;
9700        }
9701        Ok(())
9702    }
9703
9704    /// RoPE NEOX with per-dim freq factors (gemma4 global layers, rope_freqs.weight [n_dims/2]).
9705    pub fn rope_neox_ff(
9706        &self,
9707        x: &mut CudaSlice<f32>,
9708        pos: &CudaSlice<i32>,
9709        head_dim: usize,
9710        n_dims: usize,
9711        n_heads: usize,
9712        n_tokens: usize,
9713        freq_base: f32,
9714        freq_scale: f32,
9715        ff: &CudaSlice<f32>,
9716    ) -> Result<(), Box<dyn std::error::Error>> {
9717        let f = self.func("rope_neox_ff_f32");
9718        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
9719        let grid = (n_heads * n_tokens) as u32;
9720        let cfg = LaunchConfig {
9721            grid_dim: (grid, 1, 1),
9722            block_dim: ((head_dim / 2) as u32, 1, 1),
9723            shared_mem_bytes: 0,
9724        };
9725        let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
9726        let __s_b = self.gpu.stream();
9727        let mut b = __s_b.launch_builder(&f);
9728        b.arg(x)
9729            .arg(pos)
9730            .arg(&hd)
9731            .arg(&nd)
9732            .arg(&nh)
9733            .arg(&theta_scale)
9734            .arg(&freq_scale)
9735            .arg(ff);
9736        unsafe {
9737            b.launch(cfg)?;
9738        }
9739        Ok(())
9740    }
9741
9742    /// gemma4: rope q and k in one launch (per-row chain = rope_neox / rope_neox_ff verbatim).
9743    #[allow(clippy::too_many_arguments)]
9744    pub fn rope_neox2(
9745        &self,
9746        q: &mut CudaSlice<f32>,
9747        k: &mut CudaSlice<f32>,
9748        pos: &CudaSlice<i32>,
9749        head_dim: usize,
9750        n_dims: usize,
9751        nh_q: usize,
9752        nh_k: usize,
9753        n_tokens: usize,
9754        freq_base: f32,
9755        freq_scale: f32,
9756        ff: Option<&CudaSlice<f32>>,
9757    ) -> Result<(), Box<dyn std::error::Error>> {
9758        let f = self.func("rope_neox2_f32");
9759        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
9760        let grid = ((nh_q + nh_k) * n_tokens) as u32;
9761        let cfg = LaunchConfig {
9762            grid_dim: (grid, 1, 1),
9763            block_dim: ((head_dim / 2) as u32, 1, 1),
9764            shared_mem_bytes: 0,
9765        };
9766        let (hd, nd, nq, nk, nt) = (
9767            head_dim as i32,
9768            n_dims as i32,
9769            nh_q as i32,
9770            nh_k as i32,
9771            n_tokens as i32,
9772        );
9773        let __s_b = self.gpu.stream();
9774        let mut b = __s_b.launch_builder(&f);
9775        b.arg(q)
9776            .arg(k)
9777            .arg(pos)
9778            .arg(&hd)
9779            .arg(&nd)
9780            .arg(&nq)
9781            .arg(&nk)
9782            .arg(&nt)
9783            .arg(&theta_scale)
9784            .arg(&freq_scale);
9785        match ff {
9786            Some(ffv) => {
9787                b.arg(ffv);
9788                unsafe {
9789                    b.launch(cfg)?;
9790                }
9791            }
9792            None => {
9793                let null: u64 = 0;
9794                b.arg(&null);
9795                unsafe {
9796                    b.launch(cfg)?;
9797                }
9798            }
9799        }
9800        Ok(())
9801    }
9802
9803    /// gemma4 R1: dst = GELU_tanh(gate) * up.
9804    pub fn gelu_tanh_mul(
9805        &self,
9806        gate: &CudaSlice<f32>,
9807        up: &CudaSlice<f32>,
9808        dst: &mut CudaSlice<f32>,
9809        n: usize,
9810    ) -> Result<(), Box<dyn std::error::Error>> {
9811        let f = self.func("gelu_tanh_mul_f32");
9812        let cfg = LaunchConfig::for_num_elems(n as u32);
9813        let ni = n as i32;
9814        let __s_b = self.gpu.stream();
9815        let mut b = __s_b.launch_builder(&f);
9816        b.arg(gate).arg(up).arg(dst).arg(&ni);
9817        unsafe {
9818            b.launch(cfg)?;
9819        }
9820        Ok(())
9821    }
9822
9823    pub fn silu_mul(
9824        &self,
9825        gate: &CudaSlice<f32>,
9826        up: &CudaSlice<f32>,
9827        dst: &mut CudaSlice<f32>,
9828        n: usize,
9829    ) -> Result<(), Box<dyn std::error::Error>> {
9830        let f = self.func("silu_mul_f32");
9831        // float4 kernel: one thread per 4 elements (tail handled in-kernel)
9832        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
9833        let ni = n as i32;
9834        let __s_b = self.gpu.stream();
9835        let mut b = __s_b.launch_builder(&f);
9836        b.arg(gate).arg(up).arg(dst).arg(&ni);
9837        unsafe {
9838            b.launch(cfg)?;
9839        }
9840        Ok(())
9841    }
9842
9843    /// f16out twin of `silu_mul` (task #17): the epilogue also emits the fp16 GEMM operand
9844    /// for the down projection — kills the standalone convert pass. Bit-identical class.
9845    pub fn silu_mul_f16out(
9846        &self,
9847        gate: &CudaSlice<f32>,
9848        up: &CudaSlice<f32>,
9849        dst: &mut CudaSlice<f32>,
9850        dst16: &mut CudaSlice<u8>,
9851        n: usize,
9852    ) -> Result<(), Box<dyn std::error::Error>> {
9853        let f = self.func("silu_mul_f16out_f32");
9854        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
9855        let ni = n as i32;
9856        let __s_b = self.gpu.stream();
9857        let mut b = __s_b.launch_builder(&f);
9858        b.arg(gate).arg(up).arg(dst).arg(dst16).arg(&ni);
9859        unsafe {
9860            b.launch(cfg)?;
9861        }
9862        Ok(())
9863    }
9864
9865    /// FFN SwiGLU epilogue fusion (RANK3 LEVER 2): `dst = silu(gate*gs) * (up*us)` in ONE launch,
9866    /// folding the per-tensor NVFP4 macro-scale (`gs`,`us`) that would otherwise be two separate
9867    /// `scale_inplace` launches on the gate/up matmul outputs. BIT-IDENTICAL to
9868    /// scale_inplace(gate,gs); scale_inplace(up,us); silu_mul(gate,up,dst) — identical float ops in
9869    /// identical order. For non-NVFP4 weights gs==us==1.0 -> identical to `silu_mul`. Net: -2
9870    /// launches per dense FFN layer (the gate+up post-matmul scales).
9871    pub fn silu_mul_scaled(
9872        &self,
9873        gate: &CudaSlice<f32>,
9874        up: &CudaSlice<f32>,
9875        gs: f32,
9876        us: f32,
9877        dst: &mut CudaSlice<f32>,
9878        n: usize,
9879    ) -> Result<(), Box<dyn std::error::Error>> {
9880        let f = self.func("silu_mul_scaled_f32");
9881        let cfg = LaunchConfig::for_num_elems(n as u32);
9882        let ni = n as i32;
9883        let (gsf, usf) = (gs, us);
9884        let __s_b = self.gpu.stream();
9885        let mut b = __s_b.launch_builder(&f);
9886        b.arg(gate).arg(up).arg(&gsf).arg(&usf).arg(dst).arg(&ni);
9887        unsafe {
9888            b.launch(cfg)?;
9889        }
9890        Ok(())
9891    }
9892
9893    /// swigluoai (MiniMax-M3 / GPT-OSS): clamped SwiGLU epilogue, math 1:1 vs llama.cpp
9894    /// ggml_cuda_op_swiglu_oai_single. `dst = swish_alpha(clamp(gate*gs)) * (1 + clamp(up*us))`.
9895    /// gs/us fold the NVFP4 macro-scales exactly like `silu_mul_scaled`.
9896    #[allow(clippy::too_many_arguments)]
9897    pub fn swigluoai_mul_scaled(
9898        &self,
9899        gate: &CudaSlice<f32>,
9900        up: &CudaSlice<f32>,
9901        gs: f32,
9902        us: f32,
9903        alpha: f32,
9904        limit: f32,
9905        dst: &mut CudaSlice<f32>,
9906        n: usize,
9907    ) -> Result<(), Box<dyn std::error::Error>> {
9908        let f = self.func("swigluoai_mul_scaled_f32");
9909        let cfg = LaunchConfig::for_num_elems(n as u32);
9910        let ni = n as i32;
9911        let __s_b = self.gpu.stream();
9912        let mut b = __s_b.launch_builder(&f);
9913        b.arg(gate)
9914            .arg(up)
9915            .arg(&gs)
9916            .arg(&us)
9917            .arg(&alpha)
9918            .arg(&limit)
9919            .arg(dst)
9920            .arg(&ni);
9921        unsafe {
9922            b.launch(cfg)?;
9923        }
9924        Ok(())
9925    }
9926
9927    /// RANK2 LEVER (q8_1 quant-fold): SwiGLU epilogue that EMITS the q8_1 quantization of `act`
9928    /// directly (aq int8 [n] + ad f32 [n/32]), so ffn_down's standalone `quantize_q8_1` launch is
9929    /// removed — the down-proj activation has one consumer, so the quant folds into the producer for
9930    /// free (no extra HBM read; no f32 `act` write). gs/us fold the gate/up NVFP4 macro-scales like
9931    /// `silu_mul_scaled`. BIT-IDENTICAL q8_1 to silu_mul_scaled(...) then quantize_q8_1(...). Only
9932    /// valid when ffn_down uses the q8_1 dp4a/mmvq path; the caller checks `uses_q8_1_fast(ffn_down)`.
9933    /// n must be a multiple of 32 (n_ff always is).
9934    pub fn silu_mul_scaled_q8_1(
9935        &self,
9936        gate: &CudaSlice<f32>,
9937        up: &CudaSlice<f32>,
9938        gs: f32,
9939        us: f32,
9940        n: usize,
9941    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9942        let f = self.func("silu_mul_scaled_q8_1");
9943        let nblk = n / 32;
9944        let mut aq = self.alloc_uninit::<i8>(n)?; // full-overwrite output
9945        let mut ad = self.alloc_uninit::<f32>(nblk)?; // full-overwrite output
9946        // WARP-PER-BLOCK kernel: one warp (32 lanes) per 32-block -> n threads total.
9947        let cfg = LaunchConfig::for_num_elems(n as u32);
9948        let (gsf, usf, ni) = (gs, us, n as i32);
9949        let __s_b = self.gpu.stream();
9950        let mut b = __s_b.launch_builder(&f);
9951        b.arg(gate)
9952            .arg(up)
9953            .arg(&gsf)
9954            .arg(&usf)
9955            .arg(&mut aq)
9956            .arg(&mut ad)
9957            .arg(&ni);
9958        unsafe {
9959            b.launch(cfg)?;
9960        }
9961        Ok((aq, ad))
9962    }
9963
9964    pub fn add(
9965        &self,
9966        a: &CudaSlice<f32>,
9967        b_in: &CudaSlice<f32>,
9968        dst: &mut CudaSlice<f32>,
9969        n: usize,
9970    ) -> Result<(), Box<dyn std::error::Error>> {
9971        let f = self.func("add_f32");
9972        // float4 kernel: one thread per 4 elements (tail handled in-kernel)
9973        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
9974        let ni = n as i32;
9975        let __s_bld = self.gpu.stream();
9976        let mut bld = __s_bld.launch_builder(&f);
9977        bld.arg(a).arg(b_in).arg(dst).arg(&ni);
9978        unsafe {
9979            bld.launch(cfg)?;
9980        }
9981        Ok(())
9982    }
9983
9984    pub fn mul(
9985        &self,
9986        a: &CudaSlice<f32>,
9987        b_in: &CudaSlice<f32>,
9988        dst: &mut CudaSlice<f32>,
9989        n: usize,
9990    ) -> Result<(), Box<dyn std::error::Error>> {
9991        let f = self.func("mul_f32");
9992        let cfg = LaunchConfig::for_num_elems(n as u32);
9993        let ni = n as i32;
9994        let __s_bld = self.gpu.stream();
9995        let mut bld = __s_bld.launch_builder(&f);
9996        bld.arg(a).arg(b_in).arg(dst).arg(&ni);
9997        unsafe {
9998            bld.launch(cfg)?;
9999        }
10000        Ok(())
10001    }
10002
10003    /// Unified weight-tensor matmul: dispatches quant tensors to qmatvec (weights packed) and
10004    /// float tensors to cuBLASLt. y[m,out] = x[m,in] @ W[out,in]^T.
10005    pub fn matmul(
10006        &self,
10007        w: &crate::model::GpuTensor,
10008        x: &CudaSlice<f32>,
10009        m: usize,
10010    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10011        use crate::model::GpuTensor;
10012        let in_f = w.in_features();
10013        let out_f = w.out_features();
10014        // PREFILL (T>1) ROOT FIX: batched tensor-core int8 GEMM. Decodes each weight tile to int8
10015        // in smem ONCE and reuses across all tokens via mma — vs the dp4a matvec's per-token weight
10016        // re-read. Only the 4 daily-hot dtypes; m=1 decode keeps dp4a (it's bandwidth-bound, mma
10017        // gives nothing). Quantize the activation once here then call the GEMM.
10018        // m cutoff FIXED at 16: the m=4 MMA-verify A/B (2026-07-06, was MEMRA_GEMM_M) measured
10019        // NEGATIVE — the MMA tile grid starves at m=4 (BN=256 -> grid.y=1) and its FP order
10020        // shifted verify argmax at tight margins. Do not lower without re-running that battery.
10021        #[allow(non_snake_case)]
10022        // VERIFY-EXACT scope pushes the GEMM crossover out of reach (usize::MAX) — the
10023        // t>=16 dflash verify must ride the decode-exact batched class (parity law).
10024        let GEMM_M_THRESHOLD = if self.verify_exact_on() {
10025            usize::MAX
10026        } else {
10027            16usize
10028        };
10029
10030        // PREFILL GEMM (m>=16). ACCURACY-FIRST dispatch (2026-06-28, prefill-gemm-beat-research wf
10031        // wllbyo6vc step 1): the int8 W4A8 GEMM (qmatvec_gemm, q8_1 activation, s32 accumulate) is
10032        // ACCURATE (prefill logit maxdiff 0.159, < dp4a 0.55) and the default. The FP4 W4A4 mxf4 path
10033        // (try_fp4_gemm) quantizes the ACTIVATION to e2m1 4-bit (8 magnitude levels) -> maxdiff 1.0
10034        // when combined — a real accuracy loss, NOT a math bug. So FP4-W4A4 is taken ONLY under the
10035        // explicit MEMRA_FP4 opt-in AND it must come SECOND (int8 W4A8 is the correct default for NVFP4).
10036        // The workflow plan rebuilds the FP4 path (kill per-K repack, widen K, deepen pipeline, TMA) to
10037        // be both fast AND accurate; until then NVFP4 prefill defaults to the accurate int8 GEMM.
10038        // TINY-OUT_F GUARD (2026-06-28, ncu trace): the tiling GEMM's grid is (ceil(out_f/BM=64),
10039        // ceil(m/BN=256)). For tiny out_f (ssm_beta/ssm_alpha out_f=num_v_heads~32), grid.x=1 -> only
10040        // ceil(m/256) CTAs (e.g. 2 for m=512) on 82 SMs = 0.39% SM throughput, 852us EACH (measured
10041        // worst offender). The dp4a path grids (out_f, m) = far more CTAs, filling the GPU. So route
10042        // out_f < 2*BM to dp4a (skip the tiling GEMM which structurally can't fill the SMs here).
10043        const GEMM_MIN_OUT_F: usize = 128; // 2*BM; below this the GEMM grid.x starves the 82 SMs
10044        // VENDORED llama MMQ prefill GEMMs. NVFP4 W4A8 is DEFAULT-ON (2026-07-05 flip: same int8
10045        // accuracy class as the int8 GEMM below at ~1.9x pp512, rp-loader coexists with the A6
10046        // repack; MEMRA_MMQ_W4A8=0 = escape hatch). W4A4 mxf4nvf4 + Q4_K/Q5_K stay behind MEMRA_MMQ=1.
10047        // The env policy lives in mmq_supports/qmatvec_mmq. Feeds raw f32 activation `x` (the
10048        // launcher quantizes internally). out_f>=MMQ_Y/2 keeps the tile grid from starving the SMs.
10049        // FP8-ACT PREFILL (MEMRA_PP_FP8=1, probe verdict 2026-07-08): F8-E4M3-origin projections
10050        // carry their raw e4m3 device bytes (the `fp8` operand stashed at load next to the Q8_0
10051        // re-encode) — cuBLASLt FP8 TN at 620-795 TF vs 47-72 TF for this class's int8 GEMM.
10052        // Weight side EXACT (checkpoint bytes); activation rides ONE per-batch e4m3 scale
10053        // (amax/448) folded with weight_scale in-GEMM. Prefill only; decode keeps Q8_0 untouched.
10054        if m >= GEMM_M_THRESHOLD {
10055            if let Some(y) = self.try_fp8_gemm(w, x, m)? {
10056                return Ok(y);
10057            }
10058            // PER-BLOCK FP8 MMQ (lane/fp8-mmq): the block-128 class try_fp8_gemm skips (cuBLASLt
10059            // takes no block grid on sm_120). Exact per block — the checkpoint's e4m3 bytes and its
10060            // f32 grid go into the tile unchanged. TWO SOURCES, TWO DEFAULTS: the load-time stash is
10061            // opt-in (MEMRA_FP8_MMQ=1), the native-resident QT_F8_E4M3_BLK grid is DEFAULT ON
10062            // (MEMRA_FP8_MMQ=0 reverts it to dequant-per-call) — see fp8_ffi.rs for why the same
10063            // tile defaults differently by operand source.
10064            if let Some(y) = self.try_fp8_blk_mmq(w, x, m)? {
10065                return Ok(y);
10066            }
10067            // FP16-mirror prefill (MEMRA_PP_F16=1, probe 2026-07-26: 3.2-3.7x the MMQ class).
10068            // Mirror presence IS the gate (only built under the env). Decode never reaches here.
10069            if let Some(y) = self.try_f16_gemm(w, x, m)? {
10070                return Ok(y);
10071            }
10072        }
10073        // F8-E4M3 BLOCK-128 (QT_F8_E4M3_BLK, lane/fp8-blk128-decode). TWO arms, split at the SAME
10074        // m threshold the rest of this method uses:
10075        //   * m >= threshold (prefill): dequant-per-call to the ARM B' Q8_0 slab and recurse, so
10076        //     prefill keeps the floor's kernels AND the floor's bits (try_e4m3_blk_prefill).
10077        //   * m <  threshold: the native per-block GEMV — m=1 decode and the m=2..15 verify tiers.
10078        //     grid.y=m runs the exact m=1 program per (token,row), so the decode-parity law holds
10079        //     across every tier by construction with no batched twin needed.
10080        //
10081        // NOT gated on `fast`: this dtype has no dp4a twin and no Stage-A f32-dequant oracle (the
10082        // generic `deq()` switch has no block-scale input), exactly as QT_F8_E4M3 has none, so
10083        // MEMRA_FAST=0 cannot route it anywhere else. Placed before every GEMM/MMQ arm below
10084        // because gemm_supports/mmq_supports/mmvq_supports all deliberately REFUSE this qtype —
10085        // reaching the generic tail would panic rather than produce wrong numbers, and this pair of
10086        // arms is what makes sure it never gets there.
10087        if let GpuTensor::Quant { qtype, .. } = w {
10088            if *qtype == QT_F8_E4M3_BLK {
10089                if m >= GEMM_M_THRESHOLD {
10090                    if let Some(y) = self.try_e4m3_blk_prefill(w, x, m)? {
10091                        return Ok(y);
10092                    }
10093                }
10094                let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10095                if let Some(y) = self.try_e4m3_blk_pre(w, &aq, &ad, m)? {
10096                    return Ok(y);
10097                }
10098            }
10099        }
10100        if m >= GEMM_M_THRESHOLD && out_f >= GEMM_MIN_OUT_F && self.mmq_supports(w) {
10101            return self.qmatvec_mmq(w, x, m);
10102        }
10103        if m >= GEMM_M_THRESHOLD && out_f >= GEMM_MIN_OUT_F && self.gemm_supports(w) {
10104            let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10105            return self.qmatvec_gemm(w, &aq, &ad, m);
10106        }
10107        // FP4 W4A4 only as an explicit speed/accuracy tradeoff opt-in, and only if the int8 GEMM
10108        // above didn't already handle this weight (e.g. NVFP4 with in_f%64!=0, or MEMRA_NO_GEMM set).
10109        if m >= GEMM_M_THRESHOLD {
10110            if let Some(y) = self.try_fp4_gemm(w, x, m, in_f, out_f)? {
10111                return Ok(y);
10112            }
10113        }
10114        // Stage-B fast int8 dp4a is the DEFAULT since 2026-07-08 (it has been the daily path
10115        // for weeks; the old opt-in flag was a silent-slow-path landmine). MEMRA_FAST=0 reverts
10116        // to Stage-A f32-dequant (the correctness oracle path).
10117        let fast = std::env::var("MEMRA_FAST").as_deref() != Ok("0");
10118        // PERF-3 decode-GEMV: m=1 warp-per-row MMVQ (MEMRA_MMVQ). The big decode matvecs reach
10119        // `matmul` directly (ffn_down, lm_head output, wo), so route them here too — not only the
10120        // matmul_pre siblings. qmatvec_mmvq_raw quantizes the activation internally (q8_1) like the
10121        // _fast paths; the NVFP4 macro-scale is applied by the `scale != 1.0` block below.
10122        if m == 1 && fast {
10123            if let GpuTensor::Quant {
10124                bytes,
10125                qtype,
10126                row_bytes,
10127                rp,
10128                rp4,
10129                scale,
10130                ..
10131            } = w
10132            {
10133                if self.mmvq_supports(*qtype) {
10134                    // NVFP4 macro-scale rides the kernel's fused epilogue arg (one launch total);
10135                    // non-NVFP4 has scale==1.0 so qmatvec_mmvq skips scale_inplace either way.
10136                    // Q4_0 split-plane mirror (rp4): the decode arm reads it via the _rp twins.
10137                    let (bytes, rp) = match rp4 {
10138                        Some(m4) => (m4, true),
10139                        None => (bytes, *rp),
10140                    };
10141                    let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10142                    return self.qmatvec_mmvq(
10143                        bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, *scale, rp,
10144                    );
10145                }
10146            }
10147        }
10148        // BATCHED weight-resident matvec for the m=2-4 band (the MTP/verify forward's ffn_down, wo, and
10149        // lm_head `output` reach `matmul` directly at m=T=2..4). Walks the weight ONCE, dp4a vs all m
10150        // activation columns -> 1 weight read for m tokens (vs grid.y=m re-reading m times below). Quant
10151        // the activation once here (q8_1) like the _fast paths; macro-scale applied via the scale!=1.0
10152        // block below. MEMRA_NO_BATCHED -> per-m path.
10153        //
10154        // DECODE-PARITY GATE (2026-07-07, the 9B synth K=3/4/6 spec FAIL root cause): the batched
10155        // kernels are bit-identical per (token,row) to MMVQ's 32-thread warp reduce, NOT to the
10156        // dp4a kernels' 128-thread two-level reduce. Without MEMRA_MMVQ the m=1 decode chain rides
10157        // dp4a, so a verify riding batched here has a DIFFERENT FP order than the decode it must
10158        // match bit-for-bit — greedy spec flips at tight-margin tokens (the old HANDOVER "ENV LAW:
10159        // FAST+MMVQ both required" footgun, closed here). Parity law: the m>1 kernel CLASS must be
10160        // a pure function of (dtype, env) equal to the m=1 class — batched iff MMVQ. Without MMVQ
10161        // the verify falls to the per-m grid.y=m dp4a path below (each column = the exact m=1
10162        // dp4a program). MEMRA_MMVQ=1 (the daily config) is dispatch-unchanged.
10163        if (2..=16).contains(&m)
10164            && fast
10165            && std::env::var("MEMRA_NO_BATCHED").is_err()
10166            && (m <= 4 || Self::b8_enabled())
10167        {
10168            // b16 tier (2026-07-11, spec K>7): Q4_0/Q6_K have base+_rp b16 kernels; Q8_0's
10169            // b16 exists only as the split-plane _rp twin, so it joins iff the q8rp mirror
10170            // is present (rp4) — the mirror pick below then routes to the _rp family.
10171            // QT_F8_E4M3 joins unconditionally (lane/rp-on-st): its b16 IS the base kernel,
10172            // because the native e4m3 row layout is already aligned and needs no mirror.
10173            // NVFP4/Q4_K/Q8_0 all join unconditionally now (lane/rp-on-st): each has base + _rp
10174            // b16 twins, so either residency layout has its aligned form at this width. Q8_0's
10175            // old `rp4.is_some()` precondition is GONE — the mirror is a bandwidth lever, not the
10176            // exact tier's admission ticket (it was refusing FP8-ST over 23.9 MiB of ssm_beta).
10177            let m_ok = m <= 8
10178                || matches!(w, GpuTensor::Quant { qtype, .. }
10179                if *qtype == QT_Q4_0 || *qtype == QT_Q6_K || *qtype == QT_F8_E4M3
10180                    || *qtype == QT_NVFP4 || *qtype == QT_Q4_K || *qtype == QT_Q5_K || *qtype == QT_Q8_0);
10181            if m_ok {
10182                if let GpuTensor::Quant {
10183                    bytes,
10184                    qtype,
10185                    row_bytes,
10186                    rp,
10187                    rp4,
10188                    ..
10189                } = w
10190                {
10191                    if self.batched_supports(*qtype) && self.mmvq_supports(*qtype) {
10192                        let (bytes, rp) = match rp4 {
10193                            Some(m4) => (m4, true),
10194                            None => (bytes, *rp),
10195                        };
10196                        let mcols = Self::batched_mcols(m);
10197                        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10198                        let mut y = self.qmatvec_mmvq_batched(
10199                            bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, mcols, 1.0, rp,
10200                        )?;
10201                        if let GpuTensor::Quant { scale, .. } = w {
10202                            if *scale != 1.0 {
10203                                self.scale_inplace(&mut y, *scale, m * out_f)?;
10204                            }
10205                        }
10206                        return Ok(y);
10207                    }
10208                }
10209            }
10210        }
10211        // F8-E4M3 (MEMRA_ST_E4M3) catch-all for the m<16 band the arms above didn't take (m=9..15,
10212        // the K=8 verify tier; or m=2..8 under MEMRA_NO_BATCHED/MEMRA_B8=0): grid.y=m e4m3 mmvq —
10213        // the SAME per-(token,row) program as the m=1 decode launch (bit-identical by construction),
10214        // weight re-read m times (rare tier; exactness over bandwidth here). There is no _dp4a twin
10215        // for this dtype, so the generic match below must never see it under `fast`.
10216        if fast {
10217            if let GpuTensor::Quant {
10218                bytes,
10219                qtype,
10220                row_bytes,
10221                scale,
10222                ..
10223            } = w
10224            {
10225                if *qtype == QT_F8_E4M3 {
10226                    let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10227                    return self.qmatvec_mmvq(
10228                        bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, *scale, false,
10229                    );
10230                }
10231            }
10232        }
10233        let mut y = match w {
10234            GpuTensor::Quant {
10235                bytes,
10236                qtype,
10237                row_bytes,
10238                ..
10239            } if fast && *qtype == QT_Q8_0 => {
10240                self.qmatvec_q8_0_fast(bytes, x, m, in_f, out_f, *row_bytes)?
10241            }
10242            GpuTensor::Quant {
10243                bytes,
10244                qtype,
10245                row_bytes,
10246                ..
10247            } if fast && *qtype == QT_Q4_K => {
10248                self.qmatvec_q4_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
10249            }
10250            GpuTensor::Quant {
10251                bytes,
10252                qtype,
10253                row_bytes,
10254                ..
10255            } if fast && *qtype == QT_Q6_K => {
10256                self.qmatvec_q6_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
10257            }
10258            GpuTensor::Quant {
10259                bytes,
10260                qtype,
10261                row_bytes,
10262                ..
10263            } if fast && *qtype == QT_Q5_K => {
10264                self.qmatvec_q5_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
10265            }
10266            GpuTensor::Quant {
10267                bytes,
10268                qtype,
10269                row_bytes,
10270                ..
10271            } if fast && *qtype == QT_Q3_K => {
10272                self.qmatvec_q3_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
10273            }
10274            GpuTensor::Quant {
10275                bytes,
10276                qtype,
10277                row_bytes,
10278                rp,
10279                ..
10280            } if fast && *qtype == QT_NVFP4 => self.qmatvec_dp4a_named(
10281                if *rp {
10282                    "qmatvec_nvfp4_dp4a_rp"
10283                } else {
10284                    "qmatvec_nvfp4_dp4a"
10285                },
10286                bytes,
10287                x,
10288                m,
10289                in_f,
10290                out_f,
10291                *row_bytes,
10292            )?,
10293            // IQ4_XS trunk fast path — DEFAULT ON since 2026-08-02 (MEMRA_IQ_FAST=0 reverts to
10294            // Stage-A; see iq_fast_enabled). The old opt-in default was the KAT-Coder decode
10295            // anomaly (research/kat-anomaly-20260802/).
10296            GpuTensor::Quant {
10297                bytes,
10298                qtype,
10299                row_bytes,
10300                ..
10301            } if fast && *qtype == QT_IQ4_XS && Self::iq_fast_enabled() => {
10302                self.qmatvec_iq4_XS_fast(bytes, x, m, in_f, out_f, *row_bytes)?
10303            }
10304            // B3: IQ3_S uses the Stage-A f32 dequant-in-kernel path. There is NO
10305            // qmatvec_iq3_s_dp4a kernel — do NOT add a `*qtype == QT_IQ3_S` fast guard here
10306            // without first writing the matching kernel, or func() will panic
10307            // "kernel ... not in any fatbin".
10308            GpuTensor::Quant {
10309                bytes,
10310                qtype,
10311                row_bytes,
10312                rp,
10313                ..
10314            } =>
10315            // Stage-A generic: repacked NVFP4 uses the device-side split-plane tag (the
10316            // deq(row,j) form cannot address the planes; same value/product order).
10317            {
10318                self.qmatvec(
10319                    bytes,
10320                    x,
10321                    m,
10322                    in_f,
10323                    out_f,
10324                    if *rp && *qtype == QT_NVFP4 {
10325                        QT_NVFP4_RP
10326                    } else {
10327                        *qtype
10328                    },
10329                    *row_bytes,
10330                )?
10331            }
10332            GpuTensor::Float { data, .. } => self.linear(x, data, m, in_f, out_f)?,
10333            // MEMRA_FULL_PREC bf16-resident weight: dequant-on-use to f32 scratch, then the same
10334            // cuBLASLt f32 GEMV as the Float arm.
10335            GpuTensor::FloatBf16 { data, .. } => {
10336                self.linear_bf16_chunked(x, data, m, in_f, out_f, false)?
10337            }
10338        };
10339        // NVFP4 per-tensor macro-scale (post-matmul). scale==1.0 for all other quants/float -> no-op.
10340        if let GpuTensor::Quant { scale, .. } = w {
10341            if *scale != 1.0 {
10342                self.scale_inplace(&mut y, *scale, m * out_f)?;
10343            }
10344        }
10345        Ok(y)
10346    }
10347
10348    /// True when `MEMRA_FAST=0`, i.e. the Stage-A f32 oracle is the requested arithmetic and every
10349    /// `matmul_pre` call will take the raw-f32 escape rather than the q8_1 pair.
10350    ///
10351    /// WHY THIS EXISTS AS ITS OWN PREDICATE: `uses_q8_1_fast` needs a weight, but the callers that
10352    /// have to DECIDE WHETHER TO MATERIALIZE an f32 activation sit one level above any weight — the
10353    /// gemma-4 decode trunk emits a q8_1 pair per layer and hands it down, so the f32 has to be
10354    /// produced (or not) by the loop that owns the residual. Cached in a OnceLock like every other
10355    /// `*_on()` flag here: this is read once per layer per token on the decode path, and the daily
10356    /// path must not pay an env lookup for a flag that is off.
10357    pub fn stage_a_raw_needed() -> bool {
10358        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
10359        *ON.get_or_init(|| std::env::var("MEMRA_FAST").as_deref() == Ok("0"))
10360    }
10361
10362    /// True if `w` would take the int8-dp4a fast path under MEMRA_FAST (so its activation can be
10363    /// pre-quantized once and shared across sibling matmuls via `matmul_pre`).
10364    pub fn uses_q8_1_fast(&self, w: &crate::model::GpuTensor) -> bool {
10365        use crate::model::GpuTensor;
10366        if std::env::var("MEMRA_FAST").as_deref() == Ok("0") {
10367            return false;
10368        }
10369        match w {
10370            // QT_F8_E4M3_BLK is admitted for the same reason QT_F8_E4M3 is: its ONLY kernel class
10371            // takes the shared q8_1 activation, so callers may pre-quantize once and share it
10372            // across siblings. It is NOT admitted to any of the fused/dual epilogue doors those
10373            // siblings can then open (`q8_fused_params`, `e4m3_fused_params` and
10374            // `matmul_pre_dual_noscale` all match on their own qtype and refuse this one) — the
10375            // block class has no fused twin yet, so each of its projections takes its own launch.
10376            GpuTensor::Quant { qtype, .. } => {
10377                matches!(
10378                    *qtype,
10379                    QT_Q8_0
10380                        | QT_Q4_K
10381                        | QT_Q6_K
10382                        | QT_Q5_K
10383                        | QT_Q3_K
10384                        | QT_NVFP4
10385                        | QT_F8_E4M3
10386                        | QT_F8_E4M3_BLK
10387                        | QT_Q4_0
10388                ) || (*qtype == QT_IQ4_XS && Self::iq_fast_enabled())
10389            }
10390            GpuTensor::Float { .. } | GpuTensor::FloatBf16 { .. } => false,
10391        }
10392    }
10393
10394    /// matmul with a PRE-QUANTIZED q8_1 activation (aq,ad from `quantize_q8_1`). Skips the
10395    /// per-matmul re-quantize so sibling matmuls that share an input (gate+up share `z`;
10396    /// q/k/v + wqkv/gate/beta/alpha share `h`) quantize ONCE. Caller MUST have checked
10397    /// `uses_q8_1_fast(w)`; falls back to plain `matmul` otherwise (Stage-A / Float / non-fast).
10398    pub fn matmul_pre(
10399        &self,
10400        w: &crate::model::GpuTensor,
10401        aq: &CudaSlice<i8>,
10402        ad: &CudaSlice<f32>,
10403        x_fallback: &CudaSlice<f32>,
10404        m: usize,
10405    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10406        use crate::model::GpuTensor;
10407        // Every raw-f32 arm below (fp8/f16/MMQ/fp4) reads m*in_f from x_fallback. Callers that
10408        // pre-quantized and dropped the f32 input pass an EMPTY x_fallback (E4B's fusion port:
10409        // h = zeros(0)) — the length guard keeps those on the aq/ad GEMM instead of feeding a
10410        // 0-byte buffer to a convert kernel (illegal address -> cublasLt status 13; the E4B
10411        // rc=30013 dig, 2026-07-31).
10412        let x_raw_ok = x_fallback.len() >= m * w.in_features();
10413        // FP8-ACT PREFILL (MEMRA_PP_FP8=1): same arm as `matmul` — the fp8 operand needs the RAW
10414        // f32 activation (per-batch e4m3 quant differs from q8_1), so x_fallback not aq/ad.
10415        if m >= 16 && x_raw_ok && !self.verify_exact_on() {
10416            if let Some(y) = self.try_fp8_gemm(w, x_fallback, m)? {
10417                return Ok(y);
10418            }
10419            // PER-BLOCK FP8 MMQ — same arm as `matmul` (stash opt-in, native-resident default ON);
10420            // its own quantizer wants the RAW f32 activation, so x_fallback not aq/ad.
10421            if let Some(y) = self.try_fp8_blk_mmq(w, x_fallback, m)? {
10422                return Ok(y);
10423            }
10424            // FP16-mirror prefill (same arm as `matmul` — fp16 wants the RAW f32 activation).
10425            if let Some(y) = self.try_f16_gemm(w, x_fallback, m)? {
10426                return Ok(y);
10427            }
10428        }
10429        // BLOCK-128 e4m3 (QT_F8_E4M3_BLK) — the same two arms as `matmul`, split at the same m, and
10430        // placed at the same point in the order (after the prefill GEMM hooks, before every arm
10431        // that refuses this qtype). The prefill arm needs the RAW f32 activation for the Q8_0
10432        // dispatch it recurses into, so it takes x_fallback and is skipped when that is empty
10433        // (a pre-quantized caller that dropped its f32 input never runs at prefill m anyway).
10434        if m >= 16 && x_raw_ok && !self.verify_exact_on() {
10435            if let Some(y) = self.try_e4m3_blk_prefill(w, x_fallback, m)? {
10436                return Ok(y);
10437            }
10438        }
10439        if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
10440            return Ok(y);
10441        }
10442        // VENDORED llama MMQ prefill GEMMs (NVFP4 W4A8 default-on; W4A4/k-quant behind MEMRA_MMQ=1
10443        // — policy in mmq_supports) — use the RAW f32 activation (their own internal quant:
10444        // q8_1 D4 for NVFP4 W4A8, FP8/UE4M3 for W4A4, q8_1 DS4 for Q4_K/Q5_K), so x_fallback not
10445        // aq/ad.
10446        if m >= 16
10447            && w.out_features() >= 128
10448            && self.mmq_supports(w)
10449            && !self.verify_exact_on()
10450            && x_raw_ok
10451        {
10452            return self.qmatvec_mmq(w, x_fallback, m);
10453        }
10454        // Stage-C FP4 prefill (MEMRA_FP4): native mxf4 GEMM needs the f32 activation (FP4-quant differs
10455        // from q8_1), so re-quantize from x_fallback rather than reuse aq/ad. NVFP4 only, m>=16.
10456        if m >= 16 && x_raw_ok && !self.verify_exact_on() {
10457            if let Some(y) =
10458                self.try_fp4_gemm(w, x_fallback, m, w.in_features(), w.out_features())?
10459            {
10460                return Ok(y);
10461            }
10462        }
10463        // Prefill GEMM root fix: if T>1 and the dtype has a GEMM kernel, batch via tensor cores
10464        // (reuses the already-quantized aq/ad — no extra quantize). m=1 falls through to dp4a.
10465        if m >= 16 && self.gemm_supports(w) && !self.verify_exact_on() {
10466            return self.qmatvec_gemm(w, aq, ad, m);
10467        }
10468        // THE SEVENTH ARM. Every raw-f32 arm above is guarded by `x_raw_ok`; this one — the
10469        // Stage-A / Float escape — was not, and it is the ONLY one `MEMRA_FAST=0` opens. So the
10470        // 2026-07-31 E4B rc=30013 fix (add the length guard, keep empty-fallback callers off the
10471        // raw-f32 arms) protected the six arms the FAST path can reach and missed the one the
10472        // ORACLE path reaches. Consequence before this guard: the gemma-4 decode arms, which pass
10473        // `e.zeros(0)` because their trunk emits only a q8_1 pair and never materializes an f32
10474        // attn-normed activation, fell through to Stage-A `qmatvec_f32` (cu/qmatvec.cu:5442),
10475        // which reads `m * in_f` floats out of a 0-byte allocation ->
10476        // CUDA_ERROR_ILLEGAL_ADDRESS at layer 0 of the first decode token. The fault is STICKY:
10477        // it poisons the context, so every LATER request in that process fails with an unrelated
10478        // message ("cache alloc failed: ...") and the true cause appears exactly once, in the
10479        // first failure. That is what made `MEMRA_FAST=0` — the reference named in
10480        // `tools/argmax-margin-gate.sh`'s own header — return an opaque HTTP 500 on the gemma
10481        // dense artifact and left the arm with no working truth instrument.
10482        //
10483        // Refuse loudly instead of reading out of bounds. A named error at the true call site is
10484        // strictly better than an illegal address surfacing later at an unrelated sync point, and
10485        // an oracle that cannot run must say so rather than corrupt the context it runs in.
10486        if !self.uses_q8_1_fast(w) {
10487            if !x_raw_ok {
10488                return Err(format!(
10489                    "matmul_pre: q8_1-fast is off for this weight but x_fallback holds {} f32 \
10490                     (need m*in_f = {}*{} = {}). This call site pre-quantized its activation and \
10491                     dropped the f32, so there is nothing to fall back to — pass the real f32 \
10492                     activation (see Engine::rms_norm_decode, which is bit-identical to \
10493                     rms_norm_q8_1's reduction) or keep the weight on the q8_1 path.",
10494                    x_fallback.len(),
10495                    m,
10496                    w.in_features(),
10497                    m * w.in_features()
10498                )
10499                .into());
10500            }
10501            return self.matmul(w, x_fallback, m);
10502        }
10503        let in_f = w.in_features();
10504        let out_f = w.out_features();
10505        let (bytes, qtype, row_bytes, scale, rp) = match w {
10506            GpuTensor::Quant {
10507                bytes,
10508                qtype,
10509                row_bytes,
10510                scale,
10511                rp,
10512                ..
10513            } => (bytes, *qtype, *row_bytes, *scale, *rp),
10514            _ => unreachable!("uses_q8_1_fast guaranteed Quant"),
10515        };
10516        // Q4_0 split-plane mirror: only the mmvq/batched decode arms read it (the _rp twins);
10517        // the dp4a/oracle tails below keep the raw GGUF bytes.
10518        let (mbytes, mrp) = match w {
10519            GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
10520            _ => (bytes, rp),
10521        };
10522        // PERF-3 decode-GEMV: warp-per-row MMVQ for the m=1 decode arm, gated behind MEMRA_MMVQ.
10523        // Only the 4 daily-hot dtypes have an _mmvq kernel (Q8_0/Q4_K/Q6_K/NVFP4); Q5_K/Q3_K/IQ4_XS
10524        // keep _dp4a (the oracle/fallback). Bit-equivalent to _dp4a up to f32 reduction order.
10525        if m == 1 && self.mmvq_supports(qtype) {
10526            return self.qmatvec_mmvq(mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, mrp);
10527        }
10528        // BATCHED weight-resident matvec for the m=2-4 band (the MTP/verify forward: full_attn_verify
10529        // and decode_step_t run their projections at m=T=k=2..4). The plain _dp4a path below launches
10530        // grid.y=m INDEPENDENT blocks per output row -> the weight row is re-read m times from HBM/L2.
10531        // The _b2/_b4 kernels walk the weight ONCE and dp4a vs all m activation columns, so m tokens
10532        // cost ~1 weight read instead of m (decode is weight-BW-bound). BIT-IDENTICAL per (token,row)
10533        // to the _mmvq path (32-thread warp reduce — NOT the dp4a 128-thread reduce below).
10534        // m=2 -> mcols=2; m∈{3,4} -> mcols=4; m∈{5..8} -> mcols=8 (kernel guards c>=m).
10535        // MEMRA_NO_BATCHED forces the per-m grid.y=m path (the A/B reference); MEMRA_B8=0 keeps
10536        // m=5..8 on the old per-m path (b8-tier-only seam).
10537        // DECODE-PARITY GATE (2026-07-07): batched iff mmvq_supports — see matmul's parity note.
10538        // Without MEMRA_MMVQ, m=1 decode rides dp4a (the arm below at m=1); the verify must ride
10539        // the SAME class per column (grid.y=m dp4a = the exact m=1 dp4a program per column).
10540        if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
10541            && std::env::var("MEMRA_NO_BATCHED").is_err()
10542            && (m <= 4 || Self::b8_enabled())
10543            // b16 tier: every class routed here now has base + _rp b16 kernels (Q4_0/Q6_K
10544            // pre-existing; NVFP4/Q4_K/Q8_0-base/F8_E4M3 added lane/rp-on-st 2026-08-06), so
10545            // there is no mirror precondition left — `mrp` still selects the LAYOUT below.
10546            && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || qtype == QT_NVFP4
10547                || qtype == QT_Q4_K || qtype == QT_Q5_K || qtype == QT_F8_E4M3 || qtype == QT_Q8_0)
10548        {
10549            let mcols = Self::batched_mcols(m);
10550            return self.qmatvec_mmvq_batched(
10551                mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, mrp,
10552            );
10553        }
10554        // F8-E4M3 catch-all (m=9..15 / batched-disabled seams): grid.y=m e4m3 mmvq — this dtype
10555        // has NO _dp4a twin, and per (token,row) the mmvq body is the exact m=1 decode program.
10556        // Q4_0 joins the catch-all (2026-07-11): adaptive-K cap 8 makes verify t=9 reachable
10557        // for the first time (past the b8 tier) and Q4_0 has no dp4a twin either. The mirror
10558        // (mbytes/mrp) keeps the rp layout consistent with the m=1 decode program.
10559        if qtype == QT_F8_E4M3 || qtype == QT_Q4_0 {
10560            let (b2, r2) = if qtype == QT_Q4_0 {
10561                (mbytes, mrp)
10562            } else {
10563                (bytes, rp)
10564            };
10565            return self.qmatvec_mmvq(b2, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, r2);
10566        }
10567        let name = match qtype {
10568            QT_Q8_0 => "qmatvec_q8_0_dp4a",
10569            QT_Q4_K => "qmatvec_q4_K_dp4a",
10570            QT_Q6_K => "qmatvec_q6_K_dp4a",
10571            QT_Q5_K => "qmatvec_q5_K_dp4a",
10572            QT_Q3_K => "qmatvec_q3_K_dp4a",
10573            QT_NVFP4 => {
10574                if rp {
10575                    "qmatvec_nvfp4_dp4a_rp"
10576                } else {
10577                    "qmatvec_nvfp4_dp4a"
10578                }
10579            }
10580            QT_IQ4_XS => "qmatvec_iq4_XS_dp4a",
10581            _ => unreachable!(),
10582        };
10583        let f = self.func(name);
10584        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output: skip memset
10585        let cfg = LaunchConfig {
10586            grid_dim: (out_f as u32, m as u32, 1),
10587            block_dim: (128, 1, 1),
10588            shared_mem_bytes: 0,
10589        };
10590        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
10591        let __s_b = self.gpu.stream();
10592        let mut b = __s_b.launch_builder(&f);
10593        b.arg(bytes)
10594            .arg(aq)
10595            .arg(ad)
10596            .arg(&mut y)
10597            .arg(&inf)
10598            .arg(&outf)
10599            .arg(&mi)
10600            .arg(&rb);
10601        unsafe {
10602            b.launch(cfg)?;
10603        }
10604        if scale != 1.0 {
10605            self.scale_inplace(&mut y, scale, m * out_f)?;
10606        }
10607        Ok(y)
10608    }
10609
10610    /// DECODE-EXACT matmul at any m: guarantees the SAME warp-per-row (MMVQ, 32-thread) FP
10611    /// accumulation order as the T=1 decode path for EVERY token row. The spec-decode verify MUST
10612    /// use this for linear-attn projections to be bit-identical to greedy decode. The dp4a kernel
10613    /// (128 threads, two-level reduction) used by `matmul`/`matmul_pre` at m>=5 has a different
10614    /// shfl-tree shape that produces ULP differences propagating through gdn_scan into argmax flips.
10615    /// The MMVQ kernel with grid.y=m already processes each row independently (same 32-thread warp
10616    /// reduce as m=1); this method just forces that path unconditionally.
10617    pub fn matmul_decode_exact(
10618        &self,
10619        w: &crate::model::GpuTensor,
10620        x: &CudaSlice<f32>,
10621        m: usize,
10622    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10623        use crate::model::GpuTensor;
10624        // FLOAT tensors (35B ssm_beta/ssm_alpha on every linear layer, F32 ne=[2048,32]): the
10625        // generic path is cuBLASLt, whose reduction splits are n-DEPENDENT — m=1 vs m=2 col-0
10626        // outputs differ in every bit (probe 2026-07-06: 32/32 bit-diff, maxdiff 3.5e-3), which
10627        // shifted 35B verify logits 0.26-0.56 vs eager and flipped greedy at tight margins (the
10628        // p3 spec FAIL). Decode-exact contract: per-COLUMN m=1 cuBLASLt calls — each column's
10629        // reduction is the exact kernel the T=1 decode path runs, so verify==decode bit-for-bit.
10630        // m<=10 here (K+2 verify tier), so the extra launches are a handful of 4us gemvs.
10631        if let GpuTensor::Float { data, .. } = w {
10632            return self.linear_decode_exact(x, data, m, w.in_features(), w.out_features());
10633        }
10634        // MEMRA_FULL_PREC bf16-resident weight: dequant-on-use, then the per-column decode-exact
10635        // float linear (same n-independent reduction contract as the Float arm above).
10636        if let GpuTensor::FloatBf16 { data, .. } = w {
10637            let (in_f, out_f) = (w.in_features(), w.out_features());
10638            return self.linear_bf16_chunked(x, data, m, in_f, out_f, true);
10639        }
10640        if !self.uses_q8_1_fast(w) {
10641            return self.matmul(w, x, m);
10642        }
10643        let in_f = w.in_features();
10644        let out_f = w.out_features();
10645        let (bytes, qtype, row_bytes, scale, rp) = match w {
10646            GpuTensor::Quant {
10647                bytes,
10648                qtype,
10649                row_bytes,
10650                scale,
10651                rp,
10652                ..
10653            } => (bytes, *qtype, *row_bytes, *scale, *rp),
10654            _ => return self.matmul(w, x, m),
10655        };
10656        // Q4_0 split-plane mirror for the mmvq/batched arms below (dp4a tail = matmul_pre,
10657        // which does its own mirror pick).
10658        let (bytes, rp) = match w {
10659            GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
10660            _ => (bytes, rp),
10661        };
10662        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10663        // BLOCK-128 e4m3 (QT_F8_E4M3_BLK): the same single kernel every other entry dispatches, so
10664        // the decode-exact contract needs nothing special — grid.y=m runs the m=1 program per
10665        // (token,row) by construction, which is exactly what this method exists to guarantee.
10666        if let Some(y) = self.try_e4m3_blk_pre(w, &aq, &ad, m)? {
10667            return Ok(y);
10668        }
10669        // Batched weight-resident matvec for m=2-8: BIT-IDENTICAL per (token,row) to MMVQ (exact
10670        // integer dp4a, same warp reduce — kernel-check gate rel=0.00e0), one weight read for m
10671        // tokens. The dispatch the divergence fix must avoid is dp4a's 128-thread two-level
10672        // reduce, NOT this. m=5..8 is the K=4..7 spec-verify tier (b8): pre-b8 T=5 fell to the
10673        // grid.y=m per-row MMVQ below = 5 full weight reads/launch — the measured 27B K=4 cliff.
10674        // DECODE-PARITY GATE (2026-07-07): batched (MMVQ-class order) only when the m=1 decode
10675        // chain rides MMVQ too — without MEMRA_MMVQ decode is dp4a, so the exact-contract here
10676        // must be per-column dp4a (matmul_pre fallthrough), not the MMVQ order.
10677        if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
10678            && std::env::var("MEMRA_NO_BATCHED").is_err()
10679            && (m <= 4 || Self::b8_enabled())
10680            // Every b16 class has base + _rp twins after lane/rp-on-st (see matmul_pre's note):
10681            // no mirror precondition, `rp` selects the layout only.
10682            && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || qtype == QT_F8_E4M3
10683                || qtype == QT_NVFP4 || qtype == QT_Q4_K || qtype == QT_Q5_K || qtype == QT_Q8_0)
10684        {
10685            let mcols = Self::batched_mcols(m);
10686            return self.qmatvec_mmvq_batched(
10687                bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, rp,
10688            );
10689        }
10690        if self.mmvq_supports(qtype) {
10691            // MMVQ at grid.y=m: each row is processed by its own warp independently — same 32-thread
10692            // accumulation + warp_reduce_sum as m=1 decode. Bit-identical per row.
10693            return self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, scale, rp);
10694        }
10695        // Fallback for non-MMVQ quant types (Q5_K, Q3_K): use dp4a (the only available kernel).
10696        // These types are not used in the 27B's linear-attn NVFP4+Q4_K layers.
10697        self.matmul_pre(w, &aq, &ad, x, m)
10698    }
10699
10700    /// DECODE-EXACT matmul from a PRE-QUANTIZED q8_1 activation (batched-verify epilogue
10701    /// re-fuse, lane/vt-fixes fix 2, 2026-08-03): the EXACT `matmul_decode_exact` dispatch for
10702    /// q8_1-fast Quant tensors, with the caller's (aq, ad) replacing the internal
10703    /// `quantize_q8_1`. quantize_q8_1 is deterministic (same input bytes -> same q8 bytes), so
10704    /// sharing one quantize across sibling matmuls of the same activation — or consuming the
10705    /// q8 emitted by a fused epilogue (rms_norm_q8_1 / add_rms_norm_q8_1 /
10706    /// silu_mul_scaled_q8_1 / gated_rmsnorm_q8_1, all kernel-check-pinned bit-identical to
10707    /// their unfused chains) — cannot change any dispatched kernel's input bytes.
10708    /// Caller MUST guarantee `uses_q8_1_fast(w)` (the fused epilogues only exist on that path).
10709    pub fn matmul_decode_exact_pre(
10710        &self,
10711        w: &crate::model::GpuTensor,
10712        aq: &CudaSlice<i8>,
10713        ad: &CudaSlice<f32>,
10714        m: usize,
10715    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10716        use crate::model::GpuTensor;
10717        debug_assert!(
10718            self.uses_q8_1_fast(w),
10719            "matmul_decode_exact_pre: caller must guarantee q8_1-fast"
10720        );
10721        // BLOCK-128 e4m3: same single kernel, all m — see matmul_decode_exact's note.
10722        if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
10723            return Ok(y);
10724        }
10725        let in_f = w.in_features();
10726        let out_f = w.out_features();
10727        let (bytes, qtype, row_bytes, scale, rp) = match w {
10728            GpuTensor::Quant {
10729                bytes,
10730                qtype,
10731                row_bytes,
10732                scale,
10733                rp,
10734                ..
10735            } => (bytes, *qtype, *row_bytes, *scale, *rp),
10736            _ => {
10737                return Err(
10738                    "matmul_decode_exact_pre: Quant tensor required (q8_1-fast contract)".into(),
10739                );
10740            }
10741        };
10742        // Q4_0 split-plane mirror — same pick as matmul_decode_exact.
10743        let (bytes, rp) = match w {
10744            GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
10745            _ => (bytes, rp),
10746        };
10747        // Dispatch mirror of matmul_decode_exact's q8_1-fast tail, condition for condition.
10748        if (2..=16).contains(&m)
10749            && self.batched_supports(qtype)
10750            && self.mmvq_supports(qtype)
10751            && std::env::var("MEMRA_NO_BATCHED").is_err()
10752            && (m <= 4 || Self::b8_enabled())
10753            && (m <= 8
10754                || qtype == QT_Q4_0
10755                || qtype == QT_Q6_K
10756                || qtype == QT_F8_E4M3
10757                || qtype == QT_NVFP4
10758                || qtype == QT_Q4_K
10759                || qtype == QT_Q5_K
10760                || qtype == QT_Q8_0)
10761        {
10762            let mcols = Self::batched_mcols(m);
10763            return self.qmatvec_mmvq_batched(
10764                bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, rp,
10765            );
10766        }
10767        if self.mmvq_supports(qtype) {
10768            return self.qmatvec_mmvq(bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, rp);
10769        }
10770        // Non-MMVQ quant types (Q5_K/Q3_K under MEMRA_MMVQ=0): dp4a via matmul_pre — the same
10771        // fallback matmul_decode_exact takes. m <= 16 on the verify tier never reads x_fallback.
10772        let x0 = self.zeros(0)?;
10773        self.matmul_pre(w, aq, ad, &x0, m)
10774    }
10775
10776    /// DUAL gate+up batched matvec from a PRE-QUANTIZED activation, macro-scales DEFERRED
10777    /// (lane/vt-fixes fix 2): same eligibility as `matmul_decode_exact_dual`, but the caller's
10778    /// (aq, ad) replaces the internal quantize and the NVFP4 per-tensor scales are RETURNED
10779    /// instead of applied via two `scale_inplace` launches — the fused SwiGLU epilogue
10780    /// (`silu_mul_scaled_q8_1`) folds them, exactly like the m=1 decode chain does. Deferring
10781    /// is value-exact: `y[i]*s` inline in the epilogue is the same IEEE multiply scale_inplace
10782    /// would store (f32 store/load round-trips are exact). None -> caller falls back to the
10783    /// per-tensor path.
10784    pub fn matmul_decode_exact_dual_pre(
10785        &self,
10786        w0: &crate::model::GpuTensor,
10787        w1: &crate::model::GpuTensor,
10788        aq: &CudaSlice<i8>,
10789        ad: &CudaSlice<f32>,
10790        m: usize,
10791    ) -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn std::error::Error>>
10792    {
10793        use crate::model::GpuTensor;
10794        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
10795        let on = *ON.get_or_init(|| {
10796            std::env::var("MEMRA_SPEC_DUAL_T")
10797                .map(|v| v != "0")
10798                .unwrap_or(true)
10799        });
10800        if !on
10801            || !(2..=7).contains(&m)
10802            || std::env::var("MEMRA_NO_BATCHED").is_ok()
10803            || !self.uses_q8_1_fast(w0)
10804            || !self.uses_q8_1_fast(w1)
10805        {
10806            return Ok(None);
10807        }
10808        // DECODE-PARITY GATE (lane/nvfp4-strict, 2026-08-05): batched iff MMVQ — the dual
10809        // kernels are the MMVQ warp-reduce family, and without MEMRA_MMVQ the m=1 decode
10810        // chain this verify must match bit-for-bit rides dp4a (see matmul_decode_exact's
10811        // note). The singles enforce this via `mmvq_supports`; the dual door skipped it.
10812        if !self.mmvq_supports(QT_NVFP4) {
10813            return Ok(None);
10814        }
10815        let (in_f, out_f) = (w0.in_features(), w0.out_features());
10816        if w1.in_features() != in_f || w1.out_features() != out_f {
10817            return Ok(None);
10818        }
10819        let (b0, b1, row_bytes, s0, s1, rp) = match (w0, w1) {
10820            (
10821                GpuTensor::Quant {
10822                    bytes: b0,
10823                    qtype: q0,
10824                    row_bytes: rb0,
10825                    scale: s0,
10826                    rp: rp0,
10827                    rp4: None,
10828                    ..
10829                },
10830                GpuTensor::Quant {
10831                    bytes: b1,
10832                    qtype: q1,
10833                    row_bytes: rb1,
10834                    scale: s1,
10835                    rp: rp1,
10836                    rp4: None,
10837                    ..
10838                },
10839            ) if *q0 == QT_NVFP4 && *q1 == QT_NVFP4 && rb0 == rb1 && rp0 == rp1 => {
10840                (b0, b1, *rb0, *s0, *s1, *rp0)
10841            }
10842            _ => return Ok(None),
10843        };
10844        // m=5..7: only the exact-width rp duals exist (vt-fixes fix 1b); GGUF layout keeps
10845        // the singles. The b8 dual (MCOLS=8 at m=5..8) measured FLAT and stays dead.
10846        if m > 4 && !(rp && Self::b8_enabled() && std::env::var("MEMRA_B567").as_deref() != Ok("0"))
10847        {
10848            return Ok(None);
10849        }
10850        let (y0, y1) =
10851            self.qmatvec_batched_dual_raw(b0, b1, aq, ad, m, in_f, out_f, row_bytes, rp)?;
10852        Ok(Some(((y0, s0), (y1, s1))))
10853    }
10854
10855    /// DUAL gate+up BATCHED matvec at verify t=2..8 (lane/verify-economics, 2026-08-02): ONE
10856    /// launch computes both FFN projections of a verify batch — same activation, same shape,
10857    /// blockIdx.y selects the tensor. Per (tensor, token, row) the kernel body is the single
10858    /// batched program on the SAME layout (split-plane rp: b2 rp / b4 rpr2 / b8 rpr2; GGUF:
10859    /// b2 base / b4 r2 / b8 r2) -> BIT-IDENTICAL to the two single `matmul_decode_exact`
10860    /// launches (kernel-check gates bitwise on both layouts; run-spec K=1..8 arbitrates e2e).
10861    /// The one activation quantize replaces two IDENTICAL quantizes of the same `x` (same
10862    /// kernel, same input -> same q8_1 bytes), and the two independent weight streams in one
10863    /// grid restore the memory-level parallelism the two-launch form loses to tail drain +
10864    /// launch gap (m=1 dual_mr2 precedent: DRAM 40% -> 47-50% on the 27B pair).
10865    /// `Some((y0, y1))` only when both tensors are NVFP4, the SAME layout (both rp or both
10866    /// GGUF, no rp4 mirror), identical (in_f, out_f, row_bytes), q8_1-fast, and m in 2..=4
10867    /// (the b2/b4 tiers = verify T for K=1..3, the profitable-K window — the b8 dual measured
10868    /// FLAT vs the rpsc singles x3 interleaved, research/verify-economics-20260802, and was
10869    /// killed per doctrine). None -> caller runs the two singles. MEMRA_SPEC_DUAL_T=0 rollback.
10870    pub fn matmul_decode_exact_dual(
10871        &self,
10872        w0: &crate::model::GpuTensor,
10873        w1: &crate::model::GpuTensor,
10874        x: &CudaSlice<f32>,
10875        m: usize,
10876    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
10877        use crate::model::GpuTensor;
10878        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
10879        let on = *ON.get_or_init(|| {
10880            std::env::var("MEMRA_SPEC_DUAL_T")
10881                .map(|v| v != "0")
10882                .unwrap_or(true)
10883        });
10884        if !on
10885            || !(2..=4).contains(&m)
10886            || std::env::var("MEMRA_NO_BATCHED").is_ok()
10887            || !self.uses_q8_1_fast(w0)
10888            || !self.uses_q8_1_fast(w1)
10889        {
10890            return Ok(None);
10891        }
10892        // DECODE-PARITY GATE (lane/nvfp4-strict, 2026-08-05): batched iff MMVQ — same law as
10893        // the singles' `batched_supports && mmvq_supports` check in matmul_decode_exact,
10894        // which this dual door bypassed. Without MEMRA_MMVQ the m=1 decode is dp4a; the
10895        // verify must ride the per-column dp4a class, not the MMVQ-family dual.
10896        if !self.mmvq_supports(QT_NVFP4) {
10897            return Ok(None);
10898        }
10899        let (in_f, out_f) = (w0.in_features(), w0.out_features());
10900        if w1.in_features() != in_f || w1.out_features() != out_f {
10901            return Ok(None);
10902        }
10903        let (b0, b1, row_bytes, s0, s1, rp) = match (w0, w1) {
10904            (
10905                GpuTensor::Quant {
10906                    bytes: b0,
10907                    qtype: q0,
10908                    row_bytes: rb0,
10909                    scale: s0,
10910                    rp: rp0,
10911                    rp4: None,
10912                    ..
10913                },
10914                GpuTensor::Quant {
10915                    bytes: b1,
10916                    qtype: q1,
10917                    row_bytes: rb1,
10918                    scale: s1,
10919                    rp: rp1,
10920                    rp4: None,
10921                    ..
10922                },
10923            ) if *q0 == QT_NVFP4 && *q1 == QT_NVFP4 && rb0 == rb1 && rp0 == rp1 => {
10924                (b0, b1, *rb0, *s0, *s1, *rp0)
10925            }
10926            _ => return Ok(None),
10927        };
10928        // Engagement receipt (MEMRA_DEBUG=1): the first dead-arm A/B lesson — a `rp: false`
10929        // gate silently no-op'd the whole experiment; prove the arm is live in the log.
10930        if std::env::var("MEMRA_DEBUG").is_ok() {
10931            static ONCE: std::sync::Once = std::sync::Once::new();
10932            ONCE.call_once(|| eprintln!("[memra] dual gate+up batched ENGAGED (m={m} rp={rp})"));
10933        }
10934        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10935        let (y0, y1) =
10936            self.qmatvec_batched_dual_raw(b0, b1, &aq, &ad, m, in_f, out_f, row_bytes, rp)?;
10937        let mut y0 = y0;
10938        let mut y1 = y1;
10939        if s0 != 1.0 {
10940            self.scale_inplace(&mut y0, s0, m * out_f)?;
10941        }
10942        if s1 != 1.0 {
10943            self.scale_inplace(&mut y1, s1, m * out_f)?;
10944        }
10945        Ok(Some((y0, y1)))
10946    }
10947
10948    /// Launch body of the dual batched twins from raw NVFP4 weight bytes + a pre-quantized q8_1
10949    /// activation (kernel-check's bit-equivalence entry; matmul_decode_exact_dual's core).
10950    /// mcols tier = batched_mcols(m); macro-scale NOT applied. `rp` selects the split-plane
10951    /// twins (both buffers must be the repacked layout).
10952    #[allow(clippy::too_many_arguments)]
10953    pub fn qmatvec_batched_dual_raw(
10954        &self,
10955        b0: &CudaSlice<u8>,
10956        b1: &CudaSlice<u8>,
10957        aq: &CudaSlice<i8>,
10958        ad: &CudaSlice<f32>,
10959        m: usize,
10960        in_f: usize,
10961        out_f: usize,
10962        row_bytes: usize,
10963        rp: bool,
10964    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10965        const ROWS_PER_BLOCK: u32 = 4;
10966        let mcols = Self::batched_mcols(m);
10967        // EXACT-WIDTH duals at m=5..7 (vt-fixes fix 1b): rp-only; bit-identical to the two
10968        // b5/b6/b7 singles (blockIdx.y selects the tensor, same template body).
10969        let tiny_rp1 = rp
10970            && mcols == 4
10971            && out_f <= 128
10972            && std::env::var("MEMRA_NVFP4_AUX_DUAL").as_deref() != Ok("0");
10973        let (name, rows_per_block) = if tiny_rp1 {
10974            ("qmatvec_nvfp4_mmvq_dual_b4_rp", ROWS_PER_BLOCK)
10975        } else {
10976            match (mcols, rp, m) {
10977                (2, false, _) => ("qmatvec_nvfp4_mmvq_dual_b2", ROWS_PER_BLOCK),
10978                (4, false, _) => ("qmatvec_nvfp4_mmvq_dual_b4_r2", ROWS_PER_BLOCK * 2),
10979                (2, true, _) => ("qmatvec_nvfp4_mmvq_dual_b2_rp", ROWS_PER_BLOCK),
10980                (4, true, _) => ("qmatvec_nvfp4_mmvq_dual_b4_rpr2", ROWS_PER_BLOCK * 2),
10981                (8, true, 5) => ("qmatvec_nvfp4_mmvq_dual_b5_rpr2", ROWS_PER_BLOCK * 2),
10982                (8, true, 6) => ("qmatvec_nvfp4_mmvq_dual_b6_rpr2", ROWS_PER_BLOCK * 2),
10983                (8, true, 7) => ("qmatvec_nvfp4_mmvq_dual_b7_rpr2", ROWS_PER_BLOCK * 2),
10984                _ => {
10985                    return Err(
10986                        format!("qmatvec_batched_dual_raw: no dual kernel for m {m}").into(),
10987                    );
10988                }
10989            }
10990        };
10991        let f = self.func(name);
10992        let mut y0 = self.alloc_uninit::<f32>(m * out_f)?;
10993        let mut y1 = self.alloc_uninit::<f32>(m * out_f)?;
10994        let cfg = LaunchConfig {
10995            grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 2, 1),
10996            block_dim: (32, ROWS_PER_BLOCK, 1),
10997            shared_mem_bytes: 0,
10998        };
10999        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
11000        let __s_b = self.gpu.stream();
11001        let mut b = __s_b.launch_builder(&f);
11002        b.arg(b0)
11003            .arg(b1)
11004            .arg(aq)
11005            .arg(ad)
11006            .arg(&mut y0)
11007            .arg(&mut y1)
11008            .arg(&inf)
11009            .arg(&outf)
11010            .arg(&mi)
11011            .arg(&rb);
11012        unsafe {
11013            b.launch(cfg)?;
11014        }
11015        Ok((y0, y1))
11016    }
11017
11018    /// Like `matmul_pre` but RETURNS THE RAW (un-macro-scaled) matmul output together with the
11019    /// per-tensor NVFP4 scale, instead of applying `scale_inplace` internally. Used by the fused
11020    /// SwiGLU epilogue (RANK3 LEVER 2) so the gate/up scales fold into one `silu_mul_scaled` launch.
11021    /// `Some((y_raw, scale))` only on the m==1 decode fast path (mmvq / dp4a) where the scale is a
11022    /// separate post-launch op we can defer; returns `None` for every other path (prefill GEMM, FP4
11023    /// GEMM, Stage-A, Float) so the caller falls back to the scaled `matmul_pre` + `silu_mul`.
11024    /// DUAL gate+up NVFP4 matvec (mm-fusion): ONE launch computes both projections (same
11025    /// activation, same shape) — grid.y selects the tensor. Bit-identical per element to two
11026    /// mr2 launches at m=1. Returns (gate_raw, up_raw) un-scaled (caller folds the two macro
11027    /// scales into the SwiGLU epilogue, same as the matmul_pre_noscale contract). None unless
11028    /// both tensors are NVFP4 q8_1-fast with identical (in_f, out_f, row_bytes) and m==1.
11029    pub fn matmul_pre_dual_noscale(
11030        &self,
11031        w0: &crate::model::GpuTensor,
11032        w1: &crate::model::GpuTensor,
11033        aq: &CudaSlice<i8>,
11034        ad: &CudaSlice<f32>,
11035        m: usize,
11036    ) -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn std::error::Error>>
11037    {
11038        use crate::model::GpuTensor;
11039        if m != 1 || !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) {
11040            return Ok(None);
11041        }
11042        // FP-ORDER LAW (lane/nvfp4-strict, 2026-08-05): every kernel this door can dispatch
11043        // (q8_0 fused2, nvfp4 dual_mr2) is the MMVQ family — 32-thread warp reduce. Without
11044        // MEMRA_MMVQ the m=1 singles ride dp4a (128-thread two-level reduce), so fusing here
11045        // would mix dispatch families across the pair — the exact class `q8_fused_params`
11046        // already refuses for Q8_0. The NVFP4 arm lacked this check, which is why
11047        // decode-batch-gate `--mode strict`'s equalizing env (MEMRA_MMVQ=0) never pinned
11048        // NVFP4 models: decode_step_h kept riding dual_mr2 while the batched body fell to
11049        // dp4a (gate1 maxdiff 1.639e-1 / gate2 step-8 divergence at the 2026-08-05 train
11050        // HEAD, research/nvfp4-strict-20260805/). Default env (MMVQ on) is dispatch-unchanged.
11051        if !self.mmvq_supports(QT_NVFP4) {
11052            return Ok(None);
11053        }
11054        let (in_f, out_f) = (w0.in_features(), w0.out_features());
11055        if w1.in_features() != in_f || w1.out_features() != out_f {
11056            return Ok(None);
11057        }
11058        // Q8_0 ARM (lane/q27-deepdive, 2026-08-05): the dense-FFN gate+up pair on a Q8_0 trunk fell
11059        // through this NVFP4-only gate to two `matmul_pre_noscale` launches — measured 128 of the
11060        // 1015 launches/token on q27-Q8_0 decode, the single largest un-fused class in the tick
11061        // (nsys `research/q27-deepdive-20260805/nsys/`). `q8_fused2_core` already serves the same
11062        // pair shape for the shared-expert gate/up, and its kernel body is `qmatvec_q8_0_mmvq`
11063        // VERBATIM per (tensor,row) -> BIT-IDENTICAL to the two separate launches. Q8_0 carries no
11064        // macro-scale (q8_fused_params requires scale==1.0), so the noscale contract is satisfied
11065        // by returning 1.0 for both: the SwiGLU epilogue's fold becomes the identity it already is
11066        // on this dtype today. Seam: MEMRA_Q8_FFN_FUSE2=0 rolls back to the two-launch pair.
11067        // rp4 guard: with MEMRA_Q8RP the singles route to the `_rp` split-plane twin over the
11068        // mirror buffer; the fused2 kernel has no `_rp` form, so fusing there would swap
11069        // dispatch families mid-model. Bail and let the two singles run (mirror lane unchanged).
11070        let no_mirror =
11071            |w: &crate::model::GpuTensor| !matches!(w, GpuTensor::Quant { rp4: Some(_), .. });
11072        if self.q8_ffn_fuse2_on()
11073            && no_mirror(w0)
11074            && no_mirror(w1)
11075            && let Some([p0, p1]) = self.q8_fused_params(&[w0, w1])
11076        {
11077            let (y0, y1) = self.q8_fused2_core(p0.0, p1.0, aq, ad, in_f, p0.1, p1.1, p0.2)?;
11078            return Ok(Some(((y0, 1.0), (y1, 1.0))));
11079        }
11080        // F8-E4M3 ARM (lane/fp8-decode-v1, 2026-08-05): with native e4m3 residency the FFN gate+up
11081        // pair (and the ssm beta+alpha dual, which routes through this same entry) fell through
11082        // both the NVFP4 gate below and the Q8_0 arm above to two `matmul_pre_noscale` launches —
11083        // native residency was UN-FUSING the trunk relative to the Q8_0 slab it replaces. The
11084        // fused2 kernel body is `qmatvec_e4m3_mmvq` VERBATIM per (tensor,row). Contract match:
11085        // `matmul_pre_noscale` on e4m3 launches with scale 1.0 and RETURNS the per-tensor
11086        // weight_scale for the caller to fold, so we pass ws=1.0 here and return (s0,s1) — same
11087        // bits, and the two macro-scale multiplies still fold into the SwiGLU epilogue.
11088        // MEMRA_E4M3_DUAL=0 rolls back to the two-launch pair.
11089        if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
11090            let (y0, y1) =
11091                self.e4m3_fused2_core(p0.0, p1.0, aq, ad, in_f, p0.1, p1.1, p0.2, 1.0, 1.0)?;
11092            return Ok(Some(((y0, p0.3), (y1, p1.3))));
11093        }
11094        let (b0, q0, rb0, s0, rp0) = match w0 {
11095            GpuTensor::Quant {
11096                bytes,
11097                qtype,
11098                row_bytes,
11099                scale,
11100                rp,
11101                ..
11102            } => (bytes, *qtype, *row_bytes, *scale, *rp),
11103            _ => return Ok(None),
11104        };
11105        let (b1, q1, rb1, s1, rp1) = match w1 {
11106            GpuTensor::Quant {
11107                bytes,
11108                qtype,
11109                row_bytes,
11110                scale,
11111                rp,
11112                ..
11113            } => (bytes, *qtype, *row_bytes, *scale, *rp),
11114            _ => return Ok(None),
11115        };
11116        if q0 != QT_NVFP4 || q1 != QT_NVFP4 || rb0 != rb1 || rp0 != rp1 {
11117            return Ok(None);
11118        }
11119        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
11120        const RPW: u32 = 2;
11121        let rows_per_block = ROWS_PER_BLOCK * RPW;
11122        let f = self.func(if rp0 {
11123            "qmatvec_nvfp4_mmvq_dual_mr2_rp"
11124        } else {
11125            "qmatvec_nvfp4_mmvq_dual_mr2"
11126        });
11127        let mut y0 = self.alloc_uninit::<f32>(out_f)?;
11128        let mut y1 = self.alloc_uninit::<f32>(out_f)?;
11129        let cfg = LaunchConfig {
11130            grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 2, 1),
11131            block_dim: (32, ROWS_PER_BLOCK, 1),
11132            shared_mem_bytes: 0,
11133        };
11134        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, 1i32, rb0 as i64);
11135        // noscale contract: the caller folds s0/s1 into the SwiGLU epilogue — the kernel's fused
11136        // yscale args stay 1.0 here (they exist for the single-tensor callers).
11137        let one = 1.0f32;
11138        let __s_b = self.gpu.stream();
11139        let mut b = __s_b.launch_builder(&f);
11140        b.arg(b0)
11141            .arg(b1)
11142            .arg(aq)
11143            .arg(ad)
11144            .arg(&mut y0)
11145            .arg(&mut y1)
11146            .arg(&inf)
11147            .arg(&outf)
11148            .arg(&mi)
11149            .arg(&rb)
11150            .arg(&one)
11151            .arg(&one);
11152        unsafe {
11153            b.launch(cfg)?;
11154        }
11155        Ok(Some(((y0, s0), (y1, s1))))
11156    }
11157
11158    /// FUSED NVFP4 matvec TRIPLE with unequal out_f (rig-native decode increment 1,
11159    /// lane/rig-native-nvfp4): wq+wk+wv in ONE launch via the q8_0 fused2 block-offset
11160    /// recipe. Per (tensor,row,t) the kernel body is nvfp4_mmvq_multirow_rp VERBATIM ->
11161    /// bit-identical to three separate `matmul_pre` launches; yscales fold in-kernel exactly
11162    /// as the singles do. grid.y = m (the t-parallel verify rows ride the same launch).
11163    /// None when ineligible (not all rp NVFP4 / in_f mismatch / mmvq off) — callers fall
11164    /// back to the three singles.
11165    #[allow(clippy::too_many_arguments)]
11166    pub fn matmul_nvfp4_fused3(
11167        &self,
11168        w0: &crate::model::GpuTensor,
11169        w1: &crate::model::GpuTensor,
11170        w2: &crate::model::GpuTensor,
11171        aq: &CudaSlice<i8>,
11172        ad: &CudaSlice<f32>,
11173        m: usize,
11174    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
11175    {
11176        use crate::model::GpuTensor;
11177        // m == 1 rides the original fused kernel; m = 2..=8 rides the BATCHED fused twin
11178        // (lane/samplat — the fused4_b8 pattern): the seg body is nvfp4_mmvq_batched_rp_sc
11179        // verbatim, weight rows read once for all m columns, bit-identical per
11180        // (tensor,row,column) to the three bN_rpsc singles. The old "at m>1 the fused
11181        // segments would re-read the weight per row" note described the grid.y=m lift,
11182        // which this twin deliberately is NOT.
11183        if !(1..=8).contains(&m)
11184            || !self.mmvq_supports(QT_NVFP4)
11185            || !self.uses_q8_1_fast(w0)
11186            || !self.uses_q8_1_fast(w1)
11187            || !self.uses_q8_1_fast(w2)
11188        {
11189            return Ok(None);
11190        }
11191        if m > 1 {
11192            let in_f = w0.in_features();
11193            if std::env::var("MEMRA_NVFP4_FUSED3B").as_deref() == Ok("0")
11194                || !self.batched_supports(QT_NVFP4)
11195                || std::env::var("MEMRA_NO_BATCHED").is_ok()
11196                || (m > 4 && !Self::b8_enabled())
11197                || in_f % 512 != 0
11198                || in_f / 64 > 272
11199            {
11200                return Ok(None);
11201            }
11202        }
11203        let unpack = |w: &crate::model::GpuTensor| match w {
11204            GpuTensor::Quant {
11205                bytes,
11206                qtype,
11207                scale,
11208                rp,
11209                ..
11210            } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
11211            _ => None,
11212        };
11213        let (Some(p0), Some(p1), Some(p2)) = (unpack(w0), unpack(w1), unpack(w2)) else {
11214            return Ok(None);
11215        };
11216        let in_f = w0.in_features();
11217        if w1.in_features() != in_f || w2.in_features() != in_f {
11218            return Ok(None);
11219        }
11220        let (o0, o1, o2) = (w0.out_features(), w1.out_features(), w2.out_features());
11221        const ROWS_PER_BLOCK: u32 = 4; // MEMRA_MMVQ_ROWS
11222        const RPW: u32 = 2;
11223        let rows_pb = ROWS_PER_BLOCK * RPW;
11224        let nb = |o: usize| (o as u32).div_ceil(rows_pb);
11225        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
11226        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
11227        let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
11228        let (inf, oi0, oi1, oi2, mi) = (in_f as i32, o0 as i32, o1 as i32, o2 as i32, m as i32);
11229        // SAFETY: the raw pointers come straight from the &GpuTensor borrows above and are
11230        // only dereferenced for the launch-arg build inside this call.
11231        let (b0, b1, b2) = unsafe { (&*p0.0, &*p1.0, &*p2.0) };
11232        if m > 1 {
11233            // batched twin has no in-kernel scale — refuse scale carriers (GGUF trunk = 1.0).
11234            if p0.1 != 1.0 || p1.1 != 1.0 || p2.1 != 1.0 {
11235                return Ok(None);
11236            }
11237            let f = self.func("qmatvec_nvfp4_mmvq_fused3_b8_rpsc");
11238            let cfg = LaunchConfig {
11239                grid_dim: (nb(o0) + nb(o1) + nb(o2), 1, 1),
11240                block_dim: (32, ROWS_PER_BLOCK, 1),
11241                shared_mem_bytes: 0,
11242            };
11243            let __s_b = self.gpu.stream();
11244            let mut b = __s_b.launch_builder(&f);
11245            b.arg(b0)
11246                .arg(b1)
11247                .arg(b2)
11248                .arg(aq)
11249                .arg(ad)
11250                .arg(&mut y0)
11251                .arg(&mut y1)
11252                .arg(&mut y2)
11253                .arg(&inf)
11254                .arg(&oi0)
11255                .arg(&oi1)
11256                .arg(&oi2)
11257                .arg(&mi);
11258            unsafe {
11259                b.launch(cfg)?;
11260            }
11261            return Ok(Some((y0, y1, y2)));
11262        }
11263        let f = self.func("qmatvec_nvfp4_mmvq_fused3_rp");
11264        let cfg = LaunchConfig {
11265            grid_dim: (nb(o0) + nb(o1) + nb(o2), m as u32, 1),
11266            block_dim: (32, ROWS_PER_BLOCK, 1),
11267            shared_mem_bytes: 0,
11268        };
11269        let __s_b = self.gpu.stream();
11270        let mut b = __s_b.launch_builder(&f);
11271        b.arg(b0)
11272            .arg(b1)
11273            .arg(b2)
11274            .arg(aq)
11275            .arg(ad)
11276            .arg(&mut y0)
11277            .arg(&mut y1)
11278            .arg(&mut y2)
11279            .arg(&inf)
11280            .arg(&oi0)
11281            .arg(&oi1)
11282            .arg(&oi2)
11283            .arg(&mi)
11284            .arg(&p0.1)
11285            .arg(&p1.1)
11286            .arg(&p2.1);
11287        unsafe {
11288            b.launch(cfg)?;
11289        }
11290        Ok(Some((y0, y1, y2)))
11291    }
11292
11293    /// fused2 twin of `matmul_nvfp4_fused3`, for MIXED-type weight groups: the gemma4
11294    /// dense NVFP4mix recipe keeps `attn_v` and `ffn_down` at Q8_0 (full-NVFP4 was
11295    /// measured garbage on this dense class), so its q/k/v trio and gate/up/down never
11296    /// satisfy an all-NVFP4 fused3 — the pairs that ARE uniformly NVFP4 (q,k and
11297    /// gate,up) fuse here instead. m==1 only, same law as fused3/fused4: per
11298    /// (tensor,row) the kernel seg body is VERBATIM, so the fusion is bit-identical to
11299    /// two separate launches. `MEMRA_NVFP4_FUSED2=0` is the rollback seam and the
11300    /// same-binary interleaved A/B arm.
11301    pub fn matmul_nvfp4_fused2(
11302        &self,
11303        w0: &crate::model::GpuTensor,
11304        w1: &crate::model::GpuTensor,
11305        aq: &CudaSlice<i8>,
11306        ad: &CudaSlice<f32>,
11307        m: usize,
11308    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
11309        use crate::model::GpuTensor;
11310        static FUSED2_OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11311        let off =
11312            *FUSED2_OFF.get_or_init(|| std::env::var("MEMRA_NVFP4_FUSED2").as_deref() == Ok("0"));
11313        // m==1 ONLY: at m>1 the singles ride the _b16 weight-once column program (one weight
11314        // read serves all m rows); the fused segments would re-read the weight per row.
11315        if off
11316            || m != 1
11317            || !self.mmvq_supports(QT_NVFP4)
11318            || !self.uses_q8_1_fast(w0)
11319            || !self.uses_q8_1_fast(w1)
11320        {
11321            return Ok(None);
11322        }
11323        let unpack = |w: &crate::model::GpuTensor| match w {
11324            GpuTensor::Quant {
11325                bytes,
11326                qtype,
11327                scale,
11328                rp,
11329                ..
11330            } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
11331            _ => None,
11332        };
11333        let (Some(p0), Some(p1)) = (unpack(w0), unpack(w1)) else {
11334            return Ok(None);
11335        };
11336        let in_f = w0.in_features();
11337        if w1.in_features() != in_f {
11338            return Ok(None);
11339        }
11340        let (o0, o1) = (w0.out_features(), w1.out_features());
11341        const ROWS_PER_BLOCK: u32 = 4; // MEMRA_MMVQ_ROWS
11342        const RPW: u32 = 2;
11343        let rows_pb = ROWS_PER_BLOCK * RPW;
11344        let nb = |o: usize| (o as u32).div_ceil(rows_pb);
11345        let f = self.func("qmatvec_nvfp4_mmvq_fused2_rp");
11346        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
11347        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
11348        let cfg = LaunchConfig {
11349            grid_dim: (nb(o0) + nb(o1), m as u32, 1),
11350            block_dim: (32, ROWS_PER_BLOCK, 1),
11351            shared_mem_bytes: 0,
11352        };
11353        let (inf, oi0, oi1, mi) = (in_f as i32, o0 as i32, o1 as i32, m as i32);
11354        // SAFETY: the raw pointers come straight from the &GpuTensor borrows above and are
11355        // only dereferenced for the launch-arg build inside this call.
11356        let (b0, b1) = unsafe { (&*p0.0, &*p1.0) };
11357        // PDL wave-B: the kernel carries MEMRA_PDL_ENTRY — launch overlaps the producer's
11358        // drain (quantize/norm emitting aq/ad). Same math, same order, bit-identical.
11359        if Self::pdl_on() && Self::pdl_mmvq_on() && Self::pdl_nvfp4q8_on() {
11360            {
11361                use cudarc::driver::{DevicePtr, DevicePtrMut};
11362                let s = &self.gpu.stream();
11363                let (pw0, _g0) = b0.device_ptr(s);
11364                let (pw1, _g1) = b1.device_ptr(s);
11365                let (paq, _g2) = aq.device_ptr(s);
11366                let (pad, _g3) = ad.device_ptr(s);
11367                let (py0, _g4) = y0.device_ptr_mut(s);
11368                let (py1, _g5) = y1.device_ptr_mut(s);
11369                let (s0, s1) = (p0.1, p1.1);
11370                let mut ps = [
11371                    &pw0 as *const _ as *mut std::ffi::c_void,
11372                    &pw1 as *const _ as *mut _,
11373                    &paq as *const _ as *mut _,
11374                    &pad as *const _ as *mut _,
11375                    &py0 as *const _ as *mut _,
11376                    &py1 as *const _ as *mut _,
11377                    &inf as *const _ as *mut _,
11378                    &oi0 as *const _ as *mut _,
11379                    &oi1 as *const _ as *mut _,
11380                    &mi as *const _ as *mut _,
11381                    &s0 as *const _ as *mut _,
11382                    &s1 as *const _ as *mut _,
11383                ];
11384                unsafe {
11385                    self.launch_pdl(
11386                        "qmatvec_nvfp4_mmvq_fused2_rp",
11387                        cfg.grid_dim,
11388                        cfg.block_dim,
11389                        &mut ps,
11390                    )?;
11391                }
11392            }
11393            return Ok(Some((y0, y1)));
11394        }
11395        let __s_b = self.gpu.stream();
11396        let mut b = __s_b.launch_builder(&f);
11397        b.arg(b0)
11398            .arg(b1)
11399            .arg(aq)
11400            .arg(ad)
11401            .arg(&mut y0)
11402            .arg(&mut y1)
11403            .arg(&inf)
11404            .arg(&oi0)
11405            .arg(&oi1)
11406            .arg(&mi)
11407            .arg(&p0.1)
11408            .arg(&p1.1);
11409        unsafe {
11410            b.launch(cfg)?;
11411        }
11412        Ok(Some((y0, y1)))
11413    }
11414
11415    /// Slot-fed fused2 twin (alloc-free capture lane): identical launch to
11416    /// `matmul_nvfp4_fused2`, caller-owned outputs — the gemma4 dc_slotted graph body
11417    /// needs zero mem nodes, so the allocating wrapper can't serve it. Returns false
11418    /// (decline) on any non-NVFP4/rp pair; the caller chains or refuses.
11419    pub fn matmul_nvfp4_fused2_into(
11420        &self,
11421        w0: &crate::model::GpuTensor,
11422        w1: &crate::model::GpuTensor,
11423        aq: &CudaSlice<i8>,
11424        ad: &CudaSlice<f32>,
11425        y0: &mut CudaSlice<f32>,
11426        y1: &mut CudaSlice<f32>,
11427    ) -> Result<bool, Box<dyn std::error::Error>> {
11428        use crate::model::GpuTensor;
11429        static FUSED2_OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11430        let off =
11431            *FUSED2_OFF.get_or_init(|| std::env::var("MEMRA_NVFP4_FUSED2").as_deref() == Ok("0"));
11432        if off
11433            || !self.mmvq_supports(QT_NVFP4)
11434            || !self.uses_q8_1_fast(w0)
11435            || !self.uses_q8_1_fast(w1)
11436        {
11437            return Ok(false);
11438        }
11439        let unpack = |w: &crate::model::GpuTensor| match w {
11440            GpuTensor::Quant {
11441                bytes,
11442                qtype,
11443                scale,
11444                rp,
11445                ..
11446            } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
11447            _ => None,
11448        };
11449        let (Some(p0), Some(p1)) = (unpack(w0), unpack(w1)) else {
11450            return Ok(false);
11451        };
11452        let in_f = w0.in_features();
11453        if w1.in_features() != in_f {
11454            return Ok(false);
11455        }
11456        let (o0, o1) = (w0.out_features(), w1.out_features());
11457        if y0.len() < o0 || y1.len() < o1 {
11458            return Ok(false);
11459        }
11460        const ROWS_PER_BLOCK: u32 = 4; // MEMRA_MMVQ_ROWS
11461        const RPW: u32 = 2;
11462        let rows_pb = ROWS_PER_BLOCK * RPW;
11463        let nb = |o: usize| (o as u32).div_ceil(rows_pb);
11464        let f = self.func("qmatvec_nvfp4_mmvq_fused2_rp");
11465        let cfg = LaunchConfig {
11466            grid_dim: (nb(o0) + nb(o1), 1, 1),
11467            block_dim: (32, ROWS_PER_BLOCK, 1),
11468            shared_mem_bytes: 0,
11469        };
11470        let (inf, oi0, oi1, mi) = (in_f as i32, o0 as i32, o1 as i32, 1i32);
11471        // SAFETY: the raw pointers come straight from the &GpuTensor borrows above and are
11472        // only dereferenced for the launch-arg build inside this call.
11473        let (b0, b1) = unsafe { (&*p0.0, &*p1.0) };
11474        let __s_b = self.gpu.stream();
11475        let mut b = __s_b.launch_builder(&f);
11476        b.arg(b0)
11477            .arg(b1)
11478            .arg(aq)
11479            .arg(ad)
11480            .arg(&mut *y0)
11481            .arg(&mut *y1)
11482            .arg(&inf)
11483            .arg(&oi0)
11484            .arg(&oi1)
11485            .arg(&mi)
11486            .arg(&p0.1)
11487            .arg(&p1.1);
11488        unsafe {
11489            b.launch(cfg)?;
11490        }
11491        Ok(true)
11492    }
11493
11494    /// fused4 twin of `matmul_nvfp4_fused3`: the Linear-mixer projection quartet
11495    /// (wqkv + wqkv_gate + ssm_beta + ssm_alpha) in one launch, m==1 only. Per
11496    /// (tensor,row) the kernel body is nvfp4_mmvq_multirow_rp VERBATIM — bit-identical
11497    /// to four separate launches (rig-native decode increment 2, RIG-NATIVE-DECODE.md).
11498    #[allow(clippy::type_complexity)]
11499    pub fn matmul_nvfp4_fused4(
11500        &self,
11501        w0: &crate::model::GpuTensor,
11502        w1: &crate::model::GpuTensor,
11503        w2: &crate::model::GpuTensor,
11504        w3: &crate::model::GpuTensor,
11505        aq: &CudaSlice<i8>,
11506        ad: &CudaSlice<f32>,
11507        m: usize,
11508    ) -> Result<
11509        Option<(
11510            CudaSlice<f32>,
11511            CudaSlice<f32>,
11512            CudaSlice<f32>,
11513            CudaSlice<f32>,
11514        )>,
11515        Box<dyn std::error::Error>,
11516    > {
11517        use crate::model::GpuTensor;
11518        // MEMRA_NVFP4_FUSED4=0: rollback seam + the same-binary interleaved A/B arm.
11519        // m == 1 rides the original fused kernel; m = 2..=8 rides the BATCHED fused twin
11520        // (lane/samplat, 2026-08-21): same quartet-in-one-launch shape, seg body =
11521        // nvfp4_mmvq_batched_rp_sc verbatim (weight rows read once for all m columns) —
11522        // bit-identical per (tensor,row,column) to the four bN_rpsc singles it replaces.
11523        // Admission mirrors the singles' batched gates below.
11524        if !(1..=8).contains(&m)
11525            || std::env::var("MEMRA_NVFP4_FUSED4").as_deref() == Ok("0")
11526            || !self.mmvq_supports(QT_NVFP4)
11527            || !self.uses_q8_1_fast(w0)
11528            || !self.uses_q8_1_fast(w1)
11529            || !self.uses_q8_1_fast(w2)
11530            || !self.uses_q8_1_fast(w3)
11531        {
11532            return Ok(None);
11533        }
11534        if m > 1 {
11535            // the batched-twin gates: the bN_rpsc program this must stay byte-identical to
11536            // (matmul_pre's batched arm), plus the rp-sc dispatch shape requirements.
11537            let in_f = w0.in_features();
11538            if !self.batched_supports(QT_NVFP4)
11539                || std::env::var("MEMRA_NO_BATCHED").is_ok()
11540                || (m > 4 && !Self::b8_enabled())
11541                || in_f % 512 != 0
11542                || in_f / 64 > 272
11543            {
11544                return Ok(None);
11545            }
11546        }
11547        let unpack = |w: &crate::model::GpuTensor| match w {
11548            GpuTensor::Quant {
11549                bytes,
11550                qtype,
11551                scale,
11552                rp,
11553                ..
11554            } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
11555            _ => None,
11556        };
11557        let (Some(p0), Some(p1), Some(p2), Some(p3)) =
11558            (unpack(w0), unpack(w1), unpack(w2), unpack(w3))
11559        else {
11560            return Ok(None);
11561        };
11562        let in_f = w0.in_features();
11563        if w1.in_features() != in_f || w2.in_features() != in_f || w3.in_features() != in_f {
11564            return Ok(None);
11565        }
11566        let (o0, o1, o2, o3) = (
11567            w0.out_features(),
11568            w1.out_features(),
11569            w2.out_features(),
11570            w3.out_features(),
11571        );
11572        const ROWS_PER_BLOCK: u32 = 4; // MEMRA_MMVQ_ROWS
11573        const RPW: u32 = 2;
11574        let rows_pb = ROWS_PER_BLOCK * RPW;
11575        let nb = |o: usize| (o as u32).div_ceil(rows_pb);
11576        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
11577        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
11578        let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
11579        let mut y3 = self.alloc_uninit::<f32>(m * o3)?;
11580        let (inf, oi0, oi1, oi2, oi3, mi) = (
11581            in_f as i32,
11582            o0 as i32,
11583            o1 as i32,
11584            o2 as i32,
11585            o3 as i32,
11586            m as i32,
11587        );
11588        // SAFETY: the raw pointers come straight from the &GpuTensor borrows above and are
11589        // only dereferenced for the launch-arg build inside this call.
11590        let (b0, b1, b2, b3) = unsafe { (&*p0.0, &*p1.0, &*p2.0, &*p3.0) };
11591        if m > 1 {
11592            // Batched fused twin: no in-kernel scale (the bN_rpsc program has none) — refuse
11593            // scale-carrying tensors so the singles path keeps them (GGUF trunk scales are 1.0).
11594            if p0.1 != 1.0 || p1.1 != 1.0 || p2.1 != 1.0 || p3.1 != 1.0 {
11595                return Ok(None);
11596            }
11597            let f = self.func("qmatvec_nvfp4_mmvq_fused4_b8_rpsc");
11598            let cfg = LaunchConfig {
11599                grid_dim: (nb(o0) + nb(o1) + nb(o2) + nb(o3), 1, 1),
11600                block_dim: (32, ROWS_PER_BLOCK, 1),
11601                shared_mem_bytes: 0,
11602            };
11603            let __s_b = self.gpu.stream();
11604            let mut b = __s_b.launch_builder(&f);
11605            b.arg(b0)
11606                .arg(b1)
11607                .arg(b2)
11608                .arg(b3)
11609                .arg(aq)
11610                .arg(ad)
11611                .arg(&mut y0)
11612                .arg(&mut y1)
11613                .arg(&mut y2)
11614                .arg(&mut y3)
11615                .arg(&inf)
11616                .arg(&oi0)
11617                .arg(&oi1)
11618                .arg(&oi2)
11619                .arg(&oi3)
11620                .arg(&mi);
11621            unsafe {
11622                b.launch(cfg)?;
11623            }
11624            return Ok(Some((y0, y1, y2, y3)));
11625        }
11626        let f = self.func("qmatvec_nvfp4_mmvq_fused4_rp");
11627        let cfg = LaunchConfig {
11628            grid_dim: (nb(o0) + nb(o1) + nb(o2) + nb(o3), m as u32, 1),
11629            block_dim: (32, ROWS_PER_BLOCK, 1),
11630            shared_mem_bytes: 0,
11631        };
11632        let __s_b = self.gpu.stream();
11633        let mut b = __s_b.launch_builder(&f);
11634        b.arg(b0)
11635            .arg(b1)
11636            .arg(b2)
11637            .arg(b3)
11638            .arg(aq)
11639            .arg(ad)
11640            .arg(&mut y0)
11641            .arg(&mut y1)
11642            .arg(&mut y2)
11643            .arg(&mut y3)
11644            .arg(&inf)
11645            .arg(&oi0)
11646            .arg(&oi1)
11647            .arg(&oi2)
11648            .arg(&oi3)
11649            .arg(&mi)
11650            .arg(&p0.1)
11651            .arg(&p1.1)
11652            .arg(&p2.1)
11653            .arg(&p3.1);
11654        unsafe {
11655            b.launch(cfg)?;
11656        }
11657        Ok(Some((y0, y1, y2, y3)))
11658    }
11659
11660    /// FUSED Q8_0 m=1 matvec PAIR with UNEQUAL out_f (trunk launch-fusion, 2026-07-05). Folds two
11661    /// same-input q8_0 projections (35B trunk: wqkv+wqkv_gate 8192/4096, gate_shexp+up_shexp
11662    /// 512/512) into ONE launch via a block-offset split (blocks [0,nb0) -> w0, rest -> w1) — the
11663    /// dual-mr2 recipe with the same-out_f restriction lifted. Per (tensor,row) the kernel body is
11664    /// qmatvec_q8_0_mmvq VERBATIM -> BIT-IDENTICAL to two separate m=1 launches. Returns None when
11665    /// ineligible (not both Q8_0 / in_f mismatch / MEMRA_MMVQ off / MEMRA_Q8_DUAL=0) — caller falls
11666    /// back to the per-tensor path.
11667    pub fn matmul_q8_fused2(
11668        &self,
11669        w0: &crate::model::GpuTensor,
11670        w1: &crate::model::GpuTensor,
11671        aq: &CudaSlice<i8>,
11672        ad: &CudaSlice<f32>,
11673    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
11674        // e4m3 twin (lane/fp8-decode-v1): this entry is the trunk's generic m=1 pair door
11675        // (wqkv+wqkv_gate, ssm_beta+alpha, gate_shexp+up_shexp), so admitting QT_F8_E4M3 here
11676        // fuses the NATIVE-RESIDENCY FP8 trunk at every existing call site with no call-site
11677        // change. Scale is folded in-kernel per range -> the returned buffers are already scaled,
11678        // exactly like the per-tensor `matmul_pre` e4m3 dispatch this replaces.
11679        if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
11680            return Ok(Some(self.e4m3_fused2_core(
11681                p0.0,
11682                p1.0,
11683                aq,
11684                ad,
11685                w0.in_features(),
11686                p0.1,
11687                p1.1,
11688                p0.2,
11689                p0.3,
11690                p1.3,
11691            )?));
11692        }
11693        let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
11694            return Ok(None);
11695        };
11696        Ok(Some(self.q8_fused2_core(
11697            p0.0,
11698            p1.0,
11699            aq,
11700            ad,
11701            w0.in_features(),
11702            p0.1,
11703            p1.1,
11704            p0.2,
11705        )?))
11706    }
11707
11708    #[allow(clippy::too_many_arguments)]
11709    fn q8_fused2_core(
11710        &self,
11711        b0: &CudaSlice<u8>,
11712        b1: &CudaSlice<u8>,
11713        aq: &CudaSlice<i8>,
11714        ad: &CudaSlice<f32>,
11715        in_f: usize,
11716        out0: usize,
11717        out1: usize,
11718        row_bytes: usize,
11719    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11720        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
11721        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
11722        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
11723        let f = self.func("qmatvec_q8_0_mmvq_fused2");
11724        let mut y0 = self.alloc_uninit::<f32>(out0)?;
11725        let mut y1 = self.alloc_uninit::<f32>(out1)?;
11726        let cfg = LaunchConfig {
11727            grid_dim: (nb0 + nb1, 1, 1),
11728            block_dim: (32, ROWS_PER_BLOCK, 1),
11729            shared_mem_bytes: 0,
11730        };
11731        let (inf, o0, o1, rbl) = (in_f as i32, out0 as i32, out1 as i32, row_bytes as i64);
11732        let __s_b = self.gpu.stream();
11733        let mut b = __s_b.launch_builder(&f);
11734        b.arg(b0)
11735            .arg(b1)
11736            .arg(aq)
11737            .arg(ad)
11738            .arg(&mut y0)
11739            .arg(&mut y1)
11740            .arg(&inf)
11741            .arg(&o0)
11742            .arg(&o1)
11743            .arg(&rbl);
11744        unsafe {
11745            b.launch(cfg)?;
11746        }
11747        Ok((y0, y1))
11748    }
11749
11750    /// f32-activation entry for the fused2 pair: quantizes x to q8_1 ONCE then runs the fused
11751    /// launch — replaces two `matmul(w, x, 1)` calls that would each re-quantize the same x
11752    /// (35B shared-expert gate+up per MoE layer per token). Same bits: quantize_q8_1 is
11753    /// deterministic, the fused body is the MMVQ kernel verbatim. None when ineligible (the
11754    /// callers' m==1-under-MEMRA_FAST dispatch would take MMVQ; anything else falls back).
11755    pub fn matmul_q8_fused2_x(
11756        &self,
11757        w0: &crate::model::GpuTensor,
11758        w1: &crate::model::GpuTensor,
11759        x: &CudaSlice<f32>,
11760    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
11761        if !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) {
11762            return Ok(None);
11763        }
11764        if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
11765            let (aq, ad) = self.quantize_q8_1(x, 1, w0.in_features())?;
11766            return Ok(Some(self.e4m3_fused2_core(
11767                p0.0,
11768                p1.0,
11769                &aq,
11770                &ad,
11771                w0.in_features(),
11772                p0.1,
11773                p1.1,
11774                p0.2,
11775                p0.3,
11776                p1.3,
11777            )?));
11778        }
11779        let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
11780            return Ok(None);
11781        };
11782        let (aq, ad) = self.quantize_q8_1(x, 1, w0.in_features())?;
11783        Ok(Some(self.q8_fused2_core(
11784            p0.0,
11785            p1.0,
11786            &aq,
11787            &ad,
11788            w0.in_features(),
11789            p0.1,
11790            p1.1,
11791            p0.2,
11792        )?))
11793    }
11794
11795    /// Test entry for the kernel_check gate: launch the fused2 kernel from raw weight bytes,
11796    /// quantizing the f32 activation internally (mirrors qmatvec_mmvq_raw; no env gating).
11797    #[allow(clippy::too_many_arguments)]
11798    pub fn qmatvec_q8_fused2_raw(
11799        &self,
11800        b0: &CudaSlice<u8>,
11801        b1: &CudaSlice<u8>,
11802        x: &CudaSlice<f32>,
11803        in_f: usize,
11804        out0: usize,
11805        out1: usize,
11806        row_bytes: usize,
11807    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11808        let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
11809        self.q8_fused2_core(b0, b1, &aq, &ad, in_f, out0, out1, row_bytes)
11810    }
11811
11812    /// FUSED Q8_0 m=1 matvec TRIPLE (wq+wk+wv on the 35B full-attn layers: out_f 8192/512/512).
11813    /// Same block-offset recipe as `matmul_q8_fused2` with three ranges. BIT-IDENTICAL per
11814    /// (tensor,row) to three separate m=1 MMVQ launches.
11815    /// FUSED Q4_0 m=1 TRIPLE (gemma q/k/v — same quantized input; per (tensor,row) chain
11816    /// identical to the mr2 kernel). Returns None unless all three are Q4_0 with equal in_f.
11817    pub fn matmul_q4_fused3(
11818        &self,
11819        w0: &crate::model::GpuTensor,
11820        w1: &crate::model::GpuTensor,
11821        w2: &crate::model::GpuTensor,
11822        aq: &CudaSlice<i8>,
11823        ad: &CudaSlice<f32>,
11824    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
11825    {
11826        use crate::model::GpuTensor;
11827        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
11828            match w {
11829                GpuTensor::Quant {
11830                    qtype, row_bytes, ..
11831                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
11832                _ => None,
11833            }
11834        };
11835        let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (q4(w0), q4(w1), q4(w2)) else {
11836            return Ok(None);
11837        };
11838        if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
11839            return Ok(None);
11840        }
11841        // Effective (bytes, rp) per tensor: mirror (rp4) OR the in-place swap (rp flag,
11842        // bytes already split). Mixed layouts cannot share one fused launch -> fall back to
11843        // the separate matvecs (each routes its own rp).
11844        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
11845            match w {
11846                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
11847                    Some(m) => (m, true),
11848                    None => (bytes, *rp),
11849                },
11850                _ => unreachable!(),
11851            }
11852        }
11853        let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
11854        if rp0 != rp1 || rp1 != rp2 {
11855            return Ok(None);
11856        }
11857        let rp = rp0;
11858        let rpb: u32 = 4;
11859        // mr1 (one row/warp, 2026-07-14): follows the singles' MEMRA_Q40_MR default — the
11860        // fused t=1 kernels were left on mr2 when the singles flipped (DRAM-duty map:
11861        // fused3 57% / fused2 86%; small qkv segments starve under mr2's half grid).
11862        let mr1 = rp && Self::q40_mr1_on();
11863        let nb = |o: usize| {
11864            if mr1 {
11865                (o as u32).div_ceil(rpb)
11866            } else {
11867                (o as u32).div_ceil(2).div_ceil(rpb)
11868            }
11869        };
11870        let grid = nb(o0) + nb(o1) + nb(o2);
11871        let mut y0 = self.alloc_uninit::<f32>(o0)?;
11872        let mut y1 = self.alloc_uninit::<f32>(o1)?;
11873        let mut y2 = self.alloc_uninit::<f32>(o2)?;
11874        let f = self.func(if mr1 {
11875            "qmatvec_q4_0_mmvq_fused3_mr1_rp"
11876        } else if rp {
11877            "qmatvec_q4_0_mmvq_fused3_rp"
11878        } else {
11879            "qmatvec_q4_0_mmvq_fused3"
11880        });
11881        let cfg = LaunchConfig {
11882            grid_dim: (grid, 1, 1),
11883            block_dim: (32, rpb, 1),
11884            shared_mem_bytes: 0,
11885        };
11886        let inf = w0.in_features() as i32;
11887        let (oo0, oo1, oo2) = (o0 as i32, o1 as i32, o2 as i32);
11888        let (r0, r1, r2) = (rb0 as i64, rb1 as i64, rb2 as i64);
11889        // PDL wave-A (2026-07-23): the mr1 kernel carries MEMRA_PDL_ENTRY; only that
11890        // variant may take the programmatic-serialization launch.
11891        if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
11892            {
11893                use cudarc::driver::{DevicePtr, DevicePtrMut};
11894                let s = &self.gpu.stream();
11895                let (p0, _g0) = b0.device_ptr(s);
11896                let (p1, _g1) = b1.device_ptr(s);
11897                let (p2, _g2) = b2.device_ptr(s);
11898                let (paq, _g3) = aq.device_ptr(s);
11899                let (pad, _g4) = ad.device_ptr(s);
11900                let (py0, _g5) = y0.device_ptr_mut(s);
11901                let (py1, _g6) = y1.device_ptr_mut(s);
11902                let (py2, _g7) = y2.device_ptr_mut(s);
11903                let mut ps = [
11904                    &p0 as *const _ as *mut std::ffi::c_void,
11905                    &p1 as *const _ as *mut _,
11906                    &p2 as *const _ as *mut _,
11907                    &paq as *const _ as *mut _,
11908                    &pad as *const _ as *mut _,
11909                    &py0 as *const _ as *mut _,
11910                    &py1 as *const _ as *mut _,
11911                    &py2 as *const _ as *mut _,
11912                    &inf as *const _ as *mut _,
11913                    &oo0 as *const _ as *mut _,
11914                    &oo1 as *const _ as *mut _,
11915                    &oo2 as *const _ as *mut _,
11916                    &r0 as *const _ as *mut _,
11917                    &r1 as *const _ as *mut _,
11918                    &r2 as *const _ as *mut _,
11919                ];
11920                unsafe {
11921                    self.launch_pdl(
11922                        "qmatvec_q4_0_mmvq_fused3_mr1_rp",
11923                        (grid, 1, 1),
11924                        (32, rpb, 1),
11925                        &mut ps,
11926                    )?;
11927                }
11928            }
11929            return Ok(Some((y0, y1, y2)));
11930        }
11931        let __s_b = self.gpu.stream();
11932        let mut b = __s_b.launch_builder(&f);
11933        b.arg(b0)
11934            .arg(b1)
11935            .arg(b2)
11936            .arg(aq)
11937            .arg(ad)
11938            .arg(&mut y0)
11939            .arg(&mut y1)
11940            .arg(&mut y2)
11941            .arg(&inf)
11942            .arg(&oo0)
11943            .arg(&oo1)
11944            .arg(&oo2)
11945            .arg(&r0)
11946            .arg(&r1)
11947            .arg(&r2);
11948        unsafe {
11949            b.launch(cfg)?;
11950        }
11951        Ok(Some((y0, y1, y2)))
11952    }
11953
11954    /// Slot-fed fused3 twin (alloc-free capture lane): identical launch, caller-owned outputs.
11955    /// Returns Ok(false) when the fused path is unavailable (caller falls back).
11956    #[allow(clippy::too_many_arguments)]
11957    pub fn matmul_q4_fused3_into(
11958        &self,
11959        w0: &crate::model::GpuTensor,
11960        w1: &crate::model::GpuTensor,
11961        w2: &crate::model::GpuTensor,
11962        aq: &CudaSlice<i8>,
11963        ad: &CudaSlice<f32>,
11964        y0: &mut CudaSlice<f32>,
11965        y1: &mut CudaSlice<f32>,
11966        y2: &mut CudaSlice<f32>,
11967    ) -> Result<bool, Box<dyn std::error::Error>> {
11968        use crate::model::GpuTensor;
11969        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
11970            match w {
11971                GpuTensor::Quant {
11972                    qtype, row_bytes, ..
11973                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
11974                _ => None,
11975            }
11976        };
11977        let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (q4(w0), q4(w1), q4(w2)) else {
11978            return Ok(false);
11979        };
11980        if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
11981            return Ok(false);
11982        }
11983        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
11984            match w {
11985                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
11986                    Some(m) => (m, true),
11987                    None => (bytes, *rp),
11988                },
11989                _ => unreachable!(),
11990            }
11991        }
11992        let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
11993        if rp0 != rp1 || rp1 != rp2 {
11994            return Ok(false);
11995        }
11996        let rp = rp0;
11997        let rpb: u32 = 4;
11998        let mr1 = rp && Self::q40_mr1_on();
11999        let nb = |o: usize| {
12000            if mr1 {
12001                (o as u32).div_ceil(rpb)
12002            } else {
12003                (o as u32).div_ceil(2).div_ceil(rpb)
12004            }
12005        };
12006        let grid = nb(o0) + nb(o1) + nb(o2);
12007        debug_assert!(y0.len() >= o0 && y1.len() >= o1 && y2.len() >= o2);
12008        let f = self.func(if mr1 {
12009            "qmatvec_q4_0_mmvq_fused3_mr1_rp"
12010        } else if rp {
12011            "qmatvec_q4_0_mmvq_fused3_rp"
12012        } else {
12013            "qmatvec_q4_0_mmvq_fused3"
12014        });
12015        let cfg = LaunchConfig {
12016            grid_dim: (grid, 1, 1),
12017            block_dim: (32, rpb, 1),
12018            shared_mem_bytes: 0,
12019        };
12020        let inf = w0.in_features() as i32;
12021        let (oo0, oo1, oo2) = (o0 as i32, o1 as i32, o2 as i32);
12022        let (r0, r1, r2) = (rb0 as i64, rb1 as i64, rb2 as i64);
12023        // PDL wave-A: identical to the owned twin (capture-lane parity).
12024        if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
12025            use cudarc::driver::{DevicePtr, DevicePtrMut};
12026            let s = &self.gpu.stream();
12027            let (p0, _g0) = b0.device_ptr(s);
12028            let (p1, _g1) = b1.device_ptr(s);
12029            let (p2, _g2) = b2.device_ptr(s);
12030            let (paq, _g3) = aq.device_ptr(s);
12031            let (pad, _g4) = ad.device_ptr(s);
12032            let (py0, _g5) = y0.device_ptr_mut(s);
12033            let (py1, _g6) = y1.device_ptr_mut(s);
12034            let (py2, _g7) = y2.device_ptr_mut(s);
12035            let mut ps = [
12036                &p0 as *const _ as *mut std::ffi::c_void,
12037                &p1 as *const _ as *mut _,
12038                &p2 as *const _ as *mut _,
12039                &paq as *const _ as *mut _,
12040                &pad as *const _ as *mut _,
12041                &py0 as *const _ as *mut _,
12042                &py1 as *const _ as *mut _,
12043                &py2 as *const _ as *mut _,
12044                &inf as *const _ as *mut _,
12045                &oo0 as *const _ as *mut _,
12046                &oo1 as *const _ as *mut _,
12047                &oo2 as *const _ as *mut _,
12048                &r0 as *const _ as *mut _,
12049                &r1 as *const _ as *mut _,
12050                &r2 as *const _ as *mut _,
12051            ];
12052            unsafe {
12053                self.launch_pdl(
12054                    "qmatvec_q4_0_mmvq_fused3_mr1_rp",
12055                    (grid, 1, 1),
12056                    (32, rpb, 1),
12057                    &mut ps,
12058                )?;
12059            }
12060            return Ok(true);
12061        }
12062        let __s_b = self.gpu.stream();
12063        let mut b = __s_b.launch_builder(&f);
12064        b.arg(b0)
12065            .arg(b1)
12066            .arg(b2)
12067            .arg(aq)
12068            .arg(ad)
12069            .arg(&mut *y0)
12070            .arg(&mut *y1)
12071            .arg(&mut *y2)
12072            .arg(&inf)
12073            .arg(&oo0)
12074            .arg(&oo1)
12075            .arg(&oo2)
12076            .arg(&r0)
12077            .arg(&r1)
12078            .arg(&r2);
12079        unsafe {
12080            b.launch(cfg)?;
12081        }
12082        Ok(true)
12083    }
12084
12085    /// FUSED Q4_0 m=1 PAIR (gemma shared gate+up).
12086    pub fn matmul_q4_fused2(
12087        &self,
12088        w0: &crate::model::GpuTensor,
12089        w1: &crate::model::GpuTensor,
12090        aq: &CudaSlice<i8>,
12091        ad: &CudaSlice<f32>,
12092    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
12093        use crate::model::GpuTensor;
12094        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
12095            match w {
12096                GpuTensor::Quant {
12097                    qtype, row_bytes, ..
12098                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
12099                _ => None,
12100            }
12101        };
12102        let (Some((rb0, o0)), Some((rb1, o1))) = (q4(w0), q4(w1)) else {
12103            return Ok(None);
12104        };
12105        if w0.in_features() != w1.in_features() {
12106            return Ok(None);
12107        }
12108        // Effective (bytes, rp) per tensor (mirror or in-place swap); mixed -> separate matvecs.
12109        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
12110            match w {
12111                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
12112                    Some(m) => (m, true),
12113                    None => (bytes, *rp),
12114                },
12115                _ => unreachable!(),
12116            }
12117        }
12118        let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
12119        if rp0 != rp1 {
12120            return Ok(None);
12121        }
12122        let rp = rp0;
12123        let rpb: u32 = 4;
12124        // mr1 twin — see matmul_q4_fused3.
12125        let mr1 = rp && Self::q40_mr1_on();
12126        let nb = |o: usize| {
12127            if mr1 {
12128                (o as u32).div_ceil(rpb)
12129            } else {
12130                (o as u32).div_ceil(2).div_ceil(rpb)
12131            }
12132        };
12133        let grid = nb(o0) + nb(o1);
12134        let mut y0 = self.alloc_uninit::<f32>(o0)?;
12135        let mut y1 = self.alloc_uninit::<f32>(o1)?;
12136        let f = self.func(if mr1 {
12137            "qmatvec_q4_0_mmvq_fused2_mr1_rp"
12138        } else if rp {
12139            "qmatvec_q4_0_mmvq_fused2_rp"
12140        } else {
12141            "qmatvec_q4_0_mmvq_fused2"
12142        });
12143        let cfg = LaunchConfig {
12144            grid_dim: (grid, 1, 1),
12145            block_dim: (32, rpb, 1),
12146            shared_mem_bytes: 0,
12147        };
12148        let inf = w0.in_features() as i32;
12149        let (oo0, oo1) = (o0 as i32, o1 as i32);
12150        let (r0, r1) = (rb0 as i64, rb1 as i64);
12151        // PDL wave-A: mr1 kernel carries MEMRA_PDL_ENTRY.
12152        if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
12153            {
12154                use cudarc::driver::{DevicePtr, DevicePtrMut};
12155                let s = &self.gpu.stream();
12156                let (p0, _g0) = b0.device_ptr(s);
12157                let (p1, _g1) = b1.device_ptr(s);
12158                let (paq, _g2) = aq.device_ptr(s);
12159                let (pad, _g3) = ad.device_ptr(s);
12160                let (py0, _g4) = y0.device_ptr_mut(s);
12161                let (py1, _g5) = y1.device_ptr_mut(s);
12162                let mut ps = [
12163                    &p0 as *const _ as *mut std::ffi::c_void,
12164                    &p1 as *const _ as *mut _,
12165                    &paq as *const _ as *mut _,
12166                    &pad as *const _ as *mut _,
12167                    &py0 as *const _ as *mut _,
12168                    &py1 as *const _ as *mut _,
12169                    &inf as *const _ as *mut _,
12170                    &oo0 as *const _ as *mut _,
12171                    &oo1 as *const _ as *mut _,
12172                    &r0 as *const _ as *mut _,
12173                    &r1 as *const _ as *mut _,
12174                ];
12175                unsafe {
12176                    self.launch_pdl(
12177                        "qmatvec_q4_0_mmvq_fused2_mr1_rp",
12178                        (grid, 1, 1),
12179                        (32, rpb, 1),
12180                        &mut ps,
12181                    )?;
12182                }
12183            }
12184            return Ok(Some((y0, y1)));
12185        }
12186        let __s_b = self.gpu.stream();
12187        let mut b = __s_b.launch_builder(&f);
12188        b.arg(b0)
12189            .arg(b1)
12190            .arg(aq)
12191            .arg(ad)
12192            .arg(&mut y0)
12193            .arg(&mut y1)
12194            .arg(&inf)
12195            .arg(&oo0)
12196            .arg(&oo1)
12197            .arg(&r0)
12198            .arg(&r1);
12199        unsafe {
12200            b.launch(cfg)?;
12201        }
12202        Ok(Some((y0, y1)))
12203    }
12204
12205    /// Slot-fed fused2 twin (alloc-free capture lane): identical launch, caller-owned outputs.
12206    pub fn matmul_q4_fused2_into(
12207        &self,
12208        w0: &crate::model::GpuTensor,
12209        w1: &crate::model::GpuTensor,
12210        aq: &CudaSlice<i8>,
12211        ad: &CudaSlice<f32>,
12212        y0: &mut CudaSlice<f32>,
12213        y1: &mut CudaSlice<f32>,
12214    ) -> Result<bool, Box<dyn std::error::Error>> {
12215        use crate::model::GpuTensor;
12216        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
12217            match w {
12218                GpuTensor::Quant {
12219                    qtype, row_bytes, ..
12220                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
12221                _ => None,
12222            }
12223        };
12224        let (Some((rb0, o0)), Some((rb1, o1))) = (q4(w0), q4(w1)) else {
12225            return Ok(false);
12226        };
12227        if w0.in_features() != w1.in_features() {
12228            return Ok(false);
12229        }
12230        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
12231            match w {
12232                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
12233                    Some(m) => (m, true),
12234                    None => (bytes, *rp),
12235                },
12236                _ => unreachable!(),
12237            }
12238        }
12239        let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
12240        if rp0 != rp1 {
12241            return Ok(false);
12242        }
12243        let rp = rp0;
12244        let rpb: u32 = 4;
12245        let mr1 = rp && Self::q40_mr1_on();
12246        let nb = |o: usize| {
12247            if mr1 {
12248                (o as u32).div_ceil(rpb)
12249            } else {
12250                (o as u32).div_ceil(2).div_ceil(rpb)
12251            }
12252        };
12253        let grid = nb(o0) + nb(o1);
12254        debug_assert!(y0.len() >= o0 && y1.len() >= o1);
12255        let f = self.func(if mr1 {
12256            "qmatvec_q4_0_mmvq_fused2_mr1_rp"
12257        } else if rp {
12258            "qmatvec_q4_0_mmvq_fused2_rp"
12259        } else {
12260            "qmatvec_q4_0_mmvq_fused2"
12261        });
12262        let cfg = LaunchConfig {
12263            grid_dim: (grid, 1, 1),
12264            block_dim: (32, rpb, 1),
12265            shared_mem_bytes: 0,
12266        };
12267        let inf = w0.in_features() as i32;
12268        let (oo0, oo1) = (o0 as i32, o1 as i32);
12269        let (r0, r1) = (rb0 as i64, rb1 as i64);
12270        // PDL wave-A: identical to the owned twin (capture-lane parity).
12271        if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
12272            use cudarc::driver::{DevicePtr, DevicePtrMut};
12273            let s = &self.gpu.stream();
12274            let (p0, _g0) = b0.device_ptr(s);
12275            let (p1, _g1) = b1.device_ptr(s);
12276            let (paq, _g2) = aq.device_ptr(s);
12277            let (pad, _g3) = ad.device_ptr(s);
12278            let (py0, _g4) = y0.device_ptr_mut(s);
12279            let (py1, _g5) = y1.device_ptr_mut(s);
12280            let mut ps = [
12281                &p0 as *const _ as *mut std::ffi::c_void,
12282                &p1 as *const _ as *mut _,
12283                &paq as *const _ as *mut _,
12284                &pad as *const _ as *mut _,
12285                &py0 as *const _ as *mut _,
12286                &py1 as *const _ as *mut _,
12287                &inf as *const _ as *mut _,
12288                &oo0 as *const _ as *mut _,
12289                &oo1 as *const _ as *mut _,
12290                &r0 as *const _ as *mut _,
12291                &r1 as *const _ as *mut _,
12292            ];
12293            unsafe {
12294                self.launch_pdl(
12295                    "qmatvec_q4_0_mmvq_fused2_mr1_rp",
12296                    (grid, 1, 1),
12297                    (32, rpb, 1),
12298                    &mut ps,
12299                )?;
12300            }
12301            return Ok(true);
12302        }
12303        let __s_b = self.gpu.stream();
12304        let mut b = __s_b.launch_builder(&f);
12305        b.arg(b0)
12306            .arg(b1)
12307            .arg(aq)
12308            .arg(ad)
12309            .arg(&mut *y0)
12310            .arg(&mut *y1)
12311            .arg(&inf)
12312            .arg(&oo0)
12313            .arg(&oo1)
12314            .arg(&r0)
12315            .arg(&r1);
12316        unsafe {
12317            b.launch(cfg)?;
12318        }
12319        Ok(true)
12320    }
12321
12322    /// BATCHED fused2 (2026-07-13, megakernel-microcosm probe): gate+up b-tier matvecs in
12323    /// ONE segmented-grid launch — the up segment fills SMs as the gate segment drains
12324    /// (the per-launch tail waves behind the 6x-falsified b-tier plateau). Bit-identical
12325    /// per row to two mr2_rp launches. rp layout required; m in 2..=8 (b16 has no twin).
12326    pub fn matmul_q4_fused2_batched(
12327        &self,
12328        w0: &crate::model::GpuTensor,
12329        w1: &crate::model::GpuTensor,
12330        aq: &CudaSlice<i8>,
12331        ad: &CudaSlice<f32>,
12332        m: usize,
12333    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
12334        use crate::model::GpuTensor;
12335        if m < 2 || m > 8 {
12336            return Ok(None);
12337        }
12338        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
12339            match w {
12340                GpuTensor::Quant {
12341                    qtype, row_bytes, ..
12342                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
12343                _ => None,
12344            }
12345        };
12346        let (Some((rb0, o0)), Some((_rb1, o1))) = (q4(w0), q4(w1)) else {
12347            return Ok(None);
12348        };
12349        if w0.in_features() != w1.in_features() {
12350            return Ok(None);
12351        }
12352        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
12353            match w {
12354                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
12355                    Some(mr) => (mr, true),
12356                    None => (bytes, *rp),
12357                },
12358                _ => unreachable!(),
12359            }
12360        }
12361        let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
12362        if !rp0 || !rp1 {
12363            return Ok(None);
12364        }
12365        let mcols = Self::batched_mcols(m);
12366        let rpb: u32 = 4;
12367        let nb = |o: usize| (o as u32).div_ceil(2 * rpb);
12368        let grid = nb(o0) + nb(o1);
12369        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
12370        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
12371        let f = self.func(match mcols {
12372            2 => "qmatvec_q4_0_mmvq_b2_f2_rp",
12373            4 => "qmatvec_q4_0_mmvq_b4_f2_rp",
12374            _ => "qmatvec_q4_0_mmvq_b8_f2_rp",
12375        });
12376        let cfg = LaunchConfig {
12377            grid_dim: (grid, 1, 1),
12378            block_dim: (32, rpb, 1),
12379            shared_mem_bytes: 0,
12380        };
12381        let inf = w0.in_features() as i32;
12382        let (oo0, oo1, mi) = (o0 as i32, o1 as i32, m as i32);
12383        let rb = rb0 as i64;
12384        let __s_b = self.gpu.stream();
12385        let mut b = __s_b.launch_builder(&f);
12386        b.arg(b0)
12387            .arg(b1)
12388            .arg(aq)
12389            .arg(ad)
12390            .arg(&mut y0)
12391            .arg(&mut y1)
12392            .arg(&inf)
12393            .arg(&oo0)
12394            .arg(&oo1)
12395            .arg(&mi)
12396            .arg(&rb);
12397        unsafe {
12398            b.launch(cfg)?;
12399        }
12400        Ok(Some((y0, y1)))
12401    }
12402
12403    /// BATCHED fused3 (see matmul_q4_fused2_batched): three-segment single launch for the
12404    /// verify qkv triple. Same-in_f q4_0 rp tensors, m in 2..=8. Bit-identical per row.
12405    #[allow(clippy::too_many_arguments)]
12406    pub fn matmul_q4_fused3_batched(
12407        &self,
12408        w0: &crate::model::GpuTensor,
12409        w1: &crate::model::GpuTensor,
12410        w2: &crate::model::GpuTensor,
12411        aq: &CudaSlice<i8>,
12412        ad: &CudaSlice<f32>,
12413        m: usize,
12414    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
12415    {
12416        use crate::model::GpuTensor;
12417        if m < 2 || m > 8 {
12418            return Ok(None);
12419        }
12420        let q4 = |w: &GpuTensor| -> Option<usize> {
12421            match w {
12422                GpuTensor::Quant { qtype, .. } if *qtype == QT_Q4_0 => Some(w.out_features()),
12423                _ => None,
12424            }
12425        };
12426        let (Some(o0), Some(o1), Some(o2)) = (q4(w0), q4(w1), q4(w2)) else {
12427            return Ok(None);
12428        };
12429        if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
12430            return Ok(None);
12431        }
12432        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
12433            match w {
12434                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
12435                    Some(mr) => (mr, true),
12436                    None => (bytes, *rp),
12437                },
12438                _ => unreachable!(),
12439            }
12440        }
12441        let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
12442        if !rp0 || !rp1 || !rp2 {
12443            return Ok(None);
12444        }
12445        let mcols = Self::batched_mcols(m);
12446        let rpb: u32 = 4;
12447        let nb = |o: usize| (o as u32).div_ceil(2 * rpb);
12448        let grid = nb(o0) + nb(o1) + nb(o2);
12449        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
12450        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
12451        let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
12452        let f = self.func(match mcols {
12453            2 => "qmatvec_q4_0_mmvq_b2_f3_rp",
12454            4 => "qmatvec_q4_0_mmvq_b4_f3_rp",
12455            _ => "qmatvec_q4_0_mmvq_b8_f3_rp",
12456        });
12457        let cfg = LaunchConfig {
12458            grid_dim: (grid, 1, 1),
12459            block_dim: (32, rpb, 1),
12460            shared_mem_bytes: 0,
12461        };
12462        let inf = w0.in_features() as i32;
12463        let (oo0, oo1, oo2, mi) = (o0 as i32, o1 as i32, o2 as i32, m as i32);
12464        let rb = 0i64;
12465        let __s_b = self.gpu.stream();
12466        let mut b = __s_b.launch_builder(&f);
12467        b.arg(b0)
12468            .arg(b1)
12469            .arg(b2)
12470            .arg(aq)
12471            .arg(ad)
12472            .arg(&mut y0)
12473            .arg(&mut y1)
12474            .arg(&mut y2)
12475            .arg(&inf)
12476            .arg(&oo0)
12477            .arg(&oo1)
12478            .arg(&oo2)
12479            .arg(&mi)
12480            .arg(&rb);
12481        unsafe {
12482            b.launch(cfg)?;
12483        }
12484        Ok(Some((y0, y1, y2)))
12485    }
12486
12487    pub fn matmul_q8_fused3(
12488        &self,
12489        w0: &crate::model::GpuTensor,
12490        w1: &crate::model::GpuTensor,
12491        w2: &crate::model::GpuTensor,
12492        aq: &CudaSlice<i8>,
12493        ad: &CudaSlice<f32>,
12494    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
12495    {
12496        // e4m3 twin (lane/fp8-decode-v1): the full-attn wq/wk/wv triple — on the NV-27B those three
12497        // are per-tensor FP8, so native residency without this arm meant three separate launches.
12498        if let Some([p0, p1, p2]) = self.e4m3_fused_params(&[w0, w1, w2]) {
12499            return Ok(Some(self.e4m3_fused3_core(
12500                p0.0,
12501                p1.0,
12502                p2.0,
12503                aq,
12504                ad,
12505                w0.in_features(),
12506                p0.1,
12507                p1.1,
12508                p2.1,
12509                p0.2,
12510                p0.3,
12511                p1.3,
12512                p2.3,
12513            )?));
12514        }
12515        let Some([p0, p1, p2]) = self.q8_fused_params(&[w0, w1, w2]) else {
12516            return Ok(None);
12517        };
12518        Ok(Some(self.q8_fused3_core(
12519            p0.0,
12520            p1.0,
12521            p2.0,
12522            aq,
12523            ad,
12524            w0.in_features(),
12525            p0.1,
12526            p1.1,
12527            p2.1,
12528            p0.2,
12529        )?))
12530    }
12531
12532    #[allow(clippy::too_many_arguments)]
12533    fn q8_fused3_core(
12534        &self,
12535        b0: &CudaSlice<u8>,
12536        b1: &CudaSlice<u8>,
12537        b2: &CudaSlice<u8>,
12538        aq: &CudaSlice<i8>,
12539        ad: &CudaSlice<f32>,
12540        in_f: usize,
12541        out0: usize,
12542        out1: usize,
12543        out2: usize,
12544        row_bytes: usize,
12545    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12546        const ROWS_PER_BLOCK: u32 = 4;
12547        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
12548        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
12549        let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
12550        let f = self.func("qmatvec_q8_0_mmvq_fused3");
12551        let mut y0 = self.alloc_uninit::<f32>(out0)?;
12552        let mut y1 = self.alloc_uninit::<f32>(out1)?;
12553        let mut y2 = self.alloc_uninit::<f32>(out2)?;
12554        let cfg = LaunchConfig {
12555            grid_dim: (nb0 + nb1 + nb2, 1, 1),
12556            block_dim: (32, ROWS_PER_BLOCK, 1),
12557            shared_mem_bytes: 0,
12558        };
12559        let (inf, o0, o1, o2, rbl) = (
12560            in_f as i32,
12561            out0 as i32,
12562            out1 as i32,
12563            out2 as i32,
12564            row_bytes as i64,
12565        );
12566        let __s_b = self.gpu.stream();
12567        let mut b = __s_b.launch_builder(&f);
12568        b.arg(b0)
12569            .arg(b1)
12570            .arg(b2)
12571            .arg(aq)
12572            .arg(ad)
12573            .arg(&mut y0)
12574            .arg(&mut y1)
12575            .arg(&mut y2)
12576            .arg(&inf)
12577            .arg(&o0)
12578            .arg(&o1)
12579            .arg(&o2)
12580            .arg(&rbl);
12581        unsafe {
12582            b.launch(cfg)?;
12583        }
12584        Ok((y0, y1, y2))
12585    }
12586
12587    /// Test entry for the kernel_check gate: fused3 from raw weight bytes (internal q8_1 quant).
12588    #[allow(clippy::too_many_arguments)]
12589    pub fn qmatvec_q8_fused3_raw(
12590        &self,
12591        b0: &CudaSlice<u8>,
12592        b1: &CudaSlice<u8>,
12593        b2: &CudaSlice<u8>,
12594        x: &CudaSlice<f32>,
12595        in_f: usize,
12596        out0: usize,
12597        out1: usize,
12598        out2: usize,
12599        row_bytes: usize,
12600    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12601        let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
12602        self.q8_fused3_core(b0, b1, b2, &aq, &ad, in_f, out0, out1, out2, row_bytes)
12603    }
12604
12605    /// BATCHED twin of `matmul_q8_fused2` for the verify t=2-4 tier (MEMRA_SPEC_FUSED_T call
12606    /// sites, lane/close35b): ONE launch computes both same-input Q8_0 projections for m tokens.
12607    /// Per (tensor,token,row) the kernel body is q8_0_mmvq_batched VERBATIM with the identical
12608    /// row mapping (Q8_0's batched_variant is always "base") -> BIT-IDENTICAL to the two
12609    /// per-tensor _b2/_b4 launches `matmul_decode_exact` dispatches at m=2-4, with the caller's
12610    /// single shared q8_1 activation replacing two per-call re-quantizes (quantize_q8_1 is
12611    /// deterministic -> same bytes). None when ineligible (m outside 2..=4 / not both Q8_0 /
12612    /// in_f mismatch / MEMRA_MMVQ=0 / MEMRA_Q8_DUAL=0 / MEMRA_NO_BATCHED set — the last keeps
12613    /// dispatch parity: without batched kernels decode-exact runs grid.y=m MMVQ, and the fused
12614    /// twin must not introduce a batched program the reference path would not run).
12615    pub fn matmul_q8_fused2_t(
12616        &self,
12617        w0: &crate::model::GpuTensor,
12618        w1: &crate::model::GpuTensor,
12619        aq: &CudaSlice<i8>,
12620        ad: &CudaSlice<f32>,
12621        m: usize,
12622    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
12623        // m<=8 (lane/q27-deepdive, 2026-08-05): was 2..=4 (the verify tier's mcols 2/4). The
12624        // serving tick's mcols-8 tier now has its fused2_b8 wrapper, so c=5..8 batched decode
12625        // fuses too — same template body, still bit-identical to the two _b8 launches.
12626        if !(2..=8).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok() {
12627            return Ok(None);
12628        }
12629        // e4m3 twin: MEMRA_B8 parity — without it m=5..8 e4m3 decode runs the per-m grid.y=m path,
12630        // so the fused b8 launch would introduce a batched program the reference path would not run.
12631        if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
12632            if m > 4 && !Self::b8_enabled() {
12633                return Ok(None);
12634            }
12635            return Ok(Some(self.e4m3_fused2_t_core(
12636                p0.0,
12637                p1.0,
12638                aq,
12639                ad,
12640                m,
12641                w0.in_features(),
12642                p0.1,
12643                p1.1,
12644                p0.2,
12645                p0.3,
12646                p1.3,
12647            )?));
12648        }
12649        let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
12650            return Ok(None);
12651        };
12652        Ok(Some(self.q8_fused2_t_core(
12653            p0.0,
12654            p1.0,
12655            aq,
12656            ad,
12657            m,
12658            w0.in_features(),
12659            p0.1,
12660            p1.1,
12661            p0.2,
12662        )?))
12663    }
12664
12665    #[allow(clippy::too_many_arguments)]
12666    fn q8_fused2_t_core(
12667        &self,
12668        b0: &CudaSlice<u8>,
12669        b1: &CudaSlice<u8>,
12670        aq: &CudaSlice<i8>,
12671        ad: &CudaSlice<f32>,
12672        m: usize,
12673        in_f: usize,
12674        out0: usize,
12675        out1: usize,
12676        row_bytes: usize,
12677    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12678        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
12679        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
12680        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
12681        let f = self.func(match Self::batched_mcols(m) {
12682            2 => "qmatvec_q8_0_mmvq_fused2_b2",
12683            4 => "qmatvec_q8_0_mmvq_fused2_b4",
12684            // b8 = the SERVING tier (lane/q27-deepdive): c=5..8 batched decode.
12685            _ => "qmatvec_q8_0_mmvq_fused2_b8",
12686        });
12687        let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
12688        let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
12689        let cfg = LaunchConfig {
12690            grid_dim: (nb0 + nb1, 1, 1),
12691            block_dim: (32, ROWS_PER_BLOCK, 1),
12692            shared_mem_bytes: 0,
12693        };
12694        let (inf, o0, o1, mi, rbl) = (
12695            in_f as i32,
12696            out0 as i32,
12697            out1 as i32,
12698            m as i32,
12699            row_bytes as i64,
12700        );
12701        let __s_b = self.gpu.stream();
12702        let mut b = __s_b.launch_builder(&f);
12703        b.arg(b0)
12704            .arg(b1)
12705            .arg(aq)
12706            .arg(ad)
12707            .arg(&mut y0)
12708            .arg(&mut y1)
12709            .arg(&inf)
12710            .arg(&o0)
12711            .arg(&o1)
12712            .arg(&mi)
12713            .arg(&rbl);
12714        unsafe {
12715            b.launch(cfg)?;
12716        }
12717        Ok((y0, y1))
12718    }
12719
12720    /// Test entry for the kernel_check gate: fused2 batched from raw weight bytes (internal
12721    /// q8_1 quant of the [m, in_f] activation), no env gating.
12722    #[allow(clippy::too_many_arguments)]
12723    pub fn qmatvec_q8_fused2_t_raw(
12724        &self,
12725        b0: &CudaSlice<u8>,
12726        b1: &CudaSlice<u8>,
12727        x: &CudaSlice<f32>,
12728        m: usize,
12729        in_f: usize,
12730        out0: usize,
12731        out1: usize,
12732        row_bytes: usize,
12733    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12734        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
12735        self.q8_fused2_t_core(b0, b1, &aq, &ad, m, in_f, out0, out1, row_bytes)
12736    }
12737
12738    /// BATCHED twin of `matmul_q8_fused3` (wq+wk+wv at verify t=2-4). Same contract as
12739    /// `matmul_q8_fused2_t` with three ranges.
12740    #[allow(clippy::too_many_arguments)]
12741    pub fn matmul_q8_fused3_t(
12742        &self,
12743        w0: &crate::model::GpuTensor,
12744        w1: &crate::model::GpuTensor,
12745        w2: &crate::model::GpuTensor,
12746        aq: &CudaSlice<i8>,
12747        ad: &CudaSlice<f32>,
12748        m: usize,
12749    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
12750    {
12751        if !(2..=4).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok() {
12752            return Ok(None);
12753        }
12754        if let Some([p0, p1, p2]) = self.e4m3_fused_params(&[w0, w1, w2]) {
12755            return Ok(Some(self.e4m3_fused3_t_core(
12756                p0.0,
12757                p1.0,
12758                p2.0,
12759                aq,
12760                ad,
12761                m,
12762                w0.in_features(),
12763                p0.1,
12764                p1.1,
12765                p2.1,
12766                p0.2,
12767                p0.3,
12768                p1.3,
12769                p2.3,
12770            )?));
12771        }
12772        let Some([p0, p1, p2]) = self.q8_fused_params(&[w0, w1, w2]) else {
12773            return Ok(None);
12774        };
12775        Ok(Some(self.q8_fused3_t_core(
12776            p0.0,
12777            p1.0,
12778            p2.0,
12779            aq,
12780            ad,
12781            m,
12782            w0.in_features(),
12783            p0.1,
12784            p1.1,
12785            p2.1,
12786            p0.2,
12787        )?))
12788    }
12789
12790    #[allow(clippy::too_many_arguments)]
12791    fn q8_fused3_t_core(
12792        &self,
12793        b0: &CudaSlice<u8>,
12794        b1: &CudaSlice<u8>,
12795        b2: &CudaSlice<u8>,
12796        aq: &CudaSlice<i8>,
12797        ad: &CudaSlice<f32>,
12798        m: usize,
12799        in_f: usize,
12800        out0: usize,
12801        out1: usize,
12802        out2: usize,
12803        row_bytes: usize,
12804    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12805        const ROWS_PER_BLOCK: u32 = 4;
12806        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
12807        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
12808        let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
12809        let f = self.func(if Self::batched_mcols(m) == 2 {
12810            "qmatvec_q8_0_mmvq_fused3_b2"
12811        } else {
12812            "qmatvec_q8_0_mmvq_fused3_b4"
12813        });
12814        let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
12815        let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
12816        let mut y2 = self.alloc_uninit::<f32>(m * out2)?;
12817        let cfg = LaunchConfig {
12818            grid_dim: (nb0 + nb1 + nb2, 1, 1),
12819            block_dim: (32, ROWS_PER_BLOCK, 1),
12820            shared_mem_bytes: 0,
12821        };
12822        let (inf, o0, o1, o2, mi, rbl) = (
12823            in_f as i32,
12824            out0 as i32,
12825            out1 as i32,
12826            out2 as i32,
12827            m as i32,
12828            row_bytes as i64,
12829        );
12830        let __s_b = self.gpu.stream();
12831        let mut b = __s_b.launch_builder(&f);
12832        b.arg(b0)
12833            .arg(b1)
12834            .arg(b2)
12835            .arg(aq)
12836            .arg(ad)
12837            .arg(&mut y0)
12838            .arg(&mut y1)
12839            .arg(&mut y2)
12840            .arg(&inf)
12841            .arg(&o0)
12842            .arg(&o1)
12843            .arg(&o2)
12844            .arg(&mi)
12845            .arg(&rbl);
12846        unsafe {
12847            b.launch(cfg)?;
12848        }
12849        Ok((y0, y1, y2))
12850    }
12851
12852    /// Test entry for the kernel_check gate: fused3 batched from raw weight bytes.
12853    #[allow(clippy::too_many_arguments)]
12854    pub fn qmatvec_q8_fused3_t_raw(
12855        &self,
12856        b0: &CudaSlice<u8>,
12857        b1: &CudaSlice<u8>,
12858        b2: &CudaSlice<u8>,
12859        x: &CudaSlice<f32>,
12860        m: usize,
12861        in_f: usize,
12862        out0: usize,
12863        out1: usize,
12864        out2: usize,
12865        row_bytes: usize,
12866    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12867        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
12868        self.q8_fused3_t_core(b0, b1, b2, &aq, &ad, m, in_f, out0, out1, out2, row_bytes)
12869    }
12870
12871    /// Rollback seam for the Q8_0 dense-FFN gate+up fusion arm in `matmul_pre_dual_noscale`
12872    /// (lane/q27-deepdive, 2026-08-05). Default ON; `MEMRA_Q8_FFN_FUSE2=0` restores the
12873    /// two-`matmul_pre_noscale` pair. Read once — the dispatch must not vary within a run.
12874    pub fn q8_ffn_fuse2_on(&self) -> bool {
12875        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12876        *ON.get_or_init(|| std::env::var("MEMRA_Q8_FFN_FUSE2").as_deref() != Ok("0"))
12877    }
12878
12879    /// Eligibility + param extraction for the fused q8_0 launches: every tensor must be Quant Q8_0
12880    /// with macro-scale 1.0 (always true for GGUF q8_0; only NVFP4 carries scale) and share w[0]'s
12881    /// in_f (q8_0 row_bytes is a pure function of in_f, so equal in_f => equal row_bytes). MEMRA_MMVQ
12882    /// must be on: the fused body is the MMVQ kernel; without it decode m=1 runs dp4a and fusing
12883    /// would mix dispatch families (FP-order law). MEMRA_Q8_DUAL=0 = rollback seam.
12884    #[allow(clippy::type_complexity)]
12885    fn q8_fused_params<'w, const N: usize>(
12886        &self,
12887        ws: &[&'w crate::model::GpuTensor; N],
12888    ) -> Option<[(&'w CudaSlice<u8>, usize, usize); N]> {
12889        use crate::model::GpuTensor;
12890        if std::env::var("MEMRA_MMVQ").as_deref() == Ok("0") {
12891            return None;
12892        }
12893        if std::env::var("MEMRA_Q8_DUAL").is_ok_and(|v| v == "0") {
12894            return None;
12895        }
12896        let in_f = ws[0].in_features();
12897        let mut out: [Option<(&CudaSlice<u8>, usize, usize)>; N] = [None; N];
12898        for (i, w) in ws.iter().enumerate() {
12899            match w {
12900                GpuTensor::Quant {
12901                    bytes,
12902                    qtype,
12903                    row_bytes,
12904                    scale,
12905                    ..
12906                } if *qtype == QT_Q8_0 && *scale == 1.0 && w.in_features() == in_f => {
12907                    out[i] = Some((bytes, w.out_features(), *row_bytes))
12908                }
12909                _ => return None,
12910            }
12911        }
12912        Some(out.map(|o| o.unwrap()))
12913    }
12914
12915    /// Rollback seam for the F8-E4M3 launch-fusion arm (lane/fp8-decode-v1, 2026-08-05).
12916    /// Default ON; `MEMRA_E4M3_DUAL=0` restores the per-tensor m=1/batched launches.
12917    pub fn e4m3_dual_on(&self) -> bool {
12918        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12919        *ON.get_or_init(|| std::env::var("MEMRA_E4M3_DUAL").as_deref() != Ok("0"))
12920    }
12921
12922    /// Eligibility + param extraction for the FUSED e4m3 launches — the QT_F8_E4M3 twin of
12923    /// `q8_fused_params`. Differences that are inherent to the dtype, not policy:
12924    ///   * each tensor carries its OWN per-tensor `weight_scale` (returned as the 4th field);
12925    ///     Q8_0 hard-requires scale==1.0 because it has no macro-scale at all.
12926    ///   * no MEMRA_MMVQ gate: `mmvq_supports` exempts QT_F8_E4M3 (the e4m3 mmvq family is that
12927    ///     dtype's ONLY int8-act kernel class), so the per-tensor fallback these fused kernels
12928    ///     replace is ALWAYS the same mmvq body under every env — the FP-order law holds.
12929    ///   * `row_bytes == in_f` is asserted rather than derived: the native-residency load arm keeps
12930    ///     the checkpoint's raw [out_f, in_f] rows, and a re-encoded slab must never reach here.
12931    /// Rejects any split-plane mirror (`rp`/`rp4`): there is no `_rp` e4m3 fused form, so fusing
12932    /// there would swap dispatch families mid-model. MEMRA_E4M3_DUAL=0 = rollback seam.
12933    #[allow(clippy::type_complexity)]
12934    fn e4m3_fused_params<'w, const N: usize>(
12935        &self,
12936        ws: &[&'w crate::model::GpuTensor; N],
12937    ) -> Option<[(&'w CudaSlice<u8>, usize, usize, f32); N]> {
12938        use crate::model::GpuTensor;
12939        if !self.e4m3_dual_on() {
12940            return None;
12941        }
12942        let in_f = ws[0].in_features();
12943        let mut out: [Option<(&CudaSlice<u8>, usize, usize, f32)>; N] = [None; N];
12944        for (i, w) in ws.iter().enumerate() {
12945            match w {
12946                GpuTensor::Quant {
12947                    bytes,
12948                    qtype,
12949                    row_bytes,
12950                    scale,
12951                    rp,
12952                    rp4,
12953                    ..
12954                } if *qtype == QT_F8_E4M3
12955                    && w.in_features() == in_f
12956                    && *row_bytes == in_f
12957                    && !*rp
12958                    && rp4.is_none() =>
12959                {
12960                    out[i] = Some((bytes, w.out_features(), *row_bytes, *scale))
12961                }
12962                _ => return None,
12963            }
12964        }
12965        Some(out.map(|o| o.unwrap()))
12966    }
12967
12968    /// FUSED e4m3 m=1 PAIR. Block-offset split (`qmatvec_e4m3_mmvq_fused2`), per-tensor
12969    /// weight_scale folded at the write like the single-tensor `qmatvec_e4m3_mmvq` — so per
12970    /// (tensor,row) this is BIT-IDENTICAL to two separate m=1 launches, scale included.
12971    #[allow(clippy::too_many_arguments)]
12972    fn e4m3_fused2_core(
12973        &self,
12974        b0: &CudaSlice<u8>,
12975        b1: &CudaSlice<u8>,
12976        aq: &CudaSlice<i8>,
12977        ad: &CudaSlice<f32>,
12978        in_f: usize,
12979        out0: usize,
12980        out1: usize,
12981        row_bytes: usize,
12982        ws0: f32,
12983        ws1: f32,
12984    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12985        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
12986        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
12987        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
12988        let f = self.func("qmatvec_e4m3_mmvq_fused2");
12989        let mut y0 = self.alloc_uninit::<f32>(out0)?;
12990        let mut y1 = self.alloc_uninit::<f32>(out1)?;
12991        let cfg = LaunchConfig {
12992            grid_dim: (nb0 + nb1, 1, 1),
12993            block_dim: (32, ROWS_PER_BLOCK, 1),
12994            shared_mem_bytes: 0,
12995        };
12996        let (inf, o0, o1, rbl) = (in_f as i32, out0 as i32, out1 as i32, row_bytes as i64);
12997        let __s_b = self.gpu.stream();
12998        let mut b = __s_b.launch_builder(&f);
12999        b.arg(b0)
13000            .arg(b1)
13001            .arg(aq)
13002            .arg(ad)
13003            .arg(&mut y0)
13004            .arg(&mut y1)
13005            .arg(&inf)
13006            .arg(&o0)
13007            .arg(&o1)
13008            .arg(&rbl)
13009            .arg(&ws0)
13010            .arg(&ws1);
13011        unsafe {
13012            b.launch(cfg)?;
13013        }
13014        Ok((y0, y1))
13015    }
13016
13017    /// FUSED e4m3 m=1 TRIPLE (`qmatvec_e4m3_mmvq_fused3`). Same contract as the pair.
13018    #[allow(clippy::too_many_arguments)]
13019    fn e4m3_fused3_core(
13020        &self,
13021        b0: &CudaSlice<u8>,
13022        b1: &CudaSlice<u8>,
13023        b2: &CudaSlice<u8>,
13024        aq: &CudaSlice<i8>,
13025        ad: &CudaSlice<f32>,
13026        in_f: usize,
13027        out0: usize,
13028        out1: usize,
13029        out2: usize,
13030        row_bytes: usize,
13031        ws0: f32,
13032        ws1: f32,
13033        ws2: f32,
13034    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13035        const ROWS_PER_BLOCK: u32 = 4;
13036        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
13037        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
13038        let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
13039        let f = self.func("qmatvec_e4m3_mmvq_fused3");
13040        let mut y0 = self.alloc_uninit::<f32>(out0)?;
13041        let mut y1 = self.alloc_uninit::<f32>(out1)?;
13042        let mut y2 = self.alloc_uninit::<f32>(out2)?;
13043        let cfg = LaunchConfig {
13044            grid_dim: (nb0 + nb1 + nb2, 1, 1),
13045            block_dim: (32, ROWS_PER_BLOCK, 1),
13046            shared_mem_bytes: 0,
13047        };
13048        let (inf, o0, o1, o2, rbl) = (
13049            in_f as i32,
13050            out0 as i32,
13051            out1 as i32,
13052            out2 as i32,
13053            row_bytes as i64,
13054        );
13055        let __s_b = self.gpu.stream();
13056        let mut b = __s_b.launch_builder(&f);
13057        b.arg(b0)
13058            .arg(b1)
13059            .arg(b2)
13060            .arg(aq)
13061            .arg(ad)
13062            .arg(&mut y0)
13063            .arg(&mut y1)
13064            .arg(&mut y2)
13065            .arg(&inf)
13066            .arg(&o0)
13067            .arg(&o1)
13068            .arg(&o2)
13069            .arg(&rbl)
13070            .arg(&ws0)
13071            .arg(&ws1)
13072            .arg(&ws2);
13073        unsafe {
13074            b.launch(cfg)?;
13075        }
13076        Ok((y0, y1, y2))
13077    }
13078
13079    /// BATCHED FUSED e4m3 pair (m=2..8). The batched kernels carry no `ws` arg (every batched
13080    /// kernel in the tree is scale-free), so each output takes its own `scale_inplace` — the
13081    /// SAME post-op the per-tensor batched dispatch applies, hence still bit-identical.
13082    #[allow(clippy::too_many_arguments)]
13083    fn e4m3_fused2_t_core(
13084        &self,
13085        b0: &CudaSlice<u8>,
13086        b1: &CudaSlice<u8>,
13087        aq: &CudaSlice<i8>,
13088        ad: &CudaSlice<f32>,
13089        m: usize,
13090        in_f: usize,
13091        out0: usize,
13092        out1: usize,
13093        row_bytes: usize,
13094        ws0: f32,
13095        ws1: f32,
13096    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13097        const ROWS_PER_BLOCK: u32 = 4;
13098        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
13099        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
13100        let f = self.func(match Self::batched_mcols(m) {
13101            2 => "qmatvec_e4m3_mmvq_fused2_b2",
13102            4 => "qmatvec_e4m3_mmvq_fused2_b4",
13103            _ => "qmatvec_e4m3_mmvq_fused2_b8",
13104        });
13105        let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
13106        let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
13107        let cfg = LaunchConfig {
13108            grid_dim: (nb0 + nb1, 1, 1),
13109            block_dim: (32, ROWS_PER_BLOCK, 1),
13110            shared_mem_bytes: 0,
13111        };
13112        let (inf, o0, o1, mi, rbl) = (
13113            in_f as i32,
13114            out0 as i32,
13115            out1 as i32,
13116            m as i32,
13117            row_bytes as i64,
13118        );
13119        let __s_b = self.gpu.stream();
13120        let mut b = __s_b.launch_builder(&f);
13121        b.arg(b0)
13122            .arg(b1)
13123            .arg(aq)
13124            .arg(ad)
13125            .arg(&mut y0)
13126            .arg(&mut y1)
13127            .arg(&inf)
13128            .arg(&o0)
13129            .arg(&o1)
13130            .arg(&mi)
13131            .arg(&rbl);
13132        unsafe {
13133            b.launch(cfg)?;
13134        }
13135        if ws0 != 1.0 {
13136            self.scale_inplace(&mut y0, ws0, m * out0)?;
13137        }
13138        if ws1 != 1.0 {
13139            self.scale_inplace(&mut y1, ws1, m * out1)?;
13140        }
13141        Ok((y0, y1))
13142    }
13143
13144    /// BATCHED FUSED e4m3 triple (m=2..4). Same contract as the batched pair.
13145    #[allow(clippy::too_many_arguments)]
13146    fn e4m3_fused3_t_core(
13147        &self,
13148        b0: &CudaSlice<u8>,
13149        b1: &CudaSlice<u8>,
13150        b2: &CudaSlice<u8>,
13151        aq: &CudaSlice<i8>,
13152        ad: &CudaSlice<f32>,
13153        m: usize,
13154        in_f: usize,
13155        out0: usize,
13156        out1: usize,
13157        out2: usize,
13158        row_bytes: usize,
13159        ws0: f32,
13160        ws1: f32,
13161        ws2: f32,
13162    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13163        const ROWS_PER_BLOCK: u32 = 4;
13164        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
13165        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
13166        let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
13167        let f = self.func(if Self::batched_mcols(m) == 2 {
13168            "qmatvec_e4m3_mmvq_fused3_b2"
13169        } else {
13170            "qmatvec_e4m3_mmvq_fused3_b4"
13171        });
13172        let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
13173        let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
13174        let mut y2 = self.alloc_uninit::<f32>(m * out2)?;
13175        let cfg = LaunchConfig {
13176            grid_dim: (nb0 + nb1 + nb2, 1, 1),
13177            block_dim: (32, ROWS_PER_BLOCK, 1),
13178            shared_mem_bytes: 0,
13179        };
13180        let (inf, o0, o1, o2, mi, rbl) = (
13181            in_f as i32,
13182            out0 as i32,
13183            out1 as i32,
13184            out2 as i32,
13185            m as i32,
13186            row_bytes as i64,
13187        );
13188        let __s_b = self.gpu.stream();
13189        let mut b = __s_b.launch_builder(&f);
13190        b.arg(b0)
13191            .arg(b1)
13192            .arg(b2)
13193            .arg(aq)
13194            .arg(ad)
13195            .arg(&mut y0)
13196            .arg(&mut y1)
13197            .arg(&mut y2)
13198            .arg(&inf)
13199            .arg(&o0)
13200            .arg(&o1)
13201            .arg(&o2)
13202            .arg(&mi)
13203            .arg(&rbl);
13204        unsafe {
13205            b.launch(cfg)?;
13206        }
13207        if ws0 != 1.0 {
13208            self.scale_inplace(&mut y0, ws0, m * out0)?;
13209        }
13210        if ws1 != 1.0 {
13211            self.scale_inplace(&mut y1, ws1, m * out1)?;
13212        }
13213        if ws2 != 1.0 {
13214            self.scale_inplace(&mut y2, ws2, m * out2)?;
13215        }
13216        Ok((y0, y1, y2))
13217    }
13218
13219    /// BLOCK-128 e4m3 MMVQ launcher (`qmatvec_e4m3_blk_mmvq`, lane/fp8-blk128-decode 2026-08-05).
13220    /// The per-block-dequant twin of `qmatvec_mmvq`'s QT_F8_E4M3 arm: same grid/block decomposition
13221    /// (warp per output row, ROWS_PER_BLOCK warps per block, grid.y = m), same q8_1 activation, but
13222    /// the weight scale is a resident [rows, cols] f32 grid read per k128 block inside the kernel
13223    /// instead of one scalar folded at the write. It cannot share `qmatvec_mmvq`'s body because
13224    /// that launcher's arg list is fixed at (bytes, aq, ad, y, in_f, out_f, m, row_bytes [, scale]).
13225    ///
13226    /// `mr` and `rp` have no analogue here (no split-plane e4m3 layout exists), so there is exactly
13227    /// one kernel and no name table — a shape this cannot serve must be refused at LOAD, not here.
13228    pub fn qmatvec_e4m3_blk_mmvq(
13229        &self,
13230        bytes: &CudaSlice<u8>,
13231        aq: &CudaSlice<i8>,
13232        ad: &CudaSlice<f32>,
13233        scales: &CudaSlice<f32>,
13234        m: usize,
13235        in_f: usize,
13236        out_f: usize,
13237        row_bytes: usize,
13238        scale_cols: usize,
13239    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13240        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
13241        self.qmatvec_e4m3_blk_mmvq_into(
13242            bytes, aq, ad, scales, m, in_f, out_f, row_bytes, scale_cols, &mut y,
13243        )?;
13244        Ok(y)
13245    }
13246
13247    /// Slot-fed twin of `qmatvec_e4m3_blk_mmvq` (caller-owned output; the alloc-free capture lane).
13248    #[allow(clippy::too_many_arguments)]
13249    pub fn qmatvec_e4m3_blk_mmvq_into(
13250        &self,
13251        bytes: &CudaSlice<u8>,
13252        aq: &CudaSlice<i8>,
13253        ad: &CudaSlice<f32>,
13254        scales: &CudaSlice<f32>,
13255        m: usize,
13256        in_f: usize,
13257        out_f: usize,
13258        row_bytes: usize,
13259        scale_cols: usize,
13260        y: &mut CudaSlice<f32>,
13261    ) -> Result<(), Box<dyn std::error::Error>> {
13262        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
13263        let f = self.func("qmatvec_e4m3_blk_mmvq");
13264        let cfg = LaunchConfig {
13265            grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), m as u32, 1),
13266            block_dim: (32, ROWS_PER_BLOCK, 1), // warp-per-row
13267            shared_mem_bytes: 0,                // warp-only reduce
13268        };
13269        let (inf, outf, mi, rb, sc) = (
13270            in_f as i32,
13271            out_f as i32,
13272            m as i32,
13273            row_bytes as i64,
13274            scale_cols as i32,
13275        );
13276        let __s_b = self.gpu.stream();
13277        let mut b = __s_b.launch_builder(&f);
13278        b.arg(bytes)
13279            .arg(aq)
13280            .arg(ad)
13281            .arg(scales)
13282            .arg(&mut *y)
13283            .arg(&inf)
13284            .arg(&outf)
13285            .arg(&mi)
13286            .arg(&rb)
13287            .arg(&sc);
13288        unsafe {
13289            b.launch(cfg)?;
13290        }
13291        Ok(())
13292    }
13293
13294    /// BLOCK-128 e4m3 BATCHED matvec (lane/rp-on-st, 2026-08-06): the weight-read-once twin of
13295    /// `qmatvec_e4m3_blk_mmvq` for m=2..16. Per (token,row) BIT-IDENTICAL to the grid.y=m launch
13296    /// (same fmaf chain, same per-k32 `s * ad` fold, same warp reduce), so it inherits the
13297    /// decode-exactness contract while reading the weight ONCE for up to `mcols` columns instead
13298    /// of `m` times. `mcols` must be one of {2,4,8,16} and satisfy `mcols >= m`.
13299    #[allow(clippy::too_many_arguments)]
13300    pub fn qmatvec_e4m3_blk_mmvq_batched(
13301        &self,
13302        bytes: &CudaSlice<u8>,
13303        aq: &CudaSlice<i8>,
13304        ad: &CudaSlice<f32>,
13305        scales: &CudaSlice<f32>,
13306        m: usize,
13307        in_f: usize,
13308        out_f: usize,
13309        row_bytes: usize,
13310        scale_cols: usize,
13311        mcols: usize,
13312    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13313        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
13314        debug_assert!(mcols >= m, "blk batched: mcols {mcols} < m {m}");
13315        let name = match mcols {
13316            2 => "qmatvec_e4m3_blk_mmvq_b2",
13317            4 => "qmatvec_e4m3_blk_mmvq_b4",
13318            8 => "qmatvec_e4m3_blk_mmvq_b8",
13319            16 => "qmatvec_e4m3_blk_mmvq_b16",
13320            _ => {
13321                return Err(
13322                    format!("qmatvec_e4m3_blk_mmvq_batched: no kernel for mcols {mcols}").into(),
13323                );
13324            }
13325        };
13326        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
13327        let f = self.func(name);
13328        let cfg = LaunchConfig {
13329            grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
13330            block_dim: (32, ROWS_PER_BLOCK, 1),
13331            shared_mem_bytes: 0,
13332        };
13333        let (inf, outf, mi, rb, sc) = (
13334            in_f as i32,
13335            out_f as i32,
13336            m as i32,
13337            row_bytes as i64,
13338            scale_cols as i32,
13339        );
13340        let __s_b = self.gpu.stream();
13341        let mut b = __s_b.launch_builder(&f);
13342        b.arg(bytes)
13343            .arg(aq)
13344            .arg(ad)
13345            .arg(scales)
13346            .arg(&mut y)
13347            .arg(&inf)
13348            .arg(&outf)
13349            .arg(&mi)
13350            .arg(&rb)
13351            .arg(&sc);
13352        unsafe {
13353            b.launch(cfg)?;
13354        }
13355        Ok(y)
13356    }
13357
13358    /// Test entry for the kernel_check exactness gate: the block-128 e4m3 batched MMVQ from raw
13359    /// bytes with an internal q8_1 quantize (mirrors `qmatvec_batched_raw`).
13360    #[allow(clippy::too_many_arguments)]
13361    pub fn qmatvec_e4m3_blk_batched_raw(
13362        &self,
13363        bytes: &CudaSlice<u8>,
13364        x: &CudaSlice<f32>,
13365        scales: &CudaSlice<f32>,
13366        m: usize,
13367        in_f: usize,
13368        out_f: usize,
13369        row_bytes: usize,
13370        scale_cols: usize,
13371        mcols: usize,
13372    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13373        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
13374        self.qmatvec_e4m3_blk_mmvq_batched(
13375            bytes, &aq, &ad, scales, m, in_f, out_f, row_bytes, scale_cols, mcols,
13376        )
13377    }
13378
13379    /// Test entry for the kernel_check exactness gate: the block-128 e4m3 MMVQ from raw bytes with
13380    /// an internal q8_1 quantize (mirrors `qmatvec_mmvq_raw`).
13381    #[allow(clippy::too_many_arguments)]
13382    pub fn qmatvec_e4m3_blk_mmvq_raw(
13383        &self,
13384        bytes: &CudaSlice<u8>,
13385        x: &CudaSlice<f32>,
13386        scales: &CudaSlice<f32>,
13387        m: usize,
13388        in_f: usize,
13389        out_f: usize,
13390        row_bytes: usize,
13391        scale_cols: usize,
13392    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13393        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
13394        self.qmatvec_e4m3_blk_mmvq(
13395            bytes, &aq, &ad, scales, m, in_f, out_f, row_bytes, scale_cols,
13396        )
13397    }
13398
13399    /// Test entries for the kernel_check bit-parity gate: fused e4m3 launches from raw weight
13400    /// bytes with internal q8_1 quantize, no env gating (mirrors `qmatvec_q8_fused*_raw`).
13401    #[allow(clippy::too_many_arguments)]
13402    pub fn qmatvec_e4m3_fused2_raw(
13403        &self,
13404        b0: &CudaSlice<u8>,
13405        b1: &CudaSlice<u8>,
13406        x: &CudaSlice<f32>,
13407        in_f: usize,
13408        out0: usize,
13409        out1: usize,
13410        row_bytes: usize,
13411        ws0: f32,
13412        ws1: f32,
13413    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13414        let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
13415        self.e4m3_fused2_core(b0, b1, &aq, &ad, in_f, out0, out1, row_bytes, ws0, ws1)
13416    }
13417
13418    #[allow(clippy::too_many_arguments)]
13419    pub fn qmatvec_e4m3_fused3_raw(
13420        &self,
13421        b0: &CudaSlice<u8>,
13422        b1: &CudaSlice<u8>,
13423        b2: &CudaSlice<u8>,
13424        x: &CudaSlice<f32>,
13425        in_f: usize,
13426        out0: usize,
13427        out1: usize,
13428        out2: usize,
13429        row_bytes: usize,
13430        ws0: f32,
13431        ws1: f32,
13432        ws2: f32,
13433    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13434        let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
13435        self.e4m3_fused3_core(
13436            b0, b1, b2, &aq, &ad, in_f, out0, out1, out2, row_bytes, ws0, ws1, ws2,
13437        )
13438    }
13439
13440    #[allow(clippy::too_many_arguments)]
13441    pub fn qmatvec_e4m3_fused2_t_raw(
13442        &self,
13443        b0: &CudaSlice<u8>,
13444        b1: &CudaSlice<u8>,
13445        x: &CudaSlice<f32>,
13446        m: usize,
13447        in_f: usize,
13448        out0: usize,
13449        out1: usize,
13450        row_bytes: usize,
13451        ws0: f32,
13452        ws1: f32,
13453    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13454        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
13455        self.e4m3_fused2_t_core(b0, b1, &aq, &ad, m, in_f, out0, out1, row_bytes, ws0, ws1)
13456    }
13457
13458    #[allow(clippy::too_many_arguments)]
13459    pub fn qmatvec_e4m3_fused3_t_raw(
13460        &self,
13461        b0: &CudaSlice<u8>,
13462        b1: &CudaSlice<u8>,
13463        b2: &CudaSlice<u8>,
13464        x: &CudaSlice<f32>,
13465        m: usize,
13466        in_f: usize,
13467        out0: usize,
13468        out1: usize,
13469        out2: usize,
13470        row_bytes: usize,
13471        ws0: f32,
13472        ws1: f32,
13473        ws2: f32,
13474    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13475        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
13476        self.e4m3_fused3_t_core(
13477            b0, b1, b2, &aq, &ad, m, in_f, out0, out1, out2, row_bytes, ws0, ws1, ws2,
13478        )
13479    }
13480
13481    /// THE single dispatch point for `QT_F8_E4M3_BLK` from a PRE-QUANTIZED q8_1 activation
13482    /// (lane/fp8-blk128-decode). Every `matmul_pre`-family entry calls this first, so the block-128
13483    /// class has exactly ONE code path across `matmul`, `matmul_pre`, `matmul_pre_noscale`,
13484    /// `matmul_decode_exact` and `matmul_decode_exact_pre` — the same kernel at the same grid for
13485    /// every m, which is what makes verify == decode bit-for-bit at every tier for free.
13486    ///
13487    /// Returns None for any other qtype (the caller continues its normal dispatch). The `blk: Some`
13488    /// pattern is part of the match, not an unwrap: qtype and grid presence are set together in the
13489    /// one residency arm that builds this tensor, and a qtype-without-grid would be a construction
13490    /// bug — better to fall through and hit a loud refusal than to unwrap a None here.
13491    fn try_e4m3_blk_pre(
13492        &self,
13493        w: &crate::model::GpuTensor,
13494        aq: &CudaSlice<i8>,
13495        ad: &CudaSlice<f32>,
13496        m: usize,
13497    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
13498        use crate::model::GpuTensor;
13499        if let GpuTensor::Quant {
13500            bytes,
13501            qtype,
13502            row_bytes,
13503            blk: Some(g),
13504            ..
13505        } = w
13506        {
13507            if *qtype == QT_F8_E4M3_BLK {
13508                // BATCHED tier m=2..16 (lane/rp-on-st): weight read ONCE for up to mcols columns
13509                // instead of m grid.y re-reads. Bit-identical per (token,row) to the grid.y=m form
13510                // below, so the decode-exactness contract is preserved at every width. Gated by
13511                // the same seams the other batched families honor (MEMRA_NO_BATCHED, MEMRA_B8) so
13512                // one rollback door covers every dtype's batched tier.
13513                if (2..=16).contains(&m)
13514                    && std::env::var("MEMRA_NO_BATCHED").is_err()
13515                    && (m <= 4 || Self::b8_enabled())
13516                {
13517                    let mcols = Self::batched_mcols(m);
13518                    return Ok(Some(self.qmatvec_e4m3_blk_mmvq_batched(
13519                        bytes,
13520                        aq,
13521                        ad,
13522                        &g.scales,
13523                        m,
13524                        w.in_features(),
13525                        w.out_features(),
13526                        *row_bytes,
13527                        g.cols,
13528                        mcols,
13529                    )?));
13530                }
13531                return Ok(Some(self.qmatvec_e4m3_blk_mmvq(
13532                    bytes,
13533                    aq,
13534                    ad,
13535                    &g.scales,
13536                    m,
13537                    w.in_features(),
13538                    w.out_features(),
13539                    *row_bytes,
13540                    g.cols,
13541                )?));
13542            }
13543        }
13544        Ok(None)
13545    }
13546
13547    /// PREFILL (m >= GEMM_M_THRESHOLD) for `QT_F8_E4M3_BLK` — DEQUANT-PER-CALL to the Q8_0 slab
13548    /// this class's residency replaced, then the ordinary Q8_0 prefill dispatch on the transient.
13549    ///
13550    /// WHY THIS EXISTS AT ALL, i.e. the regression it prevents: the decode kernel is a warp-per-row
13551    /// GEMV. At grid.y=m it re-reads the whole weight once PER TOKEN, so letting a 512-token prefill
13552    /// chunk reach it would be a ~500x weight-traffic blowup on the single most bandwidth-bound part
13553    /// of the forward. Native residency is a DECODE win and must not be paid for in prefill, so
13554    /// prefill keeps the floor's arithmetic and the floor's kernels.
13555    ///
13556    /// WHY DEQUANT-PER-CALL rather than a second resident slab: a resident slab is dual residency —
13557    /// it gives back the entire 1.0-vs-1.0625 B/weight win this lane exists to capture (and then
13558    /// some, since the e4m3 copy stays too). The transient costs one linear device pass per
13559    /// (projection, prefill call) and frees immediately.
13560    ///
13561    /// NUMERICALLY IT IS THE FLOOR, EXACTLY: `fp8_blk_dequant_q8_0` is the merged ARM B' kernel,
13562    /// gate-proven BYTE-IDENTICAL to the host dequant+re-encode (kernel-check `fp8-blk-gpu`). So the
13563    /// slab these bytes form is bit-for-bit the slab the `MEMRA_ST_E4M3_BLK=0` arm makes resident,
13564    /// and every prefill kernel downstream sees identical input — prefill logits under this lane are
13565    /// bit-identical to prefill logits under the floor, which is what makes the decode A/B a clean
13566    /// single-variable comparison instead of a two-variable one.
13567    ///
13568    /// WHAT IT COSTS, MEASURED, AND WHY THAT COST IS MOSTLY STRUCTURAL (27B block-128 ckpt, pp512,
13569    /// this rig = RTX 5090 Laptop, ~896 GB/s GDDR7). This arm makes prefill move the weight THREE
13570    /// times instead of once: read 6.88 GB of e4m3, write 7.31 GB of Q8_0, then the MMQ reads that
13571    /// 7.31 GB back. The two extra passes are 14.19 GB = 15.8 ms at this card's roofline against a
13572    /// ~332 ms pp512, i.e. **~-4.5% pp is a floor no kernel tuning can remove** — only deleting the
13573    /// dequant can. Measured: the dequant kernel costs 27.9 ms/pass (nsys, 208 projections) after
13574    /// the 2026-08-05 vector rewrite (was 66.5 ms at one byte per thread), and e2e pp512 is
13575    /// 1451.4 vs the slab arm's 1541.6 tok/s = -5.8% (N=3 interleaved pairs). So ~1.3pp of the
13576    /// -5.8% is residual kernel inefficiency and ~4.5pp is the extra traffic itself.
13577    ///
13578    /// SO THE DEQUANT IS NO LONGER THE DEFAULT ROUTE — it is the FALLBACK. The per-block FP8 MMQ
13579    /// tile (`try_fp8_blk_mmq`) consumes the resident e4m3 bytes + grid DIRECTLY, deleting both extra
13580    /// passes, and since 2026-08-05 it runs FIRST and by default for the native-resident source
13581    /// (`fp8_blk_mmq_native_enabled`; `MEMRA_FP8_MMQ=0` is the seam back to this dequant). On paper
13582    /// the trade was unassumable — lane/fp8-mmq-v2 measured that tile at 0.85-1.09x the Q8_0 MMQ
13583    /// floor GEMM-only, so it swapped a -4.5% traffic cost for a 0-to-15% GEMM cost of unknown sign.
13584    /// Measured on the 27B (3 arms interleaved, N=3, research/fp8blk-20260805/VERDICT.md): slab
13585    /// 1540.5 / this dequant 1449.1 / the tile 1553.3 tok/s, min(tile) > max(slab). The tile wins
13586    /// because v2's denominator had its slab already resident while this class's floor must build it
13587    /// every call; same tile, opposite sign, because the question changed.
13588    ///
13589    /// THIS ARM STILL RUNS, and is not dead code: every `try_fp8_blk_mmq` precondition (in_f % 16,
13590    /// grid dims vs shape, per-tensor scale == 1.0, the e4m3-NaN scan) refuses by falling through to
13591    /// here, so a checkpoint the tile cannot take keeps exact prefill on the floor's own bits rather
13592    /// than losing the class. It is also what `MEMRA_FP8_MMQ=0` reverts to.
13593    fn try_e4m3_blk_prefill(
13594        &self,
13595        w: &crate::model::GpuTensor,
13596        x: &CudaSlice<f32>,
13597        m: usize,
13598    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
13599        use crate::model::GpuTensor;
13600        let GpuTensor::Quant {
13601            bytes,
13602            qtype,
13603            blk: Some(g),
13604            ..
13605        } = w
13606        else {
13607            return Ok(None);
13608        };
13609        if *qtype != QT_F8_E4M3_BLK {
13610            return Ok(None);
13611        }
13612        // NO-DEQUANT ROUTE, THE DEFAULT (MEMRA_FP8_MMQ=0 reverts): the per-block MMQ tile eats the
13613        // resident e4m3 bytes and grid as-is, so neither extra weight pass happens. Its own
13614        // preconditions (in_f % 16, grid dims, scale == 1.0, no e4m3 NaN code) can refuse — fall
13615        // through to the dequant below when they do, never silently produce nothing.
13616        if let Some(y) = self.try_fp8_blk_mmq(w, x, m)? {
13617            return Ok(Some(y));
13618        }
13619        let (in_f, out_f) = (w.in_features(), w.out_features());
13620        let slab = self.fp8_blk_dequant_q8_0_dev(bytes, &g.scales, out_f, in_f)?;
13621        let tmp = GpuTensor::Quant {
13622            bytes: slab,
13623            qtype: QT_Q8_0,
13624            row_bytes: in_f / 32 * 34,
13625            ne: vec![in_f as u64, out_f as u64],
13626            scale: 1.0,
13627            rp: false,
13628            #[cfg(memra_cutlass)]
13629            cutlass: None,
13630            fp8: None,
13631            blk: None,
13632            f16: None,
13633            rp4: None,
13634        };
13635        // Recursion terminates: `tmp` is QT_Q8_0 with `blk: None`, so it cannot re-enter this arm.
13636        Ok(Some(self.matmul(&tmp, x, m)?))
13637    }
13638
13639    pub fn matmul_pre_noscale(
13640        &self,
13641        w: &crate::model::GpuTensor,
13642        aq: &CudaSlice<i8>,
13643        ad: &CudaSlice<f32>,
13644        m: usize,
13645    ) -> Result<Option<(CudaSlice<f32>, f32)>, Box<dyn std::error::Error>> {
13646        use crate::model::GpuTensor;
13647        // BLOCK-128 e4m3: every scale factor is folded inside the kernel per k128, so the
13648        // "separable post-op scale" this entry exists to defer is 1.0 — return it explicitly
13649        // rather than let the tail below refuse and cost the caller a re-dispatch.
13650        if m == 1 {
13651            if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
13652                return Ok(Some((y, 1.0)));
13653            }
13654        }
13655        // Only the m==1 fast path applies the scale as a separable post-op; bail everywhere else.
13656        if m != 1 || !self.uses_q8_1_fast(w) {
13657            return Ok(None);
13658        }
13659        let in_f = w.in_features();
13660        let out_f = w.out_features();
13661        let (bytes, qtype, row_bytes, scale, rp) = match w {
13662            GpuTensor::Quant {
13663                bytes,
13664                qtype,
13665                row_bytes,
13666                scale,
13667                rp,
13668                ..
13669            } => (bytes, *qtype, *row_bytes, *scale, *rp),
13670            _ => return Ok(None),
13671        };
13672        // MMVQ warp-per-row (scale==1.0 passed -> kernel skips its internal scale; we return scale).
13673        if self.mmvq_supports(qtype) {
13674            // Q4_0 split-plane mirror (dp4a fallback below keeps the raw GGUF bytes).
13675            let (mbytes, mrp) = match w {
13676                GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
13677                _ => (bytes, rp),
13678            };
13679            let y = self.qmatvec_mmvq(
13680                mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, /*scale*/ 1.0, mrp,
13681            )?;
13682            return Ok(Some((y, scale)));
13683        }
13684        // dp4a fallback: same launch as matmul_pre but WITHOUT the post scale_inplace.
13685        let name = match qtype {
13686            QT_Q8_0 => "qmatvec_q8_0_dp4a",
13687            QT_Q4_K => "qmatvec_q4_K_dp4a",
13688            QT_Q6_K => "qmatvec_q6_K_dp4a",
13689            QT_Q5_K => "qmatvec_q5_K_dp4a",
13690            QT_Q3_K => "qmatvec_q3_K_dp4a",
13691            QT_NVFP4 => {
13692                if rp {
13693                    "qmatvec_nvfp4_dp4a_rp"
13694                } else {
13695                    "qmatvec_nvfp4_dp4a"
13696                }
13697            }
13698            QT_IQ4_XS => "qmatvec_iq4_XS_dp4a",
13699            _ => return Ok(None),
13700        };
13701        let f = self.func(name);
13702        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
13703        let cfg = LaunchConfig {
13704            grid_dim: (out_f as u32, m as u32, 1),
13705            block_dim: (128, 1, 1),
13706            shared_mem_bytes: 0,
13707        };
13708        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
13709        let __s_b = self.gpu.stream();
13710        let mut b = __s_b.launch_builder(&f);
13711        b.arg(bytes)
13712            .arg(aq)
13713            .arg(ad)
13714            .arg(&mut y)
13715            .arg(&inf)
13716            .arg(&outf)
13717            .arg(&mi)
13718            .arg(&rb);
13719        unsafe {
13720            b.launch(cfg)?;
13721        }
13722        Ok(Some((y, scale)))
13723    }
13724
13725    /// True if `qtype` has a warp-per-row MMVQ decode kernel AND MEMRA_MMVQ is set. Only the 4
13726    /// daily-hot dtypes (Q8_0, Q4_K, Q6_K, NVFP4) — others keep the _dp4a matvec (oracle/fallback).
13727    pub fn mmvq_supports(&self, qtype: i32) -> bool {
13728        // DEFAULT ON since 2026-07-08 (MEMRA_MMVQ=0 reverts to the _dp4a matvec class).
13729        // QT_F8_E4M3 is exempt from the MEMRA_MMVQ=0 escape: the e4m3 mmvq family is that dtype's
13730        // ONLY int8-act kernel class (there is no _dp4a twin), so its m=1/verify/batched dispatch
13731        // is a pure function of the dtype — the decode-parity law holds under every env.
13732        if qtype == QT_F8_E4M3 {
13733            return true;
13734        }
13735        if std::env::var("MEMRA_MMVQ").as_deref() == Ok("0") {
13736            return false;
13737        }
13738        matches!(
13739            qtype,
13740            QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q6_K | QT_NVFP4 | QT_Q4_0
13741        )
13742    }
13743
13744    /// PERF-3 warp-per-row MMVQ launcher (decode m=1 hot path). block=(32,ROWS_PER_BLOCK,1):
13745    /// one warp owns one output row, warp-only __shfl reduction (no smem barrier). Bit-equivalent
13746    /// to qmatvec_*_dp4a up to f32 reduction order. Pre-quantized q8_1 activation (aq,ad). NVFP4
13747    /// per-tensor macro-scale applied post (scale==1.0 for other dtypes -> no-op).
13748    pub fn qmatvec_mmvq(
13749        &self,
13750        bytes: &CudaSlice<u8>,
13751        aq: &CudaSlice<i8>,
13752        ad: &CudaSlice<f32>,
13753        m: usize,
13754        in_f: usize,
13755        out_f: usize,
13756        qtype: i32,
13757        row_bytes: usize,
13758        scale: f32,
13759        rp: bool,
13760    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13761        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output: skip memset
13762        self.qmatvec_mmvq_into(
13763            bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, rp, &mut y,
13764        )?;
13765        Ok(y)
13766    }
13767
13768    /// Slot-fed MMVQ twin (alloc-free capture lane): full policy body, caller-owned output.
13769    #[allow(clippy::too_many_arguments)]
13770    pub fn qmatvec_mmvq_into(
13771        &self,
13772        bytes: &CudaSlice<u8>,
13773        aq: &CudaSlice<i8>,
13774        ad: &CudaSlice<f32>,
13775        m: usize,
13776        in_f: usize,
13777        out_f: usize,
13778        qtype: i32,
13779        row_bytes: usize,
13780        scale: f32,
13781        rp: bool,
13782        y: &mut CudaSlice<f32>,
13783    ) -> Result<(), Box<dyn std::error::Error>> {
13784        debug_assert!(y.len() >= m * out_f);
13785        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
13786        // SMALL-SHAPE GRID FILL (H100 lane, 2026-07-26 microbench: attn qkv out_f=2048 =
13787        // 0.97 waves at the 4-warp block -> 66% of peak). The g2 twin (2 warps/block)
13788        // doubles the grid when the 4-warp launch would be sub-wave; per-row program
13789        // identical -> bit-identical. MEMRA_Q80_G2=0 reverts.
13790        if qtype == QT_Q8_0
13791            && rp
13792            && m == 1
13793            && out_f >= 64
13794            && (out_f as u32).div_ceil(ROWS_PER_BLOCK) < 4 * self.sm_count() as u32
13795            && {
13796                static G2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13797                *G2.get_or_init(|| std::env::var("MEMRA_Q80_G2").as_deref() != Ok("0"))
13798            }
13799        {
13800            let f = self.func("qmatvec_q8_0_mmvq_rp_g2");
13801            let cfg = LaunchConfig {
13802                grid_dim: ((out_f as u32).div_ceil(2), 1, 1),
13803                block_dim: (32, 2, 1),
13804                shared_mem_bytes: 0,
13805            };
13806            let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, 1i32, row_bytes as i64);
13807            let __s_b = self.gpu.stream();
13808            let mut b = __s_b.launch_builder(&f);
13809            b.arg(bytes)
13810                .arg(aq)
13811                .arg(ad)
13812                .arg(&mut *y)
13813                .arg(&inf)
13814                .arg(&outf)
13815                .arg(&mi)
13816                .arg(&rb);
13817            unsafe {
13818                b.launch(cfg)?;
13819            }
13820            if scale != 1.0 {
13821                self.scale_inplace(y, scale, out_f)?;
13822            }
13823            return Ok(());
13824        }
13825        // Multi-row-per-warp (mr2) policy, fixed since the 2026-07 sweeps (the MEMRA_MMVQ_MR
13826        // override + mr4 kernel were retired 2026-07-08 — mr4 regressed on register pressure and
13827        // crashed under rp; q4_K/q6_K mr2 measured flat, "no gain = no change"):
13828        //   NVFP4 m=1 -> mr2 (clean +1-2% on 9B: RPW acc chains hide the weight-load latency
13829        //     that pins the single-row kernel at 30-46% DRAM). Bit-identical per row.
13830        //   Q5_K m=1 -> mr2 (2026-07-05: the FR-Spec trimmed draft head is Q5_K 32768 rows = 8%
13831        //     of the 27B p3 spec wall; latency-bound like the other k-quants pre-fix).
13832        //   Q4_K/Q6_K m=1 -> single-row (mr2 measured +0.7% / flat — weight-bandwidth-bound).
13833        let mut mr: u32 = if m == 1 && (qtype == QT_NVFP4 || qtype == QT_Q5_K) {
13834            2
13835        } else {
13836            1
13837        };
13838        // Q4_0 mr (gemma trunk): DEFAULT 1 since 2026-07-13 (MEMRA_Q40_MR=2 reverts) — the
13839        // mr1 rp twin doubles the block count and wins the tail-quantization/latency battle
13840        // on every gemma model (E4B +3.75%: 198.9 vs 191.7; 26B +0.7%; 31B +0.9%; N=2-3
13841        // valid-window interleaved, bit-identical per row — same dot program).
13842        if m == 1 && qtype == QT_Q4_0 {
13843            static Q40MR: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
13844            // shape policy PROBED NEGATIVE (2026-07-13): tall-only mr1 197.2 vs
13845            // mr1-everywhere 198.7 — mr1 wins wide-output shapes too; arm removed.
13846            mr = *Q40MR.get_or_init(|| {
13847                std::env::var("MEMRA_Q40_MR")
13848                    .ok()
13849                    .and_then(|v| v.parse().ok())
13850                    .unwrap_or(1)
13851            });
13852        }
13853        // q5issue lane (2026-07-08): MEMRA_Q5K_ISSUE swaps the q5_K m=1 mmvq kernels for the
13854        // issue-reduced `_il` bodies (uint4 header/qh/qs loads + branchless scale decode —
13855        // cuts ~34 LDG.U16 + ~5 LDG.U8 + a warp-divergent scale branch per 32-elem group-row
13856        // to 5 LDG.128). Bit-identical per (token,row) to the reference kernels.
13857        // `1` = shape-aware policy (N=3 clock-locked micro-bench, mem P0, synthetic real shapes):
13858        //   out_f <= 65536 (trunk/frspec regime): il at the default mr — mr2_il -9.5%/-10.5%
13859        //     on 4096x4096/4096x8192, -3.1% on the 32768 frspec head vs the mr2-ref default;
13860        //   out_f > 65536 (the 248320-row 27B lm_head, already ~97% of the mem wall): mr2_il
13861        //     REGRESSES +22% there but mr1_il wins -2.1% vs the mr2-ref default -> force mr=1.
13862        // `2` = force il at the current mr for EVERY shape (A/B probe seam). Default OFF.
13863        let q5_mode = std::env::var("MEMRA_Q5K_ISSUE").ok();
13864        let q5_force = q5_mode.as_deref() == Some("2");
13865        // DEFAULT ON since 2026-07-08 (MEMRA_Q5K_ISSUE=0 reverts): +1.8% 9B plain e2e N=3
13866        // (128.2 -> 130.4), 27B flat (its big head is already at the mem wall), all gates green.
13867        let q5_il = qtype == QT_Q5_K
13868            && m == 1
13869            && (q5_force || q5_mode.as_deref().map(|v| v != "0").unwrap_or(true));
13870        if q5_il && !q5_force && out_f > 65536 {
13871            mr = 1;
13872        }
13873        // Q4_0 split-plane rp: mr2 default; MEMRA_Q40_MR=1 reaches the mr1 rp twin
13874        // (2026-07-13 — the tall-input/short-output tail-quantization probe).
13875        if qtype == QT_Q4_0 && rp && mr != 1 {
13876            mr = 2;
13877        }
13878        // Q8_0 rp (H100 lane): mr1 default — the q4_0 mr2 recipe MEASURED NEGATIVE on H100
13879        // (2026-07-26 N=3: mr1 186.2 vs mr2 171.5 tok/s; halving the grid on 132 SMs costs
13880        // more than 2-row ILP buys). mr2 kernel stays behind MEMRA_Q80_MR=2 for the corpus.
13881        if qtype == QT_Q8_0 && rp {
13882            static Q80MR: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
13883            mr = *Q80MR.get_or_init(|| {
13884                std::env::var("MEMRA_Q80_MR")
13885                    .ok()
13886                    .and_then(|v| v.parse().ok())
13887                    .unwrap_or(1)
13888            });
13889        }
13890        let name = match (qtype, mr, rp) {
13891            (QT_NVFP4, 2, false) => "qmatvec_nvfp4_mmvq_mr2",
13892            (QT_NVFP4, 2, true) => "qmatvec_nvfp4_mmvq_mr2_rp",
13893            (QT_NVFP4, _, true) => "qmatvec_nvfp4_mmvq_rp",
13894            (QT_Q4_0, 1, true) => "qmatvec_q4_0_mmvq_rp",
13895            (QT_Q4_0, _, true) => "qmatvec_q4_0_mmvq_mr2_rp",
13896            (QT_Q5_K, 2, _) => {
13897                if q5_il {
13898                    "qmatvec_q5_K_mmvq_mr2_il"
13899                } else {
13900                    "qmatvec_q5_K_mmvq_mr2"
13901                }
13902            }
13903            (QT_Q8_0, 2, true) => "qmatvec_q8_0_mmvq_mr2_rp",
13904            // rpca (cp.async-staged weight ring): MEASURED NEGATIVE on H100 for Q8_0
13905            // (2026-07-26 N=3: 181.8 vs plain rp 185.5 — the smem round-trip exceeds the
13906            // latency it hides for 8-bit direct-dp4a; the NVFP4 win case overlaps table
13907            // decode with half the bytes). OPT-IN via MEMRA_Q80_CA=1 for the corpus.
13908            (QT_Q8_0, _, true)
13909                if in_f % 1024 == 0 && {
13910                    static CA: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13911                    *CA.get_or_init(|| std::env::var("MEMRA_Q80_CA").as_deref() == Ok("1"))
13912                } =>
13913            {
13914                "qmatvec_q8_0_mmvq_rpca"
13915            }
13916            (QT_Q8_0, _, true) => "qmatvec_q8_0_mmvq_rp",
13917            (QT_Q8_0, _, _) => "qmatvec_q8_0_mmvq",
13918            // K-quant split-plane twins (H100 K-quant coalescing fix, 2026-08-01): the rp4
13919            // mirror routes here; GGUF layout keeps the plain kernels. rp bytes MUST never
13920            // reach a GGUF-layout kernel or vice versa.
13921            (QT_Q4_K, _, true) => "qmatvec_q4_K_mmvq_rp",
13922            (QT_Q6_K, _, true) => "qmatvec_q6_K_mmvq_rp",
13923            (QT_Q4_K, _, _) => "qmatvec_q4_K_mmvq",
13924            (QT_Q4_0, 2, false) => "qmatvec_q4_0_mmvq_mr2",
13925            (QT_Q4_0, _, false) => "qmatvec_q4_0_mmvq",
13926            (QT_Q5_K, _, _) => {
13927                if q5_il {
13928                    "qmatvec_q5_K_mmvq_il"
13929                } else {
13930                    "qmatvec_q5_K_mmvq"
13931                }
13932            }
13933            (QT_Q6_K, _, _) => "qmatvec_q6_K_mmvq",
13934            (QT_NVFP4, _, false) => "qmatvec_nvfp4_mmvq",
13935            (QT_F8_E4M3, _, _) => "qmatvec_e4m3_mmvq",
13936            _ => panic!("qmatvec_mmvq: qtype {qtype} has no MMVQ kernel"),
13937        };
13938        let f = self.func(name);
13939        // each block still has ROWS_PER_BLOCK warps; with mr rows/warp it covers ROWS_PER_BLOCK*mr rows.
13940        let rows_per_block = ROWS_PER_BLOCK * mr;
13941        let cfg = LaunchConfig {
13942            grid_dim: (
13943                (out_f as u32 + rows_per_block - 1) / rows_per_block,
13944                m as u32,
13945                1,
13946            ),
13947            block_dim: (32, ROWS_PER_BLOCK, 1), // warp-per-row (x mr rows each)
13948            shared_mem_bytes: 0,                // warp-only reduce at m=1
13949        };
13950        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
13951        let __s_b = self.gpu.stream();
13952        let mut b = __s_b.launch_builder(&f);
13953        // NVFP4 + e4m3 mmvq kernels take the macro-scale as a fused epilogue arg (applied at the
13954        // write — bit-identical to the old separate scale_inplace pass, minus one launch per matvec:
13955        // 53 scale launches/token on the 9B; for e4m3 the scale is the checkpoint's per-tensor f32
13956        // weight_scale). Other mmvq kernels keep the 8-arg signature.
13957        if qtype == QT_NVFP4 || qtype == QT_F8_E4M3 {
13958            // PDL wave-B: the nvfp4 mr2_rp single (gemma wo / generic rp singles) joins
13959            // the wave-A launch class — 9-arg flavor (fused macro-scale epilogue).
13960            if Self::pdl_on()
13961                && Self::pdl_mmvq_on()
13962                && Self::pdl_nvfp4q8_on()
13963                && name == "qmatvec_nvfp4_mmvq_mr2_rp"
13964            {
13965                use cudarc::driver::{DevicePtr, DevicePtrMut};
13966                let s = &self.gpu.stream();
13967                let (pw, _g0) = bytes.device_ptr(s);
13968                let (paq, _g1) = aq.device_ptr(s);
13969                let (pad, _g2) = ad.device_ptr(s);
13970                let (py, _g3) = y.device_ptr_mut(s);
13971                let mut ps = [
13972                    &pw as *const _ as *mut std::ffi::c_void,
13973                    &paq as *const _ as *mut _,
13974                    &pad as *const _ as *mut _,
13975                    &py as *const _ as *mut _,
13976                    &inf as *const _ as *mut _,
13977                    &outf as *const _ as *mut _,
13978                    &mi as *const _ as *mut _,
13979                    &rb as *const _ as *mut _,
13980                    &scale as *const _ as *mut _,
13981                ];
13982                unsafe {
13983                    self.launch_pdl(name, cfg.grid_dim, cfg.block_dim, &mut ps)?;
13984                }
13985                return Ok(());
13986            }
13987            b.arg(bytes)
13988                .arg(aq)
13989                .arg(ad)
13990                .arg(&mut *y)
13991                .arg(&inf)
13992                .arg(&outf)
13993                .arg(&mi)
13994                .arg(&rb)
13995                .arg(&scale);
13996            unsafe {
13997                b.launch(cfg)?;
13998            }
13999        } else if Self::pdl_on()
14000            && Self::pdl_mmvq_on()
14001            && (matches!(
14002                name,
14003                "qmatvec_q4_0_mmvq_rp" | "qmatvec_q6_K_mmvq" | "qmatvec_q6_K_mmvq_rp"
14004            ) || (Self::pdl_nvfp4q8_on()
14005                && matches!(name, "qmatvec_q8_0_mmvq_rp" | "qmatvec_q8_0_mmvq_mr2_rp")))
14006        {
14007            // PDL wave-A (2026-07-23): the two decode-hot single-matvec kernels carry
14008            // MEMRA_PDL_ENTRY — grid launches while the producer drains. ONLY the marked
14009            // names may take this launch (unmarked kernels would read unordered).
14010            {
14011                use cudarc::driver::{DevicePtr, DevicePtrMut};
14012                let s = &self.gpu.stream();
14013                let (pw, _g0) = bytes.device_ptr(s);
14014                let (paq, _g1) = aq.device_ptr(s);
14015                let (pad, _g2) = ad.device_ptr(s);
14016                let (py, _g3) = y.device_ptr_mut(s);
14017                let mut ps = [
14018                    &pw as *const _ as *mut std::ffi::c_void,
14019                    &paq as *const _ as *mut _,
14020                    &pad as *const _ as *mut _,
14021                    &py as *const _ as *mut _,
14022                    &inf as *const _ as *mut _,
14023                    &outf as *const _ as *mut _,
14024                    &mi as *const _ as *mut _,
14025                    &rb as *const _ as *mut _,
14026                ];
14027                unsafe {
14028                    self.launch_pdl(name, cfg.grid_dim, cfg.block_dim, &mut ps)?;
14029                }
14030            }
14031            if scale != 1.0 {
14032                self.scale_inplace(y, scale, m * out_f)?;
14033            }
14034        } else {
14035            b.arg(bytes)
14036                .arg(aq)
14037                .arg(ad)
14038                .arg(&mut *y)
14039                .arg(&inf)
14040                .arg(&outf)
14041                .arg(&mi)
14042                .arg(&rb);
14043            unsafe {
14044                b.launch(cfg)?;
14045            }
14046            if scale != 1.0 {
14047                self.scale_inplace(y, scale, m * out_f)?;
14048            }
14049        }
14050        Ok(())
14051    }
14052
14053    /// Test entry for the kernel_check bit-equivalence gate: run the warp-per-row MMVQ directly
14054    /// from raw weight bytes (quantize the f32 activation `x` to q8_1 internally). NVFP4 per-tensor
14055    /// macro-scale is NOT applied (caller compares bare, like qmatvec_*_fast). Mirrors qmatvec_gemm_raw.
14056    pub fn qmatvec_mmvq_raw(
14057        &self,
14058        bytes: &CudaSlice<u8>,
14059        x: &CudaSlice<f32>,
14060        m: usize,
14061        in_f: usize,
14062        out_f: usize,
14063        qtype: i32,
14064        row_bytes: usize,
14065        rp: bool,
14066    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14067        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
14068        self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, 1.0, rp)
14069    }
14070
14071    /// True if `qtype` has a batched weight-resident (`_b2`/`_b4`) matvec kernel. These mirror the
14072    /// `_mmvq` kernels but iterate the m token columns INSIDE one warp/row, so the weight bytes leave
14073    /// HBM/L2 once for m tokens (vs grid.y=m re-reading m times). The 5 daily-hot dtypes have them.
14074    pub fn batched_supports(&self, qtype: i32) -> bool {
14075        matches!(
14076            qtype,
14077            QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q6_K | QT_NVFP4 | QT_F8_E4M3 | QT_Q4_0
14078        )
14079    }
14080
14081    /// IQ4_XS trunk fast seam: MEMRA_IQ_FAST=0 reverts non-expert IQ4_XS matmuls to the Stage-A
14082    /// f32 oracle path. Default ON since 2026-08-02 (research/kat-anomaly-20260802/): the old
14083    /// opt-in default left every IQ4_XS-trunk artifact (KAT-Coder IQ4_XS: attn_qkv/attn_gate/
14084    /// ssm_out/shexp, ~0.52GB re-read per decode tick) on the oracle kernel — decode 106.7 ->
14085    /// 193.4 tok/s (x5 interleaved), pp512 228 -> 697, same bytes, via qmatvec_iq4_XS_dp4a. The
14086    /// supported artifacts carry IQ4_XS only in EXPERT banks (their own dispatch, not this seam),
14087    /// so this admission is dispatch-unchanged for every non-IQ4_XS-trunk model.
14088    pub fn iq_fast_enabled() -> bool {
14089        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
14090        *ON.get_or_init(|| {
14091            std::env::var("MEMRA_IQ_FAST")
14092                .map(|v| v != "0")
14093                .unwrap_or(true)
14094        })
14095    }
14096
14097    /// b8 tier seam: MEMRA_B8=0 keeps m=5..8 on the per-m grid.y=m path (m=2..4 batched dispatch
14098    /// unaffected). Default ON — the K=4..7 spec-verify weight-read-once fix.
14099    pub fn b8_enabled() -> bool {
14100        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
14101        *ON.get_or_init(|| std::env::var("MEMRA_B8").map(|v| v != "0").unwrap_or(true))
14102    }
14103
14104    /// Compile-time column batch for a runtime m: 2 -> b2, 3..4 -> b4, 5..8 -> b8.
14105    pub fn batched_mcols(m: usize) -> usize {
14106        if m == 2 {
14107            2
14108        } else if m <= 4 {
14109            4
14110        } else if m <= 8 {
14111            8
14112        } else {
14113            16
14114        }
14115    }
14116
14117    /// Kernel name for the batched matvec of `(qtype, mcols)`. mcols ∈ {2,4,8}. The b8 tier is the
14118    /// K=4..7 spec-verify fix (T=5..8): pre-b8 those T fell to grid.y=m per-row MMVQ = m full
14119    /// weight reads/launch — the measured 27B K=4 cliff (101 -> 73 tok/s at p3 despite acceptance
14120    /// holding 54%). One b8 launch reads the weight ONCE for up to 8 columns (c >= m masked).
14121    fn batched_kernel_name(qtype: i32, mcols: usize) -> Option<&'static str> {
14122        Some(match (qtype, mcols) {
14123            (QT_Q8_0, 2) => "qmatvec_q8_0_mmvq_b2",
14124            (QT_Q8_0, 4) => "qmatvec_q8_0_mmvq_b4",
14125            (QT_Q8_0, 8) => "qmatvec_q8_0_mmvq_b8",
14126            // b16 now has BOTH forms (lane/rp-on-st, 2026-08-06). It used to be rp-ONLY, which
14127            // made the q8rp mirror the exact-16 tier's admission ticket for any model carrying a
14128            // single Q8_0 matmul — measured as the FP8-ST refusal (`L0.ssm_beta qtype=0
14129            // rp4=false`, 96 t / 23.9 MiB = 0.143% of resident weight). The mirror stays a
14130            // BANDWIDTH lever on Q8_0-dominant GGUFs; it is no longer a correctness prerequisite.
14131            (QT_Q8_0, 16) => "qmatvec_q8_0_mmvq_b16",
14132            (QT_Q4_K, 2) => "qmatvec_q4_K_mmvq_b2",
14133            (QT_Q4_K, 4) => "qmatvec_q4_K_mmvq_b4",
14134            (QT_Q4_K, 8) => "qmatvec_q4_K_mmvq_b8",
14135            // b16 base + _rp (lane/rp-on-st): the 9B NVFP4 GGUF's blocker — real NVFP4 GGUFs keep
14136            // Q4_K attention next to NVFP4 MLP, and the tier's predicate is an ALL.
14137            (QT_Q4_K, 16) => "qmatvec_q4_K_mmvq_b16",
14138            (QT_Q5_K, 2) => "qmatvec_q5_K_mmvq_b2",
14139            (QT_Q5_K, 4) => "qmatvec_q5_K_mmvq_b4",
14140            (QT_Q5_K, 8) => "qmatvec_q5_K_mmvq_b8",
14141            // b16 base only (lane/rp-on-st): Q5_K has no rp twins at any width, so there is
14142            // nothing to mirror. Named by the diagnostic as `L0.wqkv_gate qtype=3` on the 9B.
14143            (QT_Q5_K, 16) => "qmatvec_q5_K_mmvq_b16",
14144            (QT_Q6_K, 2) => "qmatvec_q6_K_mmvq_b2",
14145            (QT_Q6_K, 4) => "qmatvec_q6_K_mmvq_b4",
14146            (QT_Q6_K, 8) => "qmatvec_q6_K_mmvq_b8",
14147            (QT_Q6_K, 16) => "qmatvec_q6_K_mmvq_b16",
14148            (QT_NVFP4, 2) => "qmatvec_nvfp4_mmvq_b2",
14149            (QT_NVFP4, 4) => "qmatvec_nvfp4_mmvq_b4",
14150            (QT_NVFP4, 8) => "qmatvec_nvfp4_mmvq_b8",
14151            // b16 (lane/rp-on-st): no mirror needed — NVFP4's 36 B/k32 block is already the
14152            // aligned form its own kernel walks. Unlocks the exact-16 tier for every NVFP4 model
14153            // AND for the mixed FP8-ST artifact, whose 193 NVFP4 tensors were refusing it.
14154            (QT_NVFP4, 16) => "qmatvec_nvfp4_mmvq_b16",
14155            (QT_F8_E4M3, 2) => "qmatvec_e4m3_mmvq_b2",
14156            (QT_F8_E4M3, 4) => "qmatvec_e4m3_mmvq_b4",
14157            (QT_F8_E4M3, 8) => "qmatvec_e4m3_mmvq_b8",
14158            // b16 tier (lane/rp-on-st): e4m3 needs NO split-plane mirror to reach it — its native
14159            // row-major layout is already 32B-aligned per k32 block, so the base kernel IS the
14160            // aligned form. Contrast Q8_0, whose b16 exists only as the `_rp` twin (hence q8rp).
14161            (QT_F8_E4M3, 16) => "qmatvec_e4m3_mmvq_b16",
14162            (QT_Q4_0, 2) => "qmatvec_q4_0_mmvq_b2",
14163            (QT_Q4_0, 4) => "qmatvec_q4_0_mmvq_b4",
14164            (QT_Q4_0, 8) => "qmatvec_q4_0_mmvq_b8",
14165            (QT_Q4_0, 16) => "qmatvec_q4_0_mmvq_b16",
14166            _ => return None,
14167        })
14168    }
14169
14170    /// BATCHED weight-tile-resident matvec from a PRE-QUANTIZED q8_1 activation (the m=2-8 verify/MTP
14171    /// win). One warp walks the weight row ONCE, dp4a vs all m activation columns -> weight HBM/L2
14172    /// traffic 1x for m tokens (vs grid.y=m re-reading it m times). `mcols` ∈ {2,4,8} is the
14173    /// compile-time batch; m must be <= mcols (the c >= m columns are masked in-kernel). y is
14174    /// [m, out_f] token-major. NVFP4 per-tensor macro-scale applied post
14175    /// (scale==1.0 for other dtypes -> no-op). BIT-IDENTICAL per (token,row) to qmatvec_*_mmvq.
14176    ///
14177    /// NVFP4 VARIANT DISPATCH: the batched NVFP4 kernel measured memory-LATENCY bound on the real
14178    /// 27B verify (ncu --set full, 12 steady launches: long_scoreboard 18-30 stalls/issue vs <=1.7
14179    /// for every other reason, DRAM only 41-51% active, lg_throttle 0.7, L1 hit 94% — ONE 6-LDG
14180    /// weight wavefront in flight per warp is the binding constraint, NOT bandwidth and NOT the
14181    /// column-unroll break). Two exactness-free fixes, chosen PER SHAPE from the DRAM-cold 8-copy
14182    /// msweep on all six 27B shapes (2026-07-03):
14183    ///   `pf` = next-g weight-prefetch double-buffer (48 regs, occupancy intact) — wins everywhere
14184    ///          it applies for b4 (-3..-14%), never loses;
14185    ///   `r2` = two rows/warp (67 regs -> 7 resident blocks/SM) — the bigger win (-8.5..-30%) but
14186    ///          wave-quantization-sensitive: with the grid halved to ceil(out_f/8) blocks, a
14187    ///          fractional straggler wave (waves in ~1.05-1.5) costs a full extra latency round on
14188    ///          a latency-bound kernel (27B ffn_down 640 blocks / 574 resident = 1.11 waves: +17%),
14189    ///          while <=1 wave (9B ffn_down 0.89: -30%) or >=2 waves (tail amortized; qkv 2.2:
14190    ///          -8.5%, ffn_gate 3.8: -12.5%) win. For b2, r2 wins on DEEP k-loops (in_f>=6144:
14191    ///          -8..-19%) where the 2-col body starves weight MLP hardest; pf measured negative.
14192    /// b4: r2 when waves(out_f) <= 1 (and grid fills >=half the SMs) or >= 2, else pf.
14193    /// b2: in_f>=6144 -> r2, else base.
14194    /// MEMRA_MMVQ_BV=base|pf|r2|pfr2 forces one variant everywhere (A/B + rollback seam).
14195    /// All variants BIT-IDENTICAL per (token,row): same dp4a order, scales, adg factor, reduce —
14196    /// only load issue time and the row->warp mapping change (kernel-check gates all of them).
14197    /// `rp` = the weight buffer is the A6 SPLIT-PLANE repacked layout (NVFP4 only): the same
14198    /// wave-aware auto rule applies, mapped onto the `_rp` twins (rp/rpr2/rpr2w8 mirror
14199    /// pf/r2/r2w8 — regs 44/67/64 land in the same residency classes).
14200    /// The variant the batched dispatch will pick for this (shape, m, mcols, layout) — exposed so
14201    /// gates can distinguish bit-identical variants (bit-bad==0 required) from the k-split family
14202    /// (deterministic but k-reduce-order-shifted: rel<1e-3 + run-to-run bit-identity required).
14203    /// Device SM count (cached) — grid-fill policy input.
14204    pub fn sm_count(&self) -> i32 {
14205        static SMS: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
14206        *SMS.get_or_init(|| {
14207            use cudarc::driver::sys::CUdevice_attribute_enum as A;
14208            self.gpu
14209                .ctx
14210                .attribute(A::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)
14211                .unwrap_or(82)
14212        })
14213    }
14214
14215    pub fn batched_variant(
14216        &self,
14217        _m: usize,
14218        in_f: usize,
14219        out_f: usize,
14220        qtype: i32,
14221        row_bytes: usize,
14222        mcols: usize,
14223        rp: bool,
14224    ) -> &'static str {
14225        // Q8_0 never joined the auto variant machinery (on sm_120 its only batched shapes
14226        // were tiny aux tensors). On Q8_0-trunk models the layout is the whole game: the
14227        // split-plane mirror (rp) routes to the _rp twins (H100 coalescing fix, 2026-07-26);
14228        // GGUF layout stays "base". rp bytes MUST never reach the base kernel or vice versa.
14229        if qtype == QT_Q8_0 {
14230            return if rp { "rp" } else { "base" };
14231        }
14232        static BV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
14233        let bv = *BV.get_or_init(|| match std::env::var("MEMRA_MMVQ_BV").as_deref() {
14234            Ok("base") => "base",
14235            Ok("pf") => "pf",
14236            Ok("r2") => "r2",
14237            Ok("r2w8") => "r2w8",
14238            Ok("pfr2") => "pfr2",
14239            Ok("ca") => "ca",
14240            Ok("car2") => "car2",
14241            // rp* = SPLIT-PLANE REPACKED layout kernels (A6 prototype): W must already be the
14242            // repacked buffer (msweep MSWEEP_RP harness) — never valid on GGUF-layout weights.
14243            Ok("rp") => "rp",
14244            Ok("rpr2") => "rpr2",
14245            Ok("rpr2w8") => "rpr2w8",
14246            // rpca* = cp.async software-pipelined split-plane (2026-07-05): hides the _rp
14247            // long_scoreboard load stall. rp-layout only; b4/b2 (no b8 twin).
14248            Ok("rpca") => "rpca",
14249            Ok("rpcar2") => "rpcar2",
14250            // 2026-07-06 m-small latency arc: rpsc = rpr2 + per-warp smem scale prestage (kills
14251            // the scale-plane global dependency, zero reg growth); rpms/rpmsc = m-split x2
14252            // across warp pairs (2x blocks of rpr2, column halves per warp, BIT-identical to
14253            // _rp); rpks/rpksc = k-split x2 (fastest microbench cells but k-reduce-order-shifted:
14254            // run-spec self-consistency FAILED on the 27B daily driver — verify logits must be
14255            // bit-identical to the decode path — measurement corpus ONLY, never auto).
14256            Ok("rpsc") => "rpsc",
14257            Ok("rpms") => "rpms",
14258            Ok("rpmsc") => "rpmsc",
14259            Ok("rpks") => "rpks",
14260            Ok("rpksc") => "rpksc",
14261            _ => "auto",
14262        });
14263        // cp.async ring variants need 16B-aligned rows (in_f%256==0 -> (in_f/64)*36 % 16 == 0)
14264        // and whole 32-group warp iterations (nsb%32==0 <=> in_f%1024==0). All 27B/9B trunk
14265        // shapes qualify; anything else falls back to the register variants.
14266        let ca_ok = qtype == QT_NVFP4 && (row_bytes % 16 == 0) && (in_f % 1024 == 0);
14267        // rpsc: smem scale plane fits (nsb64 <= 272) + int4-aligned staging (nsb64 % 4 == 0).
14268        // rpks/rpksc: half-plane staging alignment needs nsb64 % 8 == 0 (in_f % 512 == 0).
14269        // MEMRA_KS=0 removes the 2026-07-06 rpsc/rpks/rpksc entries from AUTO (rollback seam;
14270        // forced MEMRA_MMVQ_BV values still work).
14271        static KS_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
14272        let ks_on = *KS_ON.get_or_init(|| std::env::var("MEMRA_KS").as_deref() != Ok("0"));
14273        let sc_ok = ks_on && qtype == QT_NVFP4 && (in_f % 256 == 0) && (in_f / 64 <= 272);
14274        let ks_ok = ks_on && qtype == QT_NVFP4 && (in_f % 512 == 0) && (in_f / 64 <= 272);
14275        static SMS: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
14276        let sms = *SMS.get_or_init(|| {
14277            use cudarc::driver::sys::CUdevice_attribute_enum as A;
14278            self.gpu
14279                .ctx
14280                .attribute(A::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)
14281                .unwrap_or(82)
14282        });
14283        // k-quant r2 port (2026-07-04): q4_K/q5_K/q6_K have _r2/_r2w8 twins. ncu on the DRAM-cold
14284        // 9B msweep showed q4_K/q5_K b4 memory-latency bound like NVFP4 pre-fix (long_scoreboard
14285        // 19.6/16.4 per issue, DRAM 47.7/38.2%, L2 weight hit ~13%); q6_K lm_head is the exception
14286        // at DRAM 90-91% = wall-bound (yet r2 still wins -8%: deeper MLP raises achieved DRAM).
14287        // No _pf port (a k-quant group stages 10+ words vs NVFP4's 5 — register cost outweighs;
14288        // r2 covers the same MLP) and no rp (GGUF layout only). Q8_0 stays base: its only real
14289        // batched shapes are the tiny out_f=32 ssm_alpha/beta (8-block grids never fill one SM).
14290        // AUTO RULE = the measured winners table (differs from NVFP4's!):
14291        //   r2w8 NEVER in auto — the reg squeeze (72 -> 64 regs = stack spill) loses to unbounded
14292        //     r2 on every measured k-quant cell, incl. the wave-crossing lm_heads (q6_K 1316 vs
14293        //     r2 1258us) — kernels kept behind the force seam for the corpus;
14294        //   q4_K: r2 whenever the halved grid fills the SMs (blocks >= 4*SMs), INCLUDING the
14295        //     1.05-2.0 straggler window where NVFP4's r2 lost (qkv 1.78 waves: r2 -15% here; the
14296        //     k-quant base kernel leaves more latency on the table than a straggler wave costs);
14297        //   q5_K/q6_K: r2 only at waves >= 2 (the 248320-row lm_heads, 48+ waves: q6_K -8%, q5_K
14298        //     -2%); mid shapes measured base-or-flat (q5_K qkv 49.1 base vs 49.7 r2, attn_gate
14299        //     flat, attn_k base) — the 5/6-bit two-stream unpack makes r2's staging pricier.
14300        //   b2 same table with 8-row blocks: q4_K r2 when filled (-3..-22% all measured shapes),
14301        //     q5_K/q6_K r2 at waves >= 2 (27B lm_head -2.9%; 9B q6_K flat, harmless).
14302        let kq_r2 = matches!(qtype, QT_Q4_K | QT_Q5_K | QT_Q6_K);
14303        // MEMRA_KQ_BV=base|r2|r2w8 forces the k-quant variant WITHOUT touching the NVFP4 dispatch
14304        // (MEMRA_MMVQ_BV is global — an interleaved k-quant-only e2e A/B needs this narrower seam).
14305        static KQBV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
14306        let kq_bv = *KQBV.get_or_init(|| match std::env::var("MEMRA_KQ_BV").as_deref() {
14307            Ok("base") => "base",
14308            Ok("r2") => "r2",
14309            Ok("r2w8") => "r2w8",
14310            _ => "auto",
14311        });
14312        let variant: &'static str = if qtype == QT_Q4_0 {
14313            // Q4_0 r2 (gemma verify trunk, 2026-07-10): shared activation loads + the
14314            // row-independent ones-sum computed once per (col,group) for 2 rows. Same
14315            // fill rule as q4_K: r2 when the halved grid still fills the SMs.
14316            static Q40BV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
14317            let q40 = *Q40BV.get_or_init(|| match std::env::var("MEMRA_Q40_BV").as_deref() {
14318                // ms/sm/la = force-only measurement seams (ALL FLAT/NEGATIVE 2026-07-13,
14319                // never auto): m-split flat (nvcc keeps 72 regs); smem-slab −11% (staging
14320                // + syncs cost more than the stalls, bank-pad made no difference);
14321                // register load-ahead flat (nvcc already reorders). The b-tier limiter
14322                // is still unidentified — see the jsonl row.
14323                Ok("base") => "base",
14324                Ok("r2") => "r2",
14325                Ok("ms") => "ms",
14326                Ok("sm") => "sm",
14327                Ok("la") => "la",
14328                _ => "auto",
14329            });
14330            let v = if q40 != "auto" {
14331                q40
14332            } else if (out_f as u32).div_ceil(8) >= 4 * sms as u32 {
14333                "r2"
14334            } else {
14335                "base"
14336            };
14337            // split-plane mirror twins (2026-07-10): same fill rule, _rp names.
14338            // (m-split r2 pair twin PROBED FLAT 2026-07-13 — nvcc kept 72 regs either way
14339            // and the limiter is the per-column activation load chain (long_scoreboard
14340            // 42.5%), not occupancy; arm killed per doctrine, jsonl row is the record.)
14341            if rp {
14342                match v {
14343                    "ms" => "r2ms_rp",
14344                    "sm" => "r2sm_rp",
14345                    "la" => "r2la_rp",
14346                    "r2" => "r2_rp",
14347                    _ => "rp",
14348                }
14349            } else if matches!(v, "ms" | "sm" | "la") {
14350                "r2"
14351            } else {
14352                v
14353            }
14354        } else if qtype != QT_NVFP4 && !kq_r2 {
14355            "base"
14356        } else if kq_r2 && rp {
14357            // K-quant split-plane mirror (2026-08-01): only the plain _rp batched twins are
14358            // compiled for q4_K/q6_K — rp is a LAYOUT, it must survive every heuristic
14359            // (split-plane bytes through a GGUF-layout kernel = NaN). q5_K never mirrors.
14360            "rp"
14361        } else if kq_r2 {
14362            // k-quant r2w8 only exists at b4 (b2_r2 already 8-resident; b8 has no w8 twin) ->
14363            // mcols != 4 forced r2w8 falls to unbounded r2.
14364            if kq_bv != "auto" {
14365                if kq_bv == "r2w8" && mcols != 4 {
14366                    "r2"
14367                } else {
14368                    kq_bv
14369                }
14370            } else if bv != "auto" {
14371                match bv {
14372                    "r2" | "pfr2" | "rpr2" | "car2" => "r2",
14373                    "r2w8" | "rpr2w8" => {
14374                        if mcols != 4 {
14375                            "r2"
14376                        } else {
14377                            "r2w8"
14378                        }
14379                    }
14380                    _ => "base", // base/pf/ca/rp forced -> base (no such k-quant kernels)
14381                }
14382            } else {
14383                let blocks = (out_f + 7) / 8;
14384                let waves = blocks as f64 / (7 * sms as usize) as f64;
14385                let filled = blocks >= 4 * sms as usize;
14386                let use_r2 = if qtype == QT_Q4_K {
14387                    filled
14388                } else {
14389                    waves >= 2.0
14390                };
14391                if use_r2 { "r2" } else { "base" }
14392            }
14393        } else if bv != "auto" {
14394            // r2w8 only exists for b4/b8 (the b2_r2 kernel is already 8-blocks-resident at 60 regs).
14395            // ca/car2 need the alignment gate AND have no b8 twins; pfr2 has no b8 twin either —
14396            // unsupported (shape, mcols) combos fall back to pf/r2.
14397            // On rp buffers, forced legacy names map to their rp twins (layout law).
14398            let v = if bv == "r2w8" && mcols == 2 {
14399                "r2"
14400            } else if bv == "ca" && (!ca_ok || mcols == 8) {
14401                "pf"
14402            } else if bv == "car2" && (!ca_ok || mcols == 8) {
14403                "r2"
14404            } else if bv == "pfr2" && mcols == 8 {
14405                "r2"
14406            } else if (bv == "rpr2w8" || bv == "rpr2") && mcols == 2 {
14407                "rpr2"
14408            }
14409            // rpca* has no b8 twin (falls to rpr2w8/rpr2); needs the ca alignment gate.
14410            else if (bv == "rpca" || bv == "rpcar2") && (!ca_ok || mcols == 8) {
14411                if mcols == 8 { "rpr2w8" } else { "rpr2" }
14412            } else if bv == "rpcar2" && mcols == 2 {
14413                "rpca"
14414            }
14415            // rpsc/rpmsc/rpks* gate on smem-fit + alignment; fall to rpr2 outside it
14416            // (rpms has no smem and no alignment need — always valid on rp buffers).
14417            else if (bv == "rpsc" || bv == "rpmsc") && !sc_ok {
14418                "rpr2"
14419            } else if (bv == "rpks" || bv == "rpksc") && !ks_ok {
14420                "rpr2"
14421            } else {
14422                bv
14423            };
14424            if rp {
14425                match v {
14426                    "base" | "pf" | "ca" | "rp" => "rp",
14427                    "r2" | "pfr2" | "car2" | "rpr2" => "rpr2",
14428                    "r2w8" | "rpr2w8" => {
14429                        if mcols == 2 {
14430                            "rpr2"
14431                        } else {
14432                            "rpr2w8"
14433                        }
14434                    }
14435                    other => other, // rpca/rpcar2/rpsc/rpks/rpksc pass through (already rp-layout)
14436                }
14437            } else {
14438                v
14439            }
14440        } else if mcols == 8 {
14441            // b8 AUTO (2026-07-06 m-small latency arc, g7e DRAM-cold rp msweep m=5/6/8 all five
14442            // 27B shapes): rpsc — the rpr2w8 schedule with the warp's scale rows prestaged to
14443            // smem, leaving ONE global dependency (the quant stream) in the k-loop at zero reg
14444            // growth. BIT-identical to rpr2w8 and wins or ties EVERY b8 cell: ffn_gate m5
14445            // 50.7->46.9 m8 64.1->57.1 (-11%), qkv m8 34.6->33.0, ssm_out m8 29.7->28.8,
14446            // attn_gate m8 26.9->26.1, ffn_down m5 58.2->56.9. The faster split-grid twins are
14447            // OUT: rpksc (k-split, ffn_down m5 -21%) broke run-spec self-consistency (k-reduce
14448            // order shifts verify argmax at tie margins — verify must stay bit-identical to the
14449            // m=1 decode chain); rpmsc (m-split, bit-identical) measured NEGATIVE everywhere
14450            // (twin warp's duplicated weight stream: ffn_down m5 85.7 vs 56.9).
14451            if rp {
14452                if sc_ok { "rpsc" } else { "rpr2w8" }
14453            } else {
14454                "r2w8"
14455            }
14456        } else if mcols >= 4 {
14457            // r2 runs 7 resident blocks/SM (67 regs); its __launch_bounds__(128,8) twin `r2w8`
14458            // (64 regs) runs 8. grid = ceil(out_f/8) for both. rp twins land in the same
14459            // residency classes (rp 44 regs ~ pf-class occupancy, rpr2 67, rpr2w8 64).
14460            let blocks = (out_f + 7) / 8;
14461            let r7 = 7 * sms as usize;
14462            let r8 = 8 * sms as usize;
14463            let waves = blocks as f64 / r7 as f64;
14464            let filled = blocks >= 4 * sms as usize;
14465            // 2026-07-06 m-small latency arc: b4 keeps the wave rule (rpms/rpmsc measured
14466            // flat-to-negative at m=3/4 on every shape — the m-split twin duplicates the weight
14467            // stream; rpsc b4 also negative on r2-class picks, ffn_down m4 51.1 vs 46.5).
14468            if filled && blocks.div_ceil(r8) < blocks.div_ceil(r7) {
14469                // the extra residency drops the INTEGER wave count -> the straggler wave a
14470                // latency-bound kernel pays in full disappears (ffn_down 1.11 -> 0.98 waves:
14471                // 112.5 -> 81.6us, beats pf 90.1; qkv 2.23 -> 1.95: 58.1 -> 51.1).
14472                if rp { "rpr2w8" } else { "r2w8" }
14473            } else if waves >= 2.0 || (waves <= 1.0 && filled) {
14474                // tail amortized (>=2 waves) or single wave: unbounded r2 (no reg-squeeze tax —
14475                // gate/up 81.1 vs 83.9 bounded, attn_q 61.0 vs 63.4).
14476                if rp { "rpr2" } else { "r2" }
14477            } else {
14478                // fractional straggler-wave window with no crossing, or grid too small to fill
14479                // the SMs (tiny out_f<=1024 shapes want max row-parallelism): prefetch variant
14480                // (rp = the r1 split-plane twin — measured the attn_gate winner, 35.4 vs pf 36.4).
14481                if rp { "rp" } else { "pf" }
14482            }
14483        } else if in_f >= 6144 {
14484            // b2 deep-k (2026-07-06): every new twin measured flat-to-negative here (rpms 44.1
14485            // vs rpr2 40.8 ffn_down; rpsc 43.6; the winning rpks is banned on k-order) — rpr2
14486            // stays.
14487            if rp { "rpr2" } else { "r2" }
14488        } else if rp {
14489            // b2 shallow-k: qkv (out_f=10240, 0.97 waves at 7-resident) is the one measured cell
14490            // where the r2-schedule scale-prestage twin beats the r1 rp pick (24.7 vs 28.9us
14491            // -15%); the wider (ffn_gate 1.65 waves) and smaller (attn_gate 0.58) shapes LOSE
14492            // (41.8 vs 38.2 / 16.6 vs 14.6) — gate on the single-wave window.
14493            let waves = ((out_f + 7) / 8) as f64 / (7 * sms as usize) as f64;
14494            if sc_ok && waves >= 0.9 && waves <= 1.1 {
14495                "rpsc"
14496            } else {
14497                "rp"
14498            }
14499        } else {
14500            "base"
14501        };
14502        variant
14503    }
14504
14505    pub fn qmatvec_mmvq_batched(
14506        &self,
14507        bytes: &CudaSlice<u8>,
14508        aq: &CudaSlice<i8>,
14509        ad: &CudaSlice<f32>,
14510        m: usize,
14511        in_f: usize,
14512        out_f: usize,
14513        qtype: i32,
14514        row_bytes: usize,
14515        mcols: usize,
14516        scale: f32,
14517        rp: bool,
14518    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14519        const ROWS_PER_BLOCK: u32 = 4;
14520        // TUNE SEAM (H100 lane): MEMRA_BVAR forces the batched-variant pick for the whole
14521        // process — the auto heuristics were tuned on sm_120 (82 SMs / 858 GB/s) and the
14522        // sm_90a re-tune sweeps this seam empirically. Layout variants stay safe: an rp
14523        // weight keeps its rp-layout kernel family regardless of the override.
14524        let forced: Option<&'static str> = {
14525            static V: std::sync::OnceLock<Option<String>> = std::sync::OnceLock::new();
14526            V.get_or_init(|| std::env::var("MEMRA_BVAR").ok())
14527                .as_deref()
14528                .map(|s| Box::leak(s.to_string().into_boxed_str()) as &'static str)
14529        };
14530        let variant = match forced {
14531            Some(v) if !rp || v.contains("rp") => v,
14532            _ => self.batched_variant(m, in_f, out_f, qtype, row_bytes, mcols, rp),
14533        };
14534        let base_name = Self::batched_kernel_name(qtype, mcols).ok_or_else(|| {
14535            format!("qmatvec_mmvq_batched: no kernel for qtype {qtype} mcols {mcols}")
14536        })?;
14537        // b16 tier (t=9..16 verify): only base/_rp b16 kernels are compiled — the b2..b8
14538        // per-shape perf variants (r2/pf/...) do not apply at this width. rp is a LAYOUT,
14539        // not a perf variant: it must survive (base kernel on split-plane bytes = NaN).
14540        let variant = if mcols == 16 {
14541            if rp { "rp" } else { "base" }
14542        } else {
14543            variant
14544        };
14545        // EXACT-WIDTH b5/b6/b7 twins (lane/vt-fixes fix 1, 2026-08-03): the b8 kernels
14546        // allocate acc[WROWS][8] at ANY m, so T=5..7 verify paid the full 8-wide register
14547        // tax — the measured T=4->5 cliff. The same template at MCOLS=m runs the identical
14548        // per-(token,row) chain (columns c >= m never execute in either form) ->
14549        // BIT-IDENTICAL to the b8 launch. NVFP4 split-plane only (the sm_120 default trunk);
14550        // covers both b8 auto schedules (rpsc, rpr2w8). MEMRA_B567=0 rollback.
14551        static B567: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
14552        let b567 = *B567.get_or_init(|| std::env::var("MEMRA_B567").as_deref() != Ok("0"));
14553        if b567
14554            && qtype == QT_NVFP4
14555            && rp
14556            && mcols == 8
14557            && (5..=7).contains(&m)
14558            && matches!(variant, "rpsc" | "rpr2w8")
14559        {
14560            let f = self.func(&format!("qmatvec_nvfp4_mmvq_b{m}_{variant}"));
14561            let rows_per_block = ROWS_PER_BLOCK * 2; // r2-class schedules: 2 rows/warp
14562            let mut y = self.alloc_uninit::<f32>(m * out_f)?;
14563            let cfg = LaunchConfig {
14564                grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
14565                block_dim: (32, ROWS_PER_BLOCK, 1),
14566                shared_mem_bytes: 0,
14567            };
14568            let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
14569            let __s_b = self.gpu.stream();
14570            let mut b = __s_b.launch_builder(&f);
14571            b.arg(bytes)
14572                .arg(aq)
14573                .arg(ad)
14574                .arg(&mut y)
14575                .arg(&inf)
14576                .arg(&outf)
14577                .arg(&mi)
14578                .arg(&rb);
14579            unsafe {
14580                b.launch(cfg)?;
14581            }
14582            if scale != 1.0 {
14583                self.scale_inplace(&mut y, scale, m * out_f)?;
14584            }
14585            return Ok(y);
14586        }
14587        let (name, rows_per_block): (std::borrow::Cow<'static, str>, u32) = match variant {
14588            "base" => (base_name.into(), ROWS_PER_BLOCK),
14589            "pf" => (format!("{base_name}_pf").into(), ROWS_PER_BLOCK),
14590            "ca" => (format!("{base_name}_ca").into(), ROWS_PER_BLOCK),
14591            "rp" => (format!("{base_name}_rp").into(), ROWS_PER_BLOCK),
14592            "rpca" => (format!("{base_name}_rpca").into(), ROWS_PER_BLOCK), // 1 row/warp cp.async
14593            // split families: 2 warp-pairs x 2 rows = 4 rows/block (the k-range or column set
14594            // splits across the pair's two warps; grid.x doubles vs rpr2 at the same regs).
14595            "rpks" => (format!("{base_name}_rpks").into(), ROWS_PER_BLOCK),
14596            "rpksc" => (format!("{base_name}_rpksc").into(), ROWS_PER_BLOCK),
14597            "rpms" => (format!("{base_name}_rpms").into(), ROWS_PER_BLOCK),
14598            "rpmsc" => (format!("{base_name}_rpmsc").into(), ROWS_PER_BLOCK),
14599            "r2ms_rp" => (format!("{base_name}_r2ms_rp").into(), ROWS_PER_BLOCK),
14600            "r2sm_rp" => (format!("{base_name}_r2sm_rp").into(), ROWS_PER_BLOCK * 2),
14601            "r2la_rp" => (format!("{base_name}_r2la_rp").into(), ROWS_PER_BLOCK * 2),
14602            v => (format!("{base_name}_{v}").into(), ROWS_PER_BLOCK * 2), // r2-class: 2 rows/warp
14603        };
14604        debug_assert!(
14605            !rp || name.contains("_rp"),
14606            "rp weight dispatched to a GGUF-layout kernel"
14607        );
14608        let f = self.func(&name);
14609        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
14610        // r2sm_rp: [MCOLS][32 blk][8 int] activation slab + [MCOLS][32] f32 scales.
14611        let smem = if name.contains("_r2sm_rp") {
14612            (mcols * 32 * 9 * 4 + mcols * 32 * 4) as u32
14613        } else {
14614            0
14615        };
14616        let cfg = LaunchConfig {
14617            grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
14618            block_dim: (32, ROWS_PER_BLOCK, 1),
14619            shared_mem_bytes: smem,
14620        };
14621        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
14622        let __s_b = self.gpu.stream();
14623        let mut b = __s_b.launch_builder(&f);
14624        b.arg(bytes)
14625            .arg(aq)
14626            .arg(ad)
14627            .arg(&mut y)
14628            .arg(&inf)
14629            .arg(&outf)
14630            .arg(&mi)
14631            .arg(&rb);
14632        unsafe {
14633            b.launch(cfg)?;
14634        }
14635        if scale != 1.0 {
14636            self.scale_inplace(&mut y, scale, m * out_f)?;
14637        }
14638        Ok(y)
14639    }
14640
14641    /// BATCHED weight-tile-resident matvec from raw weight bytes (quantizes the f32 activation `x` to
14642    /// q8_1 internally; macro-scale NOT applied — caller compares bare, like qmatvec_*_fast). For the
14643    /// kernel_check bit-equivalence gate. `mcols` ∈ {2,4,8}. Works for Q8_0/Q4_K/Q5_K/Q6_K/NVFP4.
14644    pub fn qmatvec_batched_raw(
14645        &self,
14646        bytes: &CudaSlice<u8>,
14647        x: &CudaSlice<f32>,
14648        m: usize,
14649        in_f: usize,
14650        out_f: usize,
14651        qtype: i32,
14652        row_bytes: usize,
14653        mcols: usize,
14654        rp: bool,
14655    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14656        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
14657        self.qmatvec_mmvq_batched(
14658            bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, mcols, 1.0, rp,
14659        )
14660    }
14661
14662    /// Back-compat NVFP4-only batched raw launcher (used by older gates). Delegates to the generic one.
14663    pub fn qmatvec_nvfp4_batched_raw(
14664        &self,
14665        bytes: &CudaSlice<u8>,
14666        x: &CudaSlice<f32>,
14667        m: usize,
14668        in_f: usize,
14669        out_f: usize,
14670        row_bytes: usize,
14671        mcols: usize,
14672        rp: bool,
14673    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14674        self.qmatvec_batched_raw(bytes, x, m, in_f, out_f, QT_NVFP4, row_bytes, mcols, rp)
14675    }
14676
14677    /// Stage-C FP4 gate (MEMRA_FP4): if `w` is an NVFP4 weight with in_f%64==0, run the native mxf4
14678    /// block-scale GEMM and apply the per-tensor macro-scale, returning Some(y). Else None (caller
14679    /// falls through to the int8 GEMM / dp4a). Strict opt-in over the proven int8 path; m>=16 only.
14680    fn try_fp4_gemm(
14681        &self,
14682        w: &crate::model::GpuTensor,
14683        x: &CudaSlice<f32>,
14684        m: usize,
14685        in_f: usize,
14686        out_f: usize,
14687    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
14688        use crate::model::GpuTensor;
14689        if cfg!(memra_portable_cuda) {
14690            return Ok(None);
14691        }
14692        if std::env::var("MEMRA_FP4").is_err() {
14693            return Ok(None);
14694        }
14695        // CUTLASS prefill branch (m>=128 + MEMRA_FP4_CUTLASS + a repacked CutlassWeight present): route
14696        // to the CUTLASS sm120 NVFP4 GEMM, folding the per-tensor macro-scale into the epilogue alpha
14697        // (1/scale) — no post-matmul scale_inplace. Decode (m<128) and the m∈[16,128) middle band keep
14698        // the hand-roll below: CUTLASS's 128-row M-tile wastes work under 128.
14699        // The hand-roll applies the per-tensor macro-scale as a POST-matmul MULTIPLY (scale_inplace(y,
14700        // scale)); CUTLASS's epilogue does D = alpha * (A@B^T), so alpha == scale reproduces it exactly
14701        // (NOT 1/scale — the plan sketch had this inverted; the kernel_check arm gates it). scale==1.0
14702        // for the common no-macro-scale case.
14703        #[cfg(memra_cutlass)]
14704        if m >= 128 && std::env::var("MEMRA_FP4_CUTLASS").is_ok() {
14705            if let GpuTensor::Quant {
14706                bytes,
14707                qtype,
14708                scale,
14709                row_bytes,
14710                cutlass,
14711                ..
14712            } = w
14713            {
14714                if *qtype == QT_NVFP4 && in_f % 64 == 0 {
14715                    if let Some(cw) = cutlass {
14716                        // Resident fast path: load-time-repacked B + swizzled SFB (no per-call repack).
14717                        let y = self.cutlass_fp4_gemm(
14718                            &cw.b_packed,
14719                            &cw.sfb_swizzled,
14720                            x,
14721                            *scale,
14722                            m,
14723                            out_f,
14724                            in_f,
14725                        )?;
14726                        return Ok(Some(y));
14727                    } else if std::env::var("MEMRA_FP4_CUTLASS_OTF").is_ok() {
14728                        // On-the-fly repack (MEMRA_FP4_CUTLASS_OTF): de-interleave + swizzle the B operand
14729                        // from raw bytes per prefill call. No resident doubling of the NVFP4 weight VRAM
14730                        // (the load-time repack ~doubles it) — needed for models that don't fit the
14731                        // resident path (e.g. the 27B on 24GB). Slower (per-call repack) but argmax-exact.
14732                        let (b_packed, sfb_sw) =
14733                            self.build_cutlass_weight(bytes, out_f, in_f, *row_bytes)?;
14734                        let y =
14735                            self.cutlass_fp4_gemm(&b_packed, &sfb_sw, x, *scale, m, out_f, in_f)?;
14736                        return Ok(Some(y));
14737                    }
14738                }
14739            }
14740        }
14741        if let GpuTensor::Quant {
14742            bytes,
14743            qtype,
14744            row_bytes,
14745            scale,
14746            rp,
14747            ..
14748        } = w
14749        {
14750            // A6: the hand-rolled W4A4 mxf4 GEMM reads 36B GGUF blocks — no rp port (MEMRA_FP4 is
14751            // an opt-in accuracy tradeoff); repacked tensors fall through to the int8 GEMM.
14752            if *qtype == QT_NVFP4 && in_f % 64 == 0 && !*rp {
14753                let y =
14754                    self.qmatvec_gemm_nvfp4_fp4(bytes, x, m, in_f, out_f, *row_bytes, *scale)?;
14755                return Ok(Some(y));
14756            }
14757        }
14758        Ok(None)
14759    }
14760
14761    /// rms_norm + fused fp16 twin (task #14): f32 output verbatim `rms_norm` + the fp16
14762    /// copy the f16-mirror GEMM group would otherwise produce with a standalone convert
14763    /// launch. BIT-IDENTICAL end-to-end (same reduction, same __float2half values).
14764    pub fn rms_norm_f16out(
14765        &self,
14766        x: &CudaSlice<f32>,
14767        w: &CudaSlice<f32>,
14768        dst: &mut CudaSlice<f32>,
14769        dst16: &mut CudaSlice<u8>,
14770        ncols: usize,
14771        nrows: usize,
14772        eps: f32,
14773    ) -> Result<(), Box<dyn std::error::Error>> {
14774        let f = self.func("rms_norm_f16out_f32");
14775        let cfg = LaunchConfig {
14776            grid_dim: (nrows as u32, 1, 1),
14777            block_dim: (rms_block(), 1, 1),
14778            shared_mem_bytes: 0,
14779        };
14780        let (nc, e) = (ncols as i32, eps);
14781        let __s_b = self.gpu.stream();
14782        let mut b = __s_b.launch_builder(&f);
14783        b.arg(x).arg(w).arg(dst).arg(dst16).arg(&nc).arg(&e);
14784        unsafe {
14785            b.launch(cfg)?;
14786        }
14787        Ok(())
14788    }
14789
14790    /// add+norm(+f16out) fusion for the prefill trunk (round 28; add_rms_norm precedent —
14791    /// bit-identical to add_f32 -> rms_norm_f16out). block_dim matches rms_norm_f16out's.
14792    #[allow(clippy::too_many_arguments)]
14793    pub fn add_rms_norm_f16out(
14794        &self,
14795        a: &CudaSlice<f32>,
14796        b: &CudaSlice<f32>,
14797        w: &CudaSlice<f32>,
14798        res: &mut CudaSlice<f32>,
14799        dst: &mut CudaSlice<f32>,
14800        dst16: &mut CudaSlice<u8>,
14801        ncols: usize,
14802        nrows: usize,
14803        eps: f32,
14804    ) -> Result<(), Box<dyn std::error::Error>> {
14805        let f = self.func("add_rms_norm_f16out_f32");
14806        let cfg = LaunchConfig {
14807            grid_dim: (nrows as u32, 1, 1),
14808            block_dim: (rms_block(), 1, 1),
14809            shared_mem_bytes: 0,
14810        };
14811        let (nc, e) = (ncols as i32, eps);
14812        let __s_lb = self.gpu.stream();
14813        let mut lb = __s_lb.launch_builder(&f);
14814        lb.arg(a)
14815            .arg(b)
14816            .arg(w)
14817            .arg(res)
14818            .arg(dst)
14819            .arg(dst16)
14820            .arg(&nc)
14821            .arg(&e);
14822        unsafe {
14823            lb.launch(cfg)?;
14824        }
14825        Ok(())
14826    }
14827
14828    /// matmul_group with a PRE-EMITTED fp16 activation (task #14: the producer norm fused
14829    /// the convert). Mirror-less members fall back to `matmul` on the f32 activation.
14830    pub fn matmul_group_xh(
14831        &self,
14832        ws: &[&crate::model::GpuTensor],
14833        x: &CudaSlice<f32>,
14834        xh: &CudaSlice<u8>,
14835        m: usize,
14836    ) -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
14837        let mut out = Vec::with_capacity(ws.len());
14838        let in_f = ws[0].in_features();
14839        for w in ws {
14840            if w.in_features() == in_f && m >= 16 && !self.verify_exact_on() {
14841                if let Some(y) = self.try_f16_gemm_pre(w, xh, m)? {
14842                    out.push(y);
14843                    continue;
14844                }
14845            }
14846            out.push(self.matmul(w, x, m)?);
14847        }
14848        Ok(out)
14849    }
14850
14851    /// task #14 pad-proofing: zero beta/g_log at rows >= len_d[0] (pads become identity
14852    /// GDN steps). Layouts [T, H].
14853    pub fn gdn_pad_mask(
14854        &self,
14855        beta: &mut CudaSlice<f32>,
14856        g_log: &mut CudaSlice<f32>,
14857        len_d: &CudaSlice<i32>,
14858        h: usize,
14859        t: usize,
14860    ) -> Result<(), Box<dyn std::error::Error>> {
14861        let f = self.func("gdn_pad_mask_f32");
14862        let cfg = LaunchConfig::for_num_elems((t * h) as u32);
14863        let (hi, ti) = (h as i32, t as i32);
14864        let __s_b = self.gpu.stream();
14865        let mut b = __s_b.launch_builder(&f);
14866        b.arg(beta).arg(g_log).arg(len_d).arg(&hi).arg(&ti);
14867        unsafe {
14868            b.launch(cfg)?;
14869        }
14870        Ok(())
14871    }
14872
14873    /// task #14 pad-proofing: dst[ncols] = src row (len_d[0]-1) — device-indexed last-row
14874    /// gather for the padded prime graph's h_seed/hlast.
14875    pub fn row_gather_dev(
14876        &self,
14877        src: &CudaSlice<f32>,
14878        dst: &mut CudaSlice<f32>,
14879        len_d: &CudaSlice<i32>,
14880        ncols: usize,
14881    ) -> Result<(), Box<dyn std::error::Error>> {
14882        let f = self.func("row_gather_dev_f32");
14883        let cfg = LaunchConfig::for_num_elems(ncols as u32);
14884        let nc = ncols as i32;
14885        let __s_b = self.gpu.stream();
14886        let mut b = __s_b.launch_builder(&f);
14887        b.arg(src).arg(dst).arg(len_d).arg(&nc);
14888        unsafe {
14889            b.launch(cfg)?;
14890        }
14891        Ok(())
14892    }
14893
14894    /// Grouped matmul: several weights consuming ONE activation (hybrid layers: the GDN
14895    /// 4-tuple wqkv/gate/beta/alpha, attention q/k/v, ffn gate/up). Semantics identical to
14896    /// calling `matmul` per weight; the f16-mirror arm converts the activation ONCE for the
14897    /// whole group instead of once per GEMM (the standalone converts were ~250 launches/prime
14898    /// of small-kernel gap fuel — nsys 2026-07-26). Any member without a mirror (or with a
14899    /// different in_f) falls back to its own `matmul` — behavior unchanged.
14900    pub fn matmul_group(
14901        &self,
14902        ws: &[&crate::model::GpuTensor],
14903        x: &CudaSlice<f32>,
14904        m: usize,
14905    ) -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
14906        use crate::model::GpuTensor;
14907        let mut out = Vec::with_capacity(ws.len());
14908        let any_mirror = ws
14909            .iter()
14910            .any(|w| matches!(w, GpuTensor::Quant { f16: Some(_), .. }));
14911        if m >= 16 && any_mirror && !self.verify_exact_on() {
14912            let in_f = ws[0].in_features();
14913            let xh = self.f16_act(x, m * in_f, in_f)?;
14914            for w in ws {
14915                if w.in_features() == in_f {
14916                    if let Some(y) = self.try_f16_gemm_pre(w, &xh, m)? {
14917                        out.push(y);
14918                        continue;
14919                    }
14920                }
14921                out.push(self.matmul(w, x, m)?);
14922            }
14923            return Ok(out);
14924        }
14925        for w in ws {
14926            out.push(self.matmul(w, x, m)?);
14927        }
14928        Ok(out)
14929    }
14930
14931    /// Cross-request grouped matmul (task #13): run ONE projection group over the
14932    /// CONCATENATION of several sequences' activations (m = sum of per-seq rows — the
14933    /// GEMM-batch win vLLM gets from continuous batching), then split each output back
14934    /// into per-seq buffers. Zero view plumbing: gather/scatter are stream-ordered D2D
14935    /// copies (~us at prime sizes). NUMERIC CONFIG NOTE: a GEMM at m=sum tiles K
14936    /// differently than per-seq GEMMs — argmax-gated like every prefill GEMM change.
14937    pub fn matmul_group_multi(
14938        &self,
14939        ws: &[&crate::model::GpuTensor],
14940        xs: &[&CudaSlice<f32>],
14941        ms: &[usize],
14942    ) -> Result<Vec<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
14943        assert_eq!(xs.len(), ms.len());
14944        let in_f = ws[0].in_features();
14945        let total: usize = ms.iter().sum();
14946        let mut xcat = self.uninit(total * in_f)?;
14947        let mut off = 0usize;
14948        for (x, &m) in xs.iter().zip(ms) {
14949            self.copy_into(&mut xcat, off * in_f, x, m * in_f)?;
14950            off += m;
14951        }
14952        let ys = self.matmul_group(ws, &xcat, total)?;
14953        let mut out: Vec<Vec<CudaSlice<f32>>> = (0..xs.len()).map(|_| Vec::new()).collect();
14954        for (w, y) in ws.iter().zip(ys) {
14955            let out_f = w.out_features();
14956            let mut off = 0usize;
14957            for (s, &m) in ms.iter().enumerate() {
14958                let mut ys_s = self.uninit(m * out_f)?;
14959                let src = y.slice(off * out_f..(off + m) * out_f);
14960                self.gpu.stream().memcpy_dtod(&src, &mut ys_s)?;
14961                out[s].push(ys_s);
14962                off += m;
14963            }
14964        }
14965        Ok(out)
14966    }
14967
14968    /// True if `w`'s qtype has a batched tensor-core GEMM kernel (the prefill T>1 root fix).
14969    /// Only the 4 daily-hot dtypes: Q8_0, Q4_K, Q6_K, NVFP4. NVFP4 needs in_f % 64 == 0.
14970    /// DEFAULT-ON (2026-06-28): measured pp512 9B-NVFP4 = 1413 tok/s WITH this GEMM vs 298 with the
14971    /// dp4a fallback (4.7x) AND MORE accurate (prefill logit maxdiff 0.159 vs dp4a 0.55, both argmax
14972    /// MATCH). The int8 tensor-core GEMM is unconditional (its historical MEMRA_GEMM opt-in gate
14973    /// shipped with Phase 0 — mma + smem swizzle + cp.async — and was removed). Prefill-only
14974    /// (m>=GEMM_M_THRESHOLD); m=1 decode keeps dp4a/MMVQ (this returns true but matmul only calls it
14975    /// at m>=threshold). Portable CUDA targets always use the correctness fallback; on sm_120a,
14976    /// MEMRA_NO_GEMM forces that same dp4a fallback (the bit-reference).
14977    pub fn gemm_supports(&self, w: &crate::model::GpuTensor) -> bool {
14978        use crate::model::GpuTensor;
14979        if !legacy_quant_gemm_allowed(
14980            cfg!(memra_portable_cuda),
14981            cfg!(memra_hopper_mma),
14982            std::env::var_os("MEMRA_NO_GEMM").is_some(),
14983        ) {
14984            return false;
14985        }
14986        match w {
14987            GpuTensor::Quant { qtype, .. } => {
14988                matches!(*qtype, QT_Q8_0 | QT_Q4_K | QT_Q6_K | QT_Q5_K | QT_Q4_0)
14989                    || (*qtype == QT_NVFP4 && w.in_features() % 64 == 0)
14990            }
14991            GpuTensor::Float { .. } | GpuTensor::FloatBf16 { .. } => false,
14992        }
14993    }
14994
14995    /// Batched tensor-core int8 GEMM with a PRE-QUANTIZED q8_1 activation (aq,ad). The prefill
14996    /// (T>1) root fix: decode each weight 32-block to int8 in shared memory ONCE per (row-tile,
14997    /// K-step) and reuse it across all BN tokens via mma.sync.m16n8k32.s8 — amortizing the weight
14998    /// read/decode N-fold (vs the dp4a matvec's per-token re-read). s32 accumulate is exact vs
14999    /// dp4a; only the final f32 block-scale rounding differs. Caller MUST have checked
15000    /// `gemm_supports(w)`. y[m,out] token-major. NVFP4 per-tensor macro-scale applied post.
15001    pub fn qmatvec_gemm(
15002        &self,
15003        w: &crate::model::GpuTensor,
15004        aq: &CudaSlice<i8>,
15005        ad: &CudaSlice<f32>,
15006        m: usize,
15007    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15008        use crate::model::GpuTensor;
15009        let in_f = w.in_features();
15010        let out_f = w.out_features();
15011        let (bytes, qtype, row_bytes, scale, rp) = match w {
15012            GpuTensor::Quant {
15013                bytes,
15014                qtype,
15015                row_bytes,
15016                scale,
15017                rp,
15018                ..
15019            } => (bytes, *qtype, *row_bytes, *scale, *rp),
15020            _ => unreachable!("gemm_supports guaranteed Quant"),
15021        };
15022        // wgmma arm (sm_90a, task 8): the m64n64k32 warpgroup kernel reads the rp4 split-plane
15023        // mirror AS-IS (qplane rows = its A operand, the half dplane its scales) and the same
15024        // (aq, ad) activation planes. Same numeric class as the mma kernel below (exact s32 per
15025        // 32-block, one f32 scale fold per block, ascending K) — argmax/tolerance gated like
15026        // every prefill GEMM, not bit-gated. MEMRA_WGMMA=0 restores the portable kernel.
15027        if cfg!(memra_hopper_mma) && qtype == QT_Q8_0 && out_f % 64 == 0 && wgmma_gemm_enabled() {
15028            if let GpuTensor::Quant { rp4: Some(m4), .. } = w {
15029                let mut y = self.qmatvec_gemm_q8_0_wgmma_raw(m4, aq, ad, m, in_f, out_f)?;
15030                if scale != 1.0 {
15031                    self.scale_inplace(&mut y, scale, m * out_f)?;
15032                }
15033                return Ok(y);
15034            }
15035        }
15036        let name = match qtype {
15037            QT_Q8_0 => "qmatvec_gemm_q8_0",
15038            QT_Q4_K => "qmatvec_gemm_q4_K",
15039            QT_Q4_0 => {
15040                if rp {
15041                    "qmatvec_gemm_q4_0_rp"
15042                } else {
15043                    "qmatvec_gemm_q4_0"
15044                }
15045            }
15046            QT_Q5_K => "qmatvec_gemm_q5_K",
15047            QT_Q6_K => "qmatvec_gemm_q6_K",
15048            QT_NVFP4 => {
15049                if rp {
15050                    "qmatvec_gemm_nvfp4_rp"
15051                } else {
15052                    "qmatvec_gemm_nvfp4"
15053                }
15054            }
15055            _ => unreachable!(),
15056        };
15057        let f = self.func(name);
15058        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output: skip memset
15059        // CTA tile MUST match the .cu per-kernel tile. MMQ-PORT: kernel1 (Q8_0/Q4_K/Q5_K) runs llama's
15060        // 128x128 SQUARE tile (K1_BM=128 x K1_BN=128, 8 warps); kernel2 (Q6_K/NVFP4) keeps 64x256, 4 warps
15061        // (the macro BM/BN in the .cu). Grid dims are selected by qtype so each launches its own tile.
15062        let is_k1 = matches!(qtype, QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q4_0);
15063        // TUNE SEAM: MEMRA_GEMM_K1_LAUNCH overrides kernel1's launch tile to match a -D-swept fatbin.
15064        let k1_tile = if is_k1 {
15065            k1_launch_override().unwrap_or((128, 128, 8))
15066        } else {
15067            (128, 128, 8)
15068        };
15069        let (bm, bn): (u32, u32) = if is_k1 {
15070            (k1_tile.0, k1_tile.1)
15071        } else {
15072            (64, 256)
15073        };
15074        let warps: u32 = if is_k1 {
15075            k1_tile.2
15076        } else {
15077            match qtype {
15078                QT_NVFP4 => 8,
15079                _ => 4,
15080            }
15081        };
15082        let cfg = LaunchConfig {
15083            grid_dim: ((out_f as u32 + bm - 1) / bm, (m as u32 + bn - 1) / bn, 1),
15084            block_dim: (32, warps, 1),
15085            shared_mem_bytes: 0,
15086        };
15087        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
15088        let __s_b = self.gpu.stream();
15089        let mut b = __s_b.launch_builder(&f);
15090        b.arg(bytes)
15091            .arg(aq)
15092            .arg(ad)
15093            .arg(&mut y)
15094            .arg(&inf)
15095            .arg(&outf)
15096            .arg(&mi)
15097            .arg(&rb);
15098        unsafe {
15099            b.launch(cfg)?;
15100        }
15101        if scale != 1.0 {
15102            self.scale_inplace(&mut y, scale, m * out_f)?;
15103        }
15104        Ok(y)
15105    }
15106
15107    /// Test entry: run the GEMM directly from raw weight bytes + qtype (no GpuTensor). Quantizes
15108    /// the f32 activation `x` to q8_1 internally then launches the tensor-core GEMM. NVFP4 per-tensor
15109    /// macro-scale is NOT applied here (caller passes it separately, like the dp4a path). Used by
15110    /// kernel_check for the bit-equivalence gate vs qmatvec_*_dp4a.
15111    pub fn qmatvec_gemm_raw(
15112        &self,
15113        bytes: &CudaSlice<u8>,
15114        x: &CudaSlice<f32>,
15115        m: usize,
15116        in_f: usize,
15117        out_f: usize,
15118        qtype: i32,
15119        row_bytes: usize,
15120    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15121        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15122        let name = match qtype {
15123            QT_Q8_0 => "qmatvec_gemm_q8_0",
15124            QT_Q4_K => "qmatvec_gemm_q4_K",
15125            QT_Q4_0 => "qmatvec_gemm_q4_0",
15126            QT_Q5_K => "qmatvec_gemm_q5_K",
15127            QT_Q6_K => "qmatvec_gemm_q6_K",
15128            QT_NVFP4 => "qmatvec_gemm_nvfp4",
15129            QT_NVFP4_RP => "qmatvec_gemm_nvfp4_rp",
15130            _ => panic!("qmatvec_gemm_raw: qtype {qtype} has no GEMM kernel"),
15131        };
15132        let f = self.func(name);
15133        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output: skip memset
15134        // MMQ-PORT: kernel1 (Q8_0/Q4_K/Q5_K) = llama 128x128 tile, 8 warps; kernel2 (Q6_K/NVFP4) = 64x256,
15135        // 4/8 warps. Grid tile per qtype (must match the .cu K1_BM/K1_BN vs BM/BN). KEEP IN SYNC w/ qmatvec_gemm.
15136        let is_k1 = matches!(qtype, QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q4_0);
15137        // TUNE SEAM: MEMRA_GEMM_K1_LAUNCH overrides kernel1's launch tile to match a -D-swept fatbin.
15138        let k1_tile = if is_k1 {
15139            k1_launch_override().unwrap_or((128, 128, 8))
15140        } else {
15141            (128, 128, 8)
15142        };
15143        let (bm, bn): (u32, u32) = if is_k1 {
15144            (k1_tile.0, k1_tile.1)
15145        } else {
15146            (64, 256)
15147        };
15148        let warps: u32 = if is_k1 {
15149            k1_tile.2
15150        } else {
15151            match qtype {
15152                QT_NVFP4 | QT_NVFP4_RP => 8,
15153                _ => 4,
15154            }
15155        };
15156        let cfg = LaunchConfig {
15157            grid_dim: ((out_f as u32 + bm - 1) / bm, (m as u32 + bn - 1) / bn, 1),
15158            block_dim: (32, warps, 1),
15159            shared_mem_bytes: 0,
15160        };
15161        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
15162        let __s_b = self.gpu.stream();
15163        let mut b = __s_b.launch_builder(&f);
15164        b.arg(bytes)
15165            .arg(&aq)
15166            .arg(&ad)
15167            .arg(&mut y)
15168            .arg(&inf)
15169            .arg(&outf)
15170            .arg(&mi)
15171            .arg(&rb);
15172        unsafe {
15173            b.launch(cfg)?;
15174        }
15175        Ok(y)
15176    }
15177
15178    /// H100 warpgroup GEMM raw entry (task 8): launch `qmatvec_gemm_q8_0_wgmma` on an rp4
15179    /// split-plane mirror + pre-quantized (aq, ad) activation planes. One warpgroup (128 thr)
15180    /// owns a 64x64 C tile; grid (out_f/64, ceil(m/64)). out_f % 64 == 0 REQUIRED (row loads
15181    /// and dplane scale reads are unguarded); the token edge is guarded in-kernel.
15182    /// Standalone harness verdict (tools/bench_q8_gemm_wgmma.cu, 4096x4096x512): rel 1.6e-05
15183    /// vs CPU ref, 179us vs the portable mma kernel's 688us (3.84x, unpipelined).
15184    pub fn qmatvec_gemm_q8_0_wgmma_raw(
15185        &self,
15186        rp4: &CudaSlice<u8>,
15187        aq: &CudaSlice<i8>,
15188        ad: &CudaSlice<f32>,
15189        m: usize,
15190        in_f: usize,
15191        out_f: usize,
15192    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15193        assert!(
15194            out_f % 64 == 0 && in_f % 32 == 0,
15195            "wgmma GEMM needs out_f%64==0, in_f%32==0"
15196        );
15197        let f = self.func("qmatvec_gemm_q8_0_wgmma");
15198        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output
15199        let cfg = LaunchConfig {
15200            grid_dim: ((out_f / 64) as u32, (m as u32).div_ceil(64), 1),
15201            block_dim: (128, 1, 1),
15202            shared_mem_bytes: 0,
15203        };
15204        let (inf, outf, mi) = (in_f as i32, out_f as i32, m as i32);
15205        let __s_b = self.gpu.stream();
15206        let mut b = __s_b.launch_builder(&f);
15207        b.arg(rp4)
15208            .arg(aq)
15209            .arg(ad)
15210            .arg(&mut y)
15211            .arg(&inf)
15212            .arg(&outf)
15213            .arg(&mi);
15214        unsafe {
15215            b.launch(cfg)?;
15216        }
15217        Ok(y)
15218    }
15219
15220    /// y[i] *= s. NVFP4 per-tensor macro-scale broadcast over the whole output.
15221    pub fn scale_inplace(
15222        &self,
15223        y: &mut CudaSlice<f32>,
15224        s: f32,
15225        n: usize,
15226    ) -> Result<(), Box<dyn std::error::Error>> {
15227        let f = self.func("scale_f32");
15228        let cfg = LaunchConfig::for_num_elems(n as u32);
15229        let (sf, ni) = (s, n as i32);
15230        let __s_b = self.gpu.stream();
15231        let mut b = __s_b.launch_builder(&f);
15232        b.arg(y).arg(&sf).arg(&ni);
15233        unsafe {
15234            b.launch(cfg)?;
15235        }
15236        Ok(())
15237    }
15238
15239    /// MEMRA_FULL_PREC dequant-on-use: expand a bf16-resident weight (`GpuTensor::FloatBf16`, raw
15240    /// bf16 bytes) to a transient f32 scratch of `n` elements, which then feeds the existing f32
15241    /// cuBLASLt GEMV. The scratch is freed when the caller drops it, so peak VRAM = resident bf16
15242    /// weights + ONE (largest) weight's f32 expansion + activations. SLOW IS FINE (research mode).
15243    pub fn bf16_to_f32(
15244        &self,
15245        data: &cudarc::driver::CudaView<'_, u8>,
15246        n: usize,
15247    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15248        let mut out = self.alloc_uninit::<f32>(n)?;
15249        let f = self.func("bf16_to_f32");
15250        let cfg = LaunchConfig::for_num_elems(n as u32);
15251        let ni = n as i32;
15252        let __s_b = self.gpu.stream();
15253        let mut b = __s_b.launch_builder(&f);
15254        b.arg(data).arg(&mut out).arg(&ni);
15255        unsafe {
15256            b.launch(cfg)?;
15257        }
15258        Ok(out)
15259    }
15260
15261    /// Chunked bf16 linear (MEMRA_FULL_PREC): y[m,out] = x @ W_bf16^T with the f32 dequant scratch
15262    /// bounded to CHUNK_ROWS rows (256MB at in_f=4096) instead of the whole weight — the 4GB
15263    /// lm_head expansion OOM'd the 24GB budget. Row-chunking partitions OUTPUT rows; each row's
15264    /// dot is computed by the identical kernel on identical bytes, so per-(token,row) results are
15265    /// bit-identical to the unchunked form. `exact` selects linear_decode_exact (per-column m=1
15266    /// calls, the spec-verify contract) vs plain linear.
15267    fn linear_bf16_chunked(
15268        &self,
15269        x: &CudaSlice<f32>,
15270        data: &CudaSlice<u8>,
15271        m: usize,
15272        in_f: usize,
15273        out_f: usize,
15274        exact: bool,
15275    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15276        const CHUNK_BYTES: usize = 256 << 20;
15277        let chunk_rows = (CHUNK_BYTES / (in_f * 4)).max(1).min(out_f);
15278        if chunk_rows >= out_f {
15279            let wf32 = self.bf16_to_f32(&data.slice(0..in_f * out_f * 2), in_f * out_f)?;
15280            return if exact {
15281                self.linear_decode_exact(x, &wf32, m, in_f, out_f)
15282            } else {
15283                self.linear(x, &wf32, m, in_f, out_f)
15284            };
15285        }
15286        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
15287        let mut r0 = 0usize;
15288        while r0 < out_f {
15289            let rows = chunk_rows.min(out_f - r0);
15290            let wslice = data.slice(r0 * in_f * 2..(r0 + rows) * in_f * 2);
15291            let wf32 = self.bf16_to_f32(&wslice, in_f * rows)?;
15292            let yc = if exact {
15293                self.linear_decode_exact(x, &wf32, m, in_f, rows)?
15294            } else {
15295                self.linear(x, &wf32, m, in_f, rows)?
15296            };
15297            // scatter [m, rows] into y[m, out_f] at column offset r0 (m is tiny in decode/verify)
15298            for mi in 0..m {
15299                let src = yc.slice(mi * rows..(mi + 1) * rows);
15300                let mut dst = y.slice_mut(mi * out_f + r0..mi * out_f + r0 + rows);
15301                self.gpu.stream().memcpy_dtod(&src, &mut dst)?;
15302            }
15303            r0 += rows;
15304        }
15305        Ok(y)
15306    }
15307
15308    /// On-device linear: y[m,out] = x[m,in] @ W[out,in]^T, weights row-major [out,in] (ggml).
15309    /// cuBLASLt col-major mapping (see memra_runtime::Gpu::linear_f32 for the derivation).
15310    /// DECODE-EXACT float linear: per-column m=1 cuBLASLt calls. cuBLASLt's reduction split is
15311    /// n-dependent (lt_ndep probe: m=1 vs m=2 col0 differs every bit), so spec-verify batches
15312    /// must not batch float matmuls the T=1 decode chain runs at m=1. Used by the small-t MoE
15313    /// router/shexp sites and matmul_decode_exact's Float arm.
15314    pub fn linear_decode_exact(
15315        &self,
15316        x: &CudaSlice<f32>,
15317        w: &CudaSlice<f32>,
15318        m_tokens: usize,
15319        in_f: usize,
15320        out_f: usize,
15321    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15322        if m_tokens == 1 {
15323            return self.linear(x, w, 1, in_f, out_f);
15324        }
15325        let xv = self.view(x, m_tokens * in_f);
15326        let mut y = self.alloc_uninit::<f32>(m_tokens * out_f)?;
15327        for t in 0..m_tokens {
15328            let row = xv.slice(t * in_f..(t + 1) * in_f);
15329            let mut xr = self.alloc_uninit::<f32>(in_f)?;
15330            self.copy_view_into(&mut xr, 0, &row, in_f)?;
15331            let yr = self.linear(&xr, w, 1, in_f, out_f)?;
15332            self.copy_into(&mut y, t * out_f, &yr, out_f)?;
15333        }
15334        Ok(y)
15335    }
15336
15337    pub fn linear(
15338        &self,
15339        x: &CudaSlice<f32>,
15340        w: &CudaSlice<f32>,
15341        m_tokens: usize,
15342        in_f: usize,
15343        out_f: usize,
15344    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15345        use cudarc::cublaslt::{Matmul, MatmulConfig};
15346        let mut c = self.alloc_uninit::<f32>(m_tokens * out_f)?; // cuBLASLt beta=0: C fully written
15347        let cfg = MatmulConfig {
15348            transa: true,
15349            transb: false,
15350            transc: false,
15351            m: out_f as u64,
15352            n: m_tokens as u64,
15353            k: in_f as u64,
15354            alpha: 1.0,
15355            lda: in_f as i64,
15356            ldb: in_f as i64,
15357            beta: 0.0,
15358            ldc: out_f as i64,
15359            stride_a: None,
15360            stride_b: None,
15361            stride_c: None,
15362            stride_bias: None,
15363            batch_size: None,
15364        };
15365        unsafe {
15366            self.gpu.blas.matmul(cfg, w, x, &mut c, None, None)?;
15367        }
15368        Ok(c)
15369    }
15370
15371    /// Naive SDPA. Q:[head_dim,n_head,T], K/V:[head_dim,n_head_kv,T_kv] -> O:[head_dim,n_head,T].
15372    pub fn sdpa_naive(
15373        &self,
15374        q: &CudaSlice<f32>,
15375        k: &CudaSlice<f32>,
15376        v: &CudaSlice<f32>,
15377        o: &mut CudaSlice<f32>,
15378        head_dim: usize,
15379        n_head: usize,
15380        n_head_kv: usize,
15381        t: usize,
15382        t_kv: usize,
15383        scale: f32,
15384        causal: bool,
15385    ) -> Result<(), Box<dyn std::error::Error>> {
15386        let f = self.func("sdpa_naive_f32");
15387        let cfg = LaunchConfig {
15388            grid_dim: (n_head as u32, t as u32, 1),
15389            block_dim: (128, 1, 1),
15390            shared_mem_bytes: (t_kv * 4) as u32,
15391        };
15392        let (hd, nh, nhkv, ti, tkvi, cz) = (
15393            head_dim as i32,
15394            n_head as i32,
15395            n_head_kv as i32,
15396            t as i32,
15397            t_kv as i32,
15398            causal as i32,
15399        );
15400        let __s_b = self.gpu.stream();
15401        let mut b = __s_b.launch_builder(&f);
15402        b.arg(q)
15403            .arg(k)
15404            .arg(v)
15405            .arg(o)
15406            .arg(&hd)
15407            .arg(&nh)
15408            .arg(&nhkv)
15409            .arg(&ti)
15410            .arg(&tkvi)
15411            .arg(&scale)
15412            .arg(&cz);
15413        unsafe {
15414            b.launch(cfg)?;
15415        }
15416        Ok(())
15417    }
15418
15419    /// Island twin (lane/gemma-vision): causal + sliding-window attention with
15420    /// bidirectional image islands. `span_id` labels each absolute kv position
15421    /// (-1 text, >=0 island id); same-island keys are visible unconditionally,
15422    /// reproducing the reference's non-causal image batch. window 0 = no window.
15423    #[allow(clippy::too_many_arguments)]
15424    pub fn sdpa_naive_island(
15425        &self,
15426        q: &CudaSlice<f32>,
15427        k: &CudaSlice<f32>,
15428        v: &CudaSlice<f32>,
15429        o: &mut CudaSlice<f32>,
15430        span_id: &CudaSlice<i32>,
15431        head_dim: usize,
15432        n_head: usize,
15433        n_head_kv: usize,
15434        t: usize,
15435        t_kv: usize,
15436        scale: f32,
15437        window: usize,
15438    ) -> Result<(), Box<dyn std::error::Error>> {
15439        let f = self.func("sdpa_naive_island_f32");
15440        let cfg = LaunchConfig {
15441            grid_dim: (n_head as u32, t as u32, 1),
15442            block_dim: (128, 1, 1),
15443            shared_mem_bytes: (t_kv * 4) as u32,
15444        };
15445        let (hd, nh, nhkv, ti, tkvi, wi) = (
15446            head_dim as i32,
15447            n_head as i32,
15448            n_head_kv as i32,
15449            t as i32,
15450            t_kv as i32,
15451            window as i32,
15452        );
15453        let __s_b = self.gpu.stream();
15454        let mut b = __s_b.launch_builder(&f);
15455        b.arg(q)
15456            .arg(k)
15457            .arg(v)
15458            .arg(o)
15459            .arg(span_id)
15460            .arg(&hd)
15461            .arg(&nh)
15462            .arg(&nhkv)
15463            .arg(&ti)
15464            .arg(&tkvi)
15465            .arg(&scale)
15466            .arg(&wi);
15467        unsafe {
15468            b.launch(cfg)?;
15469        }
15470        Ok(())
15471    }
15472
15473    /// Windowed sdpa_naive twin (gemma4 R6): masks keys older than q_pos-(window-1).
15474    #[allow(clippy::too_many_arguments)]
15475    pub fn sdpa_naive_w(
15476        &self,
15477        q: &CudaSlice<f32>,
15478        k: &CudaSlice<f32>,
15479        v: &CudaSlice<f32>,
15480        o: &mut CudaSlice<f32>,
15481        head_dim: usize,
15482        n_head: usize,
15483        n_head_kv: usize,
15484        t: usize,
15485        t_kv: usize,
15486        scale: f32,
15487        causal: bool,
15488        window: usize,
15489    ) -> Result<(), Box<dyn std::error::Error>> {
15490        let f = self.func("sdpa_naive_w_f32");
15491        let cfg = LaunchConfig {
15492            grid_dim: (n_head as u32, t as u32, 1),
15493            block_dim: (128, 1, 1),
15494            shared_mem_bytes: (t_kv * 4) as u32,
15495        };
15496        let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
15497            head_dim as i32,
15498            n_head as i32,
15499            n_head_kv as i32,
15500            t as i32,
15501            t_kv as i32,
15502            causal as i32,
15503            window as i32,
15504        );
15505        let __s_b = self.gpu.stream();
15506        let mut b = __s_b.launch_builder(&f);
15507        b.arg(q)
15508            .arg(k)
15509            .arg(v)
15510            .arg(o)
15511            .arg(&hd)
15512            .arg(&nh)
15513            .arg(&nhkv)
15514            .arg(&ti)
15515            .arg(&tkvi)
15516            .arg(&scale)
15517            .arg(&cz)
15518            .arg(&wi);
15519        unsafe {
15520            b.launch(cfg)?;
15521        }
15522        Ok(())
15523    }
15524
15525    /// SDPA where K/V are CudaViews into a resident KV cache (decode hot path, no host round-trip).
15526    pub fn sdpa_naive_view(
15527        &self,
15528        q: &CudaSlice<f32>,
15529        k: &cudarc::driver::CudaView<f32>,
15530        v: &cudarc::driver::CudaView<f32>,
15531        o: &mut CudaSlice<f32>,
15532        head_dim: usize,
15533        n_head: usize,
15534        n_head_kv: usize,
15535        t: usize,
15536        t_kv: usize,
15537        scale: f32,
15538        causal: bool,
15539    ) -> Result<(), Box<dyn std::error::Error>> {
15540        let f = self.func("sdpa_naive_f32");
15541        let cfg = LaunchConfig {
15542            grid_dim: (n_head as u32, t as u32, 1),
15543            block_dim: (128, 1, 1),
15544            shared_mem_bytes: (t_kv * 4) as u32,
15545        };
15546        let (hd, nh, nhkv, ti, tkvi, cz) = (
15547            head_dim as i32,
15548            n_head as i32,
15549            n_head_kv as i32,
15550            t as i32,
15551            t_kv as i32,
15552            causal as i32,
15553        );
15554        let __s_b = self.gpu.stream();
15555        let mut b = __s_b.launch_builder(&f);
15556        b.arg(q)
15557            .arg(k)
15558            .arg(v)
15559            .arg(o)
15560            .arg(&hd)
15561            .arg(&nh)
15562            .arg(&nhkv)
15563            .arg(&ti)
15564            .arg(&tkvi)
15565            .arg(&scale)
15566            .arg(&cz);
15567        unsafe {
15568            b.launch(cfg)?;
15569        }
15570        Ok(())
15571    }
15572
15573    /// Correctness fallback for quantized resident K/V views. Dequantizes K and V once into f32
15574    /// workspaces, then calls `sdpa_naive`. This is an explicit API: the optimized prefill view
15575    /// dispatch remains unchanged, so callers can use it as a reference or compatibility path.
15576    /// Dequant a quantized KV view into caller-owned f32 buffers (one grid-stride launch).
15577    /// `g` picks the kf8vf8-module stamp for e4m3 caches (same flag contract as fa_decode/
15578    /// fa_prefill_view). Used by the E4B shared-KV prefill arms (2026-07-31) to feed the
15579    /// f32 fa_prefill_w / fa_prefill_hd512 twins from the target layer's quantized rows.
15580    #[allow(clippy::too_many_arguments)]
15581    pub fn fa_dequant_kv_view_f32(
15582        &self,
15583        k: &cudarc::driver::CudaView<u8>,
15584        v: &cudarc::driver::CudaView<u8>,
15585        kf: &mut CudaSlice<f32>,
15586        vf: &mut CudaSlice<f32>,
15587        kv_dim_k: usize,
15588        kv_dim_v: usize,
15589        t_kv: usize,
15590        k_tok_bytes: usize,
15591        v_tok_bytes: usize,
15592        g: bool,
15593    ) -> Result<(), Box<dyn std::error::Error>> {
15594        let f = if g {
15595            self.func_g("fa_dequant_kv_ws_f32")
15596        } else {
15597            self.func("fa_dequant_kv_ws_f32")
15598        };
15599        let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
15600        let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
15601        let cfg = LaunchConfig {
15602            grid_dim: (nblk.max(1), 1, 1),
15603            block_dim: (256, 1, 1),
15604            shared_mem_bytes: 0,
15605        };
15606        let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
15607        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
15608        let __s_b = self.gpu.stream();
15609        let mut b = __s_b.launch_builder(&f);
15610        b.arg(k)
15611            .arg(v)
15612            .arg(&mut *kf)
15613            .arg(&mut *vf)
15614            .arg(&kdk)
15615            .arg(&kdv)
15616            .arg(&tkvi)
15617            .arg(&ktb)
15618            .arg(&vtb);
15619        unsafe {
15620            b.launch(cfg)?;
15621        }
15622        Ok(())
15623    }
15624
15625    #[allow(clippy::too_many_arguments)]
15626    pub fn sdpa_naive_quantized_view(
15627        &self,
15628        q: &CudaSlice<f32>,
15629        k: &cudarc::driver::CudaView<u8>,
15630        v: &cudarc::driver::CudaView<u8>,
15631        o: &mut CudaSlice<f32>,
15632        head_dim: usize,
15633        n_head: usize,
15634        n_head_kv: usize,
15635        t: usize,
15636        t_kv: usize,
15637        scale: f32,
15638        causal: bool,
15639        k_tok_bytes: usize,
15640        v_tok_bytes: usize,
15641    ) -> Result<(), Box<dyn std::error::Error>> {
15642        let kv_dim = n_head_kv * head_dim;
15643        let mut kf = self.uninit(t_kv * kv_dim)?;
15644        let mut vf = self.uninit(t_kv * kv_dim)?;
15645        let f = self.func("fa_dequant_kv_ws_f32");
15646        let total = (2 * t_kv * kv_dim) as u64;
15647        let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
15648        let cfg = LaunchConfig {
15649            grid_dim: (nblk.max(1), 1, 1),
15650            block_dim: (256, 1, 1),
15651            shared_mem_bytes: 0,
15652        };
15653        let (kv_dim_i, t_kv_i) = (kv_dim as i32, t_kv as i32);
15654        let (k_tok_bytes_i, v_tok_bytes_i) = (k_tok_bytes as i64, v_tok_bytes as i64);
15655        let __s_b = self.gpu.stream();
15656        let mut b = __s_b.launch_builder(&f);
15657        b.arg(k)
15658            .arg(v)
15659            .arg(&mut kf)
15660            .arg(&mut vf)
15661            .arg(&kv_dim_i)
15662            .arg(&kv_dim_i)
15663            .arg(&t_kv_i)
15664            .arg(&k_tok_bytes_i)
15665            .arg(&v_tok_bytes_i);
15666        unsafe { b.launch(cfg)? };
15667        self.sdpa_naive(
15668            q, &kf, &vf, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
15669        )
15670    }
15671
15672    /// WINDOWED twin of `sdpa_naive_quantized_view` (step35 SWA prefill): dequant the KV byte
15673    /// view into f32 workspaces with the SAME `fa_dequant_kv_ws_f32` launch, then run
15674    /// `sdpa_naive_w` instead of `sdpa_naive`. `window == 0` is the unwindowed form (the kernel
15675    /// treats a non-positive window as "no window mask"), so this is a strict superset of the
15676    /// unwindowed function above and produces bit-identical output at window == 0.
15677    ///
15678    /// Why this exists: EVERY windowed FlashAttention stamp in flash_attn.cu is head_dim-256
15679    /// only (`fa_prefill_w_f32` == `fa_prefill_f32_body<256>`, and the quantized-view windowed
15680    /// twins likewise), while step35 is head_dim 128. Its SWA layers therefore have no windowed
15681    /// FA path and take this f32 floor in v0 — same cache bytes, same numeric class as the
15682    /// unwindowed quantized-view fallback, so the chunk-invariance contract holds on both.
15683    #[allow(clippy::too_many_arguments)]
15684    pub fn sdpa_naive_w_quantized_view(
15685        &self,
15686        q: &CudaSlice<f32>,
15687        k: &cudarc::driver::CudaView<u8>,
15688        v: &cudarc::driver::CudaView<u8>,
15689        o: &mut CudaSlice<f32>,
15690        head_dim: usize,
15691        n_head: usize,
15692        n_head_kv: usize,
15693        t: usize,
15694        t_kv: usize,
15695        scale: f32,
15696        causal: bool,
15697        window: usize,
15698        k_tok_bytes: usize,
15699        v_tok_bytes: usize,
15700    ) -> Result<(), Box<dyn std::error::Error>> {
15701        let kv_dim = n_head_kv * head_dim;
15702        let mut kf = self.uninit(t_kv * kv_dim)?;
15703        let mut vf = self.uninit(t_kv * kv_dim)?;
15704        let f = self.func("fa_dequant_kv_ws_f32");
15705        let total = (2 * t_kv * kv_dim) as u64;
15706        let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
15707        let cfg = LaunchConfig {
15708            grid_dim: (nblk.max(1), 1, 1),
15709            block_dim: (256, 1, 1),
15710            shared_mem_bytes: 0,
15711        };
15712        let (kv_dim_i, t_kv_i) = (kv_dim as i32, t_kv as i32);
15713        let (k_tok_bytes_i, v_tok_bytes_i) = (k_tok_bytes as i64, v_tok_bytes as i64);
15714        let __s_b = self.gpu.stream();
15715        let mut b = __s_b.launch_builder(&f);
15716        b.arg(k)
15717            .arg(v)
15718            .arg(&mut kf)
15719            .arg(&mut vf)
15720            .arg(&kv_dim_i)
15721            .arg(&kv_dim_i)
15722            .arg(&t_kv_i)
15723            .arg(&k_tok_bytes_i)
15724            .arg(&v_tok_bytes_i);
15725        unsafe { b.launch(cfg)? };
15726        self.sdpa_naive_w(
15727            q, &kf, &vf, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal, window,
15728        )
15729    }
15730
15731    /// Hand-written FlashAttention prefill (sm_120, FA-2 online softmax on validated mma.sync,
15732    /// head_dim 256 or 128 (template-stamped twins), GQA, causal). Replaces sdpa_naive for T>1.
15733    /// Q/K/V/O [head_dim, n_head(_kv), T].
15734    pub fn fa_prefill(
15735        &self,
15736        q: &CudaSlice<f32>,
15737        k: &CudaSlice<f32>,
15738        v: &CudaSlice<f32>,
15739        o: &mut CudaSlice<f32>,
15740        head_dim: usize,
15741        n_head: usize,
15742        n_head_kv: usize,
15743        t: usize,
15744        t_kv: usize,
15745        scale: f32,
15746        causal: bool,
15747    ) -> Result<(), Box<dyn std::error::Error>> {
15748        if portable_mma_gated() {
15749            return self.sdpa_naive(
15750                q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
15751            );
15752        }
15753        // FA3 v10 arm (task #20, OPT-IN MEMRA_FA3=1 — harness-proven 883us vs the shipped
15754        // kernel's 993us at T=2048): TMA-swizzled wgmma FA, fresh causal hd256 only.
15755        // NEW NUMERIC CONFIG (GDN-mma precedent): online softmax / bf16-P class — the
15756        // run-gen argmax + greedy-stream batteries arbitrate; not bit-paired.
15757        // PROMOTED default-ON hopper (2026-07-27): 3-seed 2048-prime -> 128-decode
15758        // streams MATCH vs mma, full battery green, lane interleaved 5/5 (+2.4%).
15759        // MEMRA_FA3=0 reverts; kernel-check pins the mma config regardless.
15760        let fa3_on = head_dim == 256
15761            && causal
15762            && t == t_kv
15763            && match std::env::var("MEMRA_FA3").as_deref() {
15764                Ok("0") => false,
15765                Ok("1") => true,
15766                _ => cfg!(memra_hopper_mma),
15767            };
15768        if fa3_on {
15769            let n = t * n_head * head_dim;
15770            let nkv = t * n_head_kv * head_dim;
15771            let mut q16 = self.alloc_u8_uninit(n * 2)?;
15772            let mut k16 = self.alloc_u8_uninit(nkv * 2)?;
15773            let mut v16 = self.alloc_u8_uninit(nkv * 2)?;
15774            self.f32_to_bf16_into(q, &mut q16, n)?;
15775            self.f32_to_bf16_into(k, &mut k16, nkv)?;
15776            self.f32_to_bf16_into(v, &mut v16, nkv)?;
15777            let rc = {
15778                use cudarc::driver::{DevicePtr, DevicePtrMut};
15779                let stream = self.gpu.stream();
15780                let (qp, _g1) = q16.device_ptr(&stream);
15781                let (kp, _g2) = k16.device_ptr(&stream);
15782                let (vp, _g3) = v16.device_ptr(&stream);
15783                let (op, _g4) = o.device_ptr_mut(&stream);
15784                unsafe {
15785                    memra_fa3_prefill(
15786                        qp as *const core::ffi::c_void,
15787                        kp as *const core::ffi::c_void,
15788                        vp as *const core::ffi::c_void,
15789                        op as *mut f32,
15790                        t as i32,
15791                        n_head as i32,
15792                        n_head_kv as i32,
15793                        head_dim as i32,
15794                        scale,
15795                        stream.cu_stream() as *mut core::ffi::c_void,
15796                    )
15797                }
15798            };
15799            if rc != 0 {
15800                return Err(format!("memra_fa3_prefill rc={rc}").into());
15801            }
15802            return Ok(());
15803        }
15804        // FLOOR PORT (P2+P0a+P0b+P1): 4 warps/CTA, BLOCK_Q=64 query rows, BK=32 KV tile,
15805        // Q-in-reg + register-O, grid.y=n_head_kv (4 Q-heads share staged K/V).
15806        // P1 plain arm (MEMRA_FA_P1=1 opt-in until the qwen battery): the engine-study body
15807        // (FA2 schedule + boundary split + swizzle) on the non-windowed lane. bf16 pre-convert.
15808        static FA_P1: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15809        let fa_p1 = *FA_P1.get_or_init(|| std::env::var("MEMRA_FA_P1").as_deref() == Ok("1"));
15810        if fa_p1 && head_dim == 256 && !std::env::var("MEMRA_FA_FLOOR").is_ok() {
15811            const BLOCK_Q: usize = 64;
15812            const BKX: usize = 32;
15813            let f = self.func("fa_prefill_bf16_p1");
15814            let shmem = (2 * (2 * BKX * head_dim + BLOCK_Q * BKX)
15815                + 4 * (BLOCK_Q * BKX + 2 * BLOCK_Q)) as u32;
15816            use cudarc::driver::sys::CUfunction_attribute_enum as A;
15817            f.set_attribute(
15818                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
15819                shmem as i32,
15820            )?;
15821            let cfg = LaunchConfig {
15822                grid_dim: (
15823                    (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
15824                    n_head as u32,
15825                    1,
15826                ),
15827                block_dim: (32, 4, 1),
15828                shared_mem_bytes: shmem,
15829            };
15830            let (hd, nh, nhkv, ti, tkvi, cz) = (
15831                head_dim as i32,
15832                n_head as i32,
15833                n_head_kv as i32,
15834                t as i32,
15835                t_kv as i32,
15836                causal as i32,
15837            );
15838            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
15839            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
15840            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
15841            let __s_b = self.gpu.stream();
15842            let mut b = __s_b.launch_builder(&f);
15843            b.arg(&qb)
15844                .arg(&kb)
15845                .arg(&vb)
15846                .arg(o)
15847                .arg(&hd)
15848                .arg(&nh)
15849                .arg(&nhkv)
15850                .arg(&ti)
15851                .arg(&tkvi)
15852                .arg(&scale)
15853                .arg(&cz);
15854            unsafe {
15855                b.launch(cfg)?;
15856            }
15857            return Ok(());
15858        }
15859        // Edge 5a (DEFAULT): fa_prefill_f32_pp — register-resident softmax (no sSw smem
15860        // round-trip), the FA3 softmax-GEMM overlap variant. ncu (pp512): short_scoreboard
15861        // 4.32->3.47, wait 1.99->1.45, per-call ~577us->~440us (1.31x) at flat 12.1% warps /
15862        // 255 regs / 2 CTAs (occupancy preserved). Bit-safe: 9B+27B argmax MATCH, rel 2.55e-3
15863        // vs floor 3.03e-3. MEMRA_FA_FLOOR reverts to the serialized-softmax floor kernel.
15864        const BK: usize = 32;
15865        // W2 lane (MEMRA_FA_PP_W2=1, ncu 2026-07-26): 2-warp/32-row CTA tile doubles grid.x —
15866        // bit-identical per-row math, pure coverage trade for the 6.25%-occupancy starvation.
15867        let w2 = std::env::var("MEMRA_FA_PP_W2").as_deref() == Ok("1");
15868        let (block_q, warps, w2_sfx): (usize, u32, &str) =
15869            if w2 { (32, 2, "_w2") } else { (64, 4, "") };
15870        // hd128 twins (2026-07-07): the prefill kernels are template-stamped at 256 (original
15871        // names, dispatch unchanged) and 128 (`_hd128`, the MiniMax-M3 class). Callers gate
15872        // other head_dims to sdpa_naive before reaching here.
15873        let hd_sfx = fa_hd_suffix(head_dim)?;
15874        let floor = std::env::var("MEMRA_FA_FLOOR").is_ok();
15875        // BF16-KV staging lane (2026-07-26, default ON): the kernel converts K/V to bf16
15876        // during staging anyway — pre-converting to bf16 mirrors is BIT-IDENTICAL (same
15877        // __float2bfloat16 values into the same mma) and turns the 67%-of-stalls scalar
15878        // staging into int4 vector copies. MEMRA_FA_BF16KV=0 reverts.
15879        let bf16kv = !floor && !w2 && std::env::var("MEMRA_FA_BF16KV").as_deref() != Ok("0");
15880        let (kb16, vb16) = if bf16kv {
15881            let n = t_kv * n_head_kv * head_dim;
15882            let mut kb = self.alloc_u8_uninit(n * 2)?;
15883            let mut vb = self.alloc_u8_uninit(n * 2)?;
15884            let fcv = self.func("f32_to_bf16_bulk");
15885            let ni = n as i64;
15886            let cfgc = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
15887            let __s_b = self.gpu.stream();
15888            let mut b = __s_b.launch_builder(&fcv);
15889            b.arg(k).arg(&mut kb).arg(&ni);
15890            unsafe {
15891                b.launch(cfgc)?;
15892            }
15893            let __s_b = self.gpu.stream();
15894            let mut b = __s_b.launch_builder(&fcv);
15895            b.arg(v).arg(&mut vb).arg(&ni);
15896            unsafe {
15897                b.launch(cfgc)?;
15898            }
15899            (Some(kb), Some(vb))
15900        } else {
15901            (None, None)
15902        };
15903        let f = self.func(&if bf16kv {
15904            format!("fa_prefill_bf16kv_pp{hd_sfx}")
15905        } else {
15906            format!(
15907                "fa_prefill_f32{}{}{hd_sfx}",
15908                if floor { "" } else { "_pp" },
15909                if floor { "" } else { w2_sfx }
15910            )
15911        });
15912        // persistent smem: bf16*(KV_STAGES*(sK + sV) + sP) + f32*(sS + sM + sL);
15913        // the bf16kv ring doubles the K/V stages (KV_STAGES=2).
15914        let kv_stages = if bf16kv { 2 } else { 1 };
15915        let shmem = (2 * (kv_stages * 2 * BK * head_dim + block_q * BK)
15916            + 4 * (block_q * BK + 2 * block_q)) as u32;
15917        use cudarc::driver::sys::CUfunction_attribute_enum as A;
15918        f.set_attribute(
15919            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
15920            shmem as i32,
15921        )?;
15922        let cfg = LaunchConfig {
15923            grid_dim: (
15924                (t as u32 + block_q as u32 - 1) / block_q as u32,
15925                n_head as u32,
15926                1,
15927            ),
15928            block_dim: (32, warps, 1),
15929            shared_mem_bytes: shmem,
15930        };
15931        let (hd, nh, nhkv, ti, tkvi, cz) = (
15932            head_dim as i32,
15933            n_head as i32,
15934            n_head_kv as i32,
15935            t as i32,
15936            t_kv as i32,
15937            causal as i32,
15938        );
15939        let __s_b = self.gpu.stream();
15940        let mut b = __s_b.launch_builder(&f);
15941        b.arg(q);
15942        match (&kb16, &vb16) {
15943            (Some(kb), Some(vb)) => {
15944                b.arg(kb).arg(vb);
15945            }
15946            _ => {
15947                b.arg(k).arg(v);
15948            }
15949        }
15950        b.arg(o)
15951            .arg(&hd)
15952            .arg(&nh)
15953            .arg(&nhkv)
15954            .arg(&ti)
15955            .arg(&tkvi)
15956            .arg(&scale)
15957            .arg(&cz);
15958        unsafe {
15959            b.launch(cfg)?;
15960        }
15961        Ok(())
15962    }
15963
15964    /// Windowed FA prefill (gemma4 SWA layers past the sliding window, hd256): fa_prefill's
15965    /// exact dispatch (pp default, MEMRA_FA_FLOOR seam) with the sliding-window mask + tile
15966    /// skip in-kernel. Replaces the O(T*T_kv) scalar sdpa_naive_w on the prime path.
15967    #[allow(clippy::too_many_arguments)]
15968    pub fn fa_prefill_w(
15969        &self,
15970        q: &CudaSlice<f32>,
15971        k: &CudaSlice<f32>,
15972        v: &CudaSlice<f32>,
15973        o: &mut CudaSlice<f32>,
15974        head_dim: usize,
15975        n_head: usize,
15976        n_head_kv: usize,
15977        t: usize,
15978        t_kv: usize,
15979        scale: f32,
15980        causal: bool,
15981        window: usize,
15982    ) -> Result<(), Box<dyn std::error::Error>> {
15983        // sm_90a rides the mma twins (portable_mma_gated, 2026-07-31 — the raw
15984        // portable_cuda gate was stale-conservative on Hopper; fa_prefill already flipped).
15985        if portable_mma_gated() {
15986            return self.sdpa_naive_w(
15987                q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal, window,
15988            );
15989        }
15990        // Default: bf16-prestaged twin (same treatment as hd512 — Q/K/V pre-converted once,
15991        // int4 stage copies; bit-identical, kernel_check-gated). MEMRA_FAW_STAGE=f32 reverts;
15992        // MEMRA_FA_FLOOR keeps the f32 floor stamp untouched.
15993        static FAW_F32: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15994        let faw_f32 =
15995            *FAW_F32.get_or_init(|| std::env::var("MEMRA_FAW_STAGE").as_deref() == Ok("f32"));
15996        let floor = std::env::var("MEMRA_FA_FLOOR").is_ok();
15997        self.fa_prefill_w_arm(
15998            q,
15999            k,
16000            v,
16001            o,
16002            head_dim,
16003            n_head,
16004            n_head_kv,
16005            t,
16006            t_kv,
16007            scale,
16008            causal,
16009            window,
16010            floor || faw_f32,
16011            floor,
16012        )
16013    }
16014
16015    /// Windowed FA prefill with PRE-CONVERTED bf16 operands (producer-emitted; 31B glue lane).
16016    /// Launches the P1 stamp directly — callers guarantee qb/kb/vb hold the exact bf16 of q/k/v.
16017    #[allow(clippy::too_many_arguments)]
16018    pub fn fa_prefill_w_pre(
16019        &self,
16020        qb: &CudaSlice<u8>,
16021        kb: &CudaSlice<u8>,
16022        vb: &CudaSlice<u8>,
16023        o: &mut CudaSlice<f32>,
16024        head_dim: usize,
16025        n_head: usize,
16026        n_head_kv: usize,
16027        t: usize,
16028        t_kv: usize,
16029        scale: f32,
16030        causal: bool,
16031        window: usize,
16032        v_f16: bool,
16033    ) -> Result<(), Box<dyn std::error::Error>> {
16034        const BLOCK_Q: usize = 64;
16035        const BK: usize = 32;
16036        debug_assert_eq!(head_dim, 256);
16037        let hp = fa_f16pv_on() && faw_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
16038        debug_assert!(!v_f16 || hp, "f16 V emitted but the SWA hp arm is off");
16039        if hp {
16040            const BLOCK_QH: usize = 32;
16041            // V bytes must be f16 for the h2 stamp; producer normally emits f16 (v_f16),
16042            // else re-encode through the pooled scratch (stream-ordered reuse).
16043            let mut vguard = self.fa_vf16_scratch.lock().unwrap();
16044            let vh: &CudaSlice<u8> = if v_f16 {
16045                vb
16046            } else {
16047                let n = t_kv * n_head_kv * head_dim;
16048                if vguard.as_ref().map(|b| b.len() < n * 2).unwrap_or(true) {
16049                    *vguard = Some(self.alloc_uninit::<u8>(n * 2)?);
16050                }
16051                self.bf16_to_f16_into(vb, n, vguard.as_mut().unwrap())?;
16052                vguard.as_ref().unwrap()
16053            };
16054            let f = self.func("fa_prefill_w_bf16_p1h2");
16055            let shmem = (2 * (2 * BK * head_dim + 2 * BLOCK_QH * BK) + 4 * (2 * BLOCK_QH)) as u32;
16056            use cudarc::driver::sys::CUfunction_attribute_enum as A;
16057            f.set_attribute(
16058                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16059                shmem as i32,
16060            )?;
16061            let cfg = LaunchConfig {
16062                grid_dim: ((t as u32).div_ceil(BLOCK_QH as u32), (n_head / 2) as u32, 1),
16063                block_dim: (32, 4, 1),
16064                shared_mem_bytes: shmem,
16065            };
16066            let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
16067                head_dim as i32,
16068                n_head as i32,
16069                n_head_kv as i32,
16070                t as i32,
16071                t_kv as i32,
16072                causal as i32,
16073                window as i32,
16074            );
16075            let __s_b = self.gpu.stream();
16076            let mut b = __s_b.launch_builder(&f);
16077            b.arg(qb)
16078                .arg(kb)
16079                .arg(vh)
16080                .arg(o)
16081                .arg(&hd)
16082                .arg(&nh)
16083                .arg(&nhkv)
16084                .arg(&ti)
16085                .arg(&tkvi)
16086                .arg(&scale)
16087                .arg(&cz)
16088                .arg(&wi);
16089            unsafe {
16090                b.launch(cfg)?;
16091            }
16092            return Ok(());
16093        }
16094        let f = self.func("fa_prefill_w_bf16_p1");
16095        let shmem =
16096            (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
16097        use cudarc::driver::sys::CUfunction_attribute_enum as A;
16098        f.set_attribute(
16099            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16100            shmem as i32,
16101        )?;
16102        let cfg = LaunchConfig {
16103            grid_dim: (
16104                (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
16105                n_head as u32,
16106                1,
16107            ),
16108            block_dim: (32, 4, 1),
16109            shared_mem_bytes: shmem,
16110        };
16111        let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
16112            head_dim as i32,
16113            n_head as i32,
16114            n_head_kv as i32,
16115            t as i32,
16116            t_kv as i32,
16117            causal as i32,
16118            window as i32,
16119        );
16120        let __s_b = self.gpu.stream();
16121        let mut b = __s_b.launch_builder(&f);
16122        b.arg(qb)
16123            .arg(kb)
16124            .arg(vb)
16125            .arg(o)
16126            .arg(&hd)
16127            .arg(&nh)
16128            .arg(&nhkv)
16129            .arg(&ti)
16130            .arg(&tkvi)
16131            .arg(&scale)
16132            .arg(&cz)
16133            .arg(&wi);
16134        unsafe {
16135            b.launch(cfg)?;
16136        }
16137        Ok(())
16138    }
16139
16140    /// Windowed FA prefill with the stage arm FORCED — the kernel_check bit-identity entry.
16141    #[allow(clippy::too_many_arguments)]
16142    pub fn fa_prefill_w_arm(
16143        &self,
16144        q: &CudaSlice<f32>,
16145        k: &CudaSlice<f32>,
16146        v: &CudaSlice<f32>,
16147        o: &mut CudaSlice<f32>,
16148        head_dim: usize,
16149        n_head: usize,
16150        n_head_kv: usize,
16151        t: usize,
16152        t_kv: usize,
16153        scale: f32,
16154        causal: bool,
16155        window: usize,
16156        f32_stage: bool,
16157        floor: bool,
16158    ) -> Result<(), Box<dyn std::error::Error>> {
16159        const BLOCK_Q: usize = 64;
16160        const BK: usize = 32;
16161        debug_assert_eq!(head_dim, 256, "fa_prefill_w is stamped hd256 only");
16162        // P1 (2026-07-22 engine study): per-head Br=64 stamp with the FA2 schedule (V-copy
16163        // over GEMM0, next-K over softmax+GEMM1) + boundary/interior mask split. FP order
16164        // preserved -> bit-identical (gated). MEMRA_FAW_P1=0 reverts to the g4/o2 arms.
16165        static P1_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16166        let p1 = !floor
16167            && !f32_stage
16168            && *P1_ON.get_or_init(|| {
16169                std::env::var("MEMRA_FAW_P1")
16170                    .map(|v| v != "0")
16171                    .unwrap_or(true)
16172            });
16173        let hp =
16174            p1 && fa_f16pv_on() && faw_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
16175        if hp {
16176            const BLOCK_QH: usize = 32;
16177            let f = self.func("fa_prefill_w_bf16_p1h2");
16178            let shmem = (2 * (2 * BK * head_dim + 2 * BLOCK_QH * BK) + 4 * (2 * BLOCK_QH)) as u32;
16179            use cudarc::driver::sys::CUfunction_attribute_enum as A;
16180            f.set_attribute(
16181                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16182                shmem as i32,
16183            )?;
16184            let cfg = LaunchConfig {
16185                grid_dim: ((t as u32).div_ceil(BLOCK_QH as u32), (n_head / 2) as u32, 1),
16186                block_dim: (32, 4, 1),
16187                shared_mem_bytes: shmem,
16188            };
16189            let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
16190                head_dim as i32,
16191                n_head as i32,
16192                n_head_kv as i32,
16193                t as i32,
16194                t_kv as i32,
16195                causal as i32,
16196                window as i32,
16197            );
16198            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
16199            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
16200            let vh = self.f32_to_f16(v, t_kv * n_head_kv * head_dim)?;
16201            let __s_b = self.gpu.stream();
16202            let mut b = __s_b.launch_builder(&f);
16203            b.arg(&qb)
16204                .arg(&kb)
16205                .arg(&vh)
16206                .arg(o)
16207                .arg(&hd)
16208                .arg(&nh)
16209                .arg(&nhkv)
16210                .arg(&ti)
16211                .arg(&tkvi)
16212                .arg(&scale)
16213                .arg(&cz)
16214                .arg(&wi);
16215            unsafe {
16216                b.launch(cfg)?;
16217            }
16218            return Ok(());
16219        }
16220        if p1 {
16221            let f = self.func("fa_prefill_w_bf16_p1");
16222            let shmem =
16223                (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
16224            use cudarc::driver::sys::CUfunction_attribute_enum as A;
16225            f.set_attribute(
16226                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16227                shmem as i32,
16228            )?;
16229            let cfg = LaunchConfig {
16230                grid_dim: (
16231                    (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
16232                    n_head as u32,
16233                    1,
16234                ),
16235                block_dim: (32, 4, 1),
16236                shared_mem_bytes: shmem,
16237            };
16238            let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
16239                head_dim as i32,
16240                n_head as i32,
16241                n_head_kv as i32,
16242                t as i32,
16243                t_kv as i32,
16244                causal as i32,
16245                window as i32,
16246            );
16247            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
16248            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
16249            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
16250            let __s_b = self.gpu.stream();
16251            let mut b = __s_b.launch_builder(&f);
16252            b.arg(&qb)
16253                .arg(&kb)
16254                .arg(&vb)
16255                .arg(o)
16256                .arg(&hd)
16257                .arg(&nh)
16258                .arg(&nhkv)
16259                .arg(&ti)
16260                .arg(&tkvi)
16261                .arg(&scale)
16262                .arg(&cz)
16263                .arg(&wi);
16264            unsafe {
16265                b.launch(cfg)?;
16266            }
16267            return Ok(());
16268        }
16269        // MQA head-grouping (MEMRA_FAW_G4=0 reverts): 4 heads/CTA share the staged K/V —
16270        // per-(head,row) FP chain identical to the per-head stamp -> bit-identical (gated).
16271        static G4_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16272        let g4 = !floor
16273            && !f32_stage
16274            && n_head_kv == 1
16275            && n_head % 4 == 0
16276            && *G4_ON.get_or_init(|| {
16277                std::env::var("MEMRA_FAW_G4")
16278                    .map(|v| v != "0")
16279                    .unwrap_or(true)
16280            });
16281        if g4 {
16282            const SP_M: usize = 16;
16283            // Occupancy-2 twin (MEMRA_FAW_O2=0 reverts): one shared K/V buffer inside the dead
16284            // Q-stage region -> ~36.5KB smem, 2 CTA/SM (the llama hd256 mechanism). Bit-identical.
16285            static O2_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16286            let o2 = *O2_ON.get_or_init(|| {
16287                std::env::var("MEMRA_FAW_O2")
16288                    .map(|v| v != "0")
16289                    .unwrap_or(true)
16290            });
16291            let f = self.func(if o2 {
16292                "fa_prefill_w_bf16_g4o2"
16293            } else {
16294                "fa_prefill_w_bf16_g4"
16295            });
16296            let shmem = if o2 {
16297                (2 * (4 * SP_M * head_dim + 4 * SP_M * BK) + 4 * (4 * SP_M)) as u32
16298            } else {
16299                (2 * (2 * BK * head_dim + 4 * SP_M * head_dim + 4 * SP_M * BK) + 4 * (4 * SP_M))
16300                    as u32
16301            };
16302            use cudarc::driver::sys::CUfunction_attribute_enum as A;
16303            f.set_attribute(
16304                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16305                shmem as i32,
16306            )?;
16307            let cfg = LaunchConfig {
16308                grid_dim: ((t as u32).div_ceil(SP_M as u32), (n_head / 4) as u32, 1),
16309                block_dim: (32, 4, 1),
16310                shared_mem_bytes: shmem,
16311            };
16312            let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
16313                head_dim as i32,
16314                n_head as i32,
16315                n_head_kv as i32,
16316                t as i32,
16317                t_kv as i32,
16318                causal as i32,
16319                window as i32,
16320            );
16321            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
16322            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
16323            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
16324            let __s_b = self.gpu.stream();
16325            let mut b = __s_b.launch_builder(&f);
16326            b.arg(&qb)
16327                .arg(&kb)
16328                .arg(&vb)
16329                .arg(o)
16330                .arg(&hd)
16331                .arg(&nh)
16332                .arg(&nhkv)
16333                .arg(&ti)
16334                .arg(&tkvi)
16335                .arg(&scale)
16336                .arg(&cz)
16337                .arg(&wi);
16338            unsafe {
16339                b.launch(cfg)?;
16340            }
16341            return Ok(());
16342        }
16343        let f = self.func(if floor {
16344            "fa_prefill_w_f32"
16345        } else if f32_stage {
16346            "fa_prefill_w_f32_pp"
16347        } else {
16348            "fa_prefill_w_bf16_pp"
16349        });
16350        let shmem =
16351            (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
16352        use cudarc::driver::sys::CUfunction_attribute_enum as A;
16353        f.set_attribute(
16354            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16355            shmem as i32,
16356        )?;
16357        let cfg = LaunchConfig {
16358            grid_dim: (
16359                (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
16360                n_head as u32,
16361                1,
16362            ),
16363            block_dim: (32, 4, 1),
16364            shared_mem_bytes: shmem,
16365        };
16366        let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
16367            head_dim as i32,
16368            n_head as i32,
16369            n_head_kv as i32,
16370            t as i32,
16371            t_kv as i32,
16372            causal as i32,
16373            window as i32,
16374        );
16375        if f32_stage {
16376            let __s_b = self.gpu.stream();
16377            let mut b = __s_b.launch_builder(&f);
16378            b.arg(q)
16379                .arg(k)
16380                .arg(v)
16381                .arg(o)
16382                .arg(&hd)
16383                .arg(&nh)
16384                .arg(&nhkv)
16385                .arg(&ti)
16386                .arg(&tkvi)
16387                .arg(&scale)
16388                .arg(&cz)
16389                .arg(&wi);
16390            unsafe {
16391                b.launch(cfg)?;
16392            }
16393        } else {
16394            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
16395            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
16396            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
16397            let __s_b = self.gpu.stream();
16398            let mut b = __s_b.launch_builder(&f);
16399            b.arg(&qb)
16400                .arg(&kb)
16401                .arg(&vb)
16402                .arg(o)
16403                .arg(&hd)
16404                .arg(&nh)
16405                .arg(&nhkv)
16406                .arg(&ti)
16407                .arg(&tkvi)
16408                .arg(&scale)
16409                .arg(&cz)
16410                .arg(&wi);
16411            unsafe {
16412                b.launch(cfg)?;
16413            }
16414        }
16415        Ok(())
16416    }
16417
16418    /// hd512 FA prefill (gemma4 GLOBAL layers): BLOCK_Q=32 x 2 warps, Q staged in smem,
16419    /// grid.z = 2 O-halves (each CTA computes the full 512-dim scores, accumulates half the
16420    /// V dims). Replaces the scalar sdpa_naive on the prime path's globals.
16421    #[allow(clippy::too_many_arguments)]
16422    pub fn fa_prefill_hd512(
16423        &self,
16424        q: &CudaSlice<f32>,
16425        k: &CudaSlice<f32>,
16426        v: &CudaSlice<f32>,
16427        o: &mut CudaSlice<f32>,
16428        head_dim: usize,
16429        n_head: usize,
16430        n_head_kv: usize,
16431        t: usize,
16432        t_kv: usize,
16433        scale: f32,
16434        causal: bool,
16435    ) -> Result<(), Box<dyn std::error::Error>> {
16436        // sm_90a rides the mma twins (portable_mma_gated, 2026-07-31 — same flip as _w).
16437        if portable_mma_gated() {
16438            return self.sdpa_naive(
16439                q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
16440            );
16441        }
16442        // Default: pre-convert Q/K/V to bf16 once and stage int4 (8 bf16/copy) — at 1 CTA/SM the
16443        // synchronous stage serializes with compute and MQA re-stages the same K/V per head CTA;
16444        // pre-converting halves staged bytes and cuts stage instructions 8x. BIT-IDENTICAL to the
16445        // f32-staged kernel (the converter applies the same __float2bfloat16 the stage applied;
16446        // kernel_check gates the identity). MEMRA_FA512_STAGE=f32 = rollback to the f32 kernel.
16447        static F32_STAGE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16448        let f32_stage =
16449            *F32_STAGE.get_or_init(|| std::env::var("MEMRA_FA512_STAGE").as_deref() == Ok("f32"));
16450        // Single-pass arm (MEMRA_FA512_SP=0 reverts to the z=2 bf16 kernel): GEMM0 split-K across
16451        // the 2 warps instead of recomputed per O-half CTA — the 2026-07-22 kernel-diff excess.
16452        // Own numeric config (partial-sum order) — battery-gated.
16453        static SP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16454        let sp = !f32_stage
16455            && *SP_ON.get_or_init(|| {
16456                std::env::var("MEMRA_FA512_SP")
16457                    .map(|v| v != "0")
16458                    .unwrap_or(true)
16459            });
16460        self.fa_prefill_hd512_arm(
16461            q,
16462            k,
16463            v,
16464            o,
16465            head_dim,
16466            n_head,
16467            n_head_kv,
16468            t,
16469            t_kv,
16470            scale,
16471            causal,
16472            f32_stage,
16473            sp,
16474            sp && fa_f16pv_on(),
16475        )
16476    }
16477
16478    /// hd512 single-pass FA with PRE-CONVERTED bf16 operands (producer-emitted).
16479    #[allow(clippy::too_many_arguments)]
16480    pub fn fa_prefill_hd512_pre(
16481        &self,
16482        qb: &CudaSlice<u8>,
16483        kb: &CudaSlice<u8>,
16484        vb: &CudaSlice<u8>,
16485        o: &mut CudaSlice<f32>,
16486        head_dim: usize,
16487        n_head: usize,
16488        n_head_kv: usize,
16489        t: usize,
16490        t_kv: usize,
16491        scale: f32,
16492        causal: bool,
16493        v_f16: bool,
16494    ) -> Result<(), Box<dyn std::error::Error>> {
16495        debug_assert_eq!(head_dim, 512);
16496        const SP_M: usize = 16;
16497        const BKS: usize = 32;
16498        // f16-P/V door (MEMRA_FA_F16PV=1): P and the P@V accumulation in f16 (llama's fa=1 VKQ
16499        // class); KQ/softmax/rescale-band/final-normalize stay f32. Own numeric config —
16500        // battery-gated. V bytes must be f16 for the sp16 kernel (stage/ldmatrix are typeless).
16501        let f16pv = fa_f16pv_on();
16502        let nw = if f16pv { fa512_wide_warps() } else { 2 };
16503        let hp = f16pv && fa512_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
16504        debug_assert!(!v_f16 || f16pv, "f16 V emitted without the door on");
16505        let mut vguard = self.fa_vf16_scratch.lock().unwrap();
16506        let vref: &CudaSlice<u8> = if f16pv && !v_f16 {
16507            // Fallback re-encode (producer emitted bf16); the emit lane normally hands f16.
16508            let n = t_kv * n_head_kv * head_dim;
16509            let need = n * 2;
16510            if vguard.as_ref().map(|b| b.len() < need).unwrap_or(true) {
16511                *vguard = Some(self.alloc_uninit::<u8>(need)?);
16512            }
16513            let dst = vguard.as_mut().unwrap();
16514            self.bf16_to_f16_into(vb, n, dst)?;
16515            vguard.as_ref().unwrap()
16516        } else {
16517            vb
16518        };
16519        let f = self.func(if hp {
16520            "fa_prefill_bf16_hd512_sp16h2"
16521        } else {
16522            match (f16pv, nw) {
16523                (true, 4) => "fa_prefill_bf16_hd512_sp16w4",
16524                (true, _) => "fa_prefill_bf16_hd512_sp16",
16525                _ => "fa_prefill_bf16_hd512_sp",
16526            }
16527        });
16528        let (nwarp, npart) = if hp {
16529            (4usize, 4usize)
16530        } else if nw > 2 {
16531            (nw, nw)
16532        } else {
16533            (2, 1)
16534        };
16535        // h2 drops sQ (Q register-resident) and doubles sP/sS/sL for the head pair.
16536        let shmem = if hp {
16537            (2 * (2 * BKS * head_dim + 2 * SP_M * BKS) + 4 * (2 * npart * SP_M * BKS + 2 * SP_M))
16538                as u32
16539        } else {
16540            (2 * (SP_M * head_dim + 2 * BKS * head_dim + SP_M * BKS)
16541                + 4 * (npart * SP_M * BKS + SP_M)) as u32
16542        };
16543        use cudarc::driver::sys::CUfunction_attribute_enum as A;
16544        f.set_attribute(
16545            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16546            shmem as i32,
16547        )?;
16548        let grid_y = if hp {
16549            (n_head / 2) as u32
16550        } else {
16551            n_head as u32
16552        };
16553        let cfg = LaunchConfig {
16554            grid_dim: ((t as u32).div_ceil(SP_M as u32), grid_y, 1),
16555            block_dim: (32, nwarp as u32, 1),
16556            shared_mem_bytes: shmem,
16557        };
16558        let (hd, nh, nhkv, ti, tkvi, cz) = (
16559            head_dim as i32,
16560            n_head as i32,
16561            n_head_kv as i32,
16562            t as i32,
16563            t_kv as i32,
16564            causal as i32,
16565        );
16566        let __s_b = self.gpu.stream();
16567        let mut b = __s_b.launch_builder(&f);
16568        b.arg(qb)
16569            .arg(kb)
16570            .arg(vref)
16571            .arg(o)
16572            .arg(&hd)
16573            .arg(&nh)
16574            .arg(&nhkv)
16575            .arg(&ti)
16576            .arg(&tkvi)
16577            .arg(&scale)
16578            .arg(&cz);
16579        unsafe {
16580            b.launch(cfg)?;
16581        }
16582        Ok(())
16583    }
16584
16585    /// hd512 FA prefill with the stage/sp arms FORCED — the kernel_check gate entry
16586    /// (`fa_prefill_hd512` picks the arms from MEMRA_FA512_STAGE / MEMRA_FA512_SP).
16587    #[allow(clippy::too_many_arguments)]
16588    pub fn fa_prefill_hd512_arm(
16589        &self,
16590        q: &CudaSlice<f32>,
16591        k: &CudaSlice<f32>,
16592        v: &CudaSlice<f32>,
16593        o: &mut CudaSlice<f32>,
16594        head_dim: usize,
16595        n_head: usize,
16596        n_head_kv: usize,
16597        t: usize,
16598        t_kv: usize,
16599        scale: f32,
16600        causal: bool,
16601        f32_stage: bool,
16602        sp: bool,
16603        f16pv: bool,
16604    ) -> Result<(), Box<dyn std::error::Error>> {
16605        debug_assert_eq!(head_dim, 512, "fa_prefill_hd512 is hd512 only");
16606        if sp && !f32_stage {
16607            // Single-pass: 16 q-rows/CTA, 2 warps, grid (ceil(T/16), n_head, 1).
16608            // smem: sQ[16][512] + sK[32][512] + sV[32][512] + sP[16][32] (bf16) + sS[16][32]+sL f32.
16609            // f16pv: sp16 kernel — f16 P + f16 P@V accum, V operand encoded f16.
16610            const SP_M: usize = 16;
16611            const BKS: usize = 32;
16612            let nw = if f16pv { fa512_wide_warps() } else { 2 };
16613            let hp = f16pv && fa512_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
16614            let f = self.func(if hp {
16615                "fa_prefill_bf16_hd512_sp16h2"
16616            } else {
16617                match (f16pv, nw) {
16618                    (true, 4) => "fa_prefill_bf16_hd512_sp16w4",
16619                    (true, _) => "fa_prefill_bf16_hd512_sp16",
16620                    _ => "fa_prefill_bf16_hd512_sp",
16621                }
16622            });
16623            let (nwarp, npart) = if hp {
16624                (4usize, 4usize)
16625            } else if nw > 2 {
16626                (nw, nw)
16627            } else {
16628                (2, 1)
16629            };
16630            let shmem = if hp {
16631                (2 * (2 * BKS * head_dim + 2 * SP_M * BKS)
16632                    + 4 * (2 * npart * SP_M * BKS + 2 * SP_M)) as u32
16633            } else {
16634                (2 * (SP_M * head_dim + 2 * BKS * head_dim + SP_M * BKS)
16635                    + 4 * (npart * SP_M * BKS + SP_M)) as u32
16636            };
16637            use cudarc::driver::sys::CUfunction_attribute_enum as A;
16638            f.set_attribute(
16639                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16640                shmem as i32,
16641            )?;
16642            let grid_y = if hp {
16643                (n_head / 2) as u32
16644            } else {
16645                n_head as u32
16646            };
16647            let cfg = LaunchConfig {
16648                grid_dim: ((t as u32).div_ceil(SP_M as u32), grid_y, 1),
16649                block_dim: (32, nwarp as u32, 1),
16650                shared_mem_bytes: shmem,
16651            };
16652            let (hd, nh, nhkv, ti, tkvi, cz) = (
16653                head_dim as i32,
16654                n_head as i32,
16655                n_head_kv as i32,
16656                t as i32,
16657                t_kv as i32,
16658                causal as i32,
16659            );
16660            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
16661            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
16662            let vb = if f16pv {
16663                self.f32_to_f16(v, t_kv * n_head_kv * head_dim)?
16664            } else {
16665                self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?
16666            };
16667            let __s_b = self.gpu.stream();
16668            let mut b = __s_b.launch_builder(&f);
16669            b.arg(&qb)
16670                .arg(&kb)
16671                .arg(&vb)
16672                .arg(o)
16673                .arg(&hd)
16674                .arg(&nh)
16675                .arg(&nhkv)
16676                .arg(&ti)
16677                .arg(&tkvi)
16678                .arg(&scale)
16679                .arg(&cz);
16680            unsafe {
16681                b.launch(cfg)?;
16682            }
16683            return Ok(());
16684        }
16685        const BLOCK_Q: usize = 32;
16686        const BK: usize = 32;
16687        const HALF: usize = 256;
16688        let f = self.func(if f32_stage {
16689            "fa_prefill_f32_hd512"
16690        } else {
16691            "fa_prefill_bf16_hd512"
16692        });
16693        // sQ[32][512] + sK[BK][512] + sV[BK][256] + sP[32][BK] (bf16) + sL[32] f32
16694        let shmem = (2 * (BLOCK_Q * head_dim + BK * head_dim + BK * HALF + BLOCK_Q * BK)
16695            + 4 * BLOCK_Q) as u32;
16696        use cudarc::driver::sys::CUfunction_attribute_enum as A;
16697        f.set_attribute(
16698            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16699            shmem as i32,
16700        )?;
16701        let cfg = LaunchConfig {
16702            grid_dim: (
16703                (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
16704                n_head as u32,
16705                2,
16706            ),
16707            block_dim: (32, 2, 1),
16708            shared_mem_bytes: shmem,
16709        };
16710        let (hd, nh, nhkv, ti, tkvi, cz) = (
16711            head_dim as i32,
16712            n_head as i32,
16713            n_head_kv as i32,
16714            t as i32,
16715            t_kv as i32,
16716            causal as i32,
16717        );
16718        if f32_stage {
16719            let __s_b = self.gpu.stream();
16720            let mut b = __s_b.launch_builder(&f);
16721            b.arg(q)
16722                .arg(k)
16723                .arg(v)
16724                .arg(o)
16725                .arg(&hd)
16726                .arg(&nh)
16727                .arg(&nhkv)
16728                .arg(&ti)
16729                .arg(&tkvi)
16730                .arg(&scale)
16731                .arg(&cz);
16732            unsafe {
16733                b.launch(cfg)?;
16734            }
16735        } else {
16736            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
16737            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
16738            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
16739            let __s_b = self.gpu.stream();
16740            let mut b = __s_b.launch_builder(&f);
16741            b.arg(&qb)
16742                .arg(&kb)
16743                .arg(&vb)
16744                .arg(o)
16745                .arg(&hd)
16746                .arg(&nh)
16747                .arg(&nhkv)
16748                .arg(&ti)
16749                .arg(&tkvi)
16750                .arg(&scale)
16751                .arg(&cz);
16752            unsafe {
16753                b.launch(cfg)?;
16754            }
16755        }
16756        Ok(())
16757    }
16758
16759    /// rope_neox2 with bf16 EMIT (31B glue lane): identical rope math/stores plus the post-rope
16760    /// values written as bf16 — the FA q/k operands come from this launch (bit-identical to the
16761    /// separate f32_to_bf16 the FA entries would run).
16762    #[allow(clippy::too_many_arguments)]
16763    pub fn rope_neox2_bf16e(
16764        &self,
16765        q: &mut CudaSlice<f32>,
16766        k: &mut CudaSlice<f32>,
16767        qb: &mut CudaSlice<u8>,
16768        kb: &mut CudaSlice<u8>,
16769        pos: &CudaSlice<i32>,
16770        head_dim: usize,
16771        n_dims: usize,
16772        nh_q: usize,
16773        nh_k: usize,
16774        n_tokens: usize,
16775        base: f32,
16776        freq_scale: f32,
16777        ff: Option<&CudaSlice<f32>>,
16778    ) -> Result<(), Box<dyn std::error::Error>> {
16779        let f = self.func("rope_neox2_bf16e_f32");
16780        let rows = ((nh_q + nh_k) * n_tokens) as u32;
16781        let cfg = LaunchConfig {
16782            grid_dim: (rows, 1, 1),
16783            block_dim: ((head_dim / 2) as u32, 1, 1),
16784            shared_mem_bytes: 0,
16785        };
16786        let theta_scale = base.powf(-2.0 / n_dims as f32);
16787        let (hd, nd, nhq, nhk, nt) = (
16788            head_dim as i32,
16789            n_dims as i32,
16790            nh_q as i32,
16791            nh_k as i32,
16792            n_tokens as i32,
16793        );
16794        let __s_b = self.gpu.stream();
16795        let mut b = __s_b.launch_builder(&f);
16796        match ff {
16797            Some(t) => {
16798                b.arg(&mut *q)
16799                    .arg(&mut *k)
16800                    .arg(&mut *qb)
16801                    .arg(&mut *kb)
16802                    .arg(pos)
16803                    .arg(&hd)
16804                    .arg(&nd)
16805                    .arg(&nhq)
16806                    .arg(&nhk)
16807                    .arg(&nt)
16808                    .arg(&theta_scale)
16809                    .arg(&freq_scale)
16810                    .arg(t);
16811                unsafe {
16812                    b.launch(cfg)?;
16813                }
16814            }
16815            None => {
16816                let null: u64 = 0;
16817                b.arg(&mut *q)
16818                    .arg(&mut *k)
16819                    .arg(&mut *qb)
16820                    .arg(&mut *kb)
16821                    .arg(pos)
16822                    .arg(&hd)
16823                    .arg(&nd)
16824                    .arg(&nhq)
16825                    .arg(&nhk)
16826                    .arg(&nt)
16827                    .arg(&theta_scale)
16828                    .arg(&freq_scale)
16829                    .arg(&null);
16830                unsafe {
16831                    b.launch(cfg)?;
16832                }
16833            }
16834        }
16835        Ok(())
16836    }
16837
16838    /// Flat f32 -> bf16 conversion into a fresh scratch buffer (2 bytes/elem). `n % 4 == 0`
16839    /// (float4 in, 4x bf16 out). Feeds the bf16-staged hd512 FA prefill.
16840    pub fn f32_to_bf16(
16841        &self,
16842        x: &CudaSlice<f32>,
16843        n: usize,
16844    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
16845        assert!(n % 4 == 0, "f32_to_bf16 requires n % 4 == 0, got {n}");
16846        let mut y = self.alloc_uninit::<u8>(n * 2)?;
16847        let f = self.func("f32_to_bf16_flat");
16848        let n_i = n as i64;
16849        let cfg = LaunchConfig {
16850            grid_dim: (((n / 4) as u32).div_ceil(256), 1, 1),
16851            block_dim: (256, 1, 1),
16852            shared_mem_bytes: 0,
16853        };
16854        let __s_b = self.gpu.stream();
16855        let mut b = __s_b.launch_builder(&f);
16856        b.arg(x).arg(&mut y).arg(&n_i);
16857        unsafe {
16858            b.launch(cfg)?;
16859        }
16860        Ok(y)
16861    }
16862
16863    pub fn f32_to_f16(
16864        &self,
16865        x: &CudaSlice<f32>,
16866        n: usize,
16867    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
16868        assert!(n % 4 == 0, "f32_to_f16 requires n % 4 == 0, got {n}");
16869        let mut y = self.alloc_uninit::<u8>(n * 2)?;
16870        let f = self.func("f32_to_f16_flat");
16871        let n_i = n as i64;
16872        let cfg = LaunchConfig {
16873            grid_dim: (((n / 4) as u32).div_ceil(256), 1, 1),
16874            block_dim: (256, 1, 1),
16875            shared_mem_bytes: 0,
16876        };
16877        let __s_b = self.gpu.stream();
16878        let mut b = __s_b.launch_builder(&f);
16879        b.arg(x).arg(&mut y).arg(&n_i);
16880        unsafe {
16881            b.launch(cfg)?;
16882        }
16883        Ok(y)
16884    }
16885
16886    /// bf16 bytes -> f16 bytes, n elements (the f16-P/V door's V re-encode on the emit lane).
16887    pub fn bf16_to_f16(
16888        &self,
16889        xb: &CudaSlice<u8>,
16890        n: usize,
16891    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
16892        let mut y = self.alloc_uninit::<u8>(n * 2)?;
16893        self.bf16_to_f16_into(xb, n, &mut y)?;
16894        Ok(y)
16895    }
16896
16897    /// Same conversion into a caller-owned (pooled) buffer; `y.len() >= n*2`.
16898    pub fn bf16_to_f16_into(
16899        &self,
16900        xb: &CudaSlice<u8>,
16901        n: usize,
16902        y: &mut CudaSlice<u8>,
16903    ) -> Result<(), Box<dyn std::error::Error>> {
16904        assert!(n % 2 == 0, "bf16_to_f16 requires n % 2 == 0, got {n}");
16905        assert!(y.len() >= n * 2);
16906        let f = self.func("bf16_to_f16_flat");
16907        let n2 = (n / 2) as i64;
16908        let cfg = LaunchConfig {
16909            grid_dim: (((n / 2) as u32).div_ceil(256), 1, 1),
16910            block_dim: (256, 1, 1),
16911            shared_mem_bytes: 0,
16912        };
16913        let __s_b = self.gpu.stream();
16914        let mut b = __s_b.launch_builder(&f);
16915        b.arg(xb).arg(y).arg(&n2);
16916        unsafe {
16917            b.launch(cfg)?;
16918        }
16919        Ok(())
16920    }
16921
16922    /// task #18 (attn side): varlen FA — bf16 K/V mirrors (2 launches) + ONE
16923    /// fa_prefill_bf16kv launch for every fresh sequence. Same per-block math as the
16924    /// per-seq path (bit-gateable). Caller guarantees: fresh causal (T_kv == T),
16925    /// head_dim in {256, 128}, bf16kv lane on.
16926    #[allow(clippy::too_many_arguments)]
16927    pub fn fa_prefill_vl8(
16928        &self,
16929        seqs: &[FaSeqVl],
16930        head_dim: usize,
16931        n_head: usize,
16932        n_head_kv: usize,
16933        scale: f32,
16934    ) -> Result<(), Box<dyn std::error::Error>> {
16935        const BK: usize = 32;
16936        let b = seqs.len();
16937        assert!(b >= 1 && b <= 8);
16938        let mut packed = [FaSeqVl::default(); 8];
16939        packed[..b].copy_from_slice(seqs);
16940        let v = FaVl8(packed);
16941        let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
16942        let ept = (n_head_kv * head_dim) as i32;
16943        {
16944            let f = self.func("fa_mirror_vl");
16945            let max_n = (max_t as i64) * ept as i64;
16946            let blocks = ((max_n as u32).div_ceil(4)).div_ceil(256);
16947            for which in 0..2i32 {
16948                let cfg = LaunchConfig {
16949                    grid_dim: (blocks, 1, b as u32),
16950                    block_dim: (256, 1, 1),
16951                    shared_mem_bytes: 0,
16952                };
16953                let __s_lb = self.gpu.stream();
16954                let mut lb = __s_lb.launch_builder(&f);
16955                lb.arg(&v).arg(&ept).arg(&which);
16956                unsafe {
16957                    lb.launch(cfg)?;
16958                }
16959            }
16960        }
16961        let hd_sfx = fa_hd_suffix(head_dim)?;
16962        let f = self.func(&format!("fa_prefill_bf16kv_vl{hd_sfx}"));
16963        let block_q = 64usize;
16964        let kv_stages = 2usize;
16965        let shmem = (2 * (kv_stages * 2 * BK * head_dim + block_q * BK)
16966            + 4 * (block_q * BK + 2 * block_q)) as u32;
16967        use cudarc::driver::sys::CUfunction_attribute_enum as A;
16968        f.set_attribute(
16969            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16970            shmem as i32,
16971        )?;
16972        let cfg = LaunchConfig {
16973            grid_dim: (max_t.div_ceil(block_q as u32), n_head as u32, b as u32),
16974            block_dim: (32, 4, 1),
16975            shared_mem_bytes: shmem,
16976        };
16977        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
16978        let __s_lb = self.gpu.stream();
16979        let mut lb = __s_lb.launch_builder(&f);
16980        lb.arg(&v).arg(&hd).arg(&nh).arg(&nhkv).arg(&scale);
16981        unsafe {
16982            lb.launch(cfg)?;
16983        }
16984        Ok(())
16985    }
16986
16987    /// task #18 (attn pre-FA): varlen split + QK-norm + RoPE + KV-append — FOUR launches
16988    /// for every fresh sequence (was 6 x B, plus the q/k/v split copies which the view
16989    /// inputs remove entirely). Fresh-only (append at t0=0, RoPE pos = token index).
16990    #[allow(clippy::too_many_arguments)]
16991    pub fn attn_pre_vl8(
16992        &self,
16993        seqs: &[AttnPreVl],
16994        wq: &CudaSlice<f32>,
16995        wk: &CudaSlice<f32>,
16996        head_dim: usize,
16997        rope_dims: usize,
16998        n_head: usize,
16999        n_head_kv: usize,
17000        eps: f32,
17001        freq_base: f32,
17002        freq_scale: f32,
17003        kv_dim_k: usize,
17004        kv_dim_v: usize,
17005        k_tok_bytes: usize,
17006        v_tok_bytes: usize,
17007    ) -> Result<(), Box<dyn std::error::Error>> {
17008        let b = seqs.len();
17009        assert!(b >= 1 && b <= 8);
17010        let mut packed = [AttnPreVl::default(); 8];
17011        packed[..b].copy_from_slice(seqs);
17012        let v = AttnPreVl8(packed);
17013        let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
17014        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
17015        {
17016            let f = self.func("q_gate_split_vl");
17017            let n = max_t * (n_head * head_dim) as u32;
17018            let cfg = LaunchConfig {
17019                grid_dim: (n.div_ceil(256), 1, b as u32),
17020                block_dim: (256, 1, 1),
17021                shared_mem_bytes: 0,
17022            };
17023            let __s_lb = self.gpu.stream();
17024            let mut lb = __s_lb.launch_builder(&f);
17025            lb.arg(&v).arg(&hd).arg(&nh);
17026            unsafe {
17027                lb.launch(cfg)?;
17028            }
17029        }
17030        {
17031            let f = self.func("attn_rms_vl");
17032            let cfg = LaunchConfig {
17033                grid_dim: (max_t * n_head as u32, 2, b as u32),
17034                block_dim: (rms_block(), 1, 1),
17035                shared_mem_bytes: 0,
17036            };
17037            let __s_lb = self.gpu.stream();
17038            let mut lb = __s_lb.launch_builder(&f);
17039            lb.arg(&v)
17040                .arg(wq)
17041                .arg(wk)
17042                .arg(&hd)
17043                .arg(&nh)
17044                .arg(&nhkv)
17045                .arg(&eps);
17046            unsafe {
17047                lb.launch(cfg)?;
17048            }
17049        }
17050        {
17051            let f = self.func("attn_rope_vl");
17052            let theta_scale = freq_base.powf(-2.0 / rope_dims as f32);
17053            let nd = rope_dims as i32;
17054            let cfg = LaunchConfig {
17055                grid_dim: (max_t * n_head as u32, 2, b as u32),
17056                block_dim: ((head_dim / 2) as u32, 1, 1),
17057                shared_mem_bytes: 0,
17058            };
17059            let __s_lb = self.gpu.stream();
17060            let mut lb = __s_lb.launch_builder(&f);
17061            lb.arg(&v)
17062                .arg(&hd)
17063                .arg(&nd)
17064                .arg(&nh)
17065                .arg(&nhkv)
17066                .arg(&theta_scale)
17067                .arg(&freq_scale);
17068            unsafe {
17069                lb.launch(cfg)?;
17070            }
17071        }
17072        {
17073            let f = self.func("append_kv_vl");
17074            let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
17075            let cfg = LaunchConfig {
17076                grid_dim: (nblk, max_t, b as u32),
17077                block_dim: (32, 1, 1),
17078                shared_mem_bytes: 0,
17079            };
17080            let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
17081            let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
17082            let __s_lb = self.gpu.stream();
17083            let mut lb = __s_lb.launch_builder(&f);
17084            lb.arg(&v).arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
17085            unsafe {
17086                lb.launch(cfg)?;
17087            }
17088        }
17089        Ok(())
17090    }
17091
17092    /// FA prefill where K/V are QUANTIZED CudaViews into the resident byte KV cache (the T=K verify
17093    /// path, MTP-PLAN §D.3). Uses `fa_prefill_q` (inline-dequant during stage-to-smem). The view's
17094    /// base+offset pointer is honored; the kernel reads [0..t_kv*tok_bytes). Q is the T fresh query
17095    /// rows; t = T, t_kv = cache len. k_tok_bytes/v_tok_bytes are the per-token byte strides.
17096    pub fn fa_prefill_view(
17097        &self,
17098        q: &CudaSlice<f32>,
17099        k: &cudarc::driver::CudaView<u8>,
17100        v: &cudarc::driver::CudaView<u8>,
17101        o: &mut CudaSlice<f32>,
17102        head_dim: usize,
17103        n_head: usize,
17104        n_head_kv: usize,
17105        t: usize,
17106        t_kv: usize,
17107        scale: f32,
17108        causal: bool,
17109        k_tok_bytes: usize,
17110        v_tok_bytes: usize,
17111        g: bool,
17112    ) -> Result<(), Box<dyn std::error::Error>> {
17113        if portable_mma_gated() {
17114            return self.sdpa_naive_quantized_view(
17115                q,
17116                k,
17117                v,
17118                o,
17119                head_dim,
17120                n_head,
17121                n_head_kv,
17122                t,
17123                t_kv,
17124                scale,
17125                causal,
17126                k_tok_bytes,
17127                v_tok_bytes,
17128            );
17129        }
17130        const BLOCK_Q: usize = 64;
17131        const BK: usize = 32;
17132        // g = e4m3 cache: the kernel parses via DQ_K_ELEM/DQ_V_ELEM (format macros) — the
17133        // kf8vf8-module stamp reads fp8 with the identical MMA/softmax/PV body.
17134        let name = format!("fa_prefill_q{}", fa_hd_suffix(head_dim)?);
17135        let f = if g {
17136            self.func_g(&name)
17137        } else {
17138            self.func(&name)
17139        };
17140        let shmem =
17141            (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
17142        use cudarc::driver::sys::CUfunction_attribute_enum as A;
17143        f.set_attribute(
17144            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
17145            shmem as i32,
17146        )?;
17147        let cfg = LaunchConfig {
17148            grid_dim: (
17149                (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
17150                n_head as u32,
17151                1,
17152            ),
17153            block_dim: (32, 4, 1),
17154            shared_mem_bytes: shmem,
17155        };
17156        let (hd, nh, nhkv, ti, tkvi, cz) = (
17157            head_dim as i32,
17158            n_head as i32,
17159            n_head_kv as i32,
17160            t as i32,
17161            t_kv as i32,
17162            causal as i32,
17163        );
17164        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
17165        let __s_b = self.gpu.stream();
17166        let mut b = __s_b.launch_builder(&f);
17167        b.arg(q)
17168            .arg(k)
17169            .arg(v)
17170            .arg(o)
17171            .arg(&hd)
17172            .arg(&nh)
17173            .arg(&nhkv)
17174            .arg(&ti)
17175            .arg(&tkvi)
17176            .arg(&scale)
17177            .arg(&cz)
17178            .arg(&ktb)
17179            .arg(&vtb);
17180        unsafe {
17181            b.launch(cfg)?;
17182        }
17183        Ok(())
17184    }
17185
17186    /// ARC B (2026-07-05): dequant-once chunk-prime FA. Same contract as `fa_prefill_view`, but
17187    /// instead of every (q-block, head) CTA re-dequanting the whole quantized KV stream inline
17188    /// (T/64 x n_head redundant at chunk prime — 30.5% of the 32k prime wall), dequant the full
17189    /// [t_kv, kv_dim] K and V ONCE into a resident bf16 workspace (fa_dequant_kv_ws_bf16), then
17190    /// run `fa_prefill_qw` (the bf16-workspace twin) over it. EXACT: the workspace holds the same
17191    /// __float2bfloat16(dq_*_elem(...)) values fa_prefill_q stages to smem, and the twin's MMA/
17192    /// softmax/PV code is byte-identical -> bit-identical O (kernel_check pins bitdiff=0).
17193    /// The workspace allocation is REUSED across layers/chunks (grown to the largest shape);
17194    /// contents are rewritten per call. MEMRA_PRIME_DEQW=0 falls back to fa_prefill_view (callers gate).
17195    #[allow(clippy::too_many_arguments)]
17196    pub fn fa_prefill_view_ws(
17197        &self,
17198        q: &CudaSlice<f32>,
17199        k: &cudarc::driver::CudaView<u8>,
17200        v: &cudarc::driver::CudaView<u8>,
17201        o: &mut CudaSlice<f32>,
17202        head_dim: usize,
17203        n_head: usize,
17204        n_head_kv: usize,
17205        t: usize,
17206        t_kv: usize,
17207        scale: f32,
17208        causal: bool,
17209        k_tok_bytes: usize,
17210        v_tok_bytes: usize,
17211        g: bool,
17212    ) -> Result<(), Box<dyn std::error::Error>> {
17213        if portable_mma_gated() {
17214            return self.sdpa_naive_quantized_view(
17215                q,
17216                k,
17217                v,
17218                o,
17219                head_dim,
17220                n_head,
17221                n_head_kv,
17222                t,
17223                t_kv,
17224                scale,
17225                causal,
17226                k_tok_bytes,
17227                v_tok_bytes,
17228            );
17229        }
17230        const BLOCK_Q: usize = 64;
17231        const BK: usize = 32;
17232        let kv_dim_k = n_head_kv * head_dim;
17233        let kv_dim_v = n_head_kv * head_dim;
17234        let k_ws_bytes = t_kv * kv_dim_k * 2; // bf16
17235        let v_ws_bytes = t_kv * kv_dim_v * 2;
17236        // Lock held across BOTH launches: enqueue-only (µs), all compute serializes on gpu.stream.
17237        let mut guard = self.prime_deqw_ws.lock().unwrap();
17238        let need_grow = match guard.as_ref() {
17239            Some((kw, vw)) => kw.len() < k_ws_bytes || vw.len() < v_ws_bytes,
17240            None => true,
17241        };
17242        if need_grow {
17243            let grow = |cur: usize, need: usize| if cur >= need { cur } else { need };
17244            let (ck, cv) = guard
17245                .as_ref()
17246                .map(|(a, b)| (a.len(), b.len()))
17247                .unwrap_or((0, 0));
17248            *guard = Some((
17249                self.alloc_u8(grow(ck, k_ws_bytes))?,
17250                self.alloc_u8(grow(cv, v_ws_bytes))?,
17251            ));
17252        }
17253        let (kw, vw) = guard.as_mut().unwrap();
17254        // pass 1: dequant K+V once into the bf16 workspace (grid-stride, 1 thread/elem)
17255        {
17256            // only THIS pass parses KV bytes — pass 2 reads the bf16 workspace (format-free).
17257            let f = if g {
17258                self.func_g("fa_dequant_kv_ws_bf16")
17259            } else {
17260                self.func("fa_dequant_kv_ws_bf16")
17261            };
17262            let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
17263            let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
17264            let cfg = LaunchConfig {
17265                grid_dim: (nblk.max(1), 1, 1),
17266                block_dim: (256, 1, 1),
17267                shared_mem_bytes: 0,
17268            };
17269            let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
17270            let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
17271            let __s_b = self.gpu.stream();
17272            let mut b = __s_b.launch_builder(&f);
17273            b.arg(k)
17274                .arg(v)
17275                .arg(&mut *kw)
17276                .arg(&mut *vw)
17277                .arg(&kdk)
17278                .arg(&kdv)
17279                .arg(&tkvi)
17280                .arg(&ktb)
17281                .arg(&vtb);
17282            unsafe {
17283                b.launch(cfg)?;
17284            }
17285        }
17286        // pass 2: the bf16-workspace prefill twin (same tile sizes/loop structure as fa_prefill_q).
17287        // DEFAULT: cp.async double-buffered staging twin (fa_prefill_qw_db, +32KB smem for the
17288        // second K/V tile pair, 1 CTA/SM): overlaps tile n+1's L2->smem copy with tile n's MMA.
17289        // Bit-identical output (staging is a pure byte copy; kernel_check pins bitdiff=0 under
17290        // both twins). A/B (27B g7e, N=3): 32k prime 17.10->16.51s, 16k 9.09->8.65s — the copy
17291        // latency hides behind the MMA pipe and beats the 2-CTA/SM occupancy of the sync twin.
17292        // MEMRA_PRIME_DEQW_DB=0 falls back to the single-buffer twin.
17293        let db = std::env::var("MEMRA_PRIME_DEQW_DB")
17294            .map(|v| v != "0")
17295            .unwrap_or(true);
17296        {
17297            let hd_sfx = fa_hd_suffix(head_dim)?;
17298            let f = self.func(&format!(
17299                "fa_prefill_qw{}{hd_sfx}",
17300                if db { "_db" } else { "" }
17301            ));
17302            let shmem = if db {
17303                // 4x KV tile buffers (bf16) + sP (bf16) + sL (f32)
17304                (2 * (4 * BK * head_dim + BLOCK_Q * BK) + 4 * BLOCK_Q) as u32
17305            } else {
17306                (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32
17307            };
17308            use cudarc::driver::sys::CUfunction_attribute_enum as A;
17309            f.set_attribute(
17310                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
17311                shmem as i32,
17312            )?;
17313            let cfg = LaunchConfig {
17314                grid_dim: (
17315                    (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
17316                    n_head as u32,
17317                    1,
17318                ),
17319                block_dim: (32, 4, 1),
17320                shared_mem_bytes: shmem,
17321            };
17322            let (hd, nh, nhkv, ti, tkvi, cz) = (
17323                head_dim as i32,
17324                n_head as i32,
17325                n_head_kv as i32,
17326                t as i32,
17327                t_kv as i32,
17328                causal as i32,
17329            );
17330            let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
17331            let __s_b = self.gpu.stream();
17332            let mut b = __s_b.launch_builder(&f);
17333            b.arg(q)
17334                .arg(&*kw)
17335                .arg(&*vw)
17336                .arg(o)
17337                .arg(&hd)
17338                .arg(&nh)
17339                .arg(&nhkv)
17340                .arg(&ti)
17341                .arg(&tkvi)
17342                .arg(&scale)
17343                .arg(&cz)
17344                .arg(&kdk)
17345                .arg(&kdv);
17346            unsafe {
17347                b.launch(cfg)?;
17348            }
17349        }
17350        Ok(())
17351    }
17352
17353    /// WINDOWED `fa_prefill_view_ws` twin at head_dim 128 (lane/pp-prefill 2026-08-07):
17354    /// step35's SWA prefill (win=512, 33 of 45 layers) previously had NO windowed FA prefill
17355    /// stamp — every windowed twin was hd256-only — and took `sdpa_naive_w_quantized_view`,
17356    /// the f32 floor, at 565 ms/layer on a pp4096 where the hd128 FA family does the harder
17357    /// causal-4096 in 3.3 ms (41% of the whole prime; research/pp-prefill-20260807 anatomy).
17358    /// Same two-pass shape as the unwindowed function: dequant K/V ONCE into the resident
17359    /// bf16 workspace, then the windowed qw kernel (`fa_prefill_qw_db_w_hd128`, cp.async
17360    /// double-buffered; MEMRA_PRIME_DEQW_DB=0 selects the single-buffer twin). The window
17361    /// mask is `fa_prefill_f32_body`'s exact predicate; `window == 0` is bit-identical to
17362    /// `fa_prefill_view_ws` by construction (default-arg body). NEW NUMERIC CLASS vs the
17363    /// f32 floor on SWA rows (bf16 MMA online-softmax vs f32 serial softmax) — adoption is
17364    /// gated by the full battery, and the class must change UNIFORMLY for a whole request
17365    /// (kernel selection keys on seq_end, never per chunk — the chunkfix law).
17366    /// hd128-only deliberately: the only windowed-prefill consumer at another head_dim is
17367    /// gemma4 (hd256), which already has `fa_prefill_w_f32`.
17368    #[allow(clippy::too_many_arguments)]
17369    pub fn fa_prefill_view_ws_w_hd128(
17370        &self,
17371        q: &CudaSlice<f32>,
17372        k: &cudarc::driver::CudaView<u8>,
17373        v: &cudarc::driver::CudaView<u8>,
17374        o: &mut CudaSlice<f32>,
17375        head_dim: usize,
17376        n_head: usize,
17377        n_head_kv: usize,
17378        t: usize,
17379        t_kv: usize,
17380        scale: f32,
17381        causal: bool,
17382        window: usize,
17383        k_tok_bytes: usize,
17384        v_tok_bytes: usize,
17385    ) -> Result<(), Box<dyn std::error::Error>> {
17386        assert_eq!(
17387            head_dim, 128,
17388            "fa_prefill_view_ws_w_hd128: only the hd128 twin is stamped"
17389        );
17390        if portable_mma_gated() {
17391            return self.sdpa_naive_w_quantized_view(
17392                q,
17393                k,
17394                v,
17395                o,
17396                head_dim,
17397                n_head,
17398                n_head_kv,
17399                t,
17400                t_kv,
17401                scale,
17402                causal,
17403                window,
17404                k_tok_bytes,
17405                v_tok_bytes,
17406            );
17407        }
17408        const BLOCK_Q: usize = 64;
17409        const BK: usize = 32;
17410        let kv_dim_k = n_head_kv * head_dim;
17411        let kv_dim_v = n_head_kv * head_dim;
17412        let k_ws_bytes = t_kv * kv_dim_k * 2; // bf16
17413        let v_ws_bytes = t_kv * kv_dim_v * 2;
17414        let mut guard = self.prime_deqw_ws.lock().unwrap();
17415        let need_grow = match guard.as_ref() {
17416            Some((kw, vw)) => kw.len() < k_ws_bytes || vw.len() < v_ws_bytes,
17417            None => true,
17418        };
17419        if need_grow {
17420            let grow = |cur: usize, need: usize| if cur >= need { cur } else { need };
17421            let (ck, cv) = guard
17422                .as_ref()
17423                .map(|(a, b)| (a.len(), b.len()))
17424                .unwrap_or((0, 0));
17425            *guard = Some((
17426                self.alloc_u8(grow(ck, k_ws_bytes))?,
17427                self.alloc_u8(grow(cv, v_ws_bytes))?,
17428            ));
17429        }
17430        let (kw, vw) = guard.as_mut().unwrap();
17431        // pass 1: dequant K+V once into the bf16 workspace (identical to fa_prefill_view_ws —
17432        // the workspace bytes are the SAME __float2bfloat16(dq(...)) values either way).
17433        {
17434            let f = self.func("fa_dequant_kv_ws_bf16");
17435            let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
17436            let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
17437            let cfg = LaunchConfig {
17438                grid_dim: (nblk.max(1), 1, 1),
17439                block_dim: (256, 1, 1),
17440                shared_mem_bytes: 0,
17441            };
17442            let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
17443            let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
17444            let __s_b = self.gpu.stream();
17445            let mut b = __s_b.launch_builder(&f);
17446            b.arg(k)
17447                .arg(v)
17448                .arg(&mut *kw)
17449                .arg(&mut *vw)
17450                .arg(&kdk)
17451                .arg(&kdv)
17452                .arg(&tkvi)
17453                .arg(&ktb)
17454                .arg(&vtb);
17455            unsafe {
17456                b.launch(cfg)?;
17457            }
17458        }
17459        // pass 2: the WINDOWED qw twin (db default, same as the unwindowed wrapper).
17460        let db = std::env::var("MEMRA_PRIME_DEQW_DB")
17461            .map(|v| v != "0")
17462            .unwrap_or(true);
17463        {
17464            let f = self.func(if db {
17465                "fa_prefill_qw_db_w_hd128"
17466            } else {
17467                "fa_prefill_qw_w_hd128"
17468            });
17469            let shmem = if db {
17470                (2 * (4 * BK * head_dim + BLOCK_Q * BK) + 4 * BLOCK_Q) as u32
17471            } else {
17472                (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32
17473            };
17474            use cudarc::driver::sys::CUfunction_attribute_enum as A;
17475            f.set_attribute(
17476                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
17477                shmem as i32,
17478            )?;
17479            let cfg = LaunchConfig {
17480                grid_dim: (
17481                    (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
17482                    n_head as u32,
17483                    1,
17484                ),
17485                block_dim: (32, 4, 1),
17486                shared_mem_bytes: shmem,
17487            };
17488            let (hd, nh, nhkv, ti, tkvi, cz) = (
17489                head_dim as i32,
17490                n_head as i32,
17491                n_head_kv as i32,
17492                t as i32,
17493                t_kv as i32,
17494                causal as i32,
17495            );
17496            let (kdk, kdv, wnd) = (kv_dim_k as i32, kv_dim_v as i32, window as i32);
17497            let __s_b = self.gpu.stream();
17498            let mut b = __s_b.launch_builder(&f);
17499            b.arg(q)
17500                .arg(&*kw)
17501                .arg(&*vw)
17502                .arg(o)
17503                .arg(&hd)
17504                .arg(&nh)
17505                .arg(&nhkv)
17506                .arg(&ti)
17507                .arg(&tkvi)
17508                .arg(&scale)
17509                .arg(&cz)
17510                .arg(&kdk)
17511                .arg(&kdv)
17512                .arg(&wnd);
17513            unsafe {
17514                b.launch(cfg)?;
17515            }
17516        }
17517        Ok(())
17518    }
17519
17520    /// FA decode (T=1 split-K) over the resident QUANTIZED KV cache (q8_0 K / q5_1 V) as u8 views.
17521    /// Replaces sdpa_naive_view for decode; inline-dequants per element. k_tok_bytes/v_tok_bytes are
17522    /// the per-token byte strides (differ: q8_0=34*nblk, q5_1=24*nblk per token).
17523    pub fn fa_decode(
17524        &self,
17525        q: &CudaSlice<f32>,
17526        k: &cudarc::driver::CudaView<u8>,
17527        v: &cudarc::driver::CudaView<u8>,
17528        o: &mut CudaSlice<f32>,
17529        head_dim: usize,
17530        n_head: usize,
17531        n_head_kv: usize,
17532        t_kv: usize,
17533        scale: f32,
17534        k_tok_bytes: usize,
17535        v_tok_bytes: usize,
17536    ) -> Result<(), Box<dyn std::error::Error>> {
17537        self.fa_decode_kvmod(
17538            q,
17539            k,
17540            v,
17541            o,
17542            head_dim,
17543            n_head,
17544            n_head_kv,
17545            t_kv,
17546            scale,
17547            k_tok_bytes,
17548            v_tok_bytes,
17549            false,
17550        )
17551    }
17552
17553    /// `fa_decode` with an explicit fp8-module flag (`g`): gemma windowed layers under
17554    /// MEMRA_GEMMA_WKV read an e4m3 cache — every kernel must come from the kf8vf8 module
17555    /// and the v4 lane (q8_0-hardcoded staging) is excluded.
17556    #[allow(clippy::too_many_arguments)]
17557    /// UNIFIED scalar decode launch (fa_decode_f32, nullable-ctr): ONE symbol for host-len
17558    /// (kvmod eager) and device-len (graph/stream) callers — the textually-identical f32_dc
17559    /// twin compiled apart and its ULP drift flipped 31B verify argmaxes (2026-07-12).
17560    #[allow(clippy::too_many_arguments)]
17561    #[allow(clippy::too_many_arguments)]
17562    fn fa_decode_scalar_unified(
17563        &self,
17564        q: &cudarc::driver::CudaView<f32>,
17565        k: &cudarc::driver::CudaView<u8>,
17566        v: &cudarc::driver::CudaView<u8>,
17567        o: &mut cudarc::driver::CudaViewMut<f32>,
17568        head_dim: usize,
17569        n_head: usize,
17570        n_head_kv: usize,
17571        t_kv_host: usize,
17572        t_kv_dev: Option<&CudaSlice<i32>>,
17573        scale: f32,
17574        n_splits: usize,
17575        split_keys: usize,
17576        k_tok_bytes: usize,
17577        v_tok_bytes: usize,
17578        g: bool,
17579        part_o: &mut CudaSlice<f32>,
17580        part_m: &mut CudaSlice<f32>,
17581        part_l: &mut CudaSlice<f32>,
17582        q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
17583    ) -> Result<(), Box<dyn std::error::Error>> {
17584        let f = if g {
17585            self.func_g("fa_decode_f32")
17586        } else {
17587            self.fa_func("fa_decode_f32", head_dim)
17588        };
17589        let cfg = LaunchConfig {
17590            grid_dim: (n_head as u32, n_splits as u32, 1),
17591            block_dim: (head_dim as u32, 1, 1),
17592            shared_mem_bytes: (4 * (head_dim + 32)) as u32,
17593        };
17594        let (hd, nh, nhkv, nsp) = (
17595            head_dim as i32,
17596            n_head as i32,
17597            n_head_kv as i32,
17598            n_splits as i32,
17599        );
17600        let (ktb, vtb, tkvi, ski) = (
17601            k_tok_bytes as i64,
17602            v_tok_bytes as i64,
17603            t_kv_host as i32,
17604            split_keys as i32,
17605        );
17606        let __s_b = self.gpu.stream();
17607        let mut b = __s_b.launch_builder(&f);
17608        match t_kv_dev {
17609            Some(d) => {
17610                b.arg(q)
17611                    .arg(k)
17612                    .arg(v)
17613                    .arg(&mut *part_o)
17614                    .arg(&mut *part_m)
17615                    .arg(&mut *part_l)
17616                    .arg(&hd)
17617                    .arg(&nh)
17618                    .arg(&nhkv)
17619                    .arg(&tkvi)
17620                    .arg(d)
17621                    .arg(&scale)
17622                    .arg(&nsp)
17623                    .arg(&ski)
17624                    .arg(&ktb)
17625                    .arg(&vtb);
17626                unsafe {
17627                    b.launch(cfg)?;
17628                }
17629            }
17630            None => {
17631                let null: u64 = 0;
17632                b.arg(q)
17633                    .arg(k)
17634                    .arg(v)
17635                    .arg(&mut *part_o)
17636                    .arg(&mut *part_m)
17637                    .arg(&mut *part_l)
17638                    .arg(&hd)
17639                    .arg(&nh)
17640                    .arg(&nhkv)
17641                    .arg(&tkvi)
17642                    .arg(&null)
17643                    .arg(&scale)
17644                    .arg(&nsp)
17645                    .arg(&ski)
17646                    .arg(&ktb)
17647                    .arg(&vtb);
17648                unsafe {
17649                    b.launch(cfg)?;
17650                }
17651            }
17652        }
17653        let cfg2 = LaunchConfig {
17654            grid_dim: (n_head as u32, 1, 1),
17655            block_dim: (head_dim as u32, 1, 1),
17656            shared_mem_bytes: 0,
17657        };
17658        if let Some((oq, od)) = q8_out {
17659            // wave-5b: q8-emitting combine — the wo matmul_pre consumes the pair directly.
17660            let fc = if g {
17661                self.func_g("fa_decode_combine_q8_1")
17662            } else {
17663                self.fa_func("fa_decode_combine_q8_1", head_dim)
17664            };
17665            let __s_b2 = self.gpu.stream();
17666            let mut b2 = __s_b2.launch_builder(&fc);
17667            b2.arg(&*part_o)
17668                .arg(&*part_m)
17669                .arg(&*part_l)
17670                .arg(oq)
17671                .arg(od)
17672                .arg(&hd)
17673                .arg(&nh)
17674                .arg(&nsp);
17675            unsafe {
17676                b2.launch(cfg2)?;
17677            }
17678            return Ok(());
17679        }
17680        let fc = if g {
17681            self.func_g("fa_decode_combine_f32")
17682        } else {
17683            self.fa_func("fa_decode_combine_f32", head_dim)
17684        };
17685        let __s_b2 = self.gpu.stream();
17686        let mut b2 = __s_b2.launch_builder(&fc);
17687        b2.arg(&*part_o)
17688            .arg(&*part_m)
17689            .arg(&*part_l)
17690            .arg(o)
17691            .arg(&hd)
17692            .arg(&nh)
17693            .arg(&nsp);
17694        unsafe {
17695            b2.launch(cfg2)?;
17696        }
17697        Ok(())
17698    }
17699
17700    pub fn fa_decode_kvmod(
17701        &self,
17702        q: &CudaSlice<f32>,
17703        k: &cudarc::driver::CudaView<u8>,
17704        v: &cudarc::driver::CudaView<u8>,
17705        o: &mut CudaSlice<f32>,
17706        head_dim: usize,
17707        n_head: usize,
17708        n_head_kv: usize,
17709        t_kv: usize,
17710        scale: f32,
17711        k_tok_bytes: usize,
17712        v_tok_bytes: usize,
17713        g: bool,
17714    ) -> Result<(), Box<dyn std::error::Error>> {
17715        let q_view = q.as_view();
17716        let mut o_view = o.as_view_mut();
17717        self.fa_decode_kvmod_view(
17718            &q_view,
17719            k,
17720            v,
17721            &mut o_view,
17722            head_dim,
17723            n_head,
17724            n_head_kv,
17725            t_kv,
17726            scale,
17727            k_tok_bytes,
17728            v_tok_bytes,
17729            g,
17730        )
17731    }
17732
17733    /// Row-view entry into `fa_decode_kvmod`. The kernel sees the selected Q/output rows as its
17734    /// base pointers, so the launch geometry and arithmetic are identical to the owned-slice entry.
17735    /// Batched fallback callers use this to avoid materializing rows around an otherwise unchanged
17736    /// per-session KV view and FA launch.
17737    #[allow(clippy::too_many_arguments)]
17738    pub fn fa_decode_kvmod_view(
17739        &self,
17740        q: &cudarc::driver::CudaView<f32>,
17741        k: &cudarc::driver::CudaView<u8>,
17742        v: &cudarc::driver::CudaView<u8>,
17743        o: &mut cudarc::driver::CudaViewMut<f32>,
17744        head_dim: usize,
17745        n_head: usize,
17746        n_head_kv: usize,
17747        t_kv: usize,
17748        scale: f32,
17749        k_tok_bytes: usize,
17750        v_tok_bytes: usize,
17751        g: bool,
17752    ) -> Result<(), Box<dyn std::error::Error>> {
17753        // PERF-4: the warp-per-token vec path replaces the scalar element-per-thread fa_decode_f32 —
17754        // warp-per-token fa_decode_vec_q (grid=(n_head_kv,n_splits), block=(32,gqa_ratio)).
17755        // The block dequants each KV tile ONCE into smem (bf16) and broadcasts to all gqa Q-head
17756        // warps -> each KV byte leaves HBM/L2 ~1x/group (vs 4x). ARGS identical; func/grid/block/
17757        // smem/n_splits differ. fa_decode_f32 stays the bit-reference fallback. Combine is shared.
17758        //
17759        // SPLIT-K: the scalar path has grid.x=n_head (32) blocks; the vec path only has
17760        // grid.x=n_head_kv (8). To avoid starving the GPU at mid ctx, the vec path splits MORE
17761        // aggressively (64 keys/split vs 256) so grid.y rises and 8*n_splits fills the SMs.
17762        // At VERY short ctx (t_kv<96) even 1 split can't fill the GPU from 8 KV heads, so the
17763        // broadcast can't beat the scalar path's 4x-more-blocks latency hiding — fall back to
17764        // scalar there (measured crossover: vec 0.68x at t_kv=64, 1.23x at t_kv=96, 2.2x at 256).
17765        // DEFAULT-ON (2026-06-28): clean clock-locked sweep proved vec beats scalar at every
17766        // t_kv>=96 and the gain WIDENS with ctx (graph decode: +9.5% @128, +11.6% @512, +11.8%
17767        // @2048) — the KV-byte-broadcast (4x fewer HBM reads/group) compounds as attention grows.
17768        // MEMRA_NO_FA_VEC forces the scalar bit-reference. Below FA_VEC_MIN_TKV the scalar path's
17769        // 4x-more-blocks (grid.x=n_head=32 vs n_head_kv=8) hides latency better, so keep scalar there.
17770        // g + no-v4: the g-module REGISTER twin mis-decodes the gemma windowed shape
17771        // (root-cause open, jsonl) — only reachable by forcing v4 off (MEMRA_FA_V4_MAX);
17772        // fall to the exact scalar there instead of the broken register arm.
17773        let mut fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && t_kv >= fa_vec_min_tkv();
17774        // hd256 under g: v4-or-scalar. hd128/other under g ride the REGISTER g-lane (the
17775        // dq_K_lane/dq_V_lane macros are format-aware; v3 is format-gated off, v2/smem
17776        // arms are excluded under g). Mirrored in kvmod / fa_decode_dc / fa_geom_eager.
17777        if g && head_dim == 256 && !fa_v4_at(t_kv) {
17778            fa_vec = false;
17779        }
17780        let sp = fa_split_keys(t_kv, n_head_kv);
17781        let n_splits = if fa_vec {
17782            ((t_kv + sp - 1) / sp).max(1)
17783        } else {
17784            ((t_kv + 255) / 256).max(1)
17785        };
17786        let o_len = n_head * n_splits * head_dim;
17787        let ml_len = n_head * n_splits;
17788        let mut part_guard = self.fa_part_pool.lock().unwrap();
17789        if part_guard
17790            .as_ref()
17791            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
17792            .unwrap_or(true)
17793        {
17794            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
17795            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
17796            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
17797            // later live allocations land at those addresses, and the next graph REPLAY writes
17798            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
17799            // output corruption began the burst after the trunk's t_kv growth first realloc'd
17800            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
17801            // the baked addresses alive (single-stream: eager writes the new buffers, replays
17802            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
17803            // (total retired < final size).
17804            let old = part_guard.take();
17805            let (co, cm) = old
17806                .as_ref()
17807                .map(|pp| (pp.0.len(), pp.1.len()))
17808                .unwrap_or((0, 0));
17809            if let Some(old) = old {
17810                self.fa_part_retired.lock().unwrap().push(old);
17811            }
17812            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
17813                eprintln!(
17814                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
17815                    co, o_len, cm, ml_len
17816                );
17817            }
17818            *part_guard = Some((
17819                self.alloc_uninit::<f32>(o_len.max(2 * co))?,
17820                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
17821                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
17822            ));
17823        }
17824        let pg = part_guard.as_mut().unwrap();
17825        self.gpu
17826            .stream()
17827            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
17828        self.gpu
17829            .stream()
17830            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
17831        self.gpu
17832            .stream()
17833            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
17834        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
17835        let (part_o, part_m, part_l) = (&mut *part_o, &mut *part_m, &mut *part_l);
17836        let (hd, nh, nhkv, tkvi, nsp) = (
17837            head_dim as i32,
17838            n_head as i32,
17839            n_head_kv as i32,
17840            t_kv as i32,
17841            n_splits as i32,
17842        );
17843        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
17844        // The vec kernel holds head_dim/32 register accumulators (FA_DEC_MAX_DPL=8 -> head_dim<=256).
17845        // All shipped models use head_dim=256; fall back to scalar for anything wider rather than
17846        // silently truncating the accumulator.
17847        let fa_vec = fa_vec && head_dim <= 512 && head_dim % 32 == 0;
17848        // hd-512 vec crossover (MEMRA_FA512_MIN, default 512): the DPL16 twin wins at depth
17849        // (82.5 -> vec at 1736) but the scalar's more-blocks latency hiding wins at tiny t_kv
17850        // (the same scalar-floor physics as hd256's old 96 floor; short-ctx plain regressed
17851        // 178.4 -> 173.7 when 512 rode vec unconditionally).
17852        let fa512_min = fa512_min_tkv();
17853        // FA-DEEP pick (bit-identical twins, see fa_deep_at): default module only — the
17854        // g-module keeps the v4 pick (its class is not the depth-decay class).
17855        let deep = fa_vec
17856            && head_dim == 256
17857            && fa_v4_at(t_kv)
17858            && !g
17859            && fa_deep_at(t_kv)
17860            && !matches!(fa_v4_mode(), "noB3" | "stage");
17861        let (f, cfg) = if fa_vec && head_dim == 512 && t_kv >= fa512_min {
17862            // gemma4 globals (hd 512): the DPL16 register twin (fa_decode_vec_q body with a
17863            // 16-slot accumulator ceiling). Scalar fallback measured 82.5us/layer at 1736 ctx.
17864            let gqa = (n_head / n_head_kv).max(1) as u32;
17865            let fv = self.fa_func("fa_decode_vec_q_dpl16", head_dim);
17866            (
17867                fv,
17868                LaunchConfig {
17869                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
17870                    block_dim: (32, gqa, 1),
17871                    shared_mem_bytes: 0,
17872                },
17873            )
17874        } else if fa_vec && head_dim <= 256 {
17875            let gqa = (n_head / n_head_kv).max(1) as u32;
17876            // DEEP-CTX smem twin (2026-07-05): the register-dequant path's GQA reuse rides L2,
17877            // which holds to ~8k ctx but dies at 40k (layer KV ~37MB) — the 4 GQA warps then
17878            // re-read every KV byte from DRAM (4x traffic). Above MEMRA_FA_SMEM_TKV (default
17879            // 1024 — the 2026-07-05 crossover re-sweep on real prompts: p3 spec 73.8->79.2 at
17880            // 2048, flat down to 512, p2 +5%, p1/9B unchanged; the ARC-A probe's synthetic
17881            // 2.1x smem-at-all-depths pointed here; 0=never) dispatch the smem-broadcast twin:
17882            // dequant each tile ONCE per block.
17883            // Bit-identical per (token,split): same bf16 round-trip, same accumulation order,
17884            // same partial layout -> same combine. Short/mid ctx keeps the register path (it won
17885            // there by 12x — latency, not bandwidth, rules small KV).
17886            static SMEM_TKV: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
17887            let smem_tkv = *SMEM_TKV.get_or_init(|| {
17888                std::env::var("MEMRA_FA_SMEM_TKV")
17889                    .ok()
17890                    .and_then(|v| v.parse().ok())
17891                    .unwrap_or_else(|| {
17892                        FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
17893                    })
17894            });
17895            if fa_v4_at(t_kv) && head_dim == 256 {
17896                // FA v4 lane (2026-07-10): key-per-lane score phase, zero shuffles per key.
17897                // NEW NUMERIC CONFIG (chunk-serial per-key dot) — battery-arbitrated.
17898                // g (fp8-windowed): the v4 staging is format-aware (2026-07-12) — kf8vf8 module.
17899                let v4name = match fa_v4_mode() {
17900                    "noB3" => "fa_decode_vec_q_v4_noB3", // phase probe (WRONG OUTPUT)
17901                    "stage" => "fa_decode_vec_q_v4_stage", // phase probe (WRONG OUTPUT)
17902                    _ if deep => "fa_decode_vec_q_v4_deep",
17903                    _ => "fa_decode_vec_q_v4",
17904                };
17905                let fv = if g {
17906                    self.func_g(v4name)
17907                } else {
17908                    self.func(v4name)
17909                };
17910                // fa_v4_smem (deep: fa_v4_deep_smem, +640B row pads) + sV (g: raw e4m3 sV
17911                // tile = 1B/elem — half the smem, 3->5 blocks/SM)
17912                let shmem = (if deep { 12160 } else { 11520 }
17913                    + 32 * head_dim * if g { 1 } else { 2 }) as u32;
17914                use cudarc::driver::sys::CUfunction_attribute_enum as A;
17915                fv.set_attribute(
17916                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
17917                    shmem as i32,
17918                )?;
17919                (
17920                    fv,
17921                    LaunchConfig {
17922                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
17923                        block_dim: (32, gqa, 1),
17924                        shared_mem_bytes: shmem,
17925                    },
17926                )
17927            } else if fa_v3_active(head_dim) {
17928                // FA v3 lane: dp4a-K hybrid (register-quantized Q, raw q8_0 K, staged-V kept).
17929                // smem = sV only (half of v2's).
17930                let fv = if g {
17931                    self.func_g("fa_decode_vec_q_v3")
17932                } else {
17933                    self.func("fa_decode_vec_q_v3")
17934                };
17935                let shmem = (32 * head_dim * 2) as u32; // sV bf16 [FA_DEC_TILE=32][hd]
17936                (
17937                    fv,
17938                    LaunchConfig {
17939                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
17940                        block_dim: (32, gqa, 1),
17941                        shared_mem_bytes: shmem,
17942                    },
17943                )
17944            } else if fa_v2_on() {
17945                // FAVENDOR lane: llama fattn-vec tile-batched softmax + wide-load staging on
17946                // OUR smem KV broadcast. Replaces BOTH per-key twins when on; same grid/block/
17947                // partials; same 32KB sK+sV tile as the smem twin.
17948                let fv = if g {
17949                    self.func_g("fa_decode_vec_q_v2")
17950                } else {
17951                    self.func("fa_decode_vec_q_v2")
17952                };
17953                let shmem = (2 * 32 * head_dim * 2) as u32; // sK+sV bf16 [FA_DEC_TILE=32][hd]
17954                (
17955                    fv,
17956                    LaunchConfig {
17957                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
17958                        block_dim: (32, gqa, 1),
17959                        shared_mem_bytes: shmem,
17960                    },
17961                )
17962            } else if smem_tkv > 0 && t_kv >= smem_tkv && !g && !(head_dim == 512 && Self::gkv_on())
17963            {
17964                // (fp8 exclusions: the smem twin's V-stage is q5_1-hardcoded — neither the wkv
17965                // windowed layers (g) nor the gkv globals (hd512) may be forced onto it via
17966                // MEMRA_FA_SMEM_TKV; they fall through to the format-clean register/scalar arms.)
17967                let fv = if g {
17968                    self.func_g("fa_decode_vec_q_smem")
17969                } else {
17970                    self.func("fa_decode_vec_q_smem")
17971                };
17972                let shmem = (2 * 32 * head_dim * 2) as u32; // sK+sV bf16 [FA_DEC_TILE=32][hd]
17973                use cudarc::driver::sys::CUfunction_attribute_enum as A;
17974                fv.set_attribute(
17975                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
17976                    shmem as i32,
17977                )?;
17978                (
17979                    fv,
17980                    LaunchConfig {
17981                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
17982                        block_dim: (32, gqa, 1),
17983                        shared_mem_bytes: shmem,
17984                    },
17985                )
17986            } else {
17987                // REGISTER-DEQUANT kernel (2026-07-03): per-warp direct q8_0/q5_1 register
17988                // dequant, zero dynamic shared memory.
17989                let fv = if g {
17990                    self.func_g("fa_decode_vec_q")
17991                } else {
17992                    self.func("fa_decode_vec_q")
17993                };
17994                (
17995                    fv,
17996                    LaunchConfig {
17997                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
17998                        block_dim: (32, gqa, 1),
17999                        shared_mem_bytes: 0,
18000                    },
18001                )
18002            }
18003        } else {
18004            // UNIFIED scalar (nullable-ctr symbol shared with graph/stream callers). The
18005            // split ladder value rides along so ns_eff reproduces THIS n_splits in-kernel.
18006            return self.fa_decode_scalar_unified(
18007                q,
18008                k,
18009                v,
18010                o,
18011                head_dim,
18012                n_head,
18013                n_head_kv,
18014                t_kv,
18015                None,
18016                scale,
18017                n_splits,
18018                if fa_vec { sp } else { 256 },
18019                k_tok_bytes,
18020                v_tok_bytes,
18021                g,
18022                part_o,
18023                part_m,
18024                part_l,
18025                None,
18026            );
18027        };
18028        let __s_b = self.gpu.stream();
18029        let mut b = __s_b.launch_builder(&f);
18030        b.arg(q)
18031            .arg(k)
18032            .arg(v)
18033            .arg(&mut *part_o)
18034            .arg(&mut *part_m)
18035            .arg(&mut *part_l)
18036            .arg(&hd)
18037            .arg(&nh)
18038            .arg(&nhkv)
18039            .arg(&tkvi)
18040            .arg(&scale)
18041            .arg(&nsp)
18042            .arg(&ktb)
18043            .arg(&vtb);
18044        unsafe {
18045            b.launch(cfg)?;
18046        }
18047        // (combine re-tile refuted in the fa-deep lane — flat/worse both shapes; the v4
18048        // combine stays for all arms. Receipts research/fa-decode-deep-20260802/.)
18049        let (fc, cfg2) = (
18050            if g {
18051                self.func_g("fa_decode_combine_f32")
18052            } else {
18053                self.fa_func("fa_decode_combine_f32", head_dim)
18054            },
18055            LaunchConfig {
18056                grid_dim: (n_head as u32, 1, 1),
18057                block_dim: (head_dim as u32, 1, 1),
18058                shared_mem_bytes: 0,
18059            },
18060        );
18061        let __s_b2 = self.gpu.stream();
18062        let mut b2 = __s_b2.launch_builder(&fc);
18063        b2.arg(&*part_o)
18064            .arg(&*part_m)
18065            .arg(&*part_l)
18066            .arg(o)
18067            .arg(&hd)
18068            .arg(&nh)
18069            .arg(&nsp);
18070        unsafe {
18071            b2.launch(cfg2)?;
18072        }
18073        Ok(())
18074    }
18075
18076    /// BATCHED-TICK increment 2: ONE fa_decode launch covering ALL B sequences of the
18077    /// batched decode step (blockIdx.z = sequence). Per-seq K/V cache bases ride a device
18078    /// pointer table (`kv_ptrs`, [2B] interleaved k0,v0,...); per-seq key bounds ride the
18079    /// tick's position table (`pos_seq`, T_kv = pos+1). v4-lane only: the CALLER
18080    /// (decode_batch) gates every row through `fa_seqs_eligible` AND one `fa_split_keys`
18081    /// rung (`split_keys`), so each sequence's split partition, key walk and combine order
18082    /// reproduce its per-seq eager v4 program exactly (kernel-check pins seqs-vs-loop bit
18083    /// identity; decode-batch-gate strict pins the whole tick vs decode_step_h).
18084    /// q is the stacked [B, n_head, head_dim] tick buffer read in place (no per-seq q
18085    /// copies); o is written [B, n_head, head_dim] in place (no per-seq a copies).
18086    #[allow(clippy::too_many_arguments)]
18087    pub fn fa_decode_batch_seqs_v4(
18088        &self,
18089        q: &CudaSlice<f32>,
18090        kv_ptrs: &cudarc::driver::CudaView<u64>,
18091        pos_seq: &CudaSlice<i32>,
18092        o: &mut CudaSlice<f32>,
18093        head_dim: usize,
18094        n_head: usize,
18095        n_head_kv: usize,
18096        b_n: usize,
18097        t_kv_max: usize,
18098        scale: f32,
18099        split_keys: usize,
18100        k_tok_bytes: usize,
18101        v_tok_bytes: usize,
18102    ) -> Result<(), Box<dyn std::error::Error>> {
18103        debug_assert!(head_dim == 256, "seqs twin is v4-stamped (hd256 only)");
18104        let n_splits_max = (t_kv_max + split_keys - 1) / split_keys;
18105        let o_len = b_n * n_head * n_splits_max * head_dim;
18106        let ml_len = b_n * n_head * n_splits_max;
18107        let mut part_guard = self.fa_part_pool.lock().unwrap();
18108        if part_guard
18109            .as_ref()
18110            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
18111            .unwrap_or(true)
18112        {
18113            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
18114            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
18115            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
18116            // later live allocations land at those addresses, and the next graph REPLAY writes
18117            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
18118            // output corruption began the burst after the trunk's t_kv growth first realloc'd
18119            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
18120            // the baked addresses alive (single-stream: eager writes the new buffers, replays
18121            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
18122            // (total retired < final size).
18123            let old = part_guard.take();
18124            let (co, cm) = old
18125                .as_ref()
18126                .map(|pp| (pp.0.len(), pp.1.len()))
18127                .unwrap_or((0, 0));
18128            if let Some(old) = old {
18129                self.fa_part_retired.lock().unwrap().push(old);
18130            }
18131            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
18132                eprintln!(
18133                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
18134                    co, o_len, cm, ml_len
18135                );
18136            }
18137            *part_guard = Some((
18138                self.alloc_uninit::<f32>(o_len.max(2 * co))?,
18139                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
18140                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
18141            ));
18142        }
18143        let pg = part_guard.as_mut().unwrap();
18144        self.gpu
18145            .stream()
18146            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
18147        self.gpu
18148            .stream()
18149            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
18150        self.gpu
18151            .stream()
18152            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
18153        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
18154        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
18155        let (nspm, spk) = (n_splits_max as i32, split_keys as i32);
18156        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
18157        let gqa = (n_head / n_head_kv).max(1) as u32;
18158        let f = self.func("fa_decode_vec_q_seqs_v4");
18159        // fa_v4_smem (11520B) + sV bf16 tile — the v4 eager arm's sizing on the default module.
18160        let shmem = (11520 + 32 * head_dim * 2) as u32;
18161        use cudarc::driver::sys::CUfunction_attribute_enum as A;
18162        f.set_attribute(
18163            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18164            shmem as i32,
18165        )?;
18166        let cfg = LaunchConfig {
18167            grid_dim: (n_head_kv as u32, n_splits_max as u32, b_n as u32),
18168            block_dim: (32, gqa, 1),
18169            shared_mem_bytes: shmem,
18170        };
18171        {
18172            let __s_b = self.gpu.stream();
18173            let mut b = __s_b.launch_builder(&f);
18174            b.arg(q)
18175                .arg(kv_ptrs)
18176                .arg(pos_seq)
18177                .arg(&mut *part_o)
18178                .arg(&mut *part_m)
18179                .arg(&mut *part_l)
18180                .arg(&hd)
18181                .arg(&nh)
18182                .arg(&nhkv)
18183                .arg(&scale)
18184                .arg(&nspm)
18185                .arg(&spk)
18186                .arg(&ktb)
18187                .arg(&vtb);
18188            unsafe {
18189                b.launch(cfg)?;
18190            }
18191        }
18192        let fc = self.func("fa_decode_combine_seqs");
18193        let cfg2 = LaunchConfig {
18194            grid_dim: (n_head as u32, b_n as u32, 1),
18195            block_dim: (head_dim as u32, 1, 1),
18196            shared_mem_bytes: 0,
18197        };
18198        let __s_b2 = self.gpu.stream();
18199        let mut b2 = __s_b2.launch_builder(&fc);
18200        b2.arg(&*part_o)
18201            .arg(&*part_m)
18202            .arg(&*part_l)
18203            .arg(o)
18204            .arg(&hd)
18205            .arg(&nh)
18206            .arg(pos_seq)
18207            .arg(&nspm)
18208            .arg(&spk);
18209        unsafe {
18210            b2.launch(cfg2)?;
18211        }
18212        Ok(())
18213    }
18214
18215    /// BATCHED-TICK increment 2: z-batched decode KV append — one launch appends this
18216    /// step's B rows, each into ITS OWN sequence cache at slot pos_seq[z], through the same
18217    /// [2B] interleaved pointer table the seqs FA reads. Each (block, z) warp executes the
18218    /// per-token appender's exact warp program on row z of the stacked [B, kv_dim] k/v —
18219    /// written cache bytes are BIT-IDENTICAL to the B per-seq calls it replaces
18220    /// (kernel-check pins the bytes). Default flash module only (callers exclude fp8-KV).
18221    #[allow(clippy::too_many_arguments)]
18222    pub fn append_kv_quantized_seqs(
18223        &self,
18224        k_rows: &CudaSlice<f32>,
18225        v_rows: &CudaSlice<f32>,
18226        kv_ptrs: &cudarc::driver::CudaView<u64>,
18227        pos_seq: &CudaSlice<i32>,
18228        b_n: usize,
18229        kv_dim_k: usize,
18230        kv_dim_v: usize,
18231        k_tok_bytes: usize,
18232        v_tok_bytes: usize,
18233    ) -> Result<(), Box<dyn std::error::Error>> {
18234        let f = self.func("append_quantize_kv_q8_0_q5_1_seqs");
18235        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
18236        let cfg = LaunchConfig {
18237            grid_dim: (nblk, b_n as u32, 1),
18238            block_dim: (32, 1, 1),
18239            shared_mem_bytes: 0,
18240        };
18241        let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
18242        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
18243        let __s_b = self.gpu.stream();
18244        let mut b = __s_b.launch_builder(&f);
18245        b.arg(k_rows)
18246            .arg(v_rows)
18247            .arg(kv_ptrs)
18248            .arg(pos_seq)
18249            .arg(&kdk)
18250            .arg(&kdv)
18251            .arg(&ktb)
18252            .arg(&vtb);
18253        unsafe {
18254            b.launch(cfg)?;
18255        }
18256        Ok(())
18257    }
18258
18259    /// True iff the MULTI-ROW verify FA (`fa_decode_rows`) is usable for a verify batch whose
18260    /// FIRST row attends `base_len + 1` keys: every row must take the SAME kernel eager decode
18261    /// would (the vec path) — mirrors fa_decode's gate exactly (MEMRA_NO_FA_VEC + FA_VEC_MIN_TKV +
18262    /// head_dim), evaluated at the MINIMUM row bound so no row could have picked scalar.
18263    /// MEMRA_FA_ROWS_OFF=1 is the A/B + fallback seam (per-row loop).
18264    pub fn fa_rows_eligible(&self, base_len: usize, head_dim: usize) -> bool {
18265        std::env::var("MEMRA_NO_FA_VEC").is_err()
18266            && std::env::var("MEMRA_FA_ROWS_OFF").is_err()
18267            && base_len + 1 >= fa_vec_min_tkv()
18268            && head_dim <= 256
18269            && head_dim % 32 == 0
18270    }
18271
18272    /// MULTI-ROW verify FA: run fa_decode_vec_q's EXACT per-row program for T causal query rows
18273    /// (row r attends keys [0..base_len+r+1)) in ONE kernel launch with grid.z = row, plus ONE
18274    /// row-batched combine. Replaces the T separate (fa_decode + combine) launches of the spec
18275    /// verify — same per-row split partition (n_splits_r = ceil(t_kv_r/split_keys), the
18276    /// fa_split_keys formula), same key-walk order, same reduce shapes => bit-identical outputs
18277    /// per row (kernel-check pins rows-vs-loop byte identity; run-spec is the end gate).
18278    /// Caller must have checked `fa_rows_eligible(base_len, head_dim)`.
18279    /// q is the verify's token-major [T, n_head, head_dim] stack; o is written [T, n_head, head_dim].
18280    #[allow(clippy::too_many_arguments)]
18281    pub fn fa_decode_rows(
18282        &self,
18283        q: &CudaSlice<f32>,
18284        k: &cudarc::driver::CudaView<u8>,
18285        v: &cudarc::driver::CudaView<u8>,
18286        o: &mut CudaSlice<f32>,
18287        head_dim: usize,
18288        n_head: usize,
18289        n_head_kv: usize,
18290        base_len: usize,
18291        t: usize,
18292        scale: f32,
18293        k_tok_bytes: usize,
18294        v_tok_bytes: usize,
18295        // hd512 dpl16 twin is DEVICE-LEN (graph arc): base_dev/plus feed the
18296        // kernel; host base_len keeps sizing the splits/partials. hd256 twins
18297        // keep the host arg. None is a bug for hd512 (asserted below).
18298        base_dev: Option<(&CudaSlice<i32>, i32)>,
18299        // K and V planes hold the same values (gemma globals, wv:=wk): pick
18300        // the _kv twin — V plane never read, value rides the q8_0 key dq.
18301        kv_shared: bool,
18302        // this layer's cache is e4m3 (gemma windowed under wkv): resolve the
18303        // hd256 rows kernel from the kf8vf8 module. PER-CALL — a global env
18304        // check here hijacked qwen/kernel-check hd256 rows (8 FAILs, 230ebbe).
18305        g: bool,
18306        // t=1 decode arm only: emit (int8, per-32 scales) from the dc combine
18307        // (hd512 path) — the standalone quantize launch folds away.
18308        mut q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
18309    ) -> Result<(), Box<dyn std::error::Error>> {
18310        debug_assert!(base_len + 1 >= fa_vec_min_tkv() && head_dim <= 512 && head_dim % 32 == 0);
18311        let t_kv_max = base_len + t; // LAST row's key bound
18312        let mut sp = fa_split_keys(t_kv_max, n_head_kv); // env/default — same value every row
18313        // hd512 split override (MEMRA_FA_SP512, 2026-07-11): gemma globals have n_head_kv=2 so
18314        // the grid is (2 x n_splits) — at depth ~29 splits = 58 blocks on 82 SMs (half idle,
18315        // rows_dpl16 8x off its byte floor). EVERY gemma hd512 caller shares THIS wrapper
18316        // (parity law), so the partition is freely tunable — verify and decode move together.
18317        if head_dim == 512 {
18318            static SP512: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
18319            // default 16 (2026-07-11 depth sweep, N=2: plain 155.4->156.5, depth spec
18320            // 236.9->250.4; 12/24/32 all worse). hd512 exists only on gemma globals.
18321            let v = *SP512.get_or_init(|| {
18322                std::env::var("MEMRA_FA_SP512")
18323                    .ok()
18324                    .and_then(|x| x.parse().ok())
18325                    .unwrap_or(0)
18326            });
18327            sp = if v >= 8 {
18328                v
18329            } else {
18330                FA_SP512_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
18331            };
18332        }
18333        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
18334        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
18335        let gqa = (n_head / n_head_kv).max(1) as u32;
18336        // LADDER-RUNG STRADDLE FIX (issue #10, 2026-07-13, g7e-proven): one sp for every row
18337        // diverges from eager decode when a split-ladder rung falls INSIDE the batch — row r's
18338        // eager twin used fa_split_keys(t_kv_r), the batch used fa_split_keys(t_kv_max), and
18339        // the different partition changes the combine's FP order (greedy tie flips at depth;
18340        // MEMRA_FA_SPLIT=64 pin -> PASS on the exact g7e failing config). Fix: group
18341        // consecutive rows by their OWN ladder value and launch once per group — each row then
18342        // executes the exact per-row program eager ran. Rungs land once per doubling, so this
18343        // is 1 launch in the common case and 2 on a crossing round. hd512 keeps one group (its
18344        // sp override is t_kv-independent by construction).
18345        let mut groups: Vec<(usize, usize, usize)> = Vec::new(); // (row0, t_g, sp_g)
18346        if head_dim == 512 || fa_split_keys(base_len + 1, n_head_kv) == sp {
18347            groups.push((0, t, sp));
18348        } else {
18349            let mut r0 = 0usize;
18350            while r0 < t {
18351                let sp_g = fa_split_keys(base_len + r0 + 1, n_head_kv);
18352                let mut r1 = r0 + 1;
18353                while r1 < t && fa_split_keys(base_len + r1 + 1, n_head_kv) == sp_g {
18354                    r1 += 1;
18355                }
18356                groups.push((r0, r1 - r0, sp_g));
18357                r0 = r1;
18358            }
18359        }
18360        // Deep-ctx smem twin for the VERIFY rows (2026-07-05): same threshold + rationale as
18361        // fa_decode's dispatch — at 40k the register path's GQA L2-reuse premise is dead and the
18362        // verify multiplies the 4x DRAM re-read by T rows. Bit-identical per (row,token,split).
18363        static SMEM_TKV_R: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
18364        let smem_tkv = *SMEM_TKV_R.get_or_init(|| {
18365            std::env::var("MEMRA_FA_SMEM_TKV")
18366                .ok()
18367                .and_then(|v| v.parse().ok())
18368                .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
18369        });
18370        let v4 = fa_v4_at(base_len + t) && head_dim == 256;
18371        let v3 = fa_v3_active(head_dim);
18372        let smem_rows =
18373            head_dim <= 256 && !v3 && !fa_v2_on() && smem_tkv > 0 && t_kv_max >= smem_tkv;
18374        // kv_shared twin RETIRED (2026-07-11 depth run-gen gate): the wv:=wk premise fails
18375        // POST-cache — cached K is k-normed+roped, cached V is not; the twin fed roped keys
18376        // in as values. Verify/decode/stream gates were blind (both sides shared the wrong
18377        // symbol — the parity law's blind spot); only prefill-vs-decode at depth caught it.
18378        let _ = kv_shared;
18379        // i2 twin: 2-key interleaved walk (MEMRA_FA_I2=0 reverts). i4 probed NEGATIVE
18380        // (157.3 vs 161.2 depth plain — register pressure past i2's sweet spot; jsonl).
18381        let i2 = head_dim == 512 && std::env::var("MEMRA_FA_I2").as_deref() != Ok("0");
18382        // v4-hd512 (MEMRA_FA_V512=1 opt-in, 2026-07-14): the v4 key-per-lane recipe on the
18383        // globals lane (depth profile: i2 ~4.6x off its byte floor — the v3-class
18384        // reduce-per-key latency signature). NEW NUMERIC CONFIG shared by every hd512
18385        // caller (decode+verify flip together); run-gen argmax + acceptance arbitrate.
18386        // T-BATCHED hd512 (DEFAULT ON 2026-07-14, MEMRA_FA_TB512=0 seam): one block per
18387        // (kv_head, split) stages its tile once and loops the rows over it — kills the
18388        // x t DRAM re-read of the full-ctx globals (depth cell +1.4%, plain flat, N=3
18389        // interleaved). FIXED absolute partition = NEW NUMERIC for the combine order,
18390        // shared by every hd512 caller through this wrapper (decode+verify flip together;
18391        // depth stream identical, acceptance unshifted, spec 256/256 x3 models).
18392        // Requires sp <= 32 (single staged tile; acc reused per row). The z-form v4_512
18393        // sibling (in-kernel dp4a port alone) probed FLAT — hd512 was DRAM-re-read-bound,
18394        // not unpack-bound; jsonl 2026-07-14.
18395        static TB512: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
18396        // gqa <= 16 = fa_v4_smem_512's q-array capacity; past it fall to the register twins.
18397        let tb512 = head_dim == 512
18398            && sp <= 32
18399            && n_head / n_head_kv.max(1) <= 16
18400            && *TB512.get_or_init(|| std::env::var("MEMRA_FA_TB512").as_deref() != Ok("0"));
18401        let fname = if tb512 {
18402            "fa_decode_vec_q_rows_v4_512_tb"
18403        } else if i2 {
18404            "fa_decode_vec_q_rows_dpl16_i2"
18405        } else if head_dim == 512 {
18406            "fa_decode_vec_q_rows_dpl16"
18407        }
18408        // gemma globals (parity law)
18409        else if v4 {
18410            "fa_decode_vec_q_rows_v4"
18411        } else if v3 {
18412            "fa_decode_vec_q_rows_v3"
18413        } else if fa_v2_on() {
18414            "fa_decode_vec_q_rows_v2"
18415        } else if smem_rows {
18416            "fa_decode_vec_q_rows_smem"
18417        } else {
18418            "fa_decode_vec_q_rows"
18419        };
18420        let f = if head_dim == 512 {
18421            self.fa_func(fname, head_dim)
18422        } else if g {
18423            // FP8-WINDOWED: hd256 rows over an e4m3 cache — kf8vf8 module, SAME symbol
18424            // choice as decode's kvmod dispatch (parity law: excluding v4 here paired
18425            // g-module rows against decode's g-module v4 — different programs, short-VG
18426            // maxdiff 2.0 / spec stream 0/128, 2026-07-12). rows_v4 is format-aware
18427            // since fda9790; only the smem twin stays excluded (V-stage q5_1-only).
18428            // hd128 (qwen fp8-KV) lands on the base/register rows via fname — the
18429            // dq macros are format-aware.
18430            self.func_g(if smem_rows {
18431                "fa_decode_vec_q_rows"
18432            } else {
18433                fname
18434            })
18435        } else {
18436            self.func(fname)
18437        };
18438        let shmem = if tb512 {
18439            // fa_v4_smem_512 (q 9KB gqa<=16 + k tile 18KB) + sV 32*512 (e4m3 module halves it)
18440            let gk = Self::gkv_on();
18441            let sh =
18442                (8192 + 1024 + 32 * 512 + 32 * 64 + 32 * head_dim * if gk { 1 } else { 2 }) as u32;
18443            use cudarc::driver::sys::CUfunction_attribute_enum as A;
18444            f.set_attribute(
18445                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18446                sh as i32,
18447            )?;
18448            sh
18449        } else if v4 || v3 || smem_rows || fa_v2_on() {
18450            // v4: fa_v4_smem (11.5KB) + sV; v3 stages sV only; v2/smem twins stage sK+sV.
18451            let sh = (if v4 {
18452                11520 + 32 * head_dim * if g { 1 } else { 2 }
18453            } else if v3 {
18454                32 * head_dim * 2
18455            } else {
18456                2 * 32 * head_dim * 2
18457            }) as u32;
18458            use cudarc::driver::sys::CUfunction_attribute_enum as A;
18459            f.set_attribute(
18460                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18461                sh as i32,
18462            )?;
18463            sh
18464        } else {
18465            0
18466        };
18467        // Per-GROUP launches (single group in the common case — identical to the pre-fix
18468        // single launch there): each group gets its own partials (the rows kernel indexes
18469        // partials by its LOCAL grid.z row) and q/o row-offset views.
18470        for &(r0, t_g, sp_g) in &groups {
18471            let n_splits_g = (base_len + r0 + t_g).div_ceil(sp_g);
18472            let (nspm, spk) = (n_splits_g as i32, sp_g as i32);
18473            let base_i = (base_len + r0) as i32;
18474            let o_len = t_g * n_head * n_splits_g * head_dim;
18475            let ml_len = t_g * n_head * n_splits_g;
18476            let mut part_guard = self.fa_part_pool.lock().unwrap();
18477            if part_guard
18478                .as_ref()
18479                .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
18480                .unwrap_or(true)
18481            {
18482                // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
18483                // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
18484                // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
18485                // later live allocations land at those addresses, and the next graph REPLAY writes
18486                // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
18487                // output corruption began the burst after the trunk's t_kv growth first realloc'd
18488                // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
18489                // the baked addresses alive (single-stream: eager writes the new buffers, replays
18490                // touch only the old — never concurrently). Doubling growth bounds retired VRAM
18491                // (total retired < final size).
18492                let old = part_guard.take();
18493                let (co, cm) = old
18494                    .as_ref()
18495                    .map(|pp| (pp.0.len(), pp.1.len()))
18496                    .unwrap_or((0, 0));
18497                if let Some(old) = old {
18498                    self.fa_part_retired.lock().unwrap().push(old);
18499                }
18500                if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
18501                    eprintln!(
18502                        "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
18503                        co, o_len, cm, ml_len
18504                    );
18505                }
18506                *part_guard = Some((
18507                    self.alloc_uninit::<f32>(o_len.max(2 * co))?,
18508                    self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
18509                    self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
18510                ));
18511            }
18512            let pg = part_guard.as_mut().unwrap();
18513            self.gpu
18514                .stream()
18515                .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
18516            self.gpu
18517                .stream()
18518                .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
18519            self.gpu
18520                .stream()
18521                .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
18522            let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
18523            let (part_o, part_m, part_l) = (&mut *part_o, &mut *part_m, &mut *part_l);
18524            let qv = self.view(q, t * n_head * head_dim);
18525            let q_g = qv.slice(r0 * n_head * head_dim..(r0 + t_g) * n_head * head_dim);
18526            let cfg = LaunchConfig {
18527                grid_dim: (n_head_kv as u32, n_splits_g as u32, t_g as u32),
18528                block_dim: (32, gqa, 1),
18529                shared_mem_bytes: shmem,
18530            };
18531            {
18532                let __s_b = self.gpu.stream();
18533                let mut b = __s_b.launch_builder(&f);
18534                if tb512 {
18535                    // rows-inner launch: grid.z dropped, the kernel loops n_rows itself.
18536                    let (bd, plus) =
18537                        base_dev.expect("hd512 rows twin requires a device base counter");
18538                    let plus_g = plus + r0 as i32;
18539                    let nr = t_g as i32;
18540                    if Self::pdl_on() && Self::pdl_wb_on() {
18541                        // wave-B2b: flavor mirrors fa_func(fname, 512) = gkv.
18542                        use cudarc::driver::{DevicePtr, DevicePtrMut};
18543                        let s = &self.gpu.stream();
18544                        let (pq, _b0) = q_g.device_ptr(s);
18545                        let (pk, _b1) = k.device_ptr(s);
18546                        let (pv, _b2) = v.device_ptr(s);
18547                        let (po, _b3) = part_o.device_ptr_mut(s);
18548                        let (pm, _b4) = part_m.device_ptr_mut(s);
18549                        let (pl, _b5) = part_l.device_ptr_mut(s);
18550                        let (pb, _b6) = bd.device_ptr(s);
18551                        let mut ps = [
18552                            &pq as *const _ as *mut std::ffi::c_void,
18553                            &pk as *const _ as *mut _,
18554                            &pv as *const _ as *mut _,
18555                            &po as *const _ as *mut _,
18556                            &pm as *const _ as *mut _,
18557                            &pl as *const _ as *mut _,
18558                            &hd as *const _ as *mut _,
18559                            &nh as *const _ as *mut _,
18560                            &nhkv as *const _ as *mut _,
18561                            &pb as *const _ as *mut _,
18562                            &plus_g as *const _ as *mut _,
18563                            &scale as *const _ as *mut _,
18564                            &nspm as *const _ as *mut _,
18565                            &spk as *const _ as *mut _,
18566                            &ktb as *const _ as *mut _,
18567                            &vtb as *const _ as *mut _,
18568                            &nr as *const _ as *mut _,
18569                        ];
18570                        unsafe {
18571                            self.launch_pdl_flash(
18572                                Self::gkv_on(),
18573                                "fa_decode_vec_q_rows_v4_512_tb",
18574                                (n_head_kv as u32, n_splits_g as u32, 1),
18575                                (32, gqa, 1),
18576                                shmem,
18577                                &mut ps,
18578                            )?;
18579                        }
18580                    } else {
18581                        let cfg_tb = LaunchConfig {
18582                            grid_dim: (n_head_kv as u32, n_splits_g as u32, 1),
18583                            block_dim: (32, gqa, 1),
18584                            shared_mem_bytes: shmem,
18585                        };
18586                        b.arg(&q_g)
18587                            .arg(k)
18588                            .arg(v)
18589                            .arg(&mut *part_o)
18590                            .arg(&mut *part_m)
18591                            .arg(&mut *part_l)
18592                            .arg(&hd)
18593                            .arg(&nh)
18594                            .arg(&nhkv)
18595                            .arg(bd)
18596                            .arg(&plus_g)
18597                            .arg(&scale)
18598                            .arg(&nspm)
18599                            .arg(&spk)
18600                            .arg(&ktb)
18601                            .arg(&vtb)
18602                            .arg(&nr);
18603                        unsafe {
18604                            b.launch(cfg_tb)?;
18605                        }
18606                    }
18607                } else if head_dim == 512 {
18608                    let (bd, plus) =
18609                        base_dev.expect("hd512 rows twin requires a device base counter");
18610                    let plus_g = plus + r0 as i32;
18611                    b.arg(&q_g)
18612                        .arg(k)
18613                        .arg(v)
18614                        .arg(&mut *part_o)
18615                        .arg(&mut *part_m)
18616                        .arg(&mut *part_l)
18617                        .arg(&hd)
18618                        .arg(&nh)
18619                        .arg(&nhkv)
18620                        .arg(bd)
18621                        .arg(&plus_g)
18622                        .arg(&scale)
18623                        .arg(&nspm)
18624                        .arg(&spk)
18625                        .arg(&ktb)
18626                        .arg(&vtb);
18627                    unsafe {
18628                        b.launch(cfg)?;
18629                    }
18630                } else {
18631                    b.arg(&q_g)
18632                        .arg(k)
18633                        .arg(v)
18634                        .arg(&mut *part_o)
18635                        .arg(&mut *part_m)
18636                        .arg(&mut *part_l)
18637                        .arg(&hd)
18638                        .arg(&nh)
18639                        .arg(&nhkv)
18640                        .arg(&base_i)
18641                        .arg(&scale)
18642                        .arg(&nspm)
18643                        .arg(&spk)
18644                        .arg(&ktb)
18645                        .arg(&vtb);
18646                    unsafe {
18647                        b.launch(cfg)?;
18648                    }
18649                }
18650            }
18651            let cfg2 = LaunchConfig {
18652                grid_dim: (n_head as u32, t_g as u32, 1),
18653                block_dim: (head_dim as u32, 1, 1),
18654                shared_mem_bytes: 0,
18655            };
18656            let mut o_g = o.slice_mut(r0 * n_head * head_dim..(r0 + t_g) * n_head * head_dim);
18657            if head_dim == 512 {
18658                // device-len combine (shared by verify/eager/graph — parity by symbol): the
18659                // per-row n_splits derives from the SAME counter the rows kernel read.
18660                let (bd, plus) = base_dev.unwrap();
18661                let plus_g = plus + r0 as i32;
18662                if let Some((oq, od)) = q8_out.as_mut() {
18663                    // wave-5b port (2026-07-23, t=1 decode only): q8-emitting dc combine.
18664                    debug_assert!(t == 1, "rows q8 emit is a t=1 decode arm");
18665                    if Self::pdl_on() && Self::pdl_wb_on() {
18666                        // wave-B2: flavor mirrors fa_func (hd512 + gkv → kf8vf8).
18667                        use cudarc::driver::{DevicePtr, DevicePtrMut};
18668                        let s = &self.gpu.stream();
18669                        let (po, _g0) = part_o.device_ptr(s);
18670                        let (pm, _g1) = part_m.device_ptr(s);
18671                        let (pl, _g2) = part_l.device_ptr(s);
18672                        let (pq, _g3) = oq.device_ptr_mut(s);
18673                        let (pd, _g4) = od.device_ptr_mut(s);
18674                        let (pb, _g5) = bd.device_ptr(s);
18675                        let mut ps = [
18676                            &po as *const _ as *mut std::ffi::c_void,
18677                            &pm as *const _ as *mut _,
18678                            &pl as *const _ as *mut _,
18679                            &pq as *const _ as *mut _,
18680                            &pd as *const _ as *mut _,
18681                            &hd as *const _ as *mut _,
18682                            &nh as *const _ as *mut _,
18683                            &pb as *const _ as *mut _,
18684                            &plus_g as *const _ as *mut _,
18685                            &nspm as *const _ as *mut _,
18686                            &spk as *const _ as *mut _,
18687                        ];
18688                        unsafe {
18689                            self.launch_pdl_flash(
18690                                Self::gkv_on(),
18691                                "fa_decode_combine_rows_dc_q8_1",
18692                                cfg2.grid_dim,
18693                                cfg2.block_dim,
18694                                0,
18695                                &mut ps,
18696                            )?;
18697                        }
18698                        continue;
18699                    }
18700                    let fc = self.fa_func("fa_decode_combine_rows_dc_q8_1", head_dim);
18701                    let __s_b2 = self.gpu.stream();
18702                    let mut b2 = __s_b2.launch_builder(&fc);
18703                    b2.arg(&*part_o)
18704                        .arg(&*part_m)
18705                        .arg(&*part_l)
18706                        .arg(&mut **oq)
18707                        .arg(&mut **od)
18708                        .arg(&hd)
18709                        .arg(&nh)
18710                        .arg(bd)
18711                        .arg(&plus_g)
18712                        .arg(&nspm)
18713                        .arg(&spk);
18714                    unsafe {
18715                        b2.launch(cfg2)?;
18716                    }
18717                    continue;
18718                }
18719                let fc = self.fa_func("fa_decode_combine_rows_dc", head_dim);
18720                let __s_b2 = self.gpu.stream();
18721                let mut b2 = __s_b2.launch_builder(&fc);
18722                b2.arg(&*part_o)
18723                    .arg(&*part_m)
18724                    .arg(&*part_l)
18725                    .arg(&mut o_g)
18726                    .arg(&hd)
18727                    .arg(&nh)
18728                    .arg(bd)
18729                    .arg(&plus_g)
18730                    .arg(&nspm)
18731                    .arg(&spk);
18732                unsafe {
18733                    b2.launch(cfg2)?;
18734                }
18735            } else {
18736                // q8 emit is wired for the hd512 dc-combine arm only — a Some here would
18737                // leave the caller's pair unwritten (consumer would read garbage).
18738                assert!(
18739                    q8_out.is_none(),
18740                    "rows q8 emit requires the hd512 dc combine"
18741                );
18742                let fc = self.func("fa_decode_combine_rows");
18743                let __s_b2 = self.gpu.stream();
18744                let mut b2 = __s_b2.launch_builder(&fc);
18745                b2.arg(&*part_o)
18746                    .arg(&*part_m)
18747                    .arg(&*part_l)
18748                    .arg(&mut o_g)
18749                    .arg(&hd)
18750                    .arg(&nh)
18751                    .arg(&base_i)
18752                    .arg(&nspm)
18753                    .arg(&spk);
18754                unsafe {
18755                    b2.launch(cfg2)?;
18756                }
18757            }
18758        }
18759        Ok(())
18760    }
18761
18762    /// WINDOWED verify rows (gemma R6 deep-ctx): every row attends exactly `window` keys —
18763    /// bit-identical per row to the T=1 decode's fa_decode over the window VIEW. Caller gates
18764    /// base_len + 1 >= window (no under-window rows) and head_dim == 256 (v4 stamp).
18765    #[allow(clippy::too_many_arguments)]
18766    pub fn fa_decode_rows_w(
18767        &self,
18768        q: &CudaSlice<f32>,
18769        k: &cudarc::driver::CudaView<u8>,
18770        v: &cudarc::driver::CudaView<u8>,
18771        o: &mut CudaSlice<f32>,
18772        head_dim: usize,
18773        n_head: usize,
18774        n_head_kv: usize,
18775        base_dev: &CudaSlice<i32>,
18776        base_plus: i32,
18777        t: usize,
18778        scale: f32,
18779        window: usize,
18780        k_tok_bytes: usize,
18781        v_tok_bytes: usize,
18782        q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
18783    ) -> Result<(), Box<dyn std::error::Error>> {
18784        // DEVICE-LEN (graph arc step 1, 2026-07-11): the causal base rides an i32 counter
18785        // (kernel T_kv = dev[0] + base_plus + r + 1) so depth graphs can replay with len
18786        // advancing on-device. dc paths pass kvl.len_d with plus=-1; verify/eager sync the
18787        // counter with one async set_i32_one first. Partials/splits size from `window` (host).
18788        debug_assert!(head_dim == 256);
18789        // windowed split (MEMRA_FA_SPW, default 32 — re-swept 2026-07-12 under the raw-e4m3 sV
18790        // occupancy ceiling (4 blocks/SM): t=1 decode is GRID-limited (win/sp splits x nkv
18791        // blocks), so smaller splits fill the ceiling — 1.7k 174.4/174.0 vs 48's 170.7/170.3,
18792        // 4.9k 159.8 vs 157.4 (N=2 interleaved, stable window). Spec serving prefers 64
18793        // (verify t=K+1 fills the grid via grid.z=t; depth K=7 281.3 vs 249.3 at 32) — set
18794        // MEMRA_FA_SPW=64 there, same config law as MEMRA_GEMMA_GKV=0. MUST be one value for
18795        // ALL widths: a t-keyed probe broke decode-vs-verify combine order (stream 9/128).
18796        let sp = {
18797            static SPW: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
18798            let v = *SPW.get_or_init(|| {
18799                std::env::var("MEMRA_FA_SPW")
18800                    .ok()
18801                    .and_then(|x| x.parse().ok())
18802                    .unwrap_or(0)
18803            });
18804            if v >= 8 {
18805                v
18806            } else {
18807                FA_SPW_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
18808            }
18809        };
18810        let n_splits_max = (window + sp - 1) / sp;
18811        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
18812        let (nspm, spk, wini) = (n_splits_max as i32, sp as i32, window as i32);
18813        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
18814        let gqa = (n_head / n_head_kv).max(1) as u32;
18815        let o_len = t * n_head * n_splits_max * head_dim;
18816        let ml_len = t * n_head * n_splits_max;
18817        let mut part_guard = self.fa_part_pool.lock().unwrap();
18818        if part_guard
18819            .as_ref()
18820            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
18821            .unwrap_or(true)
18822        {
18823            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
18824            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
18825            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
18826            // later live allocations land at those addresses, and the next graph REPLAY writes
18827            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
18828            // output corruption began the burst after the trunk's t_kv growth first realloc'd
18829            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
18830            // the baked addresses alive (single-stream: eager writes the new buffers, replays
18831            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
18832            // (total retired < final size).
18833            let old = part_guard.take();
18834            let (co, cm) = old
18835                .as_ref()
18836                .map(|pp| (pp.0.len(), pp.1.len()))
18837                .unwrap_or((0, 0));
18838            if let Some(old) = old {
18839                self.fa_part_retired.lock().unwrap().push(old);
18840            }
18841            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
18842                eprintln!(
18843                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
18844                    co, o_len, cm, ml_len
18845                );
18846            }
18847            *part_guard = Some((
18848                self.alloc_uninit::<f32>(o_len.max(2 * co))?,
18849                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
18850                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
18851            ));
18852        }
18853        let pg = part_guard.as_mut().unwrap();
18854        self.gpu
18855            .stream()
18856            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
18857        self.gpu
18858            .stream()
18859            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
18860        self.gpu
18861            .stream()
18862            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
18863        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
18864        // Lane pick: decode AND verify both land here in the windowed regime (parity law —
18865        // hybrid_forward verify_attn), so the pick only needs internal consistency, not
18866        // clone-of-decode bit fidelity (SASS-proven impossible for textually identical
18867        // kernels, jsonl 2026-07-10). v4 under the threshold; smem twin at/above the smem
18868        // floor (deep-ctx broadcast win); register twin between.
18869        static SMEM_TKV_W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
18870        let smem_tkv = *SMEM_TKV_W.get_or_init(|| {
18871            std::env::var("MEMRA_FA_SMEM_TKV")
18872                .ok()
18873                .and_then(|v| v.parse().ok())
18874                .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
18875        });
18876        // MULTI-ROW v4: resurrected 2026-07-14 (the '33 tok/s collapse' was a paired-map
18877        // partial-write bug, not the mechanism) and falsified HONESTLY at gqa 2: bit-exact
18878        // but −1.7% on the 31B depth cell — the sp helper warp already hides staging
18879        // in-block, and mr trades L2-cheap redundant bytes for serialized per-warp gqa
18880        // score/B3 chains. Arm deleted; jsonl row 2026-07-14 is the record.
18881        use cudarc::driver::sys::CUfunction_attribute_enum as A;
18882        // FP8-WINDOWED (wkv): the v4 family is format-aware (2026-07-12 KFMT/VFMT staging
18883        // arms) — wkv rides the SAME lane logic, resolved from the kf8vf8 module. One symbol
18884        // per (lane, format-module) keeps parity structural; the old register-i2 detour
18885        // (-33%) is retired.
18886        let wg = Self::wkv_on();
18887        // STAGING-PARALLEL v4 (MEMRA_FA_SPW2, default ON at gqa==1): warp 1 = staging helper
18888        // (v4 is 61% staging); score phases identical to v4_w. Same symbol all t.
18889        let sp2 =
18890            gqa <= 4 && fa_v4_at(window) && std::env::var("MEMRA_FA_SPW2").as_deref() != Ok("0");
18891        if sp2 {
18892            let sh = (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32;
18893            if Self::pdl_on() && Self::pdl_wb_on() {
18894                // wave-B2b: flavor mirrors wg.
18895                use cudarc::driver::{DevicePtr, DevicePtrMut};
18896                let s = &self.gpu.stream();
18897                let (pq, _b0) = q.device_ptr(s);
18898                let (pk, _b1) = k.device_ptr(s);
18899                let (pv, _b2) = v.device_ptr(s);
18900                let (po, _b3) = part_o.device_ptr_mut(s);
18901                let (pm, _b4) = part_m.device_ptr_mut(s);
18902                let (pl, _b5) = part_l.device_ptr_mut(s);
18903                let (pb, _b6) = base_dev.device_ptr(s);
18904                let mut ps = [
18905                    &pq as *const _ as *mut std::ffi::c_void,
18906                    &pk as *const _ as *mut _,
18907                    &pv as *const _ as *mut _,
18908                    &po as *const _ as *mut _,
18909                    &pm as *const _ as *mut _,
18910                    &pl as *const _ as *mut _,
18911                    &hd as *const _ as *mut _,
18912                    &nh as *const _ as *mut _,
18913                    &nhkv as *const _ as *mut _,
18914                    &pb as *const _ as *mut _,
18915                    &base_plus as *const _ as *mut _,
18916                    &scale as *const _ as *mut _,
18917                    &nspm as *const _ as *mut _,
18918                    &spk as *const _ as *mut _,
18919                    &ktb as *const _ as *mut _,
18920                    &vtb as *const _ as *mut _,
18921                    &wini as *const _ as *mut _,
18922                ];
18923                unsafe {
18924                    self.launch_pdl_flash(
18925                        wg,
18926                        "fa_decode_vec_q_rows_v4_w_sp",
18927                        (n_head_kv as u32, n_splits_max as u32, t as u32),
18928                        (32, gqa + 1, 1),
18929                        sh,
18930                        &mut ps,
18931                    )?;
18932                }
18933            } else {
18934                let f = if wg {
18935                    self.func_g("fa_decode_vec_q_rows_v4_w_sp")
18936                } else {
18937                    self.func("fa_decode_vec_q_rows_v4_w_sp")
18938                };
18939                f.set_attribute(
18940                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18941                    sh as i32,
18942                )?;
18943                let cfg = LaunchConfig {
18944                    grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
18945                    block_dim: (32, gqa + 1, 1),
18946                    shared_mem_bytes: sh,
18947                };
18948                let __s_b = self.gpu.stream();
18949                let mut b = __s_b.launch_builder(&f);
18950                b.arg(q)
18951                    .arg(k)
18952                    .arg(v)
18953                    .arg(&mut *part_o)
18954                    .arg(&mut *part_m)
18955                    .arg(&mut *part_l)
18956                    .arg(&hd)
18957                    .arg(&nh)
18958                    .arg(&nhkv)
18959                    .arg(base_dev)
18960                    .arg(&base_plus)
18961                    .arg(&scale)
18962                    .arg(&nspm)
18963                    .arg(&spk)
18964                    .arg(&ktb)
18965                    .arg(&vtb)
18966                    .arg(&wini);
18967                unsafe {
18968                    b.launch(cfg)?;
18969                }
18970            }
18971        } else {
18972            if fa_v4_at(window) && Self::pdl_on() && Self::pdl_wb_on() {
18973                // wave-B2b: the v4_w pick only (smem/reg twins stay builder-launched).
18974                let sh = (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32;
18975                use cudarc::driver::{DevicePtr, DevicePtrMut};
18976                let s = &self.gpu.stream();
18977                let (pq, _b0) = q.device_ptr(s);
18978                let (pk, _b1) = k.device_ptr(s);
18979                let (pv, _b2) = v.device_ptr(s);
18980                let (po, _b3) = part_o.device_ptr_mut(s);
18981                let (pm, _b4) = part_m.device_ptr_mut(s);
18982                let (pl, _b5) = part_l.device_ptr_mut(s);
18983                let (pb, _b6) = base_dev.device_ptr(s);
18984                let mut ps = [
18985                    &pq as *const _ as *mut std::ffi::c_void,
18986                    &pk as *const _ as *mut _,
18987                    &pv as *const _ as *mut _,
18988                    &po as *const _ as *mut _,
18989                    &pm as *const _ as *mut _,
18990                    &pl as *const _ as *mut _,
18991                    &hd as *const _ as *mut _,
18992                    &nh as *const _ as *mut _,
18993                    &nhkv as *const _ as *mut _,
18994                    &pb as *const _ as *mut _,
18995                    &base_plus as *const _ as *mut _,
18996                    &scale as *const _ as *mut _,
18997                    &nspm as *const _ as *mut _,
18998                    &spk as *const _ as *mut _,
18999                    &ktb as *const _ as *mut _,
19000                    &vtb as *const _ as *mut _,
19001                    &wini as *const _ as *mut _,
19002                ];
19003                unsafe {
19004                    self.launch_pdl_flash(
19005                        wg,
19006                        "fa_decode_vec_q_rows_v4_w",
19007                        (n_head_kv as u32, n_splits_max as u32, t as u32),
19008                        (32, gqa, 1),
19009                        sh,
19010                        &mut ps,
19011                    )?;
19012                }
19013            } else {
19014                let pick = |name: &str| {
19015                    if wg {
19016                        self.func_g(name)
19017                    } else {
19018                        self.func(name)
19019                    }
19020                };
19021                let (f, sh) = if fa_v4_at(window) {
19022                    let f = pick("fa_decode_vec_q_rows_v4_w");
19023                    (f, (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32)
19024                } else if smem_tkv > 0 && window >= smem_tkv {
19025                    // NOTE: the smem twin's V-stage is still q5_1-hardcoded — unreachable under wkv
19026                    // at the gemma window (v4 covers it); revisit if the smem floor ever drops.
19027                    (
19028                        pick("fa_decode_vec_q_rows_smem_w"),
19029                        (2 * 32 * head_dim * 2) as u32,
19030                    )
19031                } else {
19032                    (pick("fa_decode_vec_q_rows_reg_w"), 0u32)
19033                };
19034                f.set_attribute(
19035                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19036                    sh as i32,
19037                )?;
19038                let cfg = LaunchConfig {
19039                    grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
19040                    block_dim: (32, gqa, 1),
19041                    shared_mem_bytes: sh,
19042                };
19043                let __s_b = self.gpu.stream();
19044                let mut b = __s_b.launch_builder(&f);
19045                b.arg(q)
19046                    .arg(k)
19047                    .arg(v)
19048                    .arg(&mut *part_o)
19049                    .arg(&mut *part_m)
19050                    .arg(&mut *part_l)
19051                    .arg(&hd)
19052                    .arg(&nh)
19053                    .arg(&nhkv)
19054                    .arg(base_dev)
19055                    .arg(&base_plus)
19056                    .arg(&scale)
19057                    .arg(&nspm)
19058                    .arg(&spk)
19059                    .arg(&ktb)
19060                    .arg(&vtb)
19061                    .arg(&wini);
19062                unsafe {
19063                    b.launch(cfg)?;
19064                }
19065            }
19066        }
19067        let cfg2 = LaunchConfig {
19068            grid_dim: (n_head as u32, t as u32, 1),
19069            block_dim: (head_dim as u32, 1, 1),
19070            shared_mem_bytes: 0,
19071        };
19072        if let Some((oq, od)) = q8_out {
19073            // wave-5b port (2026-07-23): q8-emitting combine — the t=1 decode's wo matvec
19074            // consumes the pair directly; the standalone quantize launch folds away.
19075            if Self::pdl_on() && Self::pdl_wb_on() {
19076                // wave-B2: flavor mirrors the builder's wg choice.
19077                use cudarc::driver::{DevicePtr, DevicePtrMut};
19078                let s = &self.gpu.stream();
19079                let (po, _g0) = part_o.device_ptr(s);
19080                let (pm, _g1) = part_m.device_ptr(s);
19081                let (pl, _g2) = part_l.device_ptr(s);
19082                let (pq, _g3) = oq.device_ptr_mut(s);
19083                let (pd, _g4) = od.device_ptr_mut(s);
19084                let mut ps = [
19085                    &po as *const _ as *mut std::ffi::c_void,
19086                    &pm as *const _ as *mut _,
19087                    &pl as *const _ as *mut _,
19088                    &pq as *const _ as *mut _,
19089                    &pd as *const _ as *mut _,
19090                    &hd as *const _ as *mut _,
19091                    &nh as *const _ as *mut _,
19092                    &nspm as *const _ as *mut _,
19093                    &spk as *const _ as *mut _,
19094                    &wini as *const _ as *mut _,
19095                ];
19096                unsafe {
19097                    self.launch_pdl_flash(
19098                        wg,
19099                        "fa_decode_combine_rows_w_q8_1",
19100                        cfg2.grid_dim,
19101                        cfg2.block_dim,
19102                        0,
19103                        &mut ps,
19104                    )?;
19105                }
19106                return Ok(());
19107            }
19108            let fc = if wg {
19109                self.func_g("fa_decode_combine_rows_w_q8_1")
19110            } else {
19111                self.func("fa_decode_combine_rows_w_q8_1")
19112            };
19113            let __s_b2 = self.gpu.stream();
19114            let mut b2 = __s_b2.launch_builder(&fc);
19115            b2.arg(&*part_o)
19116                .arg(&*part_m)
19117                .arg(&*part_l)
19118                .arg(oq)
19119                .arg(od)
19120                .arg(&hd)
19121                .arg(&nh)
19122                .arg(&nspm)
19123                .arg(&spk)
19124                .arg(&wini);
19125            unsafe {
19126                b2.launch(cfg2)?;
19127            }
19128            return Ok(());
19129        }
19130        let fc = if wg {
19131            self.func_g("fa_decode_combine_rows_w")
19132        } else {
19133            self.func("fa_decode_combine_rows_w")
19134        };
19135        let __s_b2 = self.gpu.stream();
19136        let mut b2 = __s_b2.launch_builder(&fc);
19137        b2.arg(&*part_o)
19138            .arg(&*part_m)
19139            .arg(&*part_l)
19140            .arg(o)
19141            .arg(&hd)
19142            .arg(&nh)
19143            .arg(&nspm)
19144            .arg(&spk)
19145            .arg(&wini);
19146        unsafe {
19147            b2.launch(cfg2)?;
19148        }
19149        Ok(())
19150    }
19151
19152    /// ROUND-STREAM stage (c): fa rows with the causal base from a device counter. Two lanes:
19153    /// v3 (qwen stream, fa_v3_active) and v4 (gemma hd256 burst — rows_v4_dc, g-module aware);
19154    /// `t_kv_upper` sizes splits/partials — the same one-sp-for-all-rows approximation class
19155    /// the host rows path already uses (battery-arbitrated); actual per-row bounds derive
19156    /// in-kernel from the counter (+ base_plus, v4 lane only — v3's kernel has no plus arg).
19157    #[allow(clippy::too_many_arguments)]
19158    pub fn fa_decode_rows_dc(
19159        &self,
19160        q: &CudaSlice<f32>,
19161        k: &cudarc::driver::CudaView<u8>,
19162        v: &cudarc::driver::CudaView<u8>,
19163        o: &mut CudaSlice<f32>,
19164        head_dim: usize,
19165        n_head: usize,
19166        n_head_kv: usize,
19167        base_dev: &CudaSlice<i32>,
19168        t_kv_upper: usize,
19169        t: usize,
19170        scale: f32,
19171        k_tok_bytes: usize,
19172        v_tok_bytes: usize,
19173        base_plus: i32,
19174        g: bool,
19175    ) -> Result<(), Box<dyn std::error::Error>> {
19176        let v4 = head_dim == 256 && fa_v4_at(t_kv_upper);
19177        assert!(
19178            v4 || fa_v3_active(head_dim),
19179            "stream fa rows requires the v3 or v4 lane"
19180        );
19181        assert!(v4 || base_plus == 0, "v3_dc kernel takes no plus arg");
19182        if v4 {
19183            let sp = fa_split_keys(t_kv_upper, n_head_kv);
19184            let n_splits_max = (t_kv_upper + sp - 1) / sp;
19185            let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
19186            let (nspm, spk) = (n_splits_max as i32, sp as i32);
19187            let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
19188            let gqa = (n_head / n_head_kv).max(1) as u32;
19189            let o_len = t * n_head * n_splits_max * head_dim;
19190            let ml_len = t * n_head * n_splits_max;
19191            let mut part_guard = self.fa_part_pool.lock().unwrap();
19192            if part_guard
19193                .as_ref()
19194                .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
19195                .unwrap_or(true)
19196            {
19197                // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
19198                // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
19199                // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
19200                // later live allocations land at those addresses, and the next graph REPLAY writes
19201                // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
19202                // output corruption began the burst after the trunk's t_kv growth first realloc'd
19203                // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
19204                // the baked addresses alive (single-stream: eager writes the new buffers, replays
19205                // touch only the old — never concurrently). Doubling growth bounds retired VRAM
19206                // (total retired < final size).
19207                let old = part_guard.take();
19208                let (co, cm) = old
19209                    .as_ref()
19210                    .map(|pp| (pp.0.len(), pp.1.len()))
19211                    .unwrap_or((0, 0));
19212                if let Some(old) = old {
19213                    self.fa_part_retired.lock().unwrap().push(old);
19214                }
19215                if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
19216                    eprintln!(
19217                        "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
19218                        co, o_len, cm, ml_len
19219                    );
19220                }
19221                *part_guard = Some((
19222                    self.alloc_uninit::<f32>(o_len.max(2 * co))?,
19223                    self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
19224                    self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
19225                ));
19226            }
19227            let pg = part_guard.as_mut().unwrap();
19228            self.gpu
19229                .stream()
19230                .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
19231            self.gpu
19232                .stream()
19233                .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
19234            self.gpu
19235                .stream()
19236                .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
19237            let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
19238            let f = if g {
19239                self.func_g("fa_decode_vec_q_rows_v4_dc")
19240            } else {
19241                self.func("fa_decode_vec_q_rows_v4_dc")
19242            };
19243            let sh = (11520 + 32 * head_dim * if g { 1 } else { 2 }) as u32;
19244            use cudarc::driver::sys::CUfunction_attribute_enum as A;
19245            f.set_attribute(
19246                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19247                sh as i32,
19248            )?;
19249            let cfg = LaunchConfig {
19250                grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
19251                block_dim: (32, gqa, 1),
19252                shared_mem_bytes: sh,
19253            };
19254            let __s_b = self.gpu.stream();
19255            let mut b = __s_b.launch_builder(&f);
19256            b.arg(q)
19257                .arg(k)
19258                .arg(v)
19259                .arg(&mut *part_o)
19260                .arg(&mut *part_m)
19261                .arg(&mut *part_l)
19262                .arg(&hd)
19263                .arg(&nh)
19264                .arg(&nhkv)
19265                .arg(base_dev)
19266                .arg(&base_plus)
19267                .arg(&scale)
19268                .arg(&nspm)
19269                .arg(&spk)
19270                .arg(&ktb)
19271                .arg(&vtb);
19272            unsafe {
19273                b.launch(cfg)?;
19274            }
19275            let fc = self.func("fa_decode_combine_rows_dc");
19276            let cfg2 = LaunchConfig {
19277                grid_dim: (n_head as u32, t as u32, 1),
19278                block_dim: (head_dim as u32, 1, 1),
19279                shared_mem_bytes: 0,
19280            };
19281            let __s_b2 = self.gpu.stream();
19282            let mut b2 = __s_b2.launch_builder(&fc);
19283            b2.arg(&*part_o)
19284                .arg(&*part_m)
19285                .arg(&*part_l)
19286                .arg(o)
19287                .arg(&hd)
19288                .arg(&nh)
19289                .arg(base_dev)
19290                .arg(&base_plus)
19291                .arg(&nspm)
19292                .arg(&spk);
19293            unsafe {
19294                b2.launch(cfg2)?;
19295            }
19296            return Ok(());
19297        }
19298        let sp = fa_split_keys(t_kv_upper, n_head_kv);
19299        let n_splits_max = (t_kv_upper + sp - 1) / sp;
19300        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
19301        let (nspm, spk) = (n_splits_max as i32, sp as i32);
19302        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
19303        let gqa = (n_head / n_head_kv).max(1) as u32;
19304        let o_len = t * n_head * n_splits_max * head_dim;
19305        let ml_len = t * n_head * n_splits_max;
19306        let mut part_guard = self.fa_part_pool.lock().unwrap();
19307        if part_guard
19308            .as_ref()
19309            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
19310            .unwrap_or(true)
19311        {
19312            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
19313            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
19314            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
19315            // later live allocations land at those addresses, and the next graph REPLAY writes
19316            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
19317            // output corruption began the burst after the trunk's t_kv growth first realloc'd
19318            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
19319            // the baked addresses alive (single-stream: eager writes the new buffers, replays
19320            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
19321            // (total retired < final size).
19322            let old = part_guard.take();
19323            let (co, cm) = old
19324                .as_ref()
19325                .map(|pp| (pp.0.len(), pp.1.len()))
19326                .unwrap_or((0, 0));
19327            if let Some(old) = old {
19328                self.fa_part_retired.lock().unwrap().push(old);
19329            }
19330            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
19331                eprintln!(
19332                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
19333                    co, o_len, cm, ml_len
19334                );
19335            }
19336            *part_guard = Some((
19337                self.alloc_uninit::<f32>(o_len.max(2 * co))?,
19338                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
19339                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
19340            ));
19341        }
19342        let pg = part_guard.as_mut().unwrap();
19343        self.gpu
19344            .stream()
19345            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
19346        self.gpu
19347            .stream()
19348            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
19349        self.gpu
19350            .stream()
19351            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
19352        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
19353        let f = self.func("fa_decode_vec_q_rows_v3_dc");
19354        let sh = (32 * head_dim * 2) as u32;
19355        use cudarc::driver::sys::CUfunction_attribute_enum as A;
19356        f.set_attribute(
19357            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19358            sh as i32,
19359        )?;
19360        let cfg = LaunchConfig {
19361            grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
19362            block_dim: (32, gqa, 1),
19363            shared_mem_bytes: sh,
19364        };
19365        let __s_b = self.gpu.stream();
19366        let mut b = __s_b.launch_builder(&f);
19367        b.arg(q)
19368            .arg(k)
19369            .arg(v)
19370            .arg(&mut *part_o)
19371            .arg(&mut *part_m)
19372            .arg(&mut *part_l)
19373            .arg(&hd)
19374            .arg(&nh)
19375            .arg(&nhkv)
19376            .arg(base_dev)
19377            .arg(&scale)
19378            .arg(&nspm)
19379            .arg(&spk)
19380            .arg(&ktb)
19381            .arg(&vtb);
19382        unsafe {
19383            b.launch(cfg)?;
19384        }
19385        let fc = self.func("fa_decode_combine_rows_dc");
19386        let cfg2 = LaunchConfig {
19387            grid_dim: (n_head as u32, t as u32, 1),
19388            block_dim: (head_dim as u32, 1, 1),
19389            shared_mem_bytes: 0,
19390        };
19391        let plus0 = 0i32;
19392        let __s_b2 = self.gpu.stream();
19393        let mut b2 = __s_b2.launch_builder(&fc);
19394        b2.arg(&*part_o)
19395            .arg(&*part_m)
19396            .arg(&*part_l)
19397            .arg(o)
19398            .arg(&hd)
19399            .arg(&nh)
19400            .arg(base_dev)
19401            .arg(&plus0)
19402            .arg(&nspm)
19403            .arg(&spk);
19404        unsafe {
19405            b2.launch(cfg2)?;
19406        }
19407        Ok(())
19408    }
19409
19410    /// Device-counter variant of `fa_decode` (CUDA-GRAPH-PLAN Phase 2). The sequence length is read
19411    /// from `t_kv_dev[0]` (resident device i32[1]) for the attention loop bound + per-split key range;
19412    /// the GRID `n_splits` is sized for `bucket_max` (the bucket's max t_kv — baked at capture time).
19413    /// Empty splits (key range beyond the actual t_kv) write an empty partial (m=NEG_INF) so the
19414    /// shared combine skips them -> bit-correct for ANY actual t_kv <= bucket_max.
19415    ///
19416    /// BIT-IDENTITY (the gate): pass `bucket_max == actual_t_kv` and this reproduces `fa_decode`
19417    /// EXACTLY (same n_splits, same per, same split boundaries, same combine) while reading t_kv from
19418    /// device. Bucketing (bucket_max > t_kv) is for the future captured path and changes split
19419    /// grouping (different but mathematically-equal log-sum-exp merge).
19420    pub fn fa_decode_dc(
19421        &self,
19422        q: &CudaSlice<f32>,
19423        k: &cudarc::driver::CudaView<u8>,
19424        v: &cudarc::driver::CudaView<u8>,
19425        o: &mut CudaSlice<f32>,
19426        head_dim: usize,
19427        n_head: usize,
19428        n_head_kv: usize,
19429        t_kv_dev: &CudaSlice<i32>,
19430        bucket_max: usize,
19431        scale: f32,
19432        k_tok_bytes: usize,
19433        v_tok_bytes: usize,
19434        g: bool,
19435    ) -> Result<(), Box<dyn std::error::Error>> {
19436        self.fa_decode_dc_q8(
19437            q,
19438            k,
19439            v,
19440            o,
19441            head_dim,
19442            n_head,
19443            n_head_kv,
19444            t_kv_dev,
19445            bucket_max,
19446            scale,
19447            k_tok_bytes,
19448            v_tok_bytes,
19449            g,
19450            None,
19451        )
19452    }
19453
19454    /// `fa_decode_dc` with an optional q8_1 sink (wave 5b): when `q8_out` is given the
19455    /// combine emits (int8, per-32 scales) for the wo matmul_pre and skips the f32 O write.
19456    #[allow(clippy::too_many_arguments)]
19457    pub fn fa_decode_dc_q8(
19458        &self,
19459        q: &CudaSlice<f32>,
19460        k: &cudarc::driver::CudaView<u8>,
19461        v: &cudarc::driver::CudaView<u8>,
19462        o: &mut CudaSlice<f32>,
19463        head_dim: usize,
19464        n_head: usize,
19465        n_head_kv: usize,
19466        t_kv_dev: &CudaSlice<i32>,
19467        bucket_max: usize,
19468        scale: f32,
19469        k_tok_bytes: usize,
19470        v_tok_bytes: usize,
19471        g: bool,
19472        q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
19473    ) -> Result<(), Box<dyn std::error::Error>> {
19474        // The fa_vec gate + n_splits are sized from bucket_max (host, fixed at capture). The kernel
19475        // reads the ACTUAL t_kv from t_kv_dev for the per-split bound. DEFAULT-ON to MATCH the eager
19476        // `fa_decode` gate above — graph capture must mirror eager's kernel choice or the graph-vs-eager
19477        // bit-identity gate breaks. MEMRA_NO_FA_VEC forces scalar on BOTH paths in lockstep.
19478        // `g` = this layer's cache is e4m3 (gemma windowed under wkv) — every pick below must
19479        // mirror fa_decode_kvmod's g-routing or the graph diverges from eager (short/mid 1/96,
19480        // 2026-07-12).
19481        let mut fa_vec =
19482            std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
19483        if g && head_dim == 256 && !fa_v4_at(bucket_max) {
19484            fa_vec = false;
19485        } // mirror kvmod/geom
19486        let sp = fa_split_keys(bucket_max, n_head_kv);
19487        let n_splits = if fa_vec {
19488            ((bucket_max + sp - 1) / sp).max(1)
19489        } else {
19490            ((bucket_max + 255) / 256).max(1)
19491        };
19492        let o_len = n_head * n_splits * head_dim;
19493        let ml_len = n_head * n_splits;
19494        let mut part_guard = self.fa_part_pool.lock().unwrap();
19495        if part_guard
19496            .as_ref()
19497            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
19498            .unwrap_or(true)
19499        {
19500            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
19501            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
19502            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
19503            // later live allocations land at those addresses, and the next graph REPLAY writes
19504            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
19505            // output corruption began the burst after the trunk's t_kv growth first realloc'd
19506            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
19507            // the baked addresses alive (single-stream: eager writes the new buffers, replays
19508            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
19509            // (total retired < final size).
19510            let old = part_guard.take();
19511            let (co, cm) = old
19512                .as_ref()
19513                .map(|pp| (pp.0.len(), pp.1.len()))
19514                .unwrap_or((0, 0));
19515            if let Some(old) = old {
19516                self.fa_part_retired.lock().unwrap().push(old);
19517            }
19518            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
19519                eprintln!(
19520                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
19521                    co, o_len, cm, ml_len
19522                );
19523            }
19524            *part_guard = Some((
19525                self.alloc_uninit::<f32>(o_len.max(2 * co))?,
19526                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
19527                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
19528            ));
19529        }
19530        let pg = part_guard.as_mut().unwrap();
19531        self.gpu
19532            .stream()
19533            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
19534        self.gpu
19535            .stream()
19536            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
19537        self.gpu
19538            .stream()
19539            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
19540        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
19541        let (hd, nh, nhkv, nsp) = (
19542            head_dim as i32,
19543            n_head as i32,
19544            n_head_kv as i32,
19545            n_splits as i32,
19546        );
19547        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
19548        let fa_vec = fa_vec && head_dim <= 512 && head_dim % 32 == 0;
19549        // FA-DEEP pick keyed on bucket_max (the fa_v4_at precedent) — bit-identical twins,
19550        // so a threshold falling between t_kv and bucket_max cannot diverge eager-vs-graph.
19551        let deep = fa_vec
19552            && head_dim == 256
19553            && fa_v4_at(bucket_max)
19554            && !g
19555            && fa_deep_at(bucket_max)
19556            && !matches!(fa_v4_mode(), "noB3" | "stage");
19557        let (f, cfg) = if fa_vec
19558            && head_dim == 512
19559            && bucket_max >= {
19560                static FA512_MIN_DC: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
19561                *FA512_MIN_DC.get_or_init(|| {
19562                    std::env::var("MEMRA_FA512_MIN")
19563                        .ok()
19564                        .and_then(|v| v.parse().ok())
19565                        .unwrap_or(512)
19566                })
19567            } {
19568            // gemma globals dc twin (mirror the eager dpl16 pick incl the crossover floor).
19569            let gqa = (n_head / n_head_kv).max(1) as u32;
19570            (
19571                self.fa_func("fa_decode_vec_q_dpl16_dc", head_dim),
19572                LaunchConfig {
19573                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
19574                    block_dim: (32, gqa, 1),
19575                    shared_mem_bytes: 0,
19576                },
19577            )
19578        } else if fa_vec && head_dim == 512 {
19579            // under the 512 floor eager runs scalar — the SAME unified symbol, ctr non-null;
19580            // ns_eff in-kernel reproduces eager's ceil(t_kv/sp) partition for the LIVE len.
19581            let q_view = q.as_view();
19582            let mut o_view = o.as_view_mut();
19583            return self.fa_decode_scalar_unified(
19584                &q_view,
19585                k,
19586                v,
19587                &mut o_view,
19588                head_dim,
19589                n_head,
19590                n_head_kv,
19591                0,
19592                Some(t_kv_dev),
19593                scale,
19594                n_splits,
19595                sp,
19596                k_tok_bytes,
19597                v_tok_bytes,
19598                g,
19599                &mut *part_o,
19600                &mut *part_m,
19601                &mut *part_l,
19602                q8_out,
19603            );
19604        } else if fa_vec && head_dim == 256 && fa_v4_at(bucket_max) {
19605            // gemma/qwen v4 dc twin (eager default lane) — capture must mirror eager's pick,
19606            // incl the g-module route + raw-e4m3 sV sizing.
19607            let gqa = (n_head / n_head_kv).max(1) as u32;
19608            let fv = if g {
19609                self.func_g("fa_decode_vec_q_v4_dc")
19610            } else if deep {
19611                self.func("fa_decode_vec_q_v4_deep_dc")
19612            } else {
19613                self.func("fa_decode_vec_q_v4_dc")
19614            };
19615            let shmem =
19616                (if deep { 12160 } else { 11520 } + 32 * head_dim * if g { 1 } else { 2 }) as u32;
19617            use cudarc::driver::sys::CUfunction_attribute_enum as A;
19618            fv.set_attribute(
19619                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19620                shmem as i32,
19621            )?;
19622            (
19623                fv,
19624                LaunchConfig {
19625                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
19626                    block_dim: (32, gqa, 1),
19627                    shared_mem_bytes: shmem,
19628                },
19629            )
19630        } else if fa_vec && fa_v3_active(head_dim) {
19631            // FA v3 lane _dc twin: the captured graph must run the SAME walk body as eager
19632            // under MEMRA_FA_V3=1 (eager, rows-verify and graph switch together).
19633            let gqa = (n_head / n_head_kv).max(1) as u32;
19634            let fv = if g {
19635                self.func_g("fa_decode_vec_q_v3_dc")
19636            } else {
19637                self.func("fa_decode_vec_q_v3_dc")
19638            };
19639            let shmem = (32 * head_dim * 2) as u32; // sV bf16 [FA_DEC_TILE=32][hd]
19640            (
19641                fv,
19642                LaunchConfig {
19643                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
19644                    block_dim: (32, gqa, 1),
19645                    shared_mem_bytes: shmem,
19646                },
19647            )
19648        } else if fa_vec && fa_v2_on() {
19649            // FAVENDOR lane: v2 _dc twin — the captured graph must run the SAME walk body as
19650            // eager under MEMRA_FA_V2=1 or graph_decode_gate's bit-identity breaks (the flag is
19651            // a numeric config; eager, rows-verify and graph all switch together).
19652            let gqa = (n_head / n_head_kv).max(1) as u32;
19653            let fv = if g {
19654                self.func_g("fa_decode_vec_q_v2_dc")
19655            } else {
19656                self.func("fa_decode_vec_q_v2_dc")
19657            };
19658            let shmem = (2 * 32 * head_dim * 2) as u32; // sK+sV bf16 [FA_DEC_TILE=32][hd]
19659            (
19660                fv,
19661                LaunchConfig {
19662                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
19663                    block_dim: (32, gqa, 1),
19664                    shared_mem_bytes: shmem,
19665                },
19666            )
19667        } else if fa_vec {
19668            let gqa = (n_head / n_head_kv).max(1) as u32;
19669            // REGISTER-DEQUANT twin: zero dynamic smem (see fa_decode above).
19670            let fv = if g {
19671                self.func_g("fa_decode_vec_q_dc")
19672            } else {
19673                self.func("fa_decode_vec_q_dc")
19674            };
19675            (
19676                fv,
19677                LaunchConfig {
19678                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
19679                    block_dim: (32, gqa, 1),
19680                    shared_mem_bytes: 0,
19681                },
19682            )
19683        } else {
19684            let q_view = q.as_view();
19685            let mut o_view = o.as_view_mut();
19686            return self.fa_decode_scalar_unified(
19687                &q_view,
19688                k,
19689                v,
19690                &mut o_view,
19691                head_dim,
19692                n_head,
19693                n_head_kv,
19694                0,
19695                Some(t_kv_dev),
19696                scale,
19697                n_splits,
19698                if fa_vec { sp } else { 256 },
19699                k_tok_bytes,
19700                v_tok_bytes,
19701                g,
19702                &mut *part_o,
19703                &mut *part_m,
19704                &mut *part_l,
19705                q8_out,
19706            );
19707        };
19708        let ski = sp as i32; // one-partition law: the twins derive ns_eff from (T_kv, ski)
19709        let __s_b = self.gpu.stream();
19710        let mut b = __s_b.launch_builder(&f);
19711        b.arg(q)
19712            .arg(k)
19713            .arg(v)
19714            .arg(&mut *part_o)
19715            .arg(&mut *part_m)
19716            .arg(&mut *part_l)
19717            .arg(&hd)
19718            .arg(&nh)
19719            .arg(&nhkv)
19720            .arg(t_kv_dev)
19721            .arg(&scale)
19722            .arg(&nsp)
19723            .arg(&ski)
19724            .arg(&ktb)
19725            .arg(&vtb);
19726        unsafe {
19727            b.launch(cfg)?;
19728        }
19729        let cfg2 = LaunchConfig {
19730            grid_dim: (n_head as u32, 1, 1),
19731            block_dim: (head_dim as u32, 1, 1),
19732            shared_mem_bytes: 0,
19733        };
19734        if let Some((oq, od)) = q8_out {
19735            let fc = if g {
19736                self.func_g("fa_decode_combine_q8_1")
19737            } else {
19738                self.fa_func("fa_decode_combine_q8_1", head_dim)
19739            };
19740            let __s_b2 = self.gpu.stream();
19741            let mut b2 = __s_b2.launch_builder(&fc);
19742            b2.arg(&*part_o)
19743                .arg(&*part_m)
19744                .arg(&*part_l)
19745                .arg(oq)
19746                .arg(od)
19747                .arg(&hd)
19748                .arg(&nh)
19749                .arg(&nsp);
19750            unsafe {
19751                b2.launch(cfg2)?;
19752            }
19753            return Ok(());
19754        }
19755        let fc = if g {
19756            self.func_g("fa_decode_combine_f32")
19757        } else {
19758            self.fa_func("fa_decode_combine_f32", head_dim)
19759        };
19760        let __s_b2 = self.gpu.stream();
19761        let mut b2 = __s_b2.launch_builder(&fc);
19762        b2.arg(&*part_o)
19763            .arg(&*part_m)
19764            .arg(&*part_l)
19765            .arg(o)
19766            .arg(&hd)
19767            .arg(&nh)
19768            .arg(&nsp);
19769        unsafe {
19770            b2.launch(cfg2)?;
19771        }
19772        Ok(())
19773    }
19774
19775    /// EAGER fa_decode geometry for a given actual `t_kv` (CUDA-GRAPH-PLAN §3.3 bucketing). Returns
19776    /// `(fa_vec, n_splits)` EXACTLY as `fa_decode` computes them so the graph-capture path can key its
19777    /// bucket on the same `(kernel, n_splits)` pair and pass a `bucket_max` that reproduces eager's
19778    /// n_splits bit-for-bit. (Per = ceil(t_kv/n_splits) is then recomputed from the DEVICE t_kv inside
19779    /// the kernel and matches eager when n_splits matches — the bit-identity contract.)
19780    pub fn fa_geom_eager(
19781        &self,
19782        t_kv: usize,
19783        head_dim: usize,
19784        n_head_kv: usize,
19785        g: bool,
19786    ) -> (bool, usize) {
19787        // MUST mirror `fa_decode` / `fa_decode_dc` (default-ON 2026-06-28). This is the bucket-key
19788        // source: if it disagrees with the actual kernel pick, the graph captures the wrong path and
19789        // replay diverges from eager. All three sites read MEMRA_NO_FA_VEC in lockstep.
19790        let fa_ok = std::env::var("MEMRA_NO_FA_VEC").is_err() && t_kv >= fa_vec_min_tkv();
19791        // hd512 dpl16 vec lane (gemma globals, 2026-07-11 graph-arc fix): the original key
19792        // hardcoded vec = hd<=256, so for hd512 it bucketed by the SCALAR 256-key splits while
19793        // the dpl16/rows_dpl16 kernels split by the ladder — n_splits changed WITHIN a bucket
19794        // (mid-ctx graph mismatch at pos 19 + partials OOB at longer runs). Mirror the real
19795        // fa_decode dispatch: vec512 above the fa512 floor, vec256 as before.
19796        let vec512 = fa_ok && head_dim == 512 && t_kv >= fa512_min_tkv();
19797        let mut fa_vec = vec512 || (fa_ok && head_dim <= 256 && head_dim % 32 == 0);
19798        // g (fp8-windowed): mirror kvmod's clamp — only the v4 lane parses e4m3 in the vec
19799        // family; everything else falls to the g-module scalar.
19800        // hd256 under g: v4-or-scalar. hd128/other under g ride the REGISTER g-lane (the
19801        // dq_K_lane/dq_V_lane macros are format-aware; v3 is format-gated off, v2/smem
19802        // arms are excluded under g). Mirrored in kvmod / fa_decode_dc / fa_geom_eager.
19803        if g && head_dim == 256 && !fa_v4_at(t_kv) {
19804            fa_vec = false;
19805        }
19806        let sp = fa_split_keys(t_kv, n_head_kv);
19807        let n_splits = if fa_vec {
19808            ((t_kv + sp - 1) / sp).max(1)
19809        } else {
19810            ((t_kv + 255) / 256).max(1)
19811        };
19812        (fa_vec, n_splits)
19813    }
19814
19815    /// `bucket_max` (host t_kv to feed `fa_decode_dc` / `full_attn_decode_dc`) that makes the _dc
19816    /// kernel pick the SAME (fa_vec, n_splits) as eager would for actual `t_kv`. Because the dc
19817    /// launcher derives both from `bucket_max` via the same formulas, we just hand it `t_kv` itself:
19818    /// the n_splits is then identical, and the per-split boundaries (computed from the DEVICE t_kv in
19819    /// the kernel) match eager exactly. The bucket KEY (for the graph HashMap) is `(fa_vec, n_splits)`.
19820    pub fn fa_bucket_key(
19821        &self,
19822        t_kv: usize,
19823        head_dim: usize,
19824        n_head_kv: usize,
19825        g: bool,
19826    ) -> (bool, usize) {
19827        self.fa_geom_eager(t_kv, head_dim, n_head_kv, g)
19828    }
19829
19830    /// CUDA-graph capture wrapper (CUDA-GRAPH-PLAN §3.2, llama.cpp warmup pattern). Runs `step`
19831    /// inline TWICE (warmup — lets the caching allocator settle to stable pointers and any one-time
19832    /// kernel attribute/JIT happen outside capture), then captures a THIRD invocation on the Engine's
19833    /// decode stream (RELAXED mode) and instantiates it into a replayable `CudaGraph`. The closure
19834    /// must enqueue ONLY device work on `e.stream()` (no dtoh / no synchronize / no host branch on
19835    /// device data) — every per-step varying scalar must come from a device counter. Returns the
19836    /// instantiated graph; `CudaGraph::launch()` replays the whole step in one dispatch.
19837    /// `capture_graph` with CAPTURE-RETAIN: every Engine allocation made during the warmups
19838    /// and the capture is kept alive in the returned keeper — hold it as long as the graph
19839    /// replays (transients returning to the pool get reused by unrelated work and corrupt
19840    /// replays; the draft-graph root cause). Model-generic, next capture reuses it.
19841    pub fn capture_graph_retained<F>(
19842        &self,
19843        step: F,
19844    ) -> Result<
19845        (
19846            cudarc::driver::CudaGraph,
19847            Vec<Box<dyn std::any::Any + Send>>,
19848        ),
19849        Box<dyn std::error::Error>,
19850    >
19851    where
19852        F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
19853    {
19854        use cudarc::driver::sys::CUgraphInstantiate_flags;
19855        self.capture_graph_retained_flags(
19856            CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
19857            step,
19858        )
19859    }
19860
19861    /// Retained capture with an explicit instantiate flag. ALLOC-FREE captured graphs
19862    /// (zero mem nodes — the gemma slotted door) should pass UPLOAD instead of
19863    /// AUTO_FREE_ON_LAUNCH: the auto-free flag's launch-time mem-pool scan was measured at
19864    /// ~0.25us/node (205us on the 826-node step) even with nothing to free.
19865    pub fn capture_graph_retained_flags<F>(
19866        &self,
19867        flags: cudarc::driver::sys::CUgraphInstantiate_flags,
19868        mut step: F,
19869    ) -> Result<
19870        (
19871            cudarc::driver::CudaGraph,
19872            Vec<Box<dyn std::any::Any + Send>>,
19873        ),
19874        Box<dyn std::error::Error>,
19875    >
19876    where
19877        F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
19878    {
19879        use cudarc::driver::sys::CUstreamCaptureMode;
19880        // KEEP scope = WARMUPS ONLY (2026-07-13): keep_if_capturing retains via
19881        // CudaSlice::clone, which is a device ALLOC + D2D COPY on the stream — clones made
19882        // while the capture region is open become dead copy NODES replayed every launch
19883        // (E4B: 1440 copies = 0.74ms/token, the whole graph-vs-eager regression). The
19884        // warmup runs allocate the same transient sequence at the same pool addresses, so
19885        // retaining the warmup clones preserves the draft-graph fix without polluting the
19886        // captured graph.
19887        self.capture_keep.lock().unwrap().clear();
19888        let was_tracking = self.gpu.ctx.is_event_tracking();
19889        if was_tracking {
19890            unsafe {
19891                self.gpu.ctx.disable_event_tracking();
19892            }
19893        }
19894        let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
19895            self.capture_keep_on
19896                .store(true, std::sync::atomic::Ordering::Relaxed);
19897            let w = (|| {
19898                step(self)?;
19899                step(self)
19900            })();
19901            self.capture_keep_on
19902                .store(false, std::sync::atomic::Ordering::Relaxed);
19903            w?;
19904            self.gpu.stream().synchronize()?;
19905            self.gpu
19906                .stream()
19907                .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
19908            let r = step(self);
19909            let g = self.gpu.stream().end_capture(flags);
19910            r?;
19911            let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
19912            graph.upload()?;
19913            Ok(graph)
19914        };
19915        let result = run();
19916        self.capture_keep_on
19917            .store(false, std::sync::atomic::Ordering::Relaxed);
19918        if was_tracking {
19919            unsafe {
19920                self.gpu.ctx.enable_event_tracking();
19921            }
19922        }
19923        let keeper = std::mem::take(&mut *self.capture_keep.lock().unwrap());
19924        Ok((result?, keeper))
19925    }
19926
19927    pub fn capture_graph<F>(
19928        &self,
19929        mut step: F,
19930    ) -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>>
19931    where
19932        F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
19933    {
19934        use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
19935        // EVENT TRACKING OFF for capture. The Engine creates a 2nd stream (copy_stream) so cudarc is in
19936        // multi-stream mode and, by default, records a CudaEvent per CudaSlice alloc/use to serialize
19937        // cross-stream access. Those per-buffer event waits issue stream ops that are NOT permitted
19938        // inside a capture region (CUDA_ERROR_STREAM_CAPTURE_UNSUPPORTED). The captured decode step is
19939        // strictly SINGLE-STREAM (every kernel on gpu.stream), so this synchronization is unnecessary
19940        // here — disable it for the whole warmup+capture, re-enable after. SAFETY: the decode-dc path
19941        // touches only gpu.stream; no buffer crosses to copy_stream during capture.
19942        let was_tracking = self.gpu.ctx.is_event_tracking();
19943        if was_tracking {
19944            unsafe {
19945                self.gpu.ctx.disable_event_tracking();
19946            }
19947        }
19948        // Q1 PROBE (MEMRA_GRAPH_IFLAG): the generic capture body's cuMemAllocAsync nodes are
19949        // EXACTLY BALANCED by in-graph free nodes (measured census q27: 1589 ALLOC / 1589
19950        // FREE), so AUTO_FREE_ON_LAUNCH has nothing to reclaim at launch — it only pays its
19951        // per-node launch-time mem-pool scan. `upload` / `none` select the alternatives to
19952        // measure that scan's real cost on the generic path. Diagnostic door only; the
19953        // default stays AUTO_FREE until a measured A/B justifies moving it.
19954        let iflag = {
19955            static F: std::sync::OnceLock<CUgraphInstantiate_flags> = std::sync::OnceLock::new();
19956            *F.get_or_init(|| match std::env::var("MEMRA_GRAPH_IFLAG").as_deref() {
19957                // UPLOAD = the gemma slotted door's zero-mem-node choice; PRIORITY = the flag
19958                // hybrid_forward.rs:5935 actually ships (both drop the auto-free launch scan).
19959                Ok("upload") => CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_UPLOAD,
19960                Ok("priority") => {
19961                    CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY
19962                }
19963                _ => CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
19964            })
19965        };
19966        // MEMRA_GRAPH_CAPTIME=1 (Q1 lane): phase-resolved capture cost. Recapture is paid at
19967        // every kernel-class crossing, so it — not steady-state decode — is the quantity a
19968        // mem-node reduction could plausibly shrink. Only `instantiate` (cuStreamEndCapture +
19969        // cuGraphInstantiateWithFlags) and `upload` scale with node count; the warmups are
19970        // eager step executions and are node-count-invariant. Printing the split bounds the
19971        // refactor's ceiling instead of assuming it.
19972        let ct = {
19973            static T: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19974            *T.get_or_init(|| std::env::var("MEMRA_GRAPH_CAPTIME").as_deref() == Ok("1"))
19975        };
19976        // MEMRA_GRAPH_WARMUPS (Q1 lane; DEFAULT 1 since lane/graph-warmups 2026-08-05): the
19977        // phase split showed the eager warmups are 80% of recapture cost (q27 27.4 of 34.4 ms
19978        // pod / 42% of 52.6 ms 5090) — 3x larger than the ENTIRE mem-node ceiling the audit
19979        // chased, and node-count-invariant, so no capture-body refactor could touch it.
19980        // Warmup 2's theorized job was async-pool ADDRESS STABILITY: warmup 1's allocs may
19981        // grow/map the pool, warmup 2 re-walks the same sequence over the freed blocks so the
19982        // captured third run bakes settled addresses. That hazard is the #68 stale-baked-
19983        // address class — which the engine now guards STRUCTURALLY rather than by re-walking:
19984        // in-body transients are captured as BALANCED in-graph alloc/free node pairs (census
19985        // 1589/1589 — replays allocate for themselves; no baked transient pointers), every
19986        // externally-referenced buffer is stable-pointer by design (fa_part_pool retires-on-
19987        // grow and never frees, resident counters/scratch, cache set in place), and the
19988        // draft-graph path additionally rides capture_graph_retained (capture_keep holds all
19989        // warmup+capture allocs alive). One warmup therefore suffices for kernel-attr
19990        // settling and pool mapping. Arbitrated adversarially, not by taste:
19991        // graph-warmup-stress (pool-growth cycles large<->small x10, overlap arm, forced
19992        // recaptures over freed blocks — bit-identity vs eager + canary teeth) is GREEN at
19993        // warmups=1 on the deployment rig, plus graph-decode-gate 256-step bit-identity,
19994        // graph-session-gate, run-spec K=1..8 (receipts research/graph-warmups-5090-20260805/
19995        // + the pod's research/graph-allocfree-20260805/). Measured: recapture -38..-42% q27 /
19996        // -41% q9, decode +~1%, capture+prime -13ms. MEMRA_GRAPH_WARMUPS=2 = the rollback
19997        // seam; tools/graph-warmup-stress-gate.sh = the gate any regression re-runs.
19998        let warmups = {
19999            static W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
20000            *W.get_or_init(|| {
20001                std::env::var("MEMRA_GRAPH_WARMUPS")
20002                    .ok()
20003                    .and_then(|v| v.parse().ok())
20004                    .filter(|n| *n >= 1)
20005                    .unwrap_or(1)
20006            })
20007        };
20008        let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
20009            let t_w = std::time::Instant::now();
20010            // warmup: inline runs (no capture) so allocator pointers + kernel attrs are stable.
20011            for _ in 0..warmups {
20012                step(self)?;
20013            }
20014            self.gpu.stream().synchronize()?;
20015            let ms_warm = t_w.elapsed().as_secs_f64() * 1e3;
20016            // capture the third run.
20017            let t_c = std::time::Instant::now();
20018            self.gpu
20019                .stream()
20020                .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
20021            // If the body errors mid-capture, end the capture before propagating so the stream isn't
20022            // left in a capturing state.
20023            let r = step(self);
20024            let ms_body = t_c.elapsed().as_secs_f64() * 1e3;
20025            let t_i = std::time::Instant::now();
20026            let g = self.gpu.stream().end_capture(iflag);
20027            let ms_inst = t_i.elapsed().as_secs_f64() * 1e3;
20028            r?;
20029            let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
20030            let t_u = std::time::Instant::now();
20031            graph.upload()?;
20032            if ct {
20033                println!(
20034                    "[graph-captime] warmup2x {ms_warm:.2} ms  capture-body {ms_body:.2} ms  \
20035                          instantiate {ms_inst:.2} ms  upload {:.2} ms",
20036                    t_u.elapsed().as_secs_f64() * 1e3
20037                );
20038            }
20039            Ok(graph)
20040        };
20041        let result = run();
20042        if was_tracking {
20043            unsafe {
20044                self.gpu.ctx.enable_event_tracking();
20045            }
20046        }
20047        result
20048    }
20049
20050    /// gdn_scan variant where state_in/out are CudaViews (resident SSM state, in-place per step).
20051    pub fn gdn_scan_s128_view(
20052        &self,
20053        q: &CudaSlice<f32>,
20054        k: &CudaSlice<f32>,
20055        v: &CudaSlice<f32>,
20056        g: &CudaSlice<f32>,
20057        beta: &CudaSlice<f32>,
20058        state_in: &cudarc::driver::CudaView<f32>,
20059        state_out: &mut cudarc::driver::CudaViewMut<f32>,
20060        o: &mut CudaSlice<f32>,
20061        n_head: usize,
20062        t: usize,
20063        scale: f32,
20064    ) -> Result<(), Box<dyn std::error::Error>> {
20065        let f = self.func("gdn_scan_s128");
20066        const S_V: u32 = 128;
20067        const WARP: u32 = 32;
20068        const COLS: u32 = 4;
20069        let cfg = LaunchConfig {
20070            grid_dim: (n_head as u32, 1, S_V / COLS),
20071            block_dim: (WARP, COLS, 1),
20072            shared_mem_bytes: 0,
20073        };
20074        let (h, ti) = (n_head as i32, t as i32);
20075        let __s_b = self.gpu.stream();
20076        let mut b = __s_b.launch_builder(&f);
20077        b.arg(q)
20078            .arg(k)
20079            .arg(v)
20080            .arg(g)
20081            .arg(beta)
20082            .arg(state_in)
20083            .arg(state_out)
20084            .arg(o)
20085            .arg(&h)
20086            .arg(&ti)
20087            .arg(&scale);
20088        unsafe {
20089            b.launch(cfg)?;
20090        }
20091        Ok(())
20092    }
20093
20094    /// conv1d where the input is a CudaView (resident conv state assembled in place).
20095    pub fn ssm_conv1d_view(
20096        &self,
20097        x: &cudarc::driver::CudaView<f32>,
20098        w: &CudaSlice<f32>,
20099        y: &mut CudaSlice<f32>,
20100        conv_dim: usize,
20101        t: usize,
20102        d_conv: usize,
20103        silu: bool,
20104    ) -> Result<(), Box<dyn std::error::Error>> {
20105        let f = self.func("ssm_conv1d_silu_f32");
20106        // grid.x = channel, grid.y = T-tiles (block 256 strides over T) — parallel over both axes.
20107        let cfg = LaunchConfig {
20108            grid_dim: (conv_dim as u32, ((t as u32 + 255) / 256).max(1), 1),
20109            block_dim: (256, 1, 1),
20110            shared_mem_bytes: 0,
20111        };
20112        let (cd, ti, dc, s) = (conv_dim as i32, t as i32, d_conv as i32, silu as i32);
20113        let __s_b = self.gpu.stream();
20114        let mut b = __s_b.launch_builder(&f);
20115        b.arg(x).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc).arg(&s);
20116        unsafe {
20117            b.launch(cfg)?;
20118        }
20119        Ok(())
20120    }
20121
20122    /// Depthwise causal conv1d + optional SiLU.
20123    /// x:[conv_dim, T+d_conv-1] channel-major (first d_conv-1 cols = carried state),
20124    /// w:[d_conv, conv_dim] kernel-major, y:[conv_dim, T] channel-major.
20125    /// FUSED prefill conv (token-major input, zero left-state): replaces
20126    /// transpose + zeros + conv_left_pad + ssm_conv1d with ONE launch reading the matmul output
20127    /// directly. Output channel-major [conv_dim, T], SiLU applied. BIT-IDENTICAL accumulation.
20128    pub fn ssm_conv1d_tm(
20129        &self,
20130        qkv_tm: &CudaSlice<f32>,
20131        w: &CudaSlice<f32>,
20132        y: &mut CudaSlice<f32>,
20133        conv_dim: usize,
20134        t: usize,
20135        d_conv: usize,
20136    ) -> Result<(), Box<dyn std::error::Error>> {
20137        let f = self.func("ssm_conv1d_tm_f32");
20138        let cfg = LaunchConfig {
20139            grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
20140            block_dim: (256, 1, 1),
20141            shared_mem_bytes: 0,
20142        };
20143        let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
20144        let __s_b = self.gpu.stream();
20145        let mut b = __s_b.launch_builder(&f);
20146        b.arg(qkv_tm).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc);
20147        unsafe {
20148            b.launch(cfg)?;
20149        }
20150        Ok(())
20151    }
20152
20153    /// BATCHED verify conv (T>1, carried state): window reads the resident conv ring for
20154    /// negative rows; separate ring-update launch afterwards. BIT-IDENTICAL per value to the
20155    /// T=1 chain. T >= pad rides the pure input-column ring update (unchanged legacy path);
20156    /// T < pad (the MEMRA_SPEC_M2 t=2 verify arm) needs old-ring sources for the roll — the
20157    /// update kernel would race reading the ring it rewrites, so that arm clones the ring
20158    /// (dtod) and rolls via ssm_conv_ring_rebuild (PURE COPIES: the ring stores raw input
20159    /// columns; the final ring == what T sequential decode ring rolls leave).
20160    pub fn ssm_conv1d_tm_state(
20161        &self,
20162        qkv_tm: &CudaSlice<f32>,
20163        conv_state: &mut CudaSlice<f32>,
20164        w: &CudaSlice<f32>,
20165        y: &mut CudaSlice<f32>,
20166        conv_dim: usize,
20167        t: usize,
20168        d_conv: usize,
20169    ) -> Result<(), Box<dyn std::error::Error>> {
20170        self.ssm_conv1d_tm_state_pad(qkv_tm, conv_state, w, y, conv_dim, t, d_conv, None)
20171    }
20172
20173    /// task #14: `pad_len` = device true length for PADDED prime graphs — the ring update
20174    /// reads rows [len-pad, len) instead of the pad tail. None = the classic host-T path.
20175    #[allow(clippy::too_many_arguments)]
20176    pub fn ssm_conv1d_tm_state_pad(
20177        &self,
20178        qkv_tm: &CudaSlice<f32>,
20179        conv_state: &mut CudaSlice<f32>,
20180        w: &CudaSlice<f32>,
20181        y: &mut CudaSlice<f32>,
20182        conv_dim: usize,
20183        t: usize,
20184        d_conv: usize,
20185        pad_len: Option<&CudaSlice<i32>>,
20186    ) -> Result<(), Box<dyn std::error::Error>> {
20187        assert!(t >= 1, "ssm_conv1d_tm_state requires T >= 1");
20188        // clone BEFORE the window kernel is issued is not required (stream-ordered: the dtod and
20189        // the window kernel both read the pre-roll ring; the roll launches after both) — but
20190        // cloning first keeps the ordering trivially correct under any future stream split.
20191        let ring_old = if t < d_conv - 1 {
20192            Some(self.clone_dtod(conv_state)?)
20193        } else {
20194            None
20195        };
20196        {
20197            let f = self.func("ssm_conv1d_tm_state_f32");
20198            let cfg = LaunchConfig {
20199                grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
20200                block_dim: (256, 1, 1),
20201                shared_mem_bytes: 0,
20202            };
20203            let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
20204            let __s_b = self.gpu.stream();
20205            let mut b = __s_b.launch_builder(&f);
20206            b.arg(qkv_tm)
20207                .arg(&*conv_state)
20208                .arg(w)
20209                .arg(y)
20210                .arg(&cd)
20211                .arg(&ti)
20212                .arg(&dc);
20213            unsafe {
20214                b.launch(cfg)?;
20215            }
20216        }
20217        match (ring_old, pad_len) {
20218            (None, Some(len_d)) => {
20219                let f = self.func("ssm_conv_ring_update_dev_f32");
20220                let n = conv_dim * (d_conv - 1);
20221                let cfg = LaunchConfig::for_num_elems(n as u32);
20222                let (cd, dc) = (conv_dim as i32, d_conv as i32);
20223                let __s_b = self.gpu.stream();
20224                let mut b = __s_b.launch_builder(&f);
20225                b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
20226                unsafe {
20227                    b.launch(cfg)?;
20228                }
20229            }
20230            (None, None) => {
20231                let f = self.func("ssm_conv_ring_update_f32");
20232                let n = conv_dim * (d_conv - 1);
20233                let cfg = LaunchConfig::for_num_elems(n as u32);
20234                let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
20235                let __s_b = self.gpu.stream();
20236                let mut b = __s_b.launch_builder(&f);
20237                b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
20238                unsafe {
20239                    b.launch(cfg)?;
20240                }
20241            }
20242            (Some(old), _) => {
20243                self.ssm_conv_ring_rebuild(qkv_tm, &old, conv_state, conv_dim, t, d_conv)?
20244            }
20245        }
20246        Ok(())
20247    }
20248
20249    /// qkv-view twin (task #16): batched prime reads the concat GEMM output directly.
20250    pub fn ssm_conv1d_tm_state_pad_v(
20251        &self,
20252        qkv_tm: &cudarc::driver::CudaView<f32>,
20253        conv_state: &mut CudaSlice<f32>,
20254        w: &CudaSlice<f32>,
20255        y: &mut CudaSlice<f32>,
20256        conv_dim: usize,
20257        t: usize,
20258        d_conv: usize,
20259        pad_len: Option<&CudaSlice<i32>>,
20260    ) -> Result<(), Box<dyn std::error::Error>> {
20261        assert!(t >= 1, "ssm_conv1d_tm_state requires T >= 1");
20262        // clone BEFORE the window kernel is issued is not required (stream-ordered: the dtod and
20263        // the window kernel both read the pre-roll ring; the roll launches after both) — but
20264        // cloning first keeps the ordering trivially correct under any future stream split.
20265        let ring_old = if t < d_conv - 1 {
20266            Some(self.clone_dtod(conv_state)?)
20267        } else {
20268            None
20269        };
20270        {
20271            let f = self.func("ssm_conv1d_tm_state_f32");
20272            let cfg = LaunchConfig {
20273                grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
20274                block_dim: (256, 1, 1),
20275                shared_mem_bytes: 0,
20276            };
20277            let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
20278            let __s_b = self.gpu.stream();
20279            let mut b = __s_b.launch_builder(&f);
20280            b.arg(qkv_tm)
20281                .arg(&*conv_state)
20282                .arg(w)
20283                .arg(y)
20284                .arg(&cd)
20285                .arg(&ti)
20286                .arg(&dc);
20287            unsafe {
20288                b.launch(cfg)?;
20289            }
20290        }
20291        match (ring_old, pad_len) {
20292            (None, Some(len_d)) => {
20293                let f = self.func("ssm_conv_ring_update_dev_f32");
20294                let n = conv_dim * (d_conv - 1);
20295                let cfg = LaunchConfig::for_num_elems(n as u32);
20296                let (cd, dc) = (conv_dim as i32, d_conv as i32);
20297                let __s_b = self.gpu.stream();
20298                let mut b = __s_b.launch_builder(&f);
20299                b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
20300                unsafe {
20301                    b.launch(cfg)?;
20302                }
20303            }
20304            (None, None) => {
20305                let f = self.func("ssm_conv_ring_update_f32");
20306                let n = conv_dim * (d_conv - 1);
20307                let cfg = LaunchConfig::for_num_elems(n as u32);
20308                let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
20309                let __s_b = self.gpu.stream();
20310                let mut b = __s_b.launch_builder(&f);
20311                b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
20312                unsafe {
20313                    b.launch(cfg)?;
20314                }
20315            }
20316            (Some(_), _) => unreachable!(
20317                "ssm_conv1d_tm_state_pad_v: T < d_conv-1 has no view path (PRIME_MIN_T gates it)"
20318            ),
20319        }
20320        Ok(())
20321    }
20322
20323    /// PREFIX conv-ring rebuild (spec REPLAY-FREE partial accept): overwrite the resident ring
20324    /// with the state a T=1 chain holds after only the FIRST `tc` columns of `qkv_tm` — the last
20325    /// `pad` entries of [ring_old | cols 0..tc-1]. PURE COPIES (the ring stores raw inputs; no
20326    /// arithmetic, cannot perturb FP order). `ring_old` = the pre-round snapshot ring.
20327    pub fn ssm_conv_ring_rebuild(
20328        &self,
20329        qkv_tm: &CudaSlice<f32>,
20330        ring_old: &CudaSlice<f32>,
20331        conv_state: &mut CudaSlice<f32>,
20332        conv_dim: usize,
20333        tc: usize,
20334        d_conv: usize,
20335    ) -> Result<(), Box<dyn std::error::Error>> {
20336        let f = self.func("ssm_conv_ring_rebuild_f32");
20337        let n = conv_dim * (d_conv - 1);
20338        let cfg = LaunchConfig::for_num_elems(n as u32);
20339        let (cd, ti, dc) = (conv_dim as i32, tc as i32, d_conv as i32);
20340        let __s_b = self.gpu.stream();
20341        let mut b = __s_b.launch_builder(&f);
20342        b.arg(qkv_tm)
20343            .arg(ring_old)
20344            .arg(conv_state)
20345            .arg(&cd)
20346            .arg(&ti)
20347            .arg(&dc);
20348        unsafe {
20349            b.launch(cfg)?;
20350        }
20351        Ok(())
20352    }
20353
20354    /// FUSED decode GDN prep (T=1): repack + q/k L2-norm + beta sigmoid + g_log in one launch.
20355    /// Replaces 5 tiny serialized kernels on the decode critical path. L2 reduce runs as a 32-lane
20356    /// warp tree (vs l2_norm_f32's 256-thread two-level tree) — same math, different FP sum order;
20357    /// the argmax + run-spec gates are the authority.
20358    #[allow(clippy::too_many_arguments)]
20359    pub fn gdn_prep_decode(
20360        &self,
20361        conv_out: &CudaSlice<f32>,
20362        beta_raw: &CudaSlice<f32>,
20363        alpha: &CudaSlice<f32>,
20364        dt_bias: &CudaSlice<f32>,
20365        a: &CudaSlice<f32>,
20366        q_l2: &mut CudaSlice<f32>,
20367        k_l2: &mut CudaSlice<f32>,
20368        v_g: &mut CudaSlice<f32>,
20369        beta: &mut CudaSlice<f32>,
20370        g_log: &mut CudaSlice<f32>,
20371        d_state: usize,
20372        num_v: usize,
20373        num_k: usize,
20374        key_dim: usize,
20375        eps: f32,
20376    ) -> Result<(), Box<dyn std::error::Error>> {
20377        let f = self.func("gdn_prep_decode_f32");
20378        let cfg = LaunchConfig {
20379            grid_dim: (num_v as u32, 1, 1),
20380            block_dim: (32, 4, 1),
20381            shared_mem_bytes: 0,
20382        };
20383        let (ds, nv, nk, kd) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
20384        let __s_b = self.gpu.stream();
20385        let mut b = __s_b.launch_builder(&f);
20386        b.arg(conv_out)
20387            .arg(beta_raw)
20388            .arg(alpha)
20389            .arg(dt_bias)
20390            .arg(a)
20391            .arg(q_l2)
20392            .arg(k_l2)
20393            .arg(v_g)
20394            .arg(beta)
20395            .arg(g_log)
20396            .arg(&ds)
20397            .arg(&nv)
20398            .arg(&nk)
20399            .arg(&kd)
20400            .arg(&eps);
20401        unsafe {
20402            b.launch(cfg)?;
20403        }
20404        Ok(())
20405    }
20406
20407    /// FUSED prefill conv + GDN repack: token-major qkv -> q_g/k_g/v_g in ONE launch (no conv_out
20408    /// materialization, no qkv_to_gdn_repack pass). BIT-IDENTICAL values; scatter matches
20409    /// qkv_to_gdn_repack's modulo head-repeat mapping exactly.
20410    #[allow(clippy::too_many_arguments)]
20411    pub fn ssm_conv1d_gdn(
20412        &self,
20413        qkv_tm: &CudaSlice<f32>,
20414        w: &CudaSlice<f32>,
20415        q_g: &mut CudaSlice<f32>,
20416        k_g: &mut CudaSlice<f32>,
20417        v_g: &mut CudaSlice<f32>,
20418        conv_dim: usize,
20419        t: usize,
20420        d_conv: usize,
20421        d_state: usize,
20422        num_v: usize,
20423        num_k: usize,
20424        key_dim: usize,
20425    ) -> Result<(), Box<dyn std::error::Error>> {
20426        let f = self.func("ssm_conv1d_gdn_f32");
20427        let cfg = LaunchConfig {
20428            grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
20429            block_dim: (256, 1, 1),
20430            shared_mem_bytes: 0,
20431        };
20432        let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
20433        let (ds, nv, nk, kd) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
20434        let __s_b = self.gpu.stream();
20435        let mut b = __s_b.launch_builder(&f);
20436        b.arg(qkv_tm)
20437            .arg(w)
20438            .arg(q_g)
20439            .arg(k_g)
20440            .arg(v_g)
20441            .arg(&cd)
20442            .arg(&ti)
20443            .arg(&dc)
20444            .arg(&ds)
20445            .arg(&nv)
20446            .arg(&nk)
20447            .arg(&kd);
20448        unsafe {
20449            b.launch(cfg)?;
20450        }
20451        Ok(())
20452    }
20453
20454    pub fn ssm_conv1d(
20455        &self,
20456        x: &CudaSlice<f32>,
20457        w: &CudaSlice<f32>,
20458        y: &mut CudaSlice<f32>,
20459        conv_dim: usize,
20460        t: usize,
20461        d_conv: usize,
20462        silu: bool,
20463    ) -> Result<(), Box<dyn std::error::Error>> {
20464        let f = self.func("ssm_conv1d_silu_f32");
20465        let cfg = LaunchConfig {
20466            grid_dim: (conv_dim as u32, ((t as u32 + 255) / 256).max(1), 1),
20467            block_dim: (256, 1, 1),
20468            shared_mem_bytes: 0,
20469        };
20470        let (cd, ti, dc, s) = (conv_dim as i32, t as i32, d_conv as i32, silu as i32);
20471        let __s_b = self.gpu.stream();
20472        let mut b = __s_b.launch_builder(&f);
20473        b.arg(x).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc).arg(&s);
20474        unsafe {
20475            b.launch(cfg)?;
20476        }
20477        Ok(())
20478    }
20479
20480    /// Gated DeltaNet scan, S_v=128. q,k,v:[128,H,T]; g,beta:[H,T]; state:[128,128,H] transposed;
20481    /// o:[128,H,T]. Single sequence.
20482    pub fn gdn_scan_s128(
20483        &self,
20484        q: &CudaSlice<f32>,
20485        k: &CudaSlice<f32>,
20486        v: &CudaSlice<f32>,
20487        g: &CudaSlice<f32>,
20488        beta: &CudaSlice<f32>,
20489        state_in: &CudaSlice<f32>,
20490        state_out: &mut CudaSlice<f32>,
20491        o: &mut CudaSlice<f32>,
20492        n_head: usize,
20493        t: usize,
20494        scale: f32,
20495    ) -> Result<(), Box<dyn std::error::Error>> {
20496        let f = self.func("gdn_scan_s128");
20497        const S_V: u32 = 128;
20498        const WARP: u32 = 32;
20499        const COLS_PER_BLOCK: u32 = 4;
20500        let cfg = LaunchConfig {
20501            grid_dim: (n_head as u32, 1, S_V / COLS_PER_BLOCK),
20502            block_dim: (WARP, COLS_PER_BLOCK, 1),
20503            shared_mem_bytes: 0,
20504        };
20505        let (h, ti) = (n_head as i32, t as i32);
20506        let __s_b = self.gpu.stream();
20507        let mut b = __s_b.launch_builder(&f);
20508        b.arg(q)
20509            .arg(k)
20510            .arg(v)
20511            .arg(g)
20512            .arg(beta)
20513            .arg(state_in)
20514            .arg(state_out)
20515            .arg(o)
20516            .arg(&h)
20517            .arg(&ti)
20518            .arg(&scale);
20519        unsafe {
20520            b.launch(cfg)?;
20521        }
20522        Ok(())
20523    }
20524
20525    // ==== B2' batched decode state ops (decode_batch.rs) ====
20526    // Per-seq state pointers ride device u64 arrays (views into the per-step pointer table).
20527    // Bodies are the single-seq kernels per sequence — bit-identical per row.
20528
20529    #[allow(clippy::too_many_arguments)]
20530    pub fn ssm_conv1d_fused_decode_b(
20531        &self,
20532        qkv_cols: &CudaSlice<f32>,
20533        conv_state_ptrs: &cudarc::driver::CudaView<u64>,
20534        w: &CudaSlice<f32>,
20535        conv_outs: &mut CudaSlice<f32>,
20536        conv_dim: usize,
20537        d_conv: usize,
20538        b_n: usize,
20539    ) -> Result<(), Box<dyn std::error::Error>> {
20540        let f = self.func("ssm_conv1d_fused_decode_b_f32");
20541        let cfg = LaunchConfig {
20542            grid_dim: (((conv_dim + 255) / 256) as u32, 1, b_n as u32),
20543            block_dim: (256, 1, 1),
20544            shared_mem_bytes: 0,
20545        };
20546        let (cd, dc) = (conv_dim as i32, d_conv as i32);
20547        let __s_b = self.gpu.stream();
20548        let mut b = __s_b.launch_builder(&f);
20549        b.arg(qkv_cols)
20550            .arg(conv_state_ptrs)
20551            .arg(w)
20552            .arg(conv_outs)
20553            .arg(&cd)
20554            .arg(&dc);
20555        unsafe {
20556            b.launch(cfg)?;
20557        }
20558        Ok(())
20559    }
20560
20561    #[allow(clippy::too_many_arguments)]
20562    pub fn gdn_prep_decode_b(
20563        &self,
20564        conv_outs: &CudaSlice<f32>,
20565        beta_raws: &CudaSlice<f32>,
20566        alphas: &CudaSlice<f32>,
20567        dt_bias: &CudaSlice<f32>,
20568        a: &CudaSlice<f32>,
20569        q_l2: &mut CudaSlice<f32>,
20570        k_l2: &mut CudaSlice<f32>,
20571        v_g: &mut CudaSlice<f32>,
20572        beta: &mut CudaSlice<f32>,
20573        g_log: &mut CudaSlice<f32>,
20574        d_state: usize,
20575        num_v: usize,
20576        num_k: usize,
20577        key_dim: usize,
20578        eps: f32,
20579        conv_dim: usize,
20580        b_n: usize,
20581    ) -> Result<(), Box<dyn std::error::Error>> {
20582        let f = self.func("gdn_prep_decode_b_f32");
20583        let cfg = LaunchConfig {
20584            grid_dim: (num_v as u32, 1, b_n as u32),
20585            block_dim: (32, 4, 1),
20586            shared_mem_bytes: 0,
20587        };
20588        let (ds, nv, nk, kd, cd) = (
20589            d_state as i32,
20590            num_v as i32,
20591            num_k as i32,
20592            key_dim as i32,
20593            conv_dim as i32,
20594        );
20595        let __s_b = self.gpu.stream();
20596        let mut b = __s_b.launch_builder(&f);
20597        b.arg(conv_outs)
20598            .arg(beta_raws)
20599            .arg(alphas)
20600            .arg(dt_bias)
20601            .arg(a)
20602            .arg(q_l2)
20603            .arg(k_l2)
20604            .arg(v_g)
20605            .arg(beta)
20606            .arg(g_log)
20607            .arg(&ds)
20608            .arg(&nv)
20609            .arg(&nk)
20610            .arg(&kd)
20611            .arg(&eps)
20612            .arg(&cd);
20613        unsafe {
20614            b.launch(cfg)?;
20615        }
20616        Ok(())
20617    }
20618
20619    #[allow(clippy::too_many_arguments)]
20620    pub fn gdn_scan_s128_batched(
20621        &self,
20622        q: &CudaSlice<f32>,
20623        k: &CudaSlice<f32>,
20624        v: &CudaSlice<f32>,
20625        g: &CudaSlice<f32>,
20626        beta: &CudaSlice<f32>,
20627        state_in_ptrs: &cudarc::driver::CudaView<u64>,
20628        state_out_ptrs: &cudarc::driver::CudaView<u64>,
20629        o: &mut CudaSlice<f32>,
20630        n_head: usize,
20631        b_n: usize,
20632        scale: f32,
20633    ) -> Result<(), Box<dyn std::error::Error>> {
20634        let f = self.func("gdn_scan_s128_b");
20635        const S_V: u32 = 128;
20636        const WARP: u32 = 32;
20637        const COLS_PER_BLOCK: u32 = 4;
20638        let cfg = LaunchConfig {
20639            grid_dim: (n_head as u32, b_n as u32, S_V / COLS_PER_BLOCK),
20640            block_dim: (WARP, COLS_PER_BLOCK, 1),
20641            shared_mem_bytes: 0,
20642        };
20643        let h = n_head as i32;
20644        let __s_b = self.gpu.stream();
20645        let mut b = __s_b.launch_builder(&f);
20646        b.arg(q)
20647            .arg(k)
20648            .arg(v)
20649            .arg(g)
20650            .arg(beta)
20651            .arg(state_in_ptrs)
20652            .arg(state_out_ptrs)
20653            .arg(o)
20654            .arg(&h)
20655            .arg(&scale);
20656        unsafe {
20657            b.launch(cfg)?;
20658        }
20659        Ok(())
20660    }
20661
20662    /// VIEW twins of the three GDN decode state wrappers (lane/verify-launchslim): identical
20663    /// launches, row args as views into the caller's packed [T, ...] buffers — the t-parallel
20664    /// verify's per-row loop passes slices instead of paying an arithmetic-free dtod per row
20665    /// per kernel (48 layers x T rows x 4 copies/round on the money path). Same kernels, same
20666    /// numeric class; only the pointer arithmetic moved host-side.
20667    #[allow(clippy::too_many_arguments)]
20668    pub fn ssm_conv1d_fused_decode_b_view(
20669        &self,
20670        qkv_cols: &cudarc::driver::CudaView<f32>,
20671        conv_state_ptrs: &cudarc::driver::CudaView<u64>,
20672        w: &CudaSlice<f32>,
20673        conv_outs: &mut CudaSlice<f32>,
20674        conv_dim: usize,
20675        d_conv: usize,
20676        b_n: usize,
20677    ) -> Result<(), Box<dyn std::error::Error>> {
20678        let f = self.func("ssm_conv1d_fused_decode_b_f32");
20679        let cfg = LaunchConfig {
20680            grid_dim: (((conv_dim + 255) / 256) as u32, 1, b_n as u32),
20681            block_dim: (256, 1, 1),
20682            shared_mem_bytes: 0,
20683        };
20684        let (cd, dc) = (conv_dim as i32, d_conv as i32);
20685        let __s_b = self.gpu.stream();
20686        let mut b = __s_b.launch_builder(&f);
20687        b.arg(qkv_cols)
20688            .arg(conv_state_ptrs)
20689            .arg(w)
20690            .arg(conv_outs)
20691            .arg(&cd)
20692            .arg(&dc);
20693        unsafe {
20694            b.launch(cfg)?;
20695        }
20696        Ok(())
20697    }
20698
20699    #[allow(clippy::too_many_arguments)]
20700    pub fn gdn_prep_decode_b_view(
20701        &self,
20702        conv_outs: &CudaSlice<f32>,
20703        beta_raws: &cudarc::driver::CudaView<f32>,
20704        alphas: &cudarc::driver::CudaView<f32>,
20705        dt_bias: &CudaSlice<f32>,
20706        a: &CudaSlice<f32>,
20707        q_l2: &mut CudaSlice<f32>,
20708        k_l2: &mut CudaSlice<f32>,
20709        v_g: &mut CudaSlice<f32>,
20710        beta: &mut CudaSlice<f32>,
20711        g_log: &mut CudaSlice<f32>,
20712        d_state: usize,
20713        num_v: usize,
20714        num_k: usize,
20715        key_dim: usize,
20716        eps: f32,
20717        conv_dim: usize,
20718        b_n: usize,
20719    ) -> Result<(), Box<dyn std::error::Error>> {
20720        let f = self.func("gdn_prep_decode_b_f32");
20721        let cfg = LaunchConfig {
20722            grid_dim: (num_v as u32, 1, b_n as u32),
20723            block_dim: (32, 4, 1),
20724            shared_mem_bytes: 0,
20725        };
20726        let (ds, nv, nk, kd, cd) = (
20727            d_state as i32,
20728            num_v as i32,
20729            num_k as i32,
20730            key_dim as i32,
20731            conv_dim as i32,
20732        );
20733        let __s_b = self.gpu.stream();
20734        let mut b = __s_b.launch_builder(&f);
20735        b.arg(conv_outs)
20736            .arg(beta_raws)
20737            .arg(alphas)
20738            .arg(dt_bias)
20739            .arg(a)
20740            .arg(q_l2)
20741            .arg(k_l2)
20742            .arg(v_g)
20743            .arg(beta)
20744            .arg(g_log)
20745            .arg(&ds)
20746            .arg(&nv)
20747            .arg(&nk)
20748            .arg(&kd)
20749            .arg(&eps)
20750            .arg(&cd);
20751        unsafe {
20752            b.launch(cfg)?;
20753        }
20754        Ok(())
20755    }
20756
20757    #[allow(clippy::too_many_arguments)]
20758    pub fn gdn_scan_s128_batched_view(
20759        &self,
20760        q: &CudaSlice<f32>,
20761        k: &CudaSlice<f32>,
20762        v: &CudaSlice<f32>,
20763        g: &CudaSlice<f32>,
20764        beta: &CudaSlice<f32>,
20765        state_in_ptrs: &cudarc::driver::CudaView<u64>,
20766        state_out_ptrs: &cudarc::driver::CudaView<u64>,
20767        o: &mut cudarc::driver::CudaViewMut<f32>,
20768        n_head: usize,
20769        b_n: usize,
20770        scale: f32,
20771    ) -> Result<(), Box<dyn std::error::Error>> {
20772        let f = self.func("gdn_scan_s128_b");
20773        const S_V: u32 = 128;
20774        const WARP: u32 = 32;
20775        const COLS_PER_BLOCK: u32 = 4;
20776        let cfg = LaunchConfig {
20777            grid_dim: (n_head as u32, b_n as u32, S_V / COLS_PER_BLOCK),
20778            block_dim: (WARP, COLS_PER_BLOCK, 1),
20779            shared_mem_bytes: 0,
20780        };
20781        let h = n_head as i32;
20782        let __s_b = self.gpu.stream();
20783        let mut b = __s_b.launch_builder(&f);
20784        b.arg(q)
20785            .arg(k)
20786            .arg(v)
20787            .arg(g)
20788            .arg(beta)
20789            .arg(state_in_ptrs)
20790            .arg(state_out_ptrs)
20791            .arg(o)
20792            .arg(&h)
20793            .arg(&scale);
20794        unsafe {
20795            b.launch(cfg)?;
20796        }
20797        Ok(())
20798    }
20799
20800    /// A4 seam: chunked WY GDN prefill. DEFAULT ON (`MEMRA_GDN_CHUNKED=0` = rollback to the
20801    /// sequential scan). Flipped 2026-07-04 with the full battery green: kernel-check ALL
20802    /// GREEN x {9B, 27B} incl the f64-truth chunk gates; run-gen argmax 82==82 both models
20803    /// on AND off (24/24 sweep runs); run-spec K={1,2,3,4,6,8} PASS x {9B synth, 9B text,
20804    /// 27B p2, 27B p3}; e2e first-16-token agreement 6/6 (full-256 drifts at index 47-125
20805    /// on 5/6 prompts — accepted cache-state-FP class, batched-prime precedent).
20806    /// PREFILL-ONLY: decode + spec verify never route here (decode==verify dispatch
20807    /// identity law); prime_cache/forward/forward_last are the only callers.
20808    pub fn gdn_chunked_enabled() -> bool {
20809        static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20810        *E.get_or_init(|| {
20811            std::env::var("MEMRA_GDN_CHUNKED")
20812                .map(|v| v != "0")
20813                .unwrap_or(true)
20814        })
20815    }
20816
20817    /// A4 chunk size (MEMRA_GDN_CHUNK, default 32 — the sweep winner: the O(T*C) chunk
20818    /// matrices grow with C while the sequential state pass is C-flat, so smaller chunks
20819    /// win; C=32/64 also get the register-history solve template). Clamped to multiples
20820    /// of 32 in [32, 128] (kernel row mappings require it).
20821    pub fn gdn_chunk_size() -> usize {
20822        static C: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
20823        *C.get_or_init(|| {
20824            let c: usize = std::env::var("MEMRA_GDN_CHUNK")
20825                .ok()
20826                .and_then(|v| v.parse().ok())
20827                .unwrap_or(32);
20828            c.clamp(32, 128) / 32 * 32
20829        })
20830    }
20831
20832    /// A4: chunked WY / blockwise-inverse GDN prefill (see cu/hybrid.cu K1-K5 header for the
20833    /// math). Same contract as `gdn_scan_s128` (layouts, state ping-pong) but chunk-parallel:
20834    /// NOT bit-identical to the sequential scan (chunked FP accumulation order); run-gen
20835    /// argmax + run-spec batteries are the accuracy authority. PREFILL callers only.
20836    #[allow(clippy::too_many_arguments)]
20837    /// task #18: K1-K3 of the chunked WY scan (shared by the per-seq path and the
20838    /// batched-prime varlen path). Returns (gcum, P, U, W); `A` is K3-internal.
20839    #[allow(clippy::too_many_arguments, clippy::type_complexity)]
20840    #[allow(clippy::too_many_arguments)]
20841    pub fn gdn_chunk_k123(
20842        &self,
20843        q: &CudaSlice<f32>,
20844        k: &CudaSlice<f32>,
20845        v: &CudaSlice<f32>,
20846        g: &CudaSlice<f32>,
20847        beta: &CudaSlice<f32>,
20848        wb16: Option<&mut CudaSlice<u8>>,
20849        n_head: usize,
20850        t: usize,
20851        c: usize,
20852        hk: usize,
20853        k2w: Option<(&CudaSlice<u8>, &CudaSlice<u8>, &mut CudaSlice<u8>)>,
20854    ) -> Result<
20855        (
20856            CudaSlice<f32>,
20857            CudaSlice<f32>,
20858            CudaSlice<f32>,
20859            CudaSlice<f32>,
20860        ),
20861        Box<dyn std::error::Error>,
20862    > {
20863        const D: usize = 128;
20864        let h = n_head;
20865        let nc = (t + c - 1) / c;
20866        let (hi, ti, ci) = (h as i32, t as i32, c as i32);
20867        let mut gcum = self.uninit(t * h)?;
20868        let mut a = self.uninit(nc * h * c * c)?;
20869        let mut p = self.uninit(nc * h * c * c)?;
20870        let mut u = self.uninit(nc * h * c * D)?;
20871        let mut w = self.uninit(nc * h * c * D)?;
20872        {
20873            // K1
20874            let f = self.func("gdn_chunk_cumgate_f32");
20875            let cfg = LaunchConfig {
20876                grid_dim: (nc as u32, h as u32, 1),
20877                block_dim: (32, 1, 1),
20878                shared_mem_bytes: 0,
20879            };
20880            let __s_b = self.gpu.stream();
20881            let mut b = __s_b.launch_builder(&f);
20882            b.arg(g).arg(&mut gcum).arg(&hi).arg(&ti).arg(&ci);
20883            unsafe {
20884                b.launch(cfg)?;
20885            }
20886        }
20887        if let Some((qb, kb, pb)) = k2w {
20888            // K2-wgmma (MEMRA_GDN_WGMMA path, c==32): A + pre-masked Pb16 in one kernel;
20889            // the P f32 buffer stays UNWRITTEN (its only wgmma-path consumer is Pb16).
20890            assert!(c == 32, "gdn_k2_wgmma is a C==32 tile");
20891            let f = self.func("gdn_k2_wgmma");
20892            let cfg = LaunchConfig {
20893                grid_dim: (nc as u32, h as u32, 1),
20894                block_dim: (128, 1, 1),
20895                shared_mem_bytes: 0,
20896            };
20897            let hki = hk as i32;
20898            let __s_b = self.gpu.stream();
20899            let mut b = __s_b.launch_builder(&f);
20900            b.arg(qb)
20901                .arg(kb)
20902                .arg(&gcum)
20903                .arg(beta)
20904                .arg(&mut a)
20905                .arg(&mut *pb)
20906                .arg(&hi)
20907                .arg(&ti)
20908                .arg(&ci)
20909                .arg(&hki);
20910            unsafe {
20911                b.launch(cfg)?;
20912            }
20913        } else if c <= 64 && !portable_mma_gated() {
20914            // K2 register-tiled (2x2 outputs/thread, whole-chunk smem k tile)
20915            let f = self.func("gdn_chunk_attn_f32");
20916            let jt = ((c + 31) / 32) as u32;
20917            let cfg = LaunchConfig {
20918                grid_dim: (nc as u32, h as u32, jt),
20919                block_dim: (256, 1, 1),
20920                shared_mem_bytes: 0,
20921            };
20922            let hki = hk as i32;
20923            let __s_b = self.gpu.stream();
20924            let mut b = __s_b.launch_builder(&f);
20925            b.arg(q)
20926                .arg(k)
20927                .arg(&gcum)
20928                .arg(beta)
20929                .arg(&mut a)
20930                .arg(&mut p)
20931                .arg(&hi)
20932                .arg(&ti)
20933                .arg(&ci)
20934                .arg(&hki);
20935            unsafe {
20936                b.launch(cfg)?;
20937            }
20938        } else {
20939            // K2 generic (C = 128, or the portable target's low-smem fallback)
20940            assert!(
20941                hk == h,
20942                "generic K2 is broadcast-only (de-broadcast rides C==32)"
20943            );
20944            let f = self.func("gdn_chunk_attn_g_f32");
20945            let cfg = LaunchConfig {
20946                grid_dim: (nc as u32, h as u32, 1),
20947                block_dim: (32, 8, 1),
20948                shared_mem_bytes: 0,
20949            };
20950            let __s_b = self.gpu.stream();
20951            let mut b = __s_b.launch_builder(&f);
20952            b.arg(q)
20953                .arg(k)
20954                .arg(&gcum)
20955                .arg(beta)
20956                .arg(&mut a)
20957                .arg(&mut p)
20958                .arg(&hi)
20959                .arg(&ti)
20960                .arg(&ci);
20961            unsafe {
20962                b.launch(cfg)?;
20963            }
20964        }
20965        {
20966            // K3 (register-history templates for C=32/64; local-memory generic otherwise)
20967            let cfg = LaunchConfig {
20968                grid_dim: (nc as u32, h as u32, 1),
20969                block_dim: (256, 1, 1),
20970                shared_mem_bytes: 0,
20971            };
20972            match c {
20973                32 | 64 => {
20974                    let f = self.func(if c == 32 {
20975                        "gdn_chunk_solve32_f32"
20976                    } else {
20977                        "gdn_chunk_solve64_f32"
20978                    });
20979                    // mirror-fold: W's bf16 twin emitted on store (0 = skip)
20980                    let wb: u64 = match wb16 {
20981                        Some(d) => self.addr_u8(d),
20982                        None => 0,
20983                    };
20984                    let hki = hk as i32;
20985                    let __s_b = self.gpu.stream();
20986                    let mut b = __s_b.launch_builder(&f);
20987                    b.arg(v)
20988                        .arg(k)
20989                        .arg(&a)
20990                        .arg(&gcum)
20991                        .arg(&mut u)
20992                        .arg(&mut w)
20993                        .arg(&wb)
20994                        .arg(&hi)
20995                        .arg(&ti)
20996                        .arg(&hki);
20997                    unsafe {
20998                        b.launch(cfg)?;
20999                    }
21000                }
21001                _ => {
21002                    assert!(hk == h, "generic K3 is broadcast-only");
21003                    let f = self.func("gdn_chunk_solve_f32");
21004                    let __s_b = self.gpu.stream();
21005                    let mut b = __s_b.launch_builder(&f);
21006                    b.arg(v)
21007                        .arg(k)
21008                        .arg(&a)
21009                        .arg(&gcum)
21010                        .arg(&mut u)
21011                        .arg(&mut w)
21012                        .arg(&hi)
21013                        .arg(&ti)
21014                        .arg(&ci);
21015                    unsafe {
21016                        b.launch(cfg)?;
21017                    }
21018                }
21019            }
21020        }
21021        Ok((gcum, p, u, w))
21022    }
21023
21024    /// task #21 de-broadcast seam: q/k stored at num_k distinct GQA heads instead of
21025    /// the num_v broadcast. MEMRA_GDN_DB=0 reverts. Only the chunked prefill path
21026    /// consumes the compact layout (hk plumbed; hk == H reproduces broadcast exactly).
21027    pub fn gdn_db_on() -> bool {
21028        std::env::var("MEMRA_GDN_DB").as_deref() != Ok("0")
21029    }
21030
21031    /// Whether the K4/K5 mma pair serves at chunk size `c` (mirrors gdn_scan_chunked's
21032    /// seam read — env re-read per call ON PURPOSE, kernel-check pins both configs).
21033    /// DEFAULT ON for sm_120a builds too (lane/moeprime-nvfp4-direct, 2026-08-21): the pair
21034    /// was qualified on 90a only and left env-opt-in elsewhere; measured on Blackwell it
21035    /// wins on BOTH rigs — one RTX PRO 6000 (ornith15 pp14715 12,036 -> 12,751/12,957,
21036    /// +6-8%, both orders) and the local 5090 (q38-27b pp6435 1,397/1,429 -> 1,427/1,446,
21037    /// both orders) — with kernel-check/run-gen/margin-gate/run-spec green under the flag.
21038    /// bf16 HMMA (m16n8k16) is sm_80-class PTX; only the wgmma nest stays Hopper-gated.
21039    pub fn gdn_mma_enabled(&self, c: usize) -> bool {
21040        !portable_mma_gated()
21041            && c == 32
21042            && match std::env::var("MEMRA_GDN_MMA").as_deref() {
21043                Ok("1") => true,
21044                Ok("0") => false,
21045                _ => gdn_mma_default_on(),
21046            }
21047    }
21048
21049    /// task #22: whether the fused K4+K5 (+K2) wgmma path serves (nested inside the
21050    /// mma config; same per-call env read discipline). HARD-gated to the Hopper build:
21051    /// the wgmma asm bodies exist only at __CUDA_ARCH__ == 900 (MEMRA_K45_REAL,
21052    /// wgmma_common.cuh) — on every other arch the kernel compiles EMPTY, so an env
21053    /// force would silently produce garbage. Required since the sm_120a mma default
21054    /// flip made MEMRA_GDN_WGMMA=1 alone reach this branch there.
21055    pub fn gdn_wgmma_on(&self, c: usize) -> bool {
21056        cfg!(memra_hopper_mma)
21057            && self.gdn_mma_enabled(c)
21058            && std::env::var("MEMRA_GDN_WGMMA").as_deref() != Ok("0")
21059    }
21060
21061    /// task #18 conv-fuse: carried-ring conv + SiLU + GDN repack in ONE pass (the
21062    /// conv_out intermediate and its transposed re-read disappear — 11.8ms of the
21063    /// T=2048 prime). Ring update stays the separate follow-up launch (pad-aware).
21064    /// BIT-IDENTICAL values to ssm_conv1d_tm_state_pad + qkv_to_gdn_repack.
21065    #[allow(clippy::too_many_arguments)]
21066    pub fn ssm_conv1d_gdn_state_pad(
21067        &self,
21068        qkv_tm: &cudarc::driver::CudaView<f32>,
21069        conv_state: &mut CudaSlice<f32>,
21070        w: &CudaSlice<f32>,
21071        q_g: &mut CudaSlice<f32>,
21072        k_g: &mut CudaSlice<f32>,
21073        v_g: &mut CudaSlice<f32>,
21074        conv_dim: usize,
21075        t: usize,
21076        d_conv: usize,
21077        d_state: usize,
21078        num_v: usize,
21079        num_k: usize,
21080        key_dim: usize,
21081        hk: usize,
21082        pad_len: Option<&CudaSlice<i32>>,
21083    ) -> Result<(), Box<dyn std::error::Error>> {
21084        assert!(
21085            t >= d_conv - 1,
21086            "fused state conv requires T >= pad (PRIME_MIN_T gates)"
21087        );
21088        {
21089            let f = self.func("ssm_conv1d_gdn_state_f32");
21090            let cfg = LaunchConfig {
21091                grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
21092                block_dim: (256, 1, 1),
21093                shared_mem_bytes: 0,
21094            };
21095            let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
21096            let (ds, nv, nk, kd, hki) = (
21097                d_state as i32,
21098                num_v as i32,
21099                num_k as i32,
21100                key_dim as i32,
21101                hk as i32,
21102            );
21103            let __s_b = self.gpu.stream();
21104            let mut b = __s_b.launch_builder(&f);
21105            b.arg(qkv_tm)
21106                .arg(&*conv_state)
21107                .arg(w)
21108                .arg(q_g)
21109                .arg(k_g)
21110                .arg(v_g)
21111                .arg(&cd)
21112                .arg(&ti)
21113                .arg(&dc)
21114                .arg(&ds)
21115                .arg(&nv)
21116                .arg(&nk)
21117                .arg(&kd)
21118                .arg(&hki);
21119            unsafe {
21120                b.launch(cfg)?;
21121            }
21122        }
21123        match pad_len {
21124            Some(len_d) => {
21125                let f = self.func("ssm_conv_ring_update_dev_f32");
21126                let n = conv_dim * (d_conv - 1);
21127                let cfg = LaunchConfig::for_num_elems(n as u32);
21128                let (cd, dc) = (conv_dim as i32, d_conv as i32);
21129                let __s_b = self.gpu.stream();
21130                let mut b = __s_b.launch_builder(&f);
21131                b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
21132                unsafe {
21133                    b.launch(cfg)?;
21134                }
21135            }
21136            None => {
21137                let f = self.func("ssm_conv_ring_update_f32");
21138                let n = conv_dim * (d_conv - 1);
21139                let cfg = LaunchConfig::for_num_elems(n as u32);
21140                let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
21141                let __s_b = self.gpu.stream();
21142                let mut b = __s_b.launch_builder(&f);
21143                b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
21144                unsafe {
21145                    b.launch(cfg)?;
21146                }
21147            }
21148        }
21149        Ok(())
21150    }
21151
21152    /// task #18 increment 2: allocate ONE sequence's chunk buffers (no launches) —
21153    /// K1-K5 all run varlen afterwards. `a`/`w` become struct members so the varlen
21154    /// K2/K3 can write them.
21155    pub fn gdn_chunk_alloc(
21156        &self,
21157        n_head: usize,
21158        t: usize,
21159        c: usize,
21160        hk: usize,
21161    ) -> Result<GdnChunkBufs, Box<dyn std::error::Error>> {
21162        const D: usize = 128;
21163        assert!(
21164            c == 32,
21165            "gdn_chunk_alloc: varlen chain is the C==32 mma pair"
21166        );
21167        let h = n_head;
21168        let nc = (t + c - 1) / c;
21169        Ok(GdnChunkBufs {
21170            gcum: self.uninit(t * h)?,
21171            a: self.uninit(nc * h * c * c)?,
21172            p: self.uninit(nc * h * c * c)?,
21173            u: self.uninit(nc * h * c * D)?,
21174            w: self.uninit(nc * h * c * D)?,
21175            kb16: self.alloc_u8_uninit(t * hk * D * 2)?,
21176            wb16: self.alloc_u8_uninit(nc * h * c * D * 2)?,
21177            y16: self.alloc_u8_uninit(nc * h * c * D * 2)?,
21178            ssnap16: self.alloc_u8_uninit(nc * h * D * D * 2)?,
21179            qb16: self.alloc_u8_uninit(t * hk * D * 2)?,
21180            pb16: self.alloc_u8_uninit(nc * h * c * c * 2)?,
21181            o: self.uninit(D * h * t)?,
21182            t,
21183            nc,
21184        })
21185    }
21186
21187    /// view-source twin of f32_to_bf16 (the batched FA3 v mirror reads a concat view).
21188    pub fn f32_to_bf16_v(
21189        &self,
21190        x: &cudarc::driver::CudaView<f32>,
21191        dst: &mut CudaSlice<u8>,
21192        n: usize,
21193    ) -> Result<(), Box<dyn std::error::Error>> {
21194        let f = self.func("f32_to_bf16_bulk");
21195        let ni = n as i64;
21196        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
21197        let __s_b = self.gpu.stream();
21198        let mut b = __s_b.launch_builder(&f);
21199        b.arg(x).arg(dst).arg(&ni);
21200        unsafe {
21201            b.launch(cfg)?;
21202        }
21203        Ok(())
21204    }
21205
21206    /// f32 -> bf16 bulk mirror into a caller buffer (the K4/K5 operand mirrors).
21207    pub fn f32_to_bf16_into(
21208        &self,
21209        x: &CudaSlice<f32>,
21210        dst: &mut CudaSlice<u8>,
21211        n: usize,
21212    ) -> Result<(), Box<dyn std::error::Error>> {
21213        let f = self.func("f32_to_bf16_bulk");
21214        let ni = n as i64;
21215        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
21216        let __s_b = self.gpu.stream();
21217        let mut b = __s_b.launch_builder(&f);
21218        b.arg(x).arg(dst).arg(&ni);
21219        unsafe {
21220            b.launch(cfg)?;
21221        }
21222        Ok(())
21223    }
21224
21225    /// task #18 increment 2: varlen K1+K2+K3 — three launches run every sequence's
21226    /// cumgate/attn/solve (per-block math identical to the per-seq kernels).
21227    pub fn gdn_chunk_k123_vl8(
21228        &self,
21229        seqs: &[GdnSeqVl],
21230        n_head: usize,
21231        hk: usize,
21232        wq: Option<&GdnWVl8>,
21233    ) -> Result<(), Box<dyn std::error::Error>> {
21234        let b = seqs.len();
21235        assert!(b >= 1 && b <= 8, "gdn_chunk_k123_vl8: 1..=8 sequences");
21236        let mut packed = [GdnSeqVl::default(); 8];
21237        packed[..b].copy_from_slice(seqs);
21238        let v = GdnVl8(packed);
21239        let (hi, ci) = (n_head as i32, 32i32);
21240        let max_nc = seqs.iter().map(|a| a.nc).max().unwrap() as u32;
21241        {
21242            let f = self.func("gdn_chunk_cumgate_vl");
21243            let cfg = LaunchConfig {
21244                grid_dim: (max_nc, n_head as u32, b as u32),
21245                block_dim: (32, 1, 1),
21246                shared_mem_bytes: 0,
21247            };
21248            let __s_lb = self.gpu.stream();
21249            let mut lb = __s_lb.launch_builder(&f);
21250            lb.arg(&v).arg(&hi).arg(&ci);
21251            unsafe {
21252                lb.launch(cfg)?;
21253            }
21254        }
21255        let hki = hk as i32;
21256        if let Some(w) = wq {
21257            // K2-wgmma vl twin (writes A + pre-masked Pb16)
21258            let f = self.func("gdn_k2_wgmma_vl");
21259            let cfg = LaunchConfig {
21260                grid_dim: (max_nc, n_head as u32, b as u32),
21261                block_dim: (128, 1, 1),
21262                shared_mem_bytes: 0,
21263            };
21264            let __s_lb = self.gpu.stream();
21265            let mut lb = __s_lb.launch_builder(&f);
21266            lb.arg(&v).arg(w).arg(&hi).arg(&ci).arg(&hki);
21267            unsafe {
21268                lb.launch(cfg)?;
21269            }
21270        } else {
21271            let f = self.func("gdn_chunk_attn_vl");
21272            let cfg = LaunchConfig {
21273                grid_dim: (max_nc, n_head as u32, b as u32),
21274                block_dim: (256, 1, 1),
21275                shared_mem_bytes: 0,
21276            };
21277            let __s_lb = self.gpu.stream();
21278            let mut lb = __s_lb.launch_builder(&f);
21279            lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
21280            unsafe {
21281                lb.launch(cfg)?;
21282            }
21283        }
21284        {
21285            let f = self.func("gdn_chunk_solve32_vl");
21286            let cfg = LaunchConfig {
21287                grid_dim: (max_nc, n_head as u32, b as u32),
21288                block_dim: (256, 1, 1),
21289                shared_mem_bytes: 0,
21290            };
21291            let __s_lb = self.gpu.stream();
21292            let mut lb = __s_lb.launch_builder(&f);
21293            lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
21294            unsafe {
21295                lb.launch(cfg)?;
21296            }
21297        }
21298        Ok(())
21299    }
21300
21301    /// task #18 increment 3: varlen PREP chain — conv(+ring) / repack / fused-l2 /
21302    /// fused gate-prep, 5 launches for every sequence (per-element math identical
21303    /// to the per-seq kernels; l2/gate fusions write disjoint outputs).
21304    #[allow(clippy::too_many_arguments)]
21305    pub fn gdn_prep_vl8(
21306        &self,
21307        seqs: &[GdnPrepVl],
21308        conv_w: &CudaSlice<f32>,
21309        dt_bias: &CudaSlice<f32>,
21310        a: &CudaSlice<f32>,
21311        conv_dim: usize,
21312        d_conv: usize,
21313        d_state: usize,
21314        num_v: usize,
21315        num_k: usize,
21316        key_dim: usize,
21317        hk: usize,
21318        eps: f32,
21319    ) -> Result<(), Box<dyn std::error::Error>> {
21320        let b = seqs.len();
21321        assert!(b >= 1 && b <= 8);
21322        let mut packed = [GdnPrepVl::default(); 8];
21323        packed[..b].copy_from_slice(seqs);
21324        let v = GdnPrepVl8(packed);
21325        let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
21326        let (cdi, dci) = (conv_dim as i32, d_conv as i32);
21327        let conv_fuse = std::env::var("MEMRA_CONV_FUSE").as_deref() != Ok("0");
21328        assert!(
21329            conv_fuse || hk == num_v,
21330            "de-broadcast requires the fused conv"
21331        );
21332        if conv_fuse {
21333            let f = self.func("ssm_conv1d_gdn_state_vl");
21334            let cfg = LaunchConfig {
21335                grid_dim: ((conv_dim as u32).div_ceil(256), max_t, b as u32),
21336                block_dim: (256, 1, 1),
21337                shared_mem_bytes: 0,
21338            };
21339            let (dsi, nvi, nki, kdi, hki) = (
21340                d_state as i32,
21341                num_v as i32,
21342                num_k as i32,
21343                key_dim as i32,
21344                hk as i32,
21345            );
21346            let __s_lb = self.gpu.stream();
21347            let mut lb = __s_lb.launch_builder(&f);
21348            lb.arg(&v)
21349                .arg(conv_w)
21350                .arg(&cdi)
21351                .arg(&dci)
21352                .arg(&dsi)
21353                .arg(&nvi)
21354                .arg(&nki)
21355                .arg(&kdi)
21356                .arg(&hki);
21357            unsafe {
21358                lb.launch(cfg)?;
21359            }
21360        } else {
21361            let f = self.func("ssm_conv1d_tm_state_vl");
21362            let cfg = LaunchConfig {
21363                grid_dim: ((conv_dim as u32).div_ceil(256), max_t, b as u32),
21364                block_dim: (256, 1, 1),
21365                shared_mem_bytes: 0,
21366            };
21367            let __s_lb = self.gpu.stream();
21368            let mut lb = __s_lb.launch_builder(&f);
21369            lb.arg(&v).arg(conv_w).arg(&cdi).arg(&dci);
21370            unsafe {
21371                lb.launch(cfg)?;
21372            }
21373        }
21374        {
21375            let f = self.func("ssm_conv_ring_update_vl");
21376            let n = (conv_dim * (d_conv - 1)) as u32;
21377            let cfg = LaunchConfig {
21378                grid_dim: (n.div_ceil(256), 1, b as u32),
21379                block_dim: (256, 1, 1),
21380                shared_mem_bytes: 0,
21381            };
21382            let __s_lb = self.gpu.stream();
21383            let mut lb = __s_lb.launch_builder(&f);
21384            lb.arg(&v).arg(&cdi).arg(&dci);
21385            unsafe {
21386                lb.launch(cfg)?;
21387            }
21388        }
21389        if !conv_fuse {
21390            let f = self.func("qkv_to_gdn_repack_vl");
21391            let n = max_t * (num_v * d_state) as u32;
21392            let cfg = LaunchConfig {
21393                grid_dim: (n.div_ceil(256), 1, b as u32),
21394                block_dim: (256, 1, 1),
21395                shared_mem_bytes: 0,
21396            };
21397            let (dsi, nvi, nki, kdi) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
21398            let __s_lb = self.gpu.stream();
21399            let mut lb = __s_lb.launch_builder(&f);
21400            lb.arg(&v).arg(&dsi).arg(&nvi).arg(&nki).arg(&kdi);
21401            unsafe {
21402                lb.launch(cfg)?;
21403            }
21404        }
21405        if Self::l2_v2_on(d_state) {
21406            let f = self.func("gdn_l2_v2_vl");
21407            let cfg = LaunchConfig {
21408                grid_dim: ((max_t * hk as u32).div_ceil(8), 2, b as u32),
21409                block_dim: (256, 1, 1),
21410                shared_mem_bytes: 0,
21411            };
21412            let (dsi, nvi) = (d_state as i32, hk as i32);
21413            let __s_lb = self.gpu.stream();
21414            let mut lb = __s_lb.launch_builder(&f);
21415            lb.arg(&v).arg(&dsi).arg(&nvi).arg(&eps);
21416            unsafe {
21417                lb.launch(cfg)?;
21418            }
21419        } else {
21420            let f = self.func("gdn_l2_vl");
21421            let cfg = LaunchConfig {
21422                grid_dim: (max_t * hk as u32, 2, b as u32),
21423                block_dim: (256, 1, 1),
21424                shared_mem_bytes: 0,
21425            };
21426            let (dsi, nvi) = (d_state as i32, hk as i32);
21427            let __s_lb = self.gpu.stream();
21428            let mut lb = __s_lb.launch_builder(&f);
21429            lb.arg(&v).arg(&dsi).arg(&nvi).arg(&eps);
21430            unsafe {
21431                lb.launch(cfg)?;
21432            }
21433        }
21434        {
21435            let f = self.func("gdn_gate_prep_vl");
21436            let n = max_t * num_v as u32;
21437            let cfg = LaunchConfig {
21438                grid_dim: (n.div_ceil(256), 1, b as u32),
21439                block_dim: (256, 1, 1),
21440                shared_mem_bytes: 0,
21441            };
21442            let nvi = num_v as i32;
21443            let __s_lb = self.gpu.stream();
21444            let mut lb = __s_lb.launch_builder(&f);
21445            lb.arg(&v).arg(dt_bias).arg(a).arg(&nvi);
21446            unsafe {
21447                lb.launch(cfg)?;
21448            }
21449        }
21450        Ok(())
21451    }
21452
21453    /// varlen bf16 mirrors over the gdnseq_t table (which: 0 = k_l2 -> kb16, 1 = w -> wb16).
21454    pub fn gdn_mirror_vl8(
21455        &self,
21456        seqs: &[GdnSeqVl],
21457        n_head: usize,
21458        which: i32,
21459        hk: usize,
21460    ) -> Result<(), Box<dyn std::error::Error>> {
21461        let b = seqs.len();
21462        assert!(b >= 1 && b <= 8);
21463        let mut packed = [GdnSeqVl::default(); 8];
21464        packed[..b].copy_from_slice(seqs);
21465        let v = GdnVl8(packed);
21466        let ept = (if which == 0 { hk } else { n_head } * 128) as i32;
21467        let max_n = seqs
21468            .iter()
21469            .map(|s| {
21470                if which == 0 {
21471                    s.t as i64 * ept as i64
21472                } else {
21473                    s.nc as i64 * ept as i64 * 32
21474                }
21475            })
21476            .max()
21477            .unwrap();
21478        let f = self.func("gdn_mirror_vl");
21479        let blocks = ((max_n as u32).div_ceil(4)).div_ceil(256);
21480        let cfg = LaunchConfig {
21481            grid_dim: (blocks, 1, b as u32),
21482            block_dim: (256, 1, 1),
21483            shared_mem_bytes: 0,
21484        };
21485        let __s_lb = self.gpu.stream();
21486        let mut lb = __s_lb.launch_builder(&f);
21487        lb.arg(&v).arg(&ept).arg(&which);
21488        unsafe {
21489            lb.launch(cfg)?;
21490        }
21491        Ok(())
21492    }
21493
21494    /// varlen gated-norm tail (+f16out) — one launch replaces B gated_rmsnorm calls.
21495    pub fn gdn_tail_vl8(
21496        &self,
21497        seqs: &[GdnPrepVl],
21498        norm_w: &CudaSlice<f32>,
21499        d_state: usize,
21500        num_v: usize,
21501        eps: f32,
21502    ) -> Result<(), Box<dyn std::error::Error>> {
21503        let b = seqs.len();
21504        assert!(b >= 1 && b <= 8);
21505        let mut packed = [GdnPrepVl::default(); 8];
21506        packed[..b].copy_from_slice(seqs);
21507        let v = GdnPrepVl8(packed);
21508        let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
21509        let f = self.func("gated_rmsnorm_f16out_vl");
21510        // block_dim MUST match gated_rmsnorm's (128): the reduction tree order pins the scale
21511        let cfg = LaunchConfig {
21512            grid_dim: (max_t * num_v as u32, 1, b as u32),
21513            block_dim: (128, 1, 1),
21514            shared_mem_bytes: 0,
21515        };
21516        let (dsi, nvi) = (d_state as i32, num_v as i32);
21517        let __s_lb = self.gpu.stream();
21518        let mut lb = __s_lb.launch_builder(&f);
21519        lb.arg(&v).arg(norm_w).arg(&dsi).arg(&nvi).arg(&eps);
21520        unsafe {
21521            lb.launch(cfg)?;
21522        }
21523        Ok(())
21524    }
21525
21526    /// Raw device address helpers for the varlen by-value arg struct (single-stream
21527    /// launches; every buffer outlives the call — the f16 FFI discipline).
21528    pub fn addr_f32(&self, x: &CudaSlice<f32>) -> u64 {
21529        use cudarc::driver::DevicePtr;
21530        let s = self.gpu.stream();
21531        let (p, _g) = x.device_ptr(&s);
21532        p as u64
21533    }
21534    pub fn addr_f32_mut(&self, x: &mut CudaSlice<f32>) -> u64 {
21535        use cudarc::driver::DevicePtrMut;
21536        let s = self.gpu.stream();
21537        let (p, _g) = x.device_ptr_mut(&s);
21538        p as u64
21539    }
21540    pub fn addr_f32v(&self, x: &cudarc::driver::CudaView<f32>) -> u64 {
21541        use cudarc::driver::DevicePtr;
21542        let s = self.gpu.stream();
21543        let (p, _g) = x.device_ptr(&s);
21544        p as u64
21545    }
21546    pub fn addr_u8(&self, x: &CudaSlice<u8>) -> u64 {
21547        use cudarc::driver::DevicePtr;
21548        let s = self.gpu.stream();
21549        let (p, _g) = x.device_ptr(&s);
21550        p as u64
21551    }
21552
21553    /// task #18: the varlen K4+K5 pair — TWO launches run every sequence's state pass
21554    /// and output pass (grid gains a seq dim; per-block math identical to the per-seq
21555    /// launches, so this is strictly bit-gateable against them).
21556    pub fn gdn_chunk_vl8(
21557        &self,
21558        seqs: &[GdnSeqVl],
21559        n_head: usize,
21560        scale: f32,
21561        hk: usize,
21562        wq: Option<&GdnWVl8>,
21563    ) -> Result<(), Box<dyn std::error::Error>> {
21564        const NSPLIT: u32 = 4;
21565        let b = seqs.len();
21566        assert!(b >= 1 && b <= 8, "gdn_chunk_vl8: 1..=8 sequences");
21567        let mut packed = [GdnSeqVl::default(); 8];
21568        packed[..b].copy_from_slice(seqs);
21569        let v = GdnVl8(packed);
21570        let (hi, ci) = (n_head as i32, 32i32);
21571        let max_nc = seqs.iter().map(|a| a.nc).max().unwrap() as u32;
21572        let hki = hk as i32;
21573        if let Some(w) = wq {
21574            // K4+K5 fused wgmma vl twin: one launch, Y/Ssnap never materialized.
21575            let f = self.func("gdn_k45_wgmma_vl");
21576            let cfg = LaunchConfig {
21577                grid_dim: (n_head as u32, NSPLIT, b as u32),
21578                block_dim: (256, 1, 1),
21579                shared_mem_bytes: 0,
21580            };
21581            let __s_lb = self.gpu.stream();
21582            let mut lb = __s_lb.launch_builder(&f);
21583            lb.arg(&v).arg(w).arg(&scale).arg(&hi).arg(&ci).arg(&hki);
21584            unsafe {
21585                lb.launch(cfg)?;
21586            }
21587            let _ = max_nc;
21588            return Ok(());
21589        }
21590        {
21591            let f = self.func("gdn_chunk_state_mma_vl");
21592            let cfg = LaunchConfig {
21593                grid_dim: (n_head as u32, NSPLIT, b as u32),
21594                block_dim: (256, 1, 1),
21595                shared_mem_bytes: 0,
21596            };
21597            let __s_lb = self.gpu.stream();
21598            let mut lb = __s_lb.launch_builder(&f);
21599            lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
21600            unsafe {
21601                lb.launch(cfg)?;
21602            }
21603        }
21604        {
21605            let f = self.func("gdn_chunk_output_mma_vl");
21606            let cfg = LaunchConfig {
21607                grid_dim: (max_nc, n_head as u32, b as u32),
21608                block_dim: (256, 1, 1),
21609                shared_mem_bytes: 0,
21610            };
21611            let __s_lb = self.gpu.stream();
21612            let mut lb = __s_lb.launch_builder(&f);
21613            lb.arg(&v).arg(&hi).arg(&ci).arg(&scale).arg(&hki);
21614            unsafe {
21615                lb.launch(cfg)?;
21616            }
21617        }
21618        Ok(())
21619    }
21620    pub fn gdn_scan_chunked(
21621        &self,
21622        q: &CudaSlice<f32>,
21623        k: &CudaSlice<f32>,
21624        v: &CudaSlice<f32>,
21625        g: &CudaSlice<f32>,
21626        beta: &CudaSlice<f32>,
21627        kb16_pre: Option<&CudaSlice<u8>>,
21628        qb16_pre: Option<&CudaSlice<u8>>,
21629        state_in: &CudaSlice<f32>,
21630        state_out: &mut CudaSlice<f32>,
21631        o: &mut CudaSlice<f32>,
21632        n_head: usize,
21633        t: usize,
21634        scale: f32,
21635        c: usize,
21636        hk: usize,
21637    ) -> Result<(), Box<dyn std::error::Error>> {
21638        const D: usize = 128;
21639        const NSPLIT: u32 = 4;
21640        assert!(c >= 1 && c <= 128, "gdn_scan_chunked: C must be in 1..=128");
21641        let h = n_head;
21642        let nc = (t + c - 1) / c;
21643        let (hi, ti, ci) = (h as i32, t as i32, c as i32);
21644        // mirror-fold (round 27): on the mma path W's bf16 twin is emitted by K3's store
21645        // (wb16 pre-allocated and threaded through k123) and k's by the producer l2 when
21646        // the caller hands `kb16_pre` — both standalone mirror passes disappear.
21647        // default mirrors gdn_mma_enabled (incl. the sm_120a-build ON, lane/moeprime-nvfp4-direct)
21648        let gdn_mma_pre = !portable_mma_gated()
21649            && c == 32
21650            && match std::env::var("MEMRA_GDN_MMA").as_deref() {
21651                Ok("1") => true,
21652                Ok("0") => false,
21653                _ => gdn_mma_default_on(),
21654            };
21655        let mut wb16_pre: Option<CudaSlice<u8>> = if gdn_mma_pre {
21656            Some(self.alloc_u8_uninit(nc * h * c * D * 2)?)
21657        } else {
21658            None
21659        };
21660        // K2-wgmma pre-work (MEMRA_GDN_WGMMA): the kb16/qb16 mirrors hoist ABOVE K123 so
21661        // K2 rides them via cp.async; K2 writes the pre-masked Pb16 directly (the
21662        // gdn_p_bf16_masked pass and the in-branch mirror builds disappear).
21663        // mirrors gdn_wgmma_on: hard Hopper-build gate (empty wgmma bodies elsewhere)
21664        let gdn_wgmma_pre = cfg!(memra_hopper_mma)
21665            && gdn_mma_pre
21666            && std::env::var("MEMRA_GDN_WGMMA").as_deref() != Ok("0");
21667        let nk = t * hk * D;
21668        let mut kb16_local: Option<CudaSlice<u8>> = None;
21669        if gdn_mma_pre && kb16_pre.is_none() {
21670            let mut kb = self.alloc_u8_uninit(nk * 2)?;
21671            let f = self.func("f32_to_bf16_bulk");
21672            let n2 = nk as i64;
21673            let cfg2 = LaunchConfig::for_num_elems((nk as u32).div_ceil(4));
21674            let __s_b = self.gpu.stream();
21675            let mut b = __s_b.launch_builder(&f);
21676            b.arg(k).arg(&mut kb).arg(&n2);
21677            unsafe {
21678                b.launch(cfg2)?;
21679            }
21680            kb16_local = Some(kb);
21681        }
21682        let kb16_ref0: Option<&CudaSlice<u8>> = kb16_local.as_ref().or(kb16_pre);
21683        if let Some(kb) = kb16_pre {
21684            assert!(kb.len() >= nk * 2, "kb16_pre too small");
21685        }
21686        let mut qb16: Option<CudaSlice<u8>> = None;
21687        let mut pb16: Option<CudaSlice<u8>> = None;
21688        if gdn_wgmma_pre {
21689            // mirror-fold (round 35): prep's l2 v2 emits qb16 in-epilogue (kb16 pattern);
21690            // the standalone bulk cvt only serves callers without the prep mirror.
21691            if qb16_pre.is_none() {
21692                let mut qb = self.alloc_u8_uninit(nk * 2)?;
21693                let f = self.func("f32_to_bf16_bulk");
21694                let n2 = nk as i64;
21695                let cfg2 = LaunchConfig::for_num_elems((nk as u32).div_ceil(4));
21696                let __s_b = self.gpu.stream();
21697                let mut b = __s_b.launch_builder(&f);
21698                b.arg(q).arg(&mut qb).arg(&n2);
21699                unsafe {
21700                    b.launch(cfg2)?;
21701                }
21702                qb16 = Some(qb);
21703            } else if let Some(qb) = qb16_pre {
21704                assert!(qb.len() >= nk * 2, "qb16_pre too small");
21705            }
21706            pb16 = Some(self.alloc_u8_uninit(nc * h * c * c * 2)?);
21707        }
21708        let qb16_ref0: Option<&CudaSlice<u8>> = qb16.as_ref().or(qb16_pre);
21709        let k2w = if gdn_wgmma_pre {
21710            Some((
21711                *qb16_ref0.as_ref().unwrap(),
21712                *kb16_ref0.as_ref().unwrap(),
21713                pb16.as_mut().unwrap(),
21714            ))
21715        } else {
21716            None
21717        };
21718        let (gcum, p, u, w) =
21719            self.gdn_chunk_k123(q, k, v, g, beta, wb16_pre.as_mut(), n_head, t, c, hk, k2w)?;
21720        let _ = &w;
21721        let mut y = self.uninit(nc * h * c * D)?;
21722        let mut ssnap = self.uninit(nc * h * D * D)?; // chunk-start state snapshots (K5 phase 1)
21723        // K4-MMA seam (MEMRA_GDN_MMA; harness verdict 1.75x — tools/bench_gdn_k4.cu, ledger
21724        // 2026-07-26): M in mma accumulator fragments, bf16 W/k mirrors through a cp.async
21725        // ring. C==32 only (the kernel's tile). PROMOTED default-ON on the Hopper lane
21726        // after the STATE-CARRY battery (2026-07-26): 2048-token prime (64 in-kernel state
21727        // carries) -> 256 greedy decode tokens IDENTICAL to f32 on 3 seeds, AND chunked-
21728        // continuation prime (MEMRA_PRIME_CHUNK=512, 4 cross-call carries via cache.recur)
21729        // IDENTICAL on 2 seeds; plus argmax MATCH, pp512 +3.5% (17286), oracle out
21730        // mean_rel ~1e-4. kernel-check pins BOTH configs (f32 tight band forced =0; mma
21731        // band 8e-2/8e-1 vs f64 truth). =0 reverts; portable stays f32. NOT read via
21732        // OnceLock ON PURPOSE: kernel-check toggles the env per call to pin both forms.
21733        // default mirrors gdn_mma_enabled (incl. the sm_120a-build ON) — all three read
21734        // sites must agree or the pre-work arms while the scan takes the scalar route.
21735        let gdn_mma = !portable_mma_gated()
21736            && c == 32
21737            && match std::env::var("MEMRA_GDN_MMA").as_deref() {
21738                Ok("1") => true,
21739                Ok("0") => false,
21740                _ => gdn_mma_default_on(),
21741            };
21742        if gdn_mma {
21743            let wb16 = wb16_pre
21744                .take()
21745                .expect("mma path pre-allocates wb16 (K3 store fold)");
21746            let kb16_ref: &CudaSlice<u8> = kb16_ref0.expect("mma path pre-builds kb16 above K123");
21747            // K4+K5 FUSED wgmma seam (MEMRA_GDN_WGMMA, task #22; harness verdict
21748            // tools/bench_gdn_wgmma.cu v5, ledger 1f08b997: in-band Y 1.07e-2 / state
21749            // 1.03e-2 / O 1.08e-2, 91.3us vs 70.4 K4-only at H=32 T=512). K5's output
21750            // pass runs inside the persistent-M kernel; Y and Ssnap are never
21751            // materialized. New numeric class (gk folds into k^T instead of ys) —
21752            // explicit opt-in until the state-carry battery promotes it. Env read per
21753            // call (kernel-check pins configs by toggling env, GDN_MMA precedent).
21754            // PROMOTED default-ON hopper (2026-07-27): full battery green — harness
21755            // in-band, argmax gate PASS, 3-seed greedy IDENTICAL after ~2k prime,
21756            // chunked-continuation IDENTICAL, kernel-check + decode-batch gates green,
21757            // official prefill lane +0.74% interleaved x5 (5/5 rounds). =0 reverts.
21758            if gdn_wgmma_pre {
21759                // qb16/pb16 pre-built above K123 (K2-wgmma wrote the masked Pb16).
21760                let qb16 = qb16_ref0.unwrap();
21761                let pb16 = pb16.as_ref().unwrap();
21762                {
21763                    let f = self.func("gdn_k45_wgmma");
21764                    let cfg = LaunchConfig {
21765                        grid_dim: (h as u32, 4, 1),
21766                        block_dim: (256, 1, 1),
21767                        shared_mem_bytes: 0,
21768                    };
21769                    let hki = hk as i32;
21770                    let __s_b = self.gpu.stream();
21771                    let mut b = __s_b.launch_builder(&f);
21772                    b.arg(kb16_ref)
21773                        .arg(&gcum)
21774                        .arg(beta)
21775                        .arg(&u)
21776                        .arg(&wb16)
21777                        .arg(qb16)
21778                        .arg(pb16)
21779                        .arg(o)
21780                        .arg(&scale)
21781                        .arg(state_in)
21782                        .arg(&mut *state_out)
21783                        .arg(&hi)
21784                        .arg(&ti)
21785                        .arg(&ci)
21786                        .arg(&hki);
21787                    unsafe {
21788                        b.launch(cfg)?;
21789                    }
21790                }
21791                return Ok(());
21792            }
21793            // COUPLED PAIR: K4-mma writes Y and Ssnap as bf16 (their only consumer is
21794            // K5-mma, which rounds to bf16 regardless — identical numerics, half the
21795            // traffic; harness K5 63.0 -> 35.3us). Fresh bf16 buffers replace the f32 ones.
21796            let mut y16 = self.alloc_u8_uninit(nc * h * c * D * 2)?;
21797            let mut ssnap16 = self.alloc_u8_uninit(nc * h * D * D * 2)?;
21798            {
21799                let f = self.func("gdn_chunk_state_mma");
21800                let cfg = LaunchConfig {
21801                    grid_dim: (h as u32, NSPLIT, 1),
21802                    block_dim: (256, 1, 1),
21803                    shared_mem_bytes: 0,
21804                };
21805                let hki = hk as i32;
21806                let __s_b = self.gpu.stream();
21807                let mut b = __s_b.launch_builder(&f);
21808                b.arg(kb16_ref)
21809                    .arg(&gcum)
21810                    .arg(beta)
21811                    .arg(&u)
21812                    .arg(&wb16)
21813                    .arg(&mut y16)
21814                    .arg(&mut ssnap16)
21815                    .arg(state_in)
21816                    .arg(&mut *state_out)
21817                    .arg(&hi)
21818                    .arg(&ti)
21819                    .arg(&ci)
21820                    .arg(&hki);
21821                unsafe {
21822                    b.launch(cfg)?;
21823                }
21824            }
21825            {
21826                // K5-mma (bf16 St/Y consumers)
21827                let f = self.func("gdn_chunk_output_mma");
21828                let jt = ((c + 31) / 32) as u32;
21829                let cfg = LaunchConfig {
21830                    grid_dim: (nc as u32, h as u32, jt),
21831                    block_dim: (256, 1, 1),
21832                    shared_mem_bytes: 0,
21833                };
21834                let hki = hk as i32;
21835                let __s_b = self.gpu.stream();
21836                let mut b = __s_b.launch_builder(&f);
21837                b.arg(q)
21838                    .arg(&gcum)
21839                    .arg(&p)
21840                    .arg(&y16)
21841                    .arg(&ssnap16)
21842                    .arg(o)
21843                    .arg(&hi)
21844                    .arg(&ti)
21845                    .arg(&ci)
21846                    .arg(&scale)
21847                    .arg(&hki);
21848                unsafe {
21849                    b.launch(cfg)?;
21850                }
21851            }
21852            return Ok(());
21853        }
21854        {
21855            // K4 (sequential over chunks inside; blocks col-partition the state)
21856            let f = self.func("gdn_chunk_state_f32");
21857            let cfg = LaunchConfig {
21858                grid_dim: (h as u32, NSPLIT, 1),
21859                block_dim: (256, 1, 1),
21860                shared_mem_bytes: 0,
21861            };
21862            let __s_b = self.gpu.stream();
21863            let mut b = __s_b.launch_builder(&f);
21864            b.arg(k)
21865                .arg(&gcum)
21866                .arg(beta)
21867                .arg(&u)
21868                .arg(&w)
21869                .arg(&mut y)
21870                .arg(&mut ssnap)
21871                .arg(state_in)
21872                .arg(&mut *state_out)
21873                .arg(&hi)
21874                .arg(&ti)
21875                .arg(&ci);
21876            unsafe {
21877                b.launch(cfg)?;
21878            }
21879        }
21880        {
21881            // K5 (j-blocked: grid.z = 32-row output blocks per chunk; writes o fully)
21882            let f = self.func("gdn_chunk_output_f32");
21883            let jt = ((c + 31) / 32) as u32;
21884            let cfg = LaunchConfig {
21885                grid_dim: (nc as u32, h as u32, jt),
21886                block_dim: (256, 1, 1),
21887                shared_mem_bytes: 0,
21888            };
21889            let __s_b = self.gpu.stream();
21890            let mut b = __s_b.launch_builder(&f);
21891            b.arg(q)
21892                .arg(&gcum)
21893                .arg(&p)
21894                .arg(&y)
21895                .arg(&ssnap)
21896                .arg(o)
21897                .arg(&hi)
21898                .arg(&ti)
21899                .arg(&ci)
21900                .arg(&scale);
21901            unsafe {
21902                b.launch(cfg)?;
21903            }
21904        }
21905        Ok(())
21906    }
21907
21908    /// PREFILL GDN scan dispatch (the A4 seam): chunked WY form when enabled and T is in the
21909    /// batched-prefill regime, else the sequential scan. Callers: hybrid_forward::linear_attn
21910    /// (forward/forward_last) + linear_attn_prime (prime_cache). Decode (T=1) and the spec
21911    /// verify call `gdn_scan_s128` DIRECTLY — the decode==verify dispatch identity is untouched.
21912    ///
21913    /// MEMRA_GDN_DIFF=1: numerical-oracle mode — runs BOTH forms on the same inputs, prints the
21914    /// per-call (== per-layer, in call order) output/state error distribution, and keeps the
21915    /// SEQUENTIAL results so the run stays on the shipped path (stage-1 prototype evidence).
21916    #[allow(clippy::too_many_arguments)]
21917    #[allow(clippy::too_many_arguments)]
21918    pub fn gdn_scan_prefill(
21919        &self,
21920        q: &CudaSlice<f32>,
21921        k: &CudaSlice<f32>,
21922        v: &CudaSlice<f32>,
21923        g: &CudaSlice<f32>,
21924        beta: &CudaSlice<f32>,
21925        kb16_pre: Option<&CudaSlice<u8>>,
21926        qb16_pre: Option<&CudaSlice<u8>>,
21927        state_in: &CudaSlice<f32>,
21928        state_out: &mut CudaSlice<f32>,
21929        o: &mut CudaSlice<f32>,
21930        n_head: usize,
21931        t: usize,
21932        scale: f32,
21933        hk: usize,
21934    ) -> Result<(), Box<dyn std::error::Error>> {
21935        if std::env::var("MEMRA_GDN_DIFF").is_ok() && t >= 16 {
21936            assert!(hk == n_head, "GDN_DIFF oracle is broadcast-only");
21937            return self.gdn_scan_diff(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale);
21938        }
21939        if Self::gdn_chunked_enabled() && t >= 16 {
21940            self.gdn_scan_chunked(
21941                q,
21942                k,
21943                v,
21944                g,
21945                beta,
21946                kb16_pre,
21947                qb16_pre,
21948                state_in,
21949                state_out,
21950                o,
21951                n_head,
21952                t,
21953                scale,
21954                Self::gdn_chunk_size(),
21955                hk,
21956            )
21957        } else {
21958            assert!(
21959                hk == n_head,
21960                "s128 scan is broadcast-only (prep guarantees by predicate)"
21961            );
21962            self.gdn_scan_s128(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale)
21963        }
21964    }
21965
21966    /// Stage-1 oracle: run sequential AND chunked, report per-call error stats, keep sequential.
21967    #[allow(clippy::too_many_arguments)]
21968    fn gdn_scan_diff(
21969        &self,
21970        q: &CudaSlice<f32>,
21971        k: &CudaSlice<f32>,
21972        v: &CudaSlice<f32>,
21973        g: &CudaSlice<f32>,
21974        beta: &CudaSlice<f32>,
21975        state_in: &CudaSlice<f32>,
21976        state_out: &mut CudaSlice<f32>,
21977        o: &mut CudaSlice<f32>,
21978        n_head: usize,
21979        t: usize,
21980        scale: f32,
21981    ) -> Result<(), Box<dyn std::error::Error>> {
21982        static CALL: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
21983        let call = CALL.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
21984        let mut o_c = self.uninit(o.len())?;
21985        let mut st_c = self.uninit(state_out.len())?;
21986        self.gdn_scan_chunked(
21987            q,
21988            k,
21989            v,
21990            g,
21991            beta,
21992            None,
21993            None,
21994            state_in,
21995            &mut st_c,
21996            &mut o_c,
21997            n_head,
21998            t,
21999            scale,
22000            Self::gdn_chunk_size(),
22001            n_head,
22002        )?;
22003        self.gdn_scan_s128(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale)?;
22004        let (oh_s, oh_c) = (self.dtoh(o)?, self.dtoh(&o_c)?);
22005        let (sh_s, sh_c) = (self.dtoh(state_out)?, self.dtoh(&st_c)?);
22006        let stats = |a: &[f32], b: &[f32]| -> (f32, f32, f64) {
22007            let mut max_abs = 0f32;
22008            let mut max_rel = 0f32;
22009            let mut sum_rel = 0f64;
22010            for (x, y) in a.iter().zip(b) {
22011                let ad = (x - y).abs();
22012                let rel = ad / x.abs().max(y.abs()).max(1e-3);
22013                if ad > max_abs {
22014                    max_abs = ad;
22015                }
22016                if rel > max_rel {
22017                    max_rel = rel;
22018                }
22019                sum_rel += rel as f64;
22020            }
22021            (max_abs, max_rel, sum_rel / a.len() as f64)
22022        };
22023        let (o_ma, o_mr, o_mean) = stats(&oh_s, &oh_c);
22024        let (s_ma, s_mr, s_mean) = stats(&sh_s, &sh_c);
22025        println!(
22026            "[gdn-diff call {call:3} T={t} C={}] out: max_abs={o_ma:.3e} max_rel={o_mr:.3e} mean_rel={o_mean:.3e} | \
22027                  state: max_abs={s_ma:.3e} max_rel={s_mr:.3e} mean_rel={s_mean:.3e}",
22028            Self::gdn_chunk_size()
22029        );
22030        Ok(())
22031    }
22032
22033    /// softplus-based g_log: g_log[h,t] = a[h] * softplus(alpha[h,t] + dt_bias[h]). a pre-negated.
22034    pub fn gdn_glog(
22035        &self,
22036        alpha: &CudaSlice<f32>,
22037        dt_bias: &CudaSlice<f32>,
22038        a: &CudaSlice<f32>,
22039        g_log: &mut CudaSlice<f32>,
22040        n_head: usize,
22041        t: usize,
22042    ) -> Result<(), Box<dyn std::error::Error>> {
22043        let f = self.func("gdn_glog_f32");
22044        let cfg = LaunchConfig::for_num_elems((n_head * t) as u32);
22045        let (h, ti) = (n_head as i32, t as i32);
22046        let __s_b = self.gpu.stream();
22047        let mut b = __s_b.launch_builder(&f);
22048        b.arg(alpha).arg(dt_bias).arg(a).arg(g_log).arg(&h).arg(&ti);
22049        unsafe {
22050            b.launch(cfg)?;
22051        }
22052        Ok(())
22053    }
22054
22055    /// view twins (task #16): the batched prime's GDN core reads the CONCAT projection
22056    /// buffers at row offsets (CudaView) — same kernels, same values, no split copies.
22057    pub fn sigmoid_v(
22058        &self,
22059        x: &cudarc::driver::CudaView<f32>,
22060        y: &mut CudaSlice<f32>,
22061        n: usize,
22062    ) -> Result<(), Box<dyn std::error::Error>> {
22063        let f = self.func("sigmoid_f32");
22064        let cfg = LaunchConfig::for_num_elems(n as u32);
22065        let ni = n as i32;
22066        let __s_b = self.gpu.stream();
22067        let mut b = __s_b.launch_builder(&f);
22068        b.arg(x).arg(y).arg(&ni);
22069        unsafe {
22070            b.launch(cfg)?;
22071        }
22072        Ok(())
22073    }
22074
22075    pub fn gdn_glog_v(
22076        &self,
22077        alpha: &cudarc::driver::CudaView<f32>,
22078        dt_bias: &CudaSlice<f32>,
22079        a: &CudaSlice<f32>,
22080        g_log: &mut CudaSlice<f32>,
22081        n_head: usize,
22082        t: usize,
22083    ) -> Result<(), Box<dyn std::error::Error>> {
22084        let f = self.func("gdn_glog_f32");
22085        let cfg = LaunchConfig::for_num_elems((n_head * t) as u32);
22086        let (h, ti) = (n_head as i32, t as i32);
22087        let __s_b = self.gpu.stream();
22088        let mut b = __s_b.launch_builder(&f);
22089        b.arg(alpha).arg(dt_bias).arg(a).arg(g_log).arg(&h).arg(&ti);
22090        unsafe {
22091            b.launch(cfg)?;
22092        }
22093        Ok(())
22094    }
22095
22096    pub fn sigmoid(
22097        &self,
22098        x: &CudaSlice<f32>,
22099        y: &mut CudaSlice<f32>,
22100        n: usize,
22101    ) -> Result<(), Box<dyn std::error::Error>> {
22102        let f = self.func("sigmoid_f32");
22103        let cfg = LaunchConfig::for_num_elems(n as u32);
22104        let ni = n as i32;
22105        let __s_b = self.gpu.stream();
22106        let mut b = __s_b.launch_builder(&f);
22107        b.arg(x).arg(y).arg(&ni);
22108        unsafe {
22109            b.launch(cfg)?;
22110        }
22111        Ok(())
22112    }
22113
22114    /// attn out-gate fused epilogue (task #17): dst = a * sigmoid(g) + fp16 twin, one launch
22115    /// (replaces sigmoid + mul + convert). Bit-identical class.
22116    pub fn sig_mul_f16out(
22117        &self,
22118        a: &CudaSlice<f32>,
22119        g: &CudaSlice<f32>,
22120        dst: &mut CudaSlice<f32>,
22121        dst16: &mut CudaSlice<u8>,
22122        n: usize,
22123    ) -> Result<(), Box<dyn std::error::Error>> {
22124        let f = self.func("sig_mul_f16out_f32");
22125        let cfg = LaunchConfig::for_num_elems(n as u32);
22126        let ni = n as i32;
22127        let __s_b = self.gpu.stream();
22128        let mut b = __s_b.launch_builder(&f);
22129        b.arg(a).arg(g).arg(dst).arg(dst16).arg(&ni);
22130        unsafe {
22131            b.launch(cfg)?;
22132        }
22133        Ok(())
22134    }
22135
22136    /// step35 (Step-3.7-Flash) SEPARATE head-wise attention gate: one scalar per query head,
22137    /// broadcast over head_dim. `dst = a * sigmoid(g)` where `a`/`dst` are `[head_dim, n_head, T]`
22138    /// (the `q_gate_split` layout) and `g` is the PRE-sigmoid `attn_gate` projection output in
22139    /// token-major `[T, n_head]`. `dst16` is the optional fp16 operand for wo (None -> skipped).
22140    ///
22141    /// NOT interchangeable with `sig_mul_f16out`, which gates FULL WIDTH (qwen35 packs one gate
22142    /// value per (head, dim) element inside wq). Using this for that, or that for this, silently
22143    /// applies the wrong number of distinct gate values.
22144    #[allow(clippy::too_many_arguments)]
22145    pub fn attn_head_gate(
22146        &self,
22147        a: &CudaSlice<f32>,
22148        g: &CudaSlice<f32>,
22149        dst: &mut CudaSlice<f32>,
22150        dst16: Option<&mut CudaSlice<u8>>,
22151        head_dim: usize,
22152        n_head: usize,
22153        t: usize,
22154    ) -> Result<(), Box<dyn std::error::Error>> {
22155        let f = self.func("attn_head_gate_f32");
22156        let cfg = LaunchConfig::for_num_elems((head_dim * n_head * t) as u32);
22157        let (hd, nh, ti) = (head_dim as i32, n_head as i32, t as i32);
22158        // nullable device pointer by value (0 = skip), same convention as `l2_norm_pp`.
22159        let d16: u64 = match dst16 {
22160            Some(d) => self.addr_u8(d),
22161            None => 0,
22162        };
22163        let __s_b = self.gpu.stream();
22164        let mut b = __s_b.launch_builder(&f);
22165        b.arg(a)
22166            .arg(g)
22167            .arg(dst)
22168            .arg(&d16)
22169            .arg(&hd)
22170            .arg(&nh)
22171            .arg(&ti);
22172        unsafe {
22173            b.launch(cfg)?;
22174        }
22175        Ok(())
22176    }
22177
22178    /// step35 CLAMPED SwiGLU: `dst = min(silu(gate*gs), limit) * clamp(up*us, +-limit)`.
22179    /// Verbatim from llama.cpp `llama-graph.cpp:2146-2165` (routed, `swiglu_clamp_exp`) and
22180    /// `:1751-1770` (shared, `swiglu_clamp_shexp`), non-DEEPSEEK4 branch.
22181    ///
22182    /// This is NOT `swigluoai_mul_scaled`: that one clamps the gate BEFORE swish and multiplies by
22183    /// `(1 + clamp(up))`. Caller MUST check `limit > 1e-6` (upstream's eps gate) and use the plain
22184    /// `silu_mul_scaled` path otherwise — at limit=0 this kernel would clamp every positive
22185    /// activation to zero. On Step-3.7-Flash only layers 43 (7.0) and 44 (16.0) have a live limit.
22186    #[allow(clippy::too_many_arguments)]
22187    pub fn swiglu_clamped_mul_scaled(
22188        &self,
22189        gate: &CudaSlice<f32>,
22190        up: &CudaSlice<f32>,
22191        gs: f32,
22192        us: f32,
22193        limit: f32,
22194        dst: &mut CudaSlice<f32>,
22195        n: usize,
22196    ) -> Result<(), Box<dyn std::error::Error>> {
22197        debug_assert!(
22198            limit > 1e-6,
22199            "swiglu_clamped needs a live limit; use silu_mul_scaled"
22200        );
22201        let f = self.func("swiglu_clamped_mul_scaled_f32");
22202        let cfg = LaunchConfig::for_num_elems(n as u32);
22203        let ni = n as i32;
22204        let __s_b = self.gpu.stream();
22205        let mut b = __s_b.launch_builder(&f);
22206        b.arg(gate)
22207            .arg(up)
22208            .arg(&gs)
22209            .arg(&us)
22210            .arg(&limit)
22211            .arg(dst)
22212            .arg(&ni);
22213        unsafe {
22214            b.launch(cfg)?;
22215        }
22216        Ok(())
22217    }
22218
22219    /// gated RMSNorm: dst = RMSNorm(o, w[ncols]) * silu(z), per row of ncols. nrows blocks.
22220    pub fn gated_rmsnorm(
22221        &self,
22222        o: &CudaSlice<f32>,
22223        w: &CudaSlice<f32>,
22224        z: &CudaSlice<f32>,
22225        dst: &mut CudaSlice<f32>,
22226        ncols: usize,
22227        nrows: usize,
22228        eps: f32,
22229    ) -> Result<(), Box<dyn std::error::Error>> {
22230        let f = self.func("gated_rmsnorm_f32");
22231        let cfg = LaunchConfig {
22232            grid_dim: (nrows as u32, 1, 1),
22233            block_dim: (128, 1, 1),
22234            shared_mem_bytes: 0,
22235        };
22236        let (nc, e) = (ncols as i32, eps);
22237        let __s_b = self.gpu.stream();
22238        let mut b = __s_b.launch_builder(&f);
22239        b.arg(o).arg(w).arg(z).arg(dst).arg(&nc).arg(&e);
22240        unsafe {
22241            b.launch(cfg)?;
22242        }
22243        Ok(())
22244    }
22245
22246    /// f16out twin of `gated_rmsnorm` (task #17): epilogue also emits the fp16 operand for
22247    /// the ssm_out GEMM. Bit-identical class (same floats + the cvt kernel's __float2half).
22248    pub fn gated_rmsnorm_f16out(
22249        &self,
22250        o: &CudaSlice<f32>,
22251        w: &CudaSlice<f32>,
22252        z: &CudaSlice<f32>,
22253        dst: &mut CudaSlice<f32>,
22254        dst16: &mut CudaSlice<u8>,
22255        ncols: usize,
22256        nrows: usize,
22257        eps: f32,
22258    ) -> Result<(), Box<dyn std::error::Error>> {
22259        let f = self.func("gated_rmsnorm_f16out_f32");
22260        // block_dim MUST match gated_rmsnorm's (128): the reduction tree order pins the scale
22261        let cfg = LaunchConfig {
22262            grid_dim: (nrows as u32, 1, 1),
22263            block_dim: (128, 1, 1),
22264            shared_mem_bytes: 0,
22265        };
22266        let (nc, e) = (ncols as i32, eps);
22267        let __s_b = self.gpu.stream();
22268        let mut b = __s_b.launch_builder(&f);
22269        b.arg(o).arg(w).arg(z).arg(dst).arg(dst16).arg(&nc).arg(&e);
22270        unsafe {
22271            b.launch(cfg)?;
22272        }
22273        Ok(())
22274    }
22275
22276    /// add+RMSNorm emitting the f32 normed row AND its q8_1 quantization in one launch (the MoE
22277    /// layer input: z feeds the router matmul as f32, the expert dp4a as q8_1). BIT-IDENTICAL to
22278    /// add_rms_norm + quantize_q8_1. Returns (q, d) alongside the caller-provided res/z buffers.
22279    #[allow(clippy::too_many_arguments)]
22280    pub fn add_rms_norm_zq8(
22281        &self,
22282        a: &CudaSlice<f32>,
22283        b_in: &CudaSlice<f32>,
22284        w: &CudaSlice<f32>,
22285        res: &mut CudaSlice<f32>,
22286        z: &mut CudaSlice<f32>,
22287        ncols: usize,
22288        nrows: usize,
22289        eps: f32,
22290    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
22291        assert!(ncols % 32 == 0);
22292        let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
22293        let mut d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
22294        let f = self.func("add_rms_norm_zq8");
22295        let cfg = LaunchConfig {
22296            grid_dim: (nrows as u32, 1, 1),
22297            block_dim: (1024, 1, 1),
22298            shared_mem_bytes: 0,
22299        };
22300        let (nc, ep) = (ncols as i32, eps);
22301        let __s_b = self.gpu.stream();
22302        let mut b = __s_b.launch_builder(&f);
22303        b.arg(a)
22304            .arg(b_in)
22305            .arg(w)
22306            .arg(res)
22307            .arg(z)
22308            .arg(&mut q)
22309            .arg(&mut d)
22310            .arg(&nc)
22311            .arg(&ep);
22312        unsafe {
22313            b.launch(cfg)?;
22314        }
22315        Ok((q, d))
22316    }
22317
22318    /// gated RMSNorm emitting q8_1 directly (fused quantize epilogue) — the ssm_out matvec input.
22319    /// BIT-IDENTICAL bytes to gated_rmsnorm + quantize_q8_1 (ncols % 32 == 0; blocks never straddle
22320    /// rows). Saves one launch per linear-attn layer (36/token on the 9B).
22321    /// z-view twins of gated_rmsnorm(+f16out) — task #16 batched-prime split removal.
22322    pub fn gated_rmsnorm_zv(
22323        &self,
22324        o: &CudaSlice<f32>,
22325        w: &CudaSlice<f32>,
22326        z: &cudarc::driver::CudaView<f32>,
22327        dst: &mut CudaSlice<f32>,
22328        ncols: usize,
22329        nrows: usize,
22330        eps: f32,
22331    ) -> Result<(), Box<dyn std::error::Error>> {
22332        let f = self.func("gated_rmsnorm_f32");
22333        let cfg = LaunchConfig {
22334            grid_dim: (nrows as u32, 1, 1),
22335            block_dim: (128, 1, 1),
22336            shared_mem_bytes: 0,
22337        };
22338        let (nc, e) = (ncols as i32, eps);
22339        let __s_b = self.gpu.stream();
22340        let mut b = __s_b.launch_builder(&f);
22341        b.arg(o).arg(w).arg(z).arg(dst).arg(&nc).arg(&e);
22342        unsafe {
22343            b.launch(cfg)?;
22344        }
22345        Ok(())
22346    }
22347
22348    pub fn gated_rmsnorm_f16out_zv(
22349        &self,
22350        o: &CudaSlice<f32>,
22351        w: &CudaSlice<f32>,
22352        z: &cudarc::driver::CudaView<f32>,
22353        dst: &mut CudaSlice<f32>,
22354        dst16: &mut CudaSlice<u8>,
22355        ncols: usize,
22356        nrows: usize,
22357        eps: f32,
22358    ) -> Result<(), Box<dyn std::error::Error>> {
22359        let f = self.func("gated_rmsnorm_f16out_f32");
22360        // block_dim MUST match gated_rmsnorm's (128): the reduction tree order pins the scale
22361        let cfg = LaunchConfig {
22362            grid_dim: (nrows as u32, 1, 1),
22363            block_dim: (128, 1, 1),
22364            shared_mem_bytes: 0,
22365        };
22366        let (nc, e) = (ncols as i32, eps);
22367        let __s_b = self.gpu.stream();
22368        let mut b = __s_b.launch_builder(&f);
22369        b.arg(o).arg(w).arg(z).arg(dst).arg(dst16).arg(&nc).arg(&e);
22370        unsafe {
22371            b.launch(cfg)?;
22372        }
22373        Ok(())
22374    }
22375
22376    pub fn gated_rmsnorm_q8_1(
22377        &self,
22378        o: &CudaSlice<f32>,
22379        w: &CudaSlice<f32>,
22380        z: &CudaSlice<f32>,
22381        ncols: usize,
22382        nrows: usize,
22383        eps: f32,
22384    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
22385        assert!(ncols % 32 == 0);
22386        let f = self.func("gated_rmsnorm_q8_1");
22387        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
22388        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
22389        let cfg = LaunchConfig {
22390            grid_dim: (nrows as u32, 1, 1),
22391            block_dim: (128, 1, 1),
22392            shared_mem_bytes: 0,
22393        };
22394        let (nc, ep) = (ncols as i32, eps);
22395        let __s_b = self.gpu.stream();
22396        let mut b = __s_b.launch_builder(&f);
22397        b.arg(o)
22398            .arg(w)
22399            .arg(z)
22400            .arg(&mut out_q)
22401            .arg(&mut out_d)
22402            .arg(&nc)
22403            .arg(&ep);
22404        unsafe {
22405            b.launch(cfg)?;
22406        }
22407        Ok((out_q, out_d))
22408    }
22409
22410    /// transpose [rows,cols] row-major -> [cols,rows] row-major.
22411    pub fn transpose(
22412        &self,
22413        inp: &CudaSlice<f32>,
22414        rows: usize,
22415        cols: usize,
22416    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22417        let f = self.func("transpose_f32");
22418        let mut out = self.zeros(rows * cols)?;
22419        let cfg = LaunchConfig::for_num_elems((rows * cols) as u32);
22420        let (r, c) = (rows as i32, cols as i32);
22421        let __s_b = self.gpu.stream();
22422        let mut b = __s_b.launch_builder(&f);
22423        b.arg(inp).arg(&mut out).arg(&r).arg(&c);
22424        unsafe {
22425            b.launch(cfg)?;
22426        }
22427        Ok(out)
22428    }
22429
22430    /// repeat-interleave heads: in[head_dim,n_in,T] -> out[head_dim,n_out,T].
22431    pub fn repeat_heads(
22432        &self,
22433        inp: &CudaSlice<f32>,
22434        out: &mut CudaSlice<f32>,
22435        head_dim: usize,
22436        n_in: usize,
22437        n_out: usize,
22438        t: usize,
22439    ) -> Result<(), Box<dyn std::error::Error>> {
22440        let f = self.func("repeat_heads_f32");
22441        let cfg = LaunchConfig::for_num_elems((head_dim * n_out * t) as u32);
22442        let (hd, ni, no, ti) = (head_dim as i32, n_in as i32, n_out as i32, t as i32);
22443        let __s_b = self.gpu.stream();
22444        let mut b = __s_b.launch_builder(&f);
22445        b.arg(inp).arg(out).arg(&hd).arg(&ni).arg(&no).arg(&ti);
22446        unsafe {
22447            b.launch(cfg)?;
22448        }
22449        Ok(())
22450    }
22451
22452    /// q|gate split (on-device). qf:[T, n_head*2*head_dim] -> q_out,gate_out:[head_dim,n_head,T].
22453    /// Replaces the dtoh->host-double-loop->htod in full_attn / full_attn_decode.
22454    ///
22455    /// Bounds-checked against `qf` before the launch: the kernel reads `2*head_dim*n_head*t`
22456    /// floats, and running it on a `wq` output that carries no fused gate reads 2x off the end
22457    /// (silently, on the device). A layout mismatch is a typed `FusedQGateExtent` here instead.
22458    pub fn q_gate_split(
22459        &self,
22460        qf: &CudaSlice<f32>,
22461        q_out: &mut CudaSlice<f32>,
22462        gate_out: &mut CudaSlice<f32>,
22463        head_dim: usize,
22464        n_head: usize,
22465        t: usize,
22466    ) -> Result<(), Box<dyn std::error::Error>> {
22467        memra_gguf::config::check_fused_q_gate_extent(qf.len(), head_dim, n_head, t)?;
22468        let out_need = head_dim * n_head * t;
22469        if q_out.len() < out_need || gate_out.len() < out_need {
22470            return Err(format!(
22471                "q_gate_split destinations too small: need {out_need} each, have q={} gate={}",
22472                q_out.len(),
22473                gate_out.len()
22474            )
22475            .into());
22476        }
22477        let f = self.func("q_gate_split_f32");
22478        let cfg = LaunchConfig::for_num_elems((head_dim * n_head * t) as u32);
22479        let (hd, nh, ti) = (head_dim as i32, n_head as i32, t as i32);
22480        let __s_b = self.gpu.stream();
22481        let mut b = __s_b.launch_builder(&f);
22482        b.arg(qf)
22483            .arg(q_out)
22484            .arg(gate_out)
22485            .arg(&hd)
22486            .arg(&nh)
22487            .arg(&ti);
22488        unsafe {
22489            b.launch(cfg)?;
22490        }
22491        Ok(())
22492    }
22493
22494    /// qkv->GDN repack (on-device). conv_out:[conv_dim,T] channel-major ->
22495    /// q_g/k_g/v_g:[d_state,num_v,T] with q/k head-repeat kh = vh % num_k (validated modulo mapping).
22496    /// Replaces the dtoh->host-q/k/v-repack->3x-htod in linear_attn / linear_attn_decode.
22497    pub fn qkv_to_gdn_repack(
22498        &self,
22499        conv_out: &CudaSlice<f32>,
22500        q_g: &mut CudaSlice<f32>,
22501        k_g: &mut CudaSlice<f32>,
22502        v_g: &mut CudaSlice<f32>,
22503        d_state: usize,
22504        num_v: usize,
22505        num_k: usize,
22506        key_dim: usize,
22507        t: usize,
22508    ) -> Result<(), Box<dyn std::error::Error>> {
22509        let f = self.func("qkv_to_gdn_repack_f32");
22510        let cfg = LaunchConfig::for_num_elems((d_state * num_v * t) as u32);
22511        let (ds, nv, nk, kd, ti) = (
22512            d_state as i32,
22513            num_v as i32,
22514            num_k as i32,
22515            key_dim as i32,
22516            t as i32,
22517        );
22518        let __s_b = self.gpu.stream();
22519        let mut b = __s_b.launch_builder(&f);
22520        b.arg(conv_out)
22521            .arg(q_g)
22522            .arg(k_g)
22523            .arg(v_g)
22524            .arg(&ds)
22525            .arg(&nv)
22526            .arg(&nk)
22527            .arg(&kd)
22528            .arg(&ti);
22529        unsafe {
22530            b.launch(cfg)?;
22531        }
22532        Ok(())
22533    }
22534
22535    /// conv left zero-pad (prefill from zero state). src:[conv_dim,T] -> dst:[conv_dim,T+pad],
22536    /// cols 0..pad = 0, cols pad..pad+T = src. `dst` MUST be pre-zeroed. No dtoh/host-loop/htod.
22537    pub fn conv_left_pad(
22538        &self,
22539        src: &CudaSlice<f32>,
22540        dst: &mut CudaSlice<f32>,
22541        conv_dim: usize,
22542        t: usize,
22543        pad: usize,
22544    ) -> Result<(), Box<dyn std::error::Error>> {
22545        let f = self.func("conv_left_pad_f32");
22546        let cfg = LaunchConfig::for_num_elems((conv_dim * t) as u32);
22547        let (cd, ti, p) = (conv_dim as i32, t as i32, pad as i32);
22548        let __s_b = self.gpu.stream();
22549        let mut b = __s_b.launch_builder(&f);
22550        b.arg(src).arg(dst).arg(&cd).arg(&ti).arg(&p);
22551        unsafe {
22552            b.launch(cfg)?;
22553        }
22554        Ok(())
22555    }
22556
22557    /// conv-state assemble + ring roll (decode T=1). conv_state:[conv_dim,pad] (resident),
22558    /// qkv_col:[conv_dim] -> conv_in:[conv_dim,pad+1]; AND rolls conv_state (keep last pad cols).
22559    /// Replaces the dtoh->host-conv-ring-assemble->ring-update->htod in linear_attn_decode.
22560    pub fn conv_assemble_and_roll(
22561        &self,
22562        qkv_col: &CudaSlice<f32>,
22563        conv_state: &mut CudaSlice<f32>,
22564        conv_in: &mut CudaSlice<f32>,
22565        conv_dim: usize,
22566        pad: usize,
22567    ) -> Result<(), Box<dyn std::error::Error>> {
22568        let f = self.func("conv_assemble_and_roll_f32");
22569        let cfg = LaunchConfig::for_num_elems(conv_dim as u32);
22570        let (cd, p) = (conv_dim as i32, pad as i32);
22571        let __s_b = self.gpu.stream();
22572        let mut b = __s_b.launch_builder(&f);
22573        b.arg(qkv_col).arg(conv_state).arg(conv_in).arg(&cd).arg(&p);
22574        unsafe {
22575            b.launch(cfg)?;
22576        }
22577        Ok(())
22578    }
22579
22580    /// RANK3 LEVER (conv fuse, T=1 DECODE): fused conv_assemble_and_roll + ssm_conv1d_silu in ONE
22581    /// launch. Assembles the conv window [conv_state | qkv_col] in registers, computes the depthwise
22582    /// causal conv + SiLU into `conv_out`, and rolls the ring — never materializing conv_in to HBM.
22583    /// Replaces e.conv_assemble_and_roll(...) + e.ssm_conv1d(...). BIT-IDENTICAL to that two-kernel
22584    /// sequence (same 8-wide accumulation order, same SiLU). `conv_out` is [conv_dim] (T=1).
22585    pub fn ssm_conv1d_fused_decode(
22586        &self,
22587        qkv_col: &CudaSlice<f32>,
22588        conv_state: &mut CudaSlice<f32>,
22589        w: &CudaSlice<f32>,
22590        conv_out: &mut CudaSlice<f32>,
22591        conv_dim: usize,
22592        d_conv: usize,
22593    ) -> Result<(), Box<dyn std::error::Error>> {
22594        let f = self.func("ssm_conv1d_fused_decode_f32");
22595        let cfg = LaunchConfig::for_num_elems(conv_dim as u32);
22596        let (cd, dc) = (conv_dim as i32, d_conv as i32);
22597        let __s_b = self.gpu.stream();
22598        let mut b = __s_b.launch_builder(&f);
22599        b.arg(qkv_col)
22600            .arg(conv_state)
22601            .arg(w)
22602            .arg(conv_out)
22603            .arg(&cd)
22604            .arg(&dc);
22605        unsafe {
22606            b.launch(cfg)?;
22607        }
22608        Ok(())
22609    }
22610
22611    /// Copy a contiguous range [start, start+len) out of src into a fresh slice (device→device via host).
22612    /// Used for qkv split views. Small/rare; not perf-critical in Stage 1.
22613    pub fn slice_range(
22614        &self,
22615        src: &CudaSlice<f32>,
22616        start: usize,
22617        len: usize,
22618    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22619        let host = self.gpu.stream().clone_dtoh(src)?;
22620        self.gpu.stream().synchronize()?;
22621        Ok(self.htod(&host[start..start + len])?)
22622    }
22623}
22624
22625#[cfg(test)]
22626mod target_dispatch_tests {
22627    use super::legacy_quant_gemm_allowed;
22628
22629    #[test]
22630    fn legacy_quant_gemm_arch_policy_honors_the_escape_hatch() {
22631        // sm_120a native lane
22632        assert!(legacy_quant_gemm_allowed(false, false, false));
22633        assert!(!legacy_quant_gemm_allowed(false, false, true));
22634        // pure portable lane (sm_89): gated
22635        assert!(!legacy_quant_gemm_allowed(true, false, false));
22636        assert!(!legacy_quant_gemm_allowed(true, false, true));
22637        // Hopper-MMA lane (sm_90a): portable build, int8-MMA GEMM re-admitted
22638        assert!(legacy_quant_gemm_allowed(true, true, false));
22639        assert!(!legacy_quant_gemm_allowed(true, true, true));
22640    }
22641
22642    #[cfg(all(memra_portable_cuda, not(memra_hopper_mma)))]
22643    #[test]
22644    fn portable_build_disables_legacy_quant_gemm_without_an_env_override() {
22645        assert!(!legacy_quant_gemm_allowed(
22646            cfg!(memra_portable_cuda),
22647            cfg!(memra_hopper_mma),
22648            false
22649        ));
22650    }
22651
22652    #[cfg(memra_hopper_mma)]
22653    #[test]
22654    fn hopper_mma_build_re_admits_legacy_quant_gemm() {
22655        assert!(legacy_quant_gemm_allowed(
22656            cfg!(memra_portable_cuda),
22657            cfg!(memra_hopper_mma),
22658            false
22659        ));
22660        assert!(super::portable_mma_gated() == false);
22661    }
22662}
22663
22664/// The memra-kv device seam (Phase D): the cache's 7 ops delegate to the engine's
22665/// inherent methods (inherent methods win name resolution, so no recursion).
22666impl memra_kv::KvDev for Engine {
22667    fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22668        Engine::zeros(self, n)
22669    }
22670    fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22671        Engine::uninit(self, n)
22672    }
22673    fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
22674        Engine::alloc_u8(self, n)
22675    }
22676    fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
22677        Engine::htod_i32(self, v)
22678    }
22679    fn clone_dtod(
22680        &self,
22681        src: &CudaSlice<f32>,
22682    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22683        Engine::clone_dtod(self, src)
22684    }
22685    fn copy_into(
22686        &self,
22687        dst: &mut CudaSlice<f32>,
22688        off: usize,
22689        src: &CudaSlice<f32>,
22690        len: usize,
22691    ) -> Result<(), Box<dyn std::error::Error>> {
22692        Engine::copy_into(self, dst, off, src, len)
22693    }
22694    fn set_i32_one(
22695        &self,
22696        d: &mut CudaSlice<i32>,
22697        v: i32,
22698    ) -> Result<(), Box<dyn std::error::Error>> {
22699        Engine::set_i32_one(self, d, v)
22700    }
22701}
22702
22703#[cfg(test)]
22704mod fused_gate_bounds_tests {
22705    use super::*;
22706
22707    /// The fused `[q|gate]` split's read-site guard, on the device.
22708    ///
22709    /// `q_gate_split_f32` reads `2*head_dim*n_head*T` floats out of `qf`. A checkpoint whose gate
22710    /// is a SEPARATE tensor produces a `wq` output of exactly half that, so before 2026-08-19 the
22711    /// kernel launched and read 2x past the end of the allocation — an out-of-bounds DEVICE read:
22712    /// no panic, no error, just whatever memory follows. The guard turns it into a typed
22713    /// `FusedQGateExtent` before the launch.
22714    ///
22715    /// Catch demonstration for this test (guard temporarily removed, then restored):
22716    /// `compute-sanitizer --tool memcheck` on the half-width case reported invalid `__global__`
22717    /// reads of size 4 in `q_gate_split_f32`; with the guard in place the same run is clean and
22718    /// the call returns `Err`. Receipt in the lane report.
22719    #[test]
22720    #[ignore = "requires a CUDA GPU"]
22721    fn q_gate_split_refuses_a_separate_gate_wq_instead_of_reading_past_it() {
22722        let e = Engine::new(0).unwrap();
22723        let (head_dim, n_head, t) = (8usize, 4usize, 2usize);
22724        let fused = 2 * head_dim * n_head * t;
22725        let out_n = head_dim * n_head * t;
22726
22727        // half-width `qf` = the separate-gate / ungated layout. MUST be refused.
22728        let narrow = e.htod(&vec![1.0f32; out_n]).unwrap();
22729        let mut q = e.uninit(out_n).unwrap();
22730        let mut gate = e.uninit(out_n).unwrap();
22731        let err = e
22732            .q_gate_split(&narrow, &mut q, &mut gate, head_dim, n_head, t)
22733            .expect_err("half-width wq must be refused, not read past")
22734            .to_string();
22735        assert!(err.contains("NO fused gate"), "{err}");
22736        assert!(err.contains(&format!("{fused}")), "{err}");
22737
22738        // full-width `qf` = a real qwen3.5 fused layout. MUST still run, and split correctly:
22739        // per head hh the block is [q(head_dim) | gate(head_dim)] at stride 2*head_dim.
22740        let host: Vec<f32> = (0..fused).map(|i| i as f32).collect();
22741        let wide = e.htod(&host).unwrap();
22742        e.q_gate_split(&wide, &mut q, &mut gate, head_dim, n_head, t)
22743            .expect("full-width wq splits");
22744        let (qh, gh) = (e.dtoh(&q).unwrap(), e.dtoh(&gate).unwrap());
22745        for tok in 0..t {
22746            for hh in 0..n_head {
22747                for d in 0..head_dim {
22748                    let base = tok * (n_head * 2 * head_dim) + hh * (2 * head_dim);
22749                    let idx = tok * (n_head * head_dim) + hh * head_dim + d;
22750                    assert_eq!(qh[idx], host[base + d], "q t{tok} h{hh} d{d}");
22751                    assert_eq!(gh[idx], host[base + head_dim + d], "gate t{tok} h{hh} d{d}");
22752                }
22753            }
22754        }
22755
22756        // undersized destinations are refused too (the other half of the extent contract)
22757        let mut small = e.uninit(out_n - 1).unwrap();
22758        assert!(
22759            e.q_gate_split(&wide, &mut small, &mut gate, head_dim, n_head, t)
22760                .is_err()
22761        );
22762    }
22763}
22764
22765/// FULL-WIDTH-ROPE CONTRACT on the fused rms_norm+qkv+rope kernels
22766/// (lane/graph-s-key-exactness-20260819, probe O-10 follow-up). CPU-only: the guard runs before
22767/// any launch, so the refusal is testable without a device.
22768#[cfg(test)]
22769mod fused_rope_width_tests {
22770    use super::Engine;
22771
22772    /// gemma-4: rotary width == head width on both classes (GGUF 256/256 and 512/512, and the
22773    /// safetensors route derives the same), which is why the fusion is legal there today.
22774    #[test]
22775    fn full_width_is_accepted() {
22776        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 256, 256).is_ok());
22777        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_cat", 512, 512).is_ok());
22778        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append", 128, 128).is_ok());
22779    }
22780
22781    /// The widths the gemma-4 31B OFFICIAL artifact declares, read from its own GGUF header
22782    /// (`gemma-4-31B-it-official-Q8_0-MTP.gguf`, box3, 2026-08-19):
22783    ///
22784    /// ```text
22785    /// attention.key_length     512   rope.dimension_count     512   (global class)
22786    /// attention.key_length_swa 256   rope.dimension_count_swa 256   (SWA class)
22787    /// ```
22788    ///
22789    /// Both classes satisfy `n_rot == head_dim`, which is why the fusion is legal for gemma and
22790    /// why `HybridModel::gemma4_rope_dims` can feed this guard without refusing what we serve.
22791    /// An artifact that ever declares otherwise gets a loud refusal at the first fused launch
22792    /// instead of a silently over-rotated head.
22793    #[test]
22794    fn gemma4_official_artifact_widths_pass() {
22795        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 512, 512).is_ok());
22796        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append_dc", 256, 256).is_ok());
22797    }
22798
22799    /// The bug this guard exists to make impossible: a partial-rotary arch fused onto a kernel
22800    /// with no `n_dims`, silently rotating the pass-through band.
22801    #[test]
22802    fn partial_rotary_is_refused_with_the_geometry_named() {
22803        // qwen3.5: n_rot 64 of head_dim 256 (the shape probe O-10 pinned in the split path).
22804        let err = Engine::full_width_rope_only("rms_norm_qkv_rope", 64, 256)
22805            .expect_err("partial rotary must refuse");
22806        let msg = err.to_string();
22807        assert!(msg.contains("PARTIAL ROTARY REFUSED"), "{msg}");
22808        assert!(msg.contains("n_rot 64"), "{msg}");
22809        assert!(msg.contains("head_dim 256"), "{msg}");
22810        assert!(
22811            msg.contains("64..256"),
22812            "names the band it would corrupt: {msg}"
22813        );
22814        // step35 full-attn: 64 of 128 (upstream halves n_rot_full).
22815        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append_dc", 64, 128).is_err());
22816        // and the reverse mismatch (a wider rope than the head) is not "close enough" either.
22817        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 256, 128).is_err());
22818    }
22819}