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        let f = self.func("filter_stats_f32");
1455        let (ni, nr, rs) = (n as i32, nrow as i32, row_stride as i64);
1456        let cfg = LaunchConfig {
1457            grid_dim: (nrow as u32, 1, 1),
1458            block_dim: (1024, 1, 1),
1459            shared_mem_bytes: 0,
1460        };
1461        let __s_b = self.gpu.stream();
1462        let mut b = __s_b.launch_builder(&f);
1463        b.arg(x)
1464            .arg(&rs)
1465            .arg(rows)
1466            .arg(&mut *out_th)
1467            .arg(&mut *out_z)
1468            .arg(&mut *out_max)
1469            .arg(&ni)
1470            .arg(&nr)
1471            .arg(&temp)
1472            .arg(&top_k)
1473            .arg(&top_p)
1474            .arg(&min_p);
1475        unsafe {
1476            b.launch(cfg)?;
1477        }
1478        Ok(())
1479    }
1480
1481    /// out[pair] = filtered-softmax prob of ids[pair] in row rows[pair] (th/z per PAIR).
1482    #[allow(clippy::too_many_arguments)]
1483    pub fn softmax_gather_filtered(
1484        &self,
1485        x: &CudaSlice<f32>,
1486        row_stride: usize,
1487        ids: &CudaSlice<u32>,
1488        rows: &CudaSlice<i32>,
1489        th: &CudaSlice<f32>,
1490        z: &CudaSlice<f32>,
1491        out: &mut CudaSlice<f32>,
1492        n: usize,
1493        npair: usize,
1494        temp: f32,
1495    ) -> Result<(), Box<dyn std::error::Error>> {
1496        let f = self.func("softmax_gather_filtered_f32");
1497        let (ni, np, rs) = (n as i32, npair as i32, row_stride as i64);
1498        let cfg = LaunchConfig {
1499            grid_dim: (npair as u32, 1, 1),
1500            block_dim: (256, 1, 1),
1501            shared_mem_bytes: 0,
1502        };
1503        let __s_b = self.gpu.stream();
1504        let mut b = __s_b.launch_builder(&f);
1505        b.arg(x)
1506            .arg(&rs)
1507            .arg(ids)
1508            .arg(rows)
1509            .arg(th)
1510            .arg(z)
1511            .arg(&mut *out)
1512            .arg(&ni)
1513            .arg(&np)
1514            .arg(&temp);
1515        unsafe {
1516            b.launch(cfg)?;
1517        }
1518        Ok(())
1519    }
1520
1521    /// Filtered residual sample: token ~ norm(max(0, fp - fq)) with fp/fq the filtered softmaxes.
1522    #[allow(clippy::too_many_arguments)]
1523    pub fn residual_sample_filtered(
1524        &self,
1525        p: &CudaSlice<f32>,
1526        q: Option<&CudaSlice<f32>>,
1527        n: usize,
1528        temp: f32,
1529        seed: u64,
1530        stream_pos: u32,
1531        p_stats: (f32, f32, f32),
1532        q_stats: (f32, f32, f32),
1533        out_tok: &mut CudaSlice<u32>,
1534    ) -> Result<(), Box<dyn std::error::Error>> {
1535        let f = self.func("residual_sample_filtered_f32");
1536        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
1537        let has_q: i32 = q.is_some() as i32;
1538        let qbuf = q.unwrap_or(p);
1539        let (pm, pth, pz) = p_stats;
1540        let (qm, qth, qz) = q_stats;
1541        let cfg = LaunchConfig {
1542            grid_dim: (1, 1, 1),
1543            block_dim: (1024, 1, 1),
1544            shared_mem_bytes: 0,
1545        };
1546        let __s_b = self.gpu.stream();
1547        let mut b = __s_b.launch_builder(&f);
1548        b.arg(p)
1549            .arg(qbuf)
1550            .arg(&has_q)
1551            .arg(&ni)
1552            .arg(&temp)
1553            .arg(&slo)
1554            .arg(&shi)
1555            .arg(&stream_pos)
1556            .arg(&pm)
1557            .arg(&pth)
1558            .arg(&pz)
1559            .arg(&qm)
1560            .arg(&qth)
1561            .arg(&qz)
1562            .arg(&mut *out_tok);
1563        unsafe {
1564            b.launch(cfg)?;
1565        }
1566        Ok(())
1567    }
1568
1569    /// Gumbel-max draw from the FILTERED distribution (masked perturb; argmax after).
1570    #[allow(clippy::too_many_arguments)]
1571    pub fn gumbel_perturb_filtered(
1572        &self,
1573        x: &CudaSlice<f32>,
1574        y: &mut CudaSlice<f32>,
1575        n: usize,
1576        seed: u64,
1577        stream_pos: u32,
1578        temp: f32,
1579        row_max: f32,
1580        th: f32,
1581    ) -> Result<(), Box<dyn std::error::Error>> {
1582        let f = self.func("gumbel_perturb_filtered_f32");
1583        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
1584        let cfg = LaunchConfig {
1585            grid_dim: (n.div_ceil(256) as u32, 1, 1),
1586            block_dim: (256, 1, 1),
1587            shared_mem_bytes: 0,
1588        };
1589        let __s_b = self.gpu.stream();
1590        let mut b = __s_b.launch_builder(&f);
1591        b.arg(x)
1592            .arg(&mut *y)
1593            .arg(&ni)
1594            .arg(&slo)
1595            .arg(&shi)
1596            .arg(&stream_pos)
1597            .arg(&temp)
1598            .arg(&row_max)
1599            .arg(&th);
1600        unsafe {
1601            b.launch(cfg)?;
1602        }
1603        Ok(())
1604    }
1605
1606    /// Keskar penalties applied IN PLACE to a logits buffer: history token ids get
1607    /// rep-divided/multiplied + freq*count + presence subtracted. Symmetric p/q usage keeps
1608    /// filtered rejection sampling exact for the penalized target.
1609    #[allow(clippy::too_many_arguments)]
1610    pub fn penalize_logits(
1611        &self,
1612        x: &mut CudaSlice<f32>,
1613        hist: &CudaSlice<u32>,
1614        n_hist: usize,
1615        rep: f32,
1616        freq: f32,
1617        present: f32,
1618        n: usize,
1619    ) -> Result<(), Box<dyn std::error::Error>> {
1620        if n_hist == 0 {
1621            return Ok(());
1622        }
1623        let f = self.func("penalize_logits_f32");
1624        let (nh, ni) = (n_hist as i32, n as i32);
1625        let cfg = LaunchConfig {
1626            grid_dim: (n_hist.div_ceil(128) as u32, 1, 1),
1627            block_dim: (128, 1, 1),
1628            shared_mem_bytes: 0,
1629        };
1630        let __s_b = self.gpu.stream();
1631        let mut b = __s_b.launch_builder(&f);
1632        b.arg(&mut *x)
1633            .arg(hist)
1634            .arg(&nh)
1635            .arg(&rep)
1636            .arg(&freq)
1637            .arg(&present)
1638            .arg(&ni);
1639        unsafe {
1640            b.launch(cfg)?;
1641        }
1642        Ok(())
1643    }
1644
1645    /// Rows variant: penalize `nrow` contiguous rows of length n in one launch.
1646    #[allow(clippy::too_many_arguments)]
1647    pub fn penalize_logits_rows(
1648        &self,
1649        x: &mut CudaSlice<f32>,
1650        hist: &CudaSlice<u32>,
1651        n_hist: usize,
1652        rep: f32,
1653        freq: f32,
1654        present: f32,
1655        n: usize,
1656        nrow: usize,
1657    ) -> Result<(), Box<dyn std::error::Error>> {
1658        if n_hist == 0 || nrow == 0 {
1659            return Ok(());
1660        }
1661        let f = self.func("penalize_logits_rows_f32");
1662        let (nh, ni, nr) = (n_hist as i32, n as i32, nrow as i32);
1663        let cfg = LaunchConfig {
1664            grid_dim: (n_hist.div_ceil(128) as u32, nrow as u32, 1),
1665            block_dim: (128, 1, 1),
1666            shared_mem_bytes: 0,
1667        };
1668        let __s_b = self.gpu.stream();
1669        let mut b = __s_b.launch_builder(&f);
1670        b.arg(&mut *x)
1671            .arg(hist)
1672            .arg(&nh)
1673            .arg(&rep)
1674            .arg(&freq)
1675            .arg(&present)
1676            .arg(&ni)
1677            .arg(&nr);
1678        unsafe {
1679            b.launch(cfg)?;
1680        }
1681        Ok(())
1682    }
1683
1684    /// WEIGHT PREFETCH (SOTA item 3, 2026-07-13, DEFAULT ON): during a bandwidth-idle
1685    /// window (the fa launch reads KV, not weights) prefetch the NEXT matvec's
1686    /// decode-plane bytes into L2 so it reads L2-warm. Value-free scheduling op — same
1687    /// class as prefetch_l2 (numerics untouched by construction). Wired only where it
1688    /// measured positive: the E4B dc attn arm (+0.65%). 26B (flat — MoE ffn dominates),
1689    /// 31B (−0.2% — decode at the DRAM wall) and the ffn gate/up cascade (−1% — 29MB/layer
1690    /// floods the fill path) all probed and NOT wired. MEMRA_WPF=0 rollback seam.
1691    pub fn wpf_level() -> u32 {
1692        static ON: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
1693        *ON.get_or_init(|| {
1694            std::env::var("MEMRA_WPF")
1695                .ok()
1696                .and_then(|v| v.parse().ok())
1697                .unwrap_or(1)
1698        })
1699    }
1700
1701    /// PDL launch arm (SOTA item 2, 2026-07-13, DEFAULT ON): the six MEMRA_PDL_ENTRY glue
1702    /// kernels launch through cuLaunchKernelEx with PROGRAMMATIC_STREAM_SERIALIZATION — the
1703    /// grid launches while the predecessor drains (~120ns/kernel back, pdl_probe), the
1704    /// kernels' entry grid-dep sync restores read order (SASS-audited: ACQBULK precedes
1705    /// every LDG in all six). Valid windows: E4B +1.0-1.2% (128 AND 384-tok gens);
1706    /// 26B/31B/qwen flat no-harm. Battery: kernel-check GREEN, run-gen tokens IDENTICAL x3
1707    /// gemma, spec 64/64 E4B K=1/4/8 + 26B/31B K=4 + qwen PASS. Works eager AND under
1708    /// capture (capture encodes native programmatic edges — the post-capture edge-REWRITE
1709    /// arm died: engine graphs hold cuMemAllocAsync alloc nodes, edge edits on those return
1710    /// CUDA_ERROR_NOT_SUPPORTED). MEMRA_PDL=0 rollback seam.
1711    /// See the `verify_exact` field. Scoped by the dflash round around its t=16 verify.
1712    pub fn set_verify_exact(&self, on: bool) {
1713        self.verify_exact
1714            .store(on, std::sync::atomic::Ordering::Relaxed);
1715    }
1716    pub(crate) fn verify_exact_on(&self) -> bool {
1717        self.verify_exact.load(std::sync::atomic::Ordering::Relaxed)
1718    }
1719
1720    /// m=1 norm+rope+append fold seam (2026-07-23): MEMRA_QKV_APPEND=0 reverts to the
1721    /// fused-norm-rope + standalone-append pair (the exact-oracle bisect arm).
1722    pub fn qkv_append_on() -> bool {
1723        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1724        *ON.get_or_init(|| {
1725            std::env::var("MEMRA_QKV_APPEND")
1726                .map(|v| v != "0")
1727                .unwrap_or(true)
1728        })
1729    }
1730
1731    /// PDL wave-B1a seam: the four dense-glue kernels (rms_norm_f32, add_rms_norm_f32,
1732    /// add_scale_rms_norm_q8_1, quantize_q8_1). MEMRA_PDL_WB=0 reverts alone.
1733    pub fn pdl_wb_on() -> bool {
1734        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1735        *ON.get_or_init(|| {
1736            std::env::var("MEMRA_PDL_WB")
1737                .map(|v| v != "0")
1738                .unwrap_or(true)
1739        })
1740    }
1741
1742    /// PDL wave-A seam: the mmvq matvec PDL launches only (the six glue kernels keep
1743    /// their own MEMRA_PDL master seam). MEMRA_PDL_MMVQ=0 reverts wave-A alone — the
1744    /// per-model no-harm bisect knob.
1745    pub fn pdl_mmvq_on() -> bool {
1746        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1747        *ON.get_or_init(|| {
1748            std::env::var("MEMRA_PDL_MMVQ")
1749                .map(|v| v != "0")
1750                .unwrap_or(true)
1751        })
1752    }
1753
1754    pub fn pdl_on() -> bool {
1755        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1756        *ON.get_or_init(|| std::env::var("MEMRA_PDL").map(|v| v != "0").unwrap_or(true))
1757    }
1758
1759    /// PDL wave-B seam (gap-diagnosis arc, GAP-DIAGNOSIS.md verdict 8): the gemma
1760    /// NVFP4mix decode chain's hot kernels — nvfp4 fused2/mr2 and the q8_0 `_rp`
1761    /// singles — join the wave-A launch class. Scheduling-only (the entry macro waits
1762    /// on the producer before any read), bit-identical by construction.
1763    /// MEMRA_PDL_NVFP4=0 reverts wave-B alone.
1764    pub fn pdl_nvfp4q8_on() -> bool {
1765        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1766        *ON.get_or_init(|| {
1767            std::env::var("MEMRA_PDL_NVFP4")
1768                .map(|v| v != "0")
1769                .unwrap_or(true)
1770        })
1771    }
1772
1773    /// Raw CUfunction for a PDL-attributed launch: the SAME kernels.fatbin loaded once more
1774    /// through the raw driver API (cudarc hides its CUfunction handles; a duplicate module
1775    /// of tiny glue kernels is free). Resolved lazily per name, cached process-wide.
1776    /// Fused t=1 q4_0 mr policy: env MEMRA_Q40_MR wins (1/2); else the per-model
1777    /// FUSED_MR1_DEFAULT (dense gemma = mr1, MoE = mr2 — see the static's doc).
1778    fn q40_mr1_on() -> bool {
1779        static Q40MR: std::sync::OnceLock<Option<u32>> = std::sync::OnceLock::new();
1780        match *Q40MR.get_or_init(|| {
1781            std::env::var("MEMRA_Q40_MR")
1782                .ok()
1783                .and_then(|v| v.parse().ok())
1784        }) {
1785            Some(v) => v == 1,
1786            None => crate::FUSED_MR1_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
1787        }
1788    }
1789
1790    /// PDL wave-B2: flash-module PDL functions. `g` selects the kf8vf8 flavor — the
1791    /// caller MUST pass the SAME flavor its builder launch would resolve (fa_func/func_g
1792    /// mirror); the flavors differ semantically (KV byte formats), a wrong-module launch
1793    /// writes wrong bytes silently.
1794    fn pdl_func_flash(
1795        &self,
1796        g: bool,
1797        name: &'static str,
1798    ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
1799        use cudarc::driver::sys as cu;
1800        // PER-CONTEXT caches (M1-PP2 cross-device fix, 8x box 2026-08-02): CUmodule and
1801        // CUfunction handles are CONTEXT-scoped, and a remote-stage Engine
1802        // (MEMRA_PP_DEVICES=a,b) lives in the other device's primary context. The old
1803        // process-wide OnceLock cache handed stage 1 the dev-a handles, so every stage-1
1804        // launch_pdl* died CUDA_ERROR_INVALID_HANDLE. Key module + function caches by
1805        // this engine's CUcontext; single-context runs behave exactly as before.
1806        static MODS: std::sync::Mutex<Option<std::collections::HashMap<(usize, bool), usize>>> =
1807            std::sync::Mutex::new(None);
1808        static FNS: std::sync::Mutex<
1809            Option<std::collections::HashMap<(usize, bool, &'static str), usize>>,
1810        > = std::sync::Mutex::new(None);
1811        let ctx_key = self.ctx().cu_ctx() as usize;
1812        if let Some(&f) = FNS
1813            .lock()
1814            .unwrap()
1815            .get_or_insert_with(Default::default)
1816            .get(&(ctx_key, g, name))
1817        {
1818            return Ok(f as cu::CUfunction);
1819        }
1820        let module = {
1821            let mut mods = MODS.lock().unwrap();
1822            let map = mods.get_or_insert_with(Default::default);
1823            match map.get(&(ctx_key, g)) {
1824                Some(&m) => m,
1825                None => {
1826                    let m = self.pdl_load_module_in_ctx(if g {
1827                        FLASH_FATBIN_KF8VF8
1828                    } else {
1829                        FLASH_FATBIN
1830                    })?;
1831                    map.insert((ctx_key, g), m);
1832                    m
1833                }
1834            }
1835        };
1836        let cname = std::ffi::CString::new(name)?;
1837        let mut f: cu::CUfunction = std::ptr::null_mut();
1838        let r = unsafe { cu::cuModuleGetFunction(&mut f, module as cu::CUmodule, cname.as_ptr()) };
1839        if r != cu::CUresult::CUDA_SUCCESS {
1840            return Err(format!("pdl_func_flash {name} (g={g}): {r:?}").into());
1841        }
1842        FNS.lock()
1843            .unwrap()
1844            .get_or_insert_with(Default::default)
1845            .insert((ctx_key, g, name), f as usize);
1846        Ok(f)
1847    }
1848
1849    /// Load a fatbin as a raw CUmodule IN THIS ENGINE'S CONTEXT. `cuModuleLoadData` binds
1850    /// the module to the thread's CURRENT context — a remote-stage engine must not
1851    /// inherit the primary's (the INVALID_HANDLE class above). Restores the caller's
1852    /// current context before returning.
1853    fn pdl_load_module_in_ctx(&self, bytes: &[u8]) -> Result<usize, Box<dyn std::error::Error>> {
1854        use cudarc::driver::sys as cu;
1855        let mut prev: cu::CUcontext = std::ptr::null_mut();
1856        unsafe {
1857            cu::cuCtxGetCurrent(&mut prev).result()?;
1858        }
1859        self.ctx().bind_to_thread()?;
1860        let mut m: cu::CUmodule = std::ptr::null_mut();
1861        let r = unsafe { cu::cuModuleLoadData(&mut m, bytes.as_ptr() as *const std::ffi::c_void) };
1862        let restore = if prev.is_null() {
1863            cu::CUresult::CUDA_SUCCESS
1864        } else {
1865            unsafe { cu::cuCtxSetCurrent(prev) }
1866        };
1867        if r != cu::CUresult::CUDA_SUCCESS {
1868            return Err(format!("pdl module load: {r:?}").into());
1869        }
1870        if restore != cu::CUresult::CUDA_SUCCESS {
1871            return Err(format!("pdl module load: ctx restore {restore:?}").into());
1872        }
1873        Ok(m as usize)
1874    }
1875
1876    fn pdl_func(
1877        &self,
1878        name: &'static str,
1879    ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
1880        use cudarc::driver::sys as cu;
1881        // PER-CONTEXT caches — same M1-PP2 cross-device fix as pdl_func_flash (handles
1882        // are context-scoped; key everything by this engine's CUcontext).
1883        static MODULES: std::sync::Mutex<Option<std::collections::HashMap<usize, usize>>> =
1884            std::sync::Mutex::new(None);
1885        // PDL wave-A: the mmvq kernels live in the qmatvec fatbin, not kernels.cu — second
1886        // duplicate module, loaded lazily on the first kernels-module miss.
1887        static QMODULES: std::sync::Mutex<Option<std::collections::HashMap<usize, usize>>> =
1888            std::sync::Mutex::new(None);
1889        static FNS: std::sync::Mutex<
1890            Option<std::collections::HashMap<(usize, &'static str), usize>>,
1891        > = std::sync::Mutex::new(None);
1892        let ctx_key = self.ctx().cu_ctx() as usize;
1893        if let Some(&f) = FNS
1894            .lock()
1895            .unwrap()
1896            .get_or_insert_with(Default::default)
1897            .get(&(ctx_key, name))
1898        {
1899            return Ok(f as cu::CUfunction);
1900        }
1901        let module = {
1902            let mut mods = MODULES.lock().unwrap();
1903            let map = mods.get_or_insert_with(Default::default);
1904            match map.get(&ctx_key) {
1905                Some(&m) => m,
1906                None => {
1907                    let m = self.pdl_load_module_in_ctx(FATBIN)?;
1908                    map.insert(ctx_key, m);
1909                    m
1910                }
1911            }
1912        };
1913        let cname = std::ffi::CString::new(name)?;
1914        let mut f: cu::CUfunction = std::ptr::null_mut();
1915        let mut r =
1916            unsafe { cu::cuModuleGetFunction(&mut f, module as cu::CUmodule, cname.as_ptr()) };
1917        if r == cu::CUresult::CUDA_ERROR_NOT_FOUND {
1918            let qmodule = {
1919                let mut mods = QMODULES.lock().unwrap();
1920                let map = mods.get_or_insert_with(Default::default);
1921                match map.get(&ctx_key) {
1922                    Some(&m) => m,
1923                    None => {
1924                        let m = self.pdl_load_module_in_ctx(QMATVEC_FATBIN)?;
1925                        map.insert(ctx_key, m);
1926                        m
1927                    }
1928                }
1929            };
1930            r = unsafe { cu::cuModuleGetFunction(&mut f, qmodule as cu::CUmodule, cname.as_ptr()) };
1931        }
1932        if r != cu::CUresult::CUDA_SUCCESS {
1933            return Err(format!("pdl_func {name}: {r:?}").into());
1934        }
1935        FNS.lock()
1936            .unwrap()
1937            .get_or_insert_with(Default::default)
1938            .insert((ctx_key, name), f as usize);
1939        Ok(f)
1940    }
1941
1942    /// cuLaunchKernelEx with CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION on the
1943    /// compute stream. ONLY legal for kernels whose entry carries MEMRA_PDL_ENTRY.
1944    ///
1945    /// # Safety
1946    /// `params` must match the kernel's exact parameter list (order, types, count) —
1947    /// a mismatch corrupts the launch silently.
1948    /// Flash-module twin of `launch_pdl` — `g` picks the kf8vf8 flavor (must mirror the
1949    /// builder path's fa_func/func_g choice exactly).
1950    ///
1951    /// # Safety
1952    /// Same contract as `launch_pdl`.
1953    unsafe fn launch_pdl_flash(
1954        &self,
1955        g: bool,
1956        name: &'static str,
1957        grid: (u32, u32, u32),
1958        block: (u32, u32, u32),
1959        smem: u32,
1960        params: &mut [*mut std::ffi::c_void],
1961    ) -> Result<(), Box<dyn std::error::Error>> {
1962        use cudarc::driver::sys as cu;
1963        let f = self.pdl_func_flash(g, name)?;
1964        if smem > 0 {
1965            // mirror the builder path's opt-in ceiling (idempotent host-side set).
1966            let r =
1967                unsafe {
1968                    cu::cuFuncSetAttribute(f,
1969                cu::CUfunction_attribute_enum::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
1970                smem as i32)
1971                };
1972            if r != cu::CUresult::CUDA_SUCCESS {
1973                return Err(format!("pdl smem attr {name}: {r:?}").into());
1974            }
1975        }
1976        let mut attr = cu::CUlaunchAttribute {
1977            id: cu::CUlaunchAttributeID::CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION,
1978            pad: [0; 4],
1979            value: cu::CUlaunchAttributeValue {
1980                programmaticStreamSerializationAllowed: 1,
1981            },
1982        };
1983        let cfg = cu::CUlaunchConfig {
1984            gridDimX: grid.0,
1985            gridDimY: grid.1,
1986            gridDimZ: grid.2,
1987            blockDimX: block.0,
1988            blockDimY: block.1,
1989            blockDimZ: block.2,
1990            sharedMemBytes: smem,
1991            hStream: self.gpu.stream().cu_stream(),
1992            attrs: &mut attr,
1993            numAttrs: 1,
1994        };
1995        let r = unsafe { cu::cuLaunchKernelEx(&cfg, f, params.as_mut_ptr(), std::ptr::null_mut()) };
1996        if r != cu::CUresult::CUDA_SUCCESS {
1997            return Err(format!("launch_pdl_flash {name}: {r:?}").into());
1998        }
1999        Ok(())
2000    }
2001
2002    unsafe fn launch_pdl(
2003        &self,
2004        name: &'static str,
2005        grid: (u32, u32, u32),
2006        block: (u32, u32, u32),
2007        params: &mut [*mut std::ffi::c_void],
2008    ) -> Result<(), Box<dyn std::error::Error>> {
2009        use cudarc::driver::sys as cu;
2010        let f = self.pdl_func(name)?;
2011        let mut attr = cu::CUlaunchAttribute {
2012            id: cu::CUlaunchAttributeID::CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION,
2013            pad: [0; 4],
2014            value: cu::CUlaunchAttributeValue {
2015                programmaticStreamSerializationAllowed: 1,
2016            },
2017        };
2018        let cfg = cu::CUlaunchConfig {
2019            gridDimX: grid.0,
2020            gridDimY: grid.1,
2021            gridDimZ: grid.2,
2022            blockDimX: block.0,
2023            blockDimY: block.1,
2024            blockDimZ: block.2,
2025            sharedMemBytes: 0,
2026            hStream: self.gpu.stream().cu_stream(),
2027            attrs: &mut attr,
2028            numAttrs: 1,
2029        };
2030        let r = unsafe { cu::cuLaunchKernelEx(&cfg, f, params.as_mut_ptr(), std::ptr::null_mut()) };
2031        if r != cu::CUresult::CUDA_SUCCESS {
2032            return Err(format!("launch_pdl {name}: {r:?}").into());
2033        }
2034        Ok(())
2035    }
2036
2037    /// L2-prefetch a quant weight's DECODE plane (the rp4 split-plane mirror when present —
2038    /// that is what the m<=8 dispatch reads — else the raw block bytes). No-op on float arms.
2039    pub fn prefetch_weight_l2(
2040        &self,
2041        w: &crate::model::GpuTensor,
2042    ) -> Result<(), Box<dyn std::error::Error>> {
2043        if let crate::model::GpuTensor::Quant { bytes, rp4, .. } = w {
2044            let p = rp4.as_ref().unwrap_or(bytes);
2045            self.prefetch_l2(p, p.len())?;
2046        }
2047        Ok(())
2048    }
2049
2050    /// DSpark markov chain ops (dflash lane): gather one bf16 row of a [V, rank] table
2051    /// by the DEVICE token id at tok[idx] into f32.
2052    pub fn gather_row_bf16(
2053        &self,
2054        table: &CudaSlice<u8>,
2055        tok: &CudaSlice<u32>,
2056        idx: usize,
2057        dst: &mut CudaSlice<f32>,
2058        ncols: usize,
2059    ) -> Result<(), Box<dyn std::error::Error>> {
2060        let f = self.func("gather_row_bf16_f32");
2061        let cfg = LaunchConfig {
2062            grid_dim: (ncols.div_ceil(256) as u32, 1, 1),
2063            block_dim: (256, 1, 1),
2064            shared_mem_bytes: 0,
2065        };
2066        let (nc, ix) = (ncols as i32, idx as i32);
2067        let __s_b = self.gpu.stream();
2068        let mut b = __s_b.launch_builder(&f);
2069        b.arg(table).arg(tok).arg(&ix).arg(dst).arg(&nc);
2070        unsafe {
2071            b.launch(cfg)?;
2072        }
2073        Ok(())
2074    }
2075
2076    /// logits[row_off .. row_off+n] += bias[0..n] (in place, one row).
2077    pub fn add_row_inplace(
2078        &self,
2079        logits: &mut CudaSlice<f32>,
2080        bias: &CudaSlice<f32>,
2081        n: usize,
2082        row_off: usize,
2083    ) -> Result<(), Box<dyn std::error::Error>> {
2084        let f = self.func("add_row_inplace_f32");
2085        let cfg = LaunchConfig {
2086            grid_dim: (n.div_ceil(256) as u32, 1, 1),
2087            block_dim: (256, 1, 1),
2088            shared_mem_bytes: 0,
2089        };
2090        let (ni, off) = (n as i32, row_off as i64);
2091        let __s_b = self.gpu.stream();
2092        let mut b = __s_b.launch_builder(&f);
2093        b.arg(logits).arg(bias).arg(&ni).arg(&off);
2094        unsafe {
2095            b.launch(cfg)?;
2096        }
2097        Ok(())
2098    }
2099
2100    /// L2 prefetch of a device byte range (latency-hiding arc; value-free scheduling op).
2101    pub fn prefetch_l2(
2102        &self,
2103        p: &CudaSlice<u8>,
2104        n: usize,
2105    ) -> Result<(), Box<dyn std::error::Error>> {
2106        let f = self.func("prefetch_l2_bytes");
2107        let lines = n.div_ceil(128);
2108        let ni = n as i64;
2109        let cfg = LaunchConfig {
2110            grid_dim: (lines.div_ceil(256) as u32, 1, 1),
2111            block_dim: (256, 1, 1),
2112            shared_mem_bytes: 0,
2113        };
2114        let __s_b = self.gpu.stream();
2115        let mut b = __s_b.launch_builder(&f);
2116        b.arg(p).arg(&ni);
2117        unsafe {
2118            b.launch(cfg)?;
2119        }
2120        Ok(())
2121    }
2122
2123    /// MoE router GEMV (MEMRA_ROUTER_KERNEL): deterministic warp-per-(expert,token) f32 dot.
2124    /// Different FP order than the cuBLAS path it replaces — battery-gated numeric config.
2125    pub fn router_gemv(
2126        &self,
2127        w: &CudaSlice<f32>,
2128        x: &CudaSlice<f32>,
2129        n_embd: usize,
2130        n_experts: usize,
2131        t: usize,
2132    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2133        // float4 v2 probed 2026-07-14: +0.25% but flips near-tie routing (new FP order,
2134        // stream differs) — too small to justify a numeric config change; deleted.
2135        // w8 twin (2026-07-31): on the 132-SM H100 the lone-warp form is 14.8% of the q35
2136        // decode step (latency-bound) — the calculus flipped. MEMRA_ROUTER_V2=0 reverts to
2137        // the warp form (rollback seam; new FP order, battery-arbitrated per model).
2138        let w8 = match std::env::var("MEMRA_ROUTER_V2").as_deref() {
2139            Ok("0") => false,
2140            Ok(_) => true,
2141            Err(_) => ROUTER_W8_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
2142        };
2143        // FAST-ROUTER batch twin (lane/fast-router, 2026-08-02): at prefill m the per-(e,tok)
2144        // w8 form re-streams both operand rows per output (GEMV program at GEMM shape — the
2145        // concat-prime exactness fix paid -10% q35 board-2048 prefill through it). The batch
2146        // twin (8x8 expert-x-token register tile) is BIT-IDENTICAL per row (same k order,
2147        // same tree, same fold — kernel-check sweeps m=1..2048 on real router weights), so
2148        // the crossover is pure perf, not a numeric config. MIN_T from the on-box sweep
2149        // (research/fast-router-20260802/crossover-router*.jsonl); decode t=1 and small-t
2150        // spec verify keep the plain w8 form. MEMRA_ROUTER_BATCH=0: rollback seam
2151        // (perf-only, bits equal).
2152        let batch = w8 && t >= ROUTER_BATCH_MIN_T && router_batch_on();
2153        self.router_gemv_form(w, x, n_embd, n_experts, t, w8, batch)
2154    }
2155
2156    /// Form-explicit router GEMV launch (kernel-check bit-identity gate + crossover bench
2157    /// force both forms; `batch` requires `w8`).
2158    pub fn router_gemv_form(
2159        &self,
2160        w: &CudaSlice<f32>,
2161        x: &CudaSlice<f32>,
2162        n_embd: usize,
2163        n_experts: usize,
2164        t: usize,
2165        w8: bool,
2166        batch: bool,
2167    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2168        debug_assert!(!batch || w8, "batch twin exists for the w8 form only");
2169        let mut y = self.alloc_uninit::<f32>(t * n_experts)?;
2170        let f = if batch {
2171            self.func("router_gemv_f32_w8_batch")
2172        } else if w8 {
2173            self.func("router_gemv_f32_w8")
2174        } else {
2175            self.func("router_gemv_f32")
2176        };
2177        let (ne, nx, ti) = (n_embd as i32, n_experts as i32, t as i32);
2178        let cfg = if batch {
2179            LaunchConfig {
2180                grid_dim: (n_experts.div_ceil(8) as u32, t.div_ceil(8) as u32, 1),
2181                block_dim: (32, 8, 1),
2182                shared_mem_bytes: 0,
2183            }
2184        } else {
2185            LaunchConfig {
2186                grid_dim: (n_experts as u32, t as u32, 1),
2187                block_dim: (32, if w8 { 8 } else { 1 }, 1),
2188                shared_mem_bytes: 0,
2189            }
2190        };
2191        let __s_b = self.gpu.stream();
2192        let mut b = __s_b.launch_builder(&f);
2193        b.arg(w).arg(x).arg(&mut y).arg(&ne).arg(&nx).arg(&ti);
2194        unsafe {
2195            b.launch(cfg)?;
2196        }
2197        Ok(y)
2198    }
2199
2200    /// f32 row permute: dst[idx[i], :] = src[i, :] (grouped-GEMM CSR -> pair-id reorder).
2201    pub fn rows_permute(
2202        &self,
2203        src: &CudaSlice<f32>,
2204        idx: &CudaSlice<i32>,
2205        nrows: usize,
2206        ncols: usize,
2207    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2208        let mut dst = self.alloc_uninit::<f32>(nrows * ncols)?;
2209        let f = self.func("rows_permute_f32");
2210        let (nc, nr) = (ncols as i32, nrows as i32);
2211        let cfg = LaunchConfig {
2212            grid_dim: (nrows as u32, 1, 1),
2213            block_dim: (256, 1, 1),
2214            shared_mem_bytes: 0,
2215        };
2216        let __s_b = self.gpu.stream();
2217        let mut b = __s_b.launch_builder(&f);
2218        b.arg(src).arg(idx).arg(&mut dst).arg(&nc).arg(&nr);
2219        unsafe {
2220            b.launch(cfg)?;
2221        }
2222        Ok(dst)
2223    }
2224
2225    /// shexp gate fused dot: g[tok] = sigmoid(dot(x[tok,:], w)) — replaces the per-layer
2226    /// cuBLASLt m=1 GEMM + separate sigmoid launch on the qwen35moe decode path (the
2227    /// splitKreduce x40/step dig, 2026-07-31). One fold order for every t, so the t=1
2228    /// decode chain and the small-t spec-verify chain match per row by construction.
2229    pub fn sigmoid_dot_rows(
2230        &self,
2231        x: &CudaSlice<f32>,
2232        w: &CudaSlice<f32>,
2233        n_embd: usize,
2234        t: usize,
2235    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2236        // MEMRA_SHEXP_DOT=0: rollback seam to the cuBLASLt linear + sigmoid pair (numeric
2237        // config; same class as MEMRA_ROUTER_V2).
2238        static OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2239        if *OFF.get_or_init(|| std::env::var("MEMRA_SHEXP_DOT").as_deref() == Ok("0")) {
2240            let gs = self.linear(x, w, t, n_embd, 1)?;
2241            let mut g = self.uninit(t)?;
2242            self.sigmoid(&gs, &mut g, t)?;
2243            return Ok(g);
2244        }
2245        // FAST-ROUTER lane note (2026-08-02): a register-tiled 8-token batch twin of this
2246        // kernel was built, proven bit-identical, and measured SLOWER at every prefill t on
2247        // the 5090 (0.62-0.89x — launch-latency-bound op, ~7us/layer at m=2048;
2248        // research/fast-router-20260802/crossover-router.jsonl). Dispatch arm killed per
2249        // flags doctrine; this per-token form serves every t.
2250        let mut g = self.alloc_uninit::<f32>(t)?;
2251        let f = self.func("sigmoid_dot_rows_f32");
2252        let (ne, ti) = (n_embd as i32, t as i32);
2253        let cfg = LaunchConfig {
2254            grid_dim: (t as u32, 1, 1),
2255            block_dim: (32, 8, 1),
2256            shared_mem_bytes: 0,
2257        };
2258        let __s_b = self.gpu.stream();
2259        let mut b = __s_b.launch_builder(&f);
2260        b.arg(x).arg(w).arg(&mut g).arg(&ne).arg(&ti);
2261        unsafe {
2262            b.launch(cfg)?;
2263        }
2264        Ok(g)
2265    }
2266
2267    /// ROUND-STREAM stream rollback: all counters <- pos_start + base + n_acc.
2268    pub fn spec_rollback_stream(
2269        &self,
2270        len_ptrs: &CudaSlice<u64>,
2271        pos_start: &CudaSlice<i32>,
2272        acc: &CudaSlice<u32>,
2273        base: usize,
2274        n_rows: usize,
2275    ) -> Result<(), Box<dyn std::error::Error>> {
2276        let f = self.func("spec_rollback_stream");
2277        let (b, nr) = (base as i32, n_rows as i32);
2278        let cfg = LaunchConfig {
2279            grid_dim: (n_rows.div_ceil(64) as u32, 1, 1),
2280            block_dim: (64, 1, 1),
2281            shared_mem_bytes: 0,
2282        };
2283        let __s_bl = self.gpu.stream();
2284        let mut bl = __s_bl.launch_builder(&f);
2285        bl.arg(len_ptrs).arg(pos_start).arg(acc).arg(&b).arg(&nr);
2286        unsafe {
2287            bl.launch(cfg)?;
2288        }
2289        Ok(())
2290    }
2291
2292    /// PLAIN-DECODE GRAPH ring store: ring[(pos_start - base) % cap] = vam[0].
2293    pub fn plain_tok_ring(
2294        &self,
2295        vam: &CudaSlice<u32>,
2296        pos_start: &CudaSlice<i32>,
2297        base: usize,
2298        ring: &mut CudaSlice<u32>,
2299    ) -> Result<(), Box<dyn std::error::Error>> {
2300        let f = self.func("plain_tok_ring");
2301        let (b, cap) = (base as i32, ring.len() as i32);
2302        let cfg = LaunchConfig {
2303            grid_dim: (1, 1, 1),
2304            block_dim: (32, 1, 1),
2305            shared_mem_bytes: 0,
2306        };
2307        let __s_bl = self.gpu.stream();
2308        let mut bl = __s_bl.launch_builder(&f);
2309        bl.arg(vam).arg(pos_start).arg(&b).arg(&mut *ring).arg(&cap);
2310        unsafe {
2311            bl.launch(cfg)?;
2312        }
2313        Ok(())
2314    }
2315
2316    /// ROUND-STREAM stage (c) 4 epilogue: ring commit + tiny counter copies.
2317    pub fn spec_ring_commit(
2318        &self,
2319        vtok: &CudaSlice<u32>,
2320        acc: &CudaSlice<u32>,
2321        brk: &CudaSlice<u32>,
2322        ring: &mut CudaSlice<u32>,
2323        pend: &mut CudaSlice<u32>,
2324    ) -> Result<(), Box<dyn std::error::Error>> {
2325        let f = self.func("spec_ring_commit");
2326        let cfg = LaunchConfig {
2327            grid_dim: (1, 1, 1),
2328            block_dim: (32, 1, 1),
2329            shared_mem_bytes: 0,
2330        };
2331        let __s_b = self.gpu.stream();
2332        let mut b = __s_b.launch_builder(&f);
2333        b.arg(vtok).arg(acc).arg(brk).arg(ring).arg(pend);
2334        unsafe {
2335            b.launch(cfg)?;
2336        }
2337        Ok(())
2338    }
2339    pub fn i32_copy_add(
2340        &self,
2341        src: &CudaSlice<i32>,
2342        dst: &mut CudaSlice<i32>,
2343        delta: i32,
2344    ) -> Result<(), Box<dyn std::error::Error>> {
2345        let f = self.func("i32_copy_add");
2346        let cfg = LaunchConfig {
2347            grid_dim: (1, 1, 1),
2348            block_dim: (32, 1, 1),
2349            shared_mem_bytes: 0,
2350        };
2351        let __s_b = self.gpu.stream();
2352        let mut b = __s_b.launch_builder(&f);
2353        b.arg(src).arg(dst).arg(&delta);
2354        unsafe {
2355            b.launch(cfg)?;
2356        }
2357        Ok(())
2358    }
2359    pub fn u32_copy(
2360        &self,
2361        src: &CudaSlice<u32>,
2362        dst: &mut CudaSlice<u32>,
2363    ) -> Result<(), Box<dyn std::error::Error>> {
2364        let f = self.func("u32_copy");
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(src).arg(dst);
2373        unsafe {
2374            b.launch(cfg)?;
2375        }
2376        Ok(())
2377    }
2378
2379    /// ROUND-GRAPH adaptive depth: brk[0] <- clamp(acc[0] + 1, floor, cap) — the host
2380    /// adaptive policy as a captured device op (policy-identical: the accept walk depth
2381    /// caps acceptance exactly like drafting fewer tokens).
2382    pub fn spec_adapt_k(
2383        &self,
2384        acc: &CudaSlice<u32>,
2385        brk: &mut CudaSlice<u32>,
2386        floor: usize,
2387        cap: usize,
2388    ) -> Result<(), Box<dyn std::error::Error>> {
2389        let f = self.func("spec_adapt_k");
2390        let (fl, cp) = (floor as i32, cap as i32);
2391        let cfg = LaunchConfig {
2392            grid_dim: (1, 1, 1),
2393            block_dim: (32, 1, 1),
2394            shared_mem_bytes: 0,
2395        };
2396        let __s_b = self.gpu.stream();
2397        let mut b = __s_b.launch_builder(&f);
2398        b.arg(acc).arg(brk).arg(&fl).arg(&cp);
2399        unsafe {
2400            b.launch(cfg)?;
2401        }
2402        Ok(())
2403    }
2404
2405    /// ROUND-STREAM stage (c) 3: accept walk fully device-driven (brk + assembled vtok).
2406    pub fn spec_accept_greedy_dc(
2407        &self,
2408        preds: &CudaSlice<u32>,
2409        vtok: &CudaSlice<u32>,
2410        last_pred: &CudaSlice<u32>,
2411        brk: &CudaSlice<u32>,
2412        out: &mut CudaSlice<u32>,
2413    ) -> Result<(), Box<dyn std::error::Error>> {
2414        let f = self.func("spec_accept_greedy_dc");
2415        let cfg = LaunchConfig {
2416            grid_dim: (1, 1, 1),
2417            block_dim: (32, 1, 1),
2418            shared_mem_bytes: 0,
2419        };
2420        let __s_b = self.gpu.stream();
2421        let mut b = __s_b.launch_builder(&f);
2422        b.arg(preds).arg(vtok).arg(last_pred).arg(brk).arg(out);
2423        unsafe {
2424            b.launch(cfg)?;
2425        }
2426        Ok(())
2427    }
2428
2429    /// ROUND-STREAM stage (c) 2: verify-chain device-pos entries.
2430    pub fn pos_iota(
2431        &self,
2432        pos0: &CudaSlice<i32>,
2433        out: &mut CudaSlice<i32>,
2434        t: usize,
2435    ) -> Result<(), Box<dyn std::error::Error>> {
2436        let f = self.func("pos_iota_i32");
2437        let ti = t as i32;
2438        let cfg = LaunchConfig {
2439            grid_dim: (1, 1, 1),
2440            block_dim: (t.max(1) as u32, 1, 1),
2441            shared_mem_bytes: 0,
2442        };
2443        let __s_b = self.gpu.stream();
2444        let mut b = __s_b.launch_builder(&f);
2445        b.arg(pos0).arg(out).arg(&ti);
2446        unsafe {
2447            b.launch(cfg)?;
2448        }
2449        Ok(())
2450    }
2451    #[allow(clippy::too_many_arguments)]
2452    pub fn append_kv_quantized_rows_dc(
2453        &self,
2454        k_rows: &CudaSlice<f32>,
2455        v_rows: &CudaSlice<f32>,
2456        kc: &mut CudaSlice<u8>,
2457        vc: &mut CudaSlice<u8>,
2458        t0_dev: &CudaSlice<i32>,
2459        t: usize,
2460        kv_dim_k: usize,
2461        kv_dim_v: usize,
2462        k_tok_bytes: usize,
2463        v_tok_bytes: usize,
2464        g: bool,
2465    ) -> Result<(), Box<dyn std::error::Error>> {
2466        let f = if g {
2467            self.func_g("append_quantize_kv_q8_0_q5_1_rows_dc")
2468        } else {
2469            self.func("append_quantize_kv_q8_0_q5_1_rows_dc")
2470        };
2471        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
2472        let cfg = LaunchConfig {
2473            grid_dim: (nblk, t as u32, 1),
2474            block_dim: (32, 1, 1),
2475            shared_mem_bytes: 0,
2476        };
2477        let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
2478        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
2479        let __s_b = self.gpu.stream();
2480        let mut b = __s_b.launch_builder(&f);
2481        b.arg(k_rows)
2482            .arg(v_rows)
2483            .arg(kc)
2484            .arg(vc)
2485            .arg(t0_dev)
2486            .arg(&kdk)
2487            .arg(&kdv)
2488            .arg(&ktb)
2489            .arg(&vtb);
2490        unsafe {
2491            b.launch(cfg)?;
2492        }
2493        Ok(())
2494    }
2495
2496    /// t=1 dc append with a FUSED len_d increment (wave 5c) — one launch replaces
2497    /// append_rows_dc + inc_seqlen. Single block (read-before-inc ordering).
2498    #[allow(clippy::too_many_arguments)]
2499    pub fn append_kv_quantized_row_dc_inc(
2500        &self,
2501        k_row: &CudaSlice<f32>,
2502        v_row: &CudaSlice<f32>,
2503        kc: &mut CudaSlice<u8>,
2504        vc: &mut CudaSlice<u8>,
2505        t0_dev: &mut CudaSlice<i32>,
2506        kv_dim_k: usize,
2507        kv_dim_v: usize,
2508        k_tok_bytes: usize,
2509        v_tok_bytes: usize,
2510        g: bool,
2511    ) -> Result<(), Box<dyn std::error::Error>> {
2512        let f = if g {
2513            self.func_g("append_quantize_kv_q8_0_q5_1_dc_inc")
2514        } else {
2515            self.func("append_quantize_kv_q8_0_q5_1_dc_inc")
2516        };
2517        let nthreads = ((kv_dim_k.max(kv_dim_v) / 32) * 32).min(1024) as u32;
2518        let cfg = LaunchConfig {
2519            grid_dim: (1, 1, 1),
2520            block_dim: (nthreads, 1, 1),
2521            shared_mem_bytes: 0,
2522        };
2523        let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
2524        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
2525        let __s_b = self.gpu.stream();
2526        let mut b = __s_b.launch_builder(&f);
2527        b.arg(k_row)
2528            .arg(v_row)
2529            .arg(kc)
2530            .arg(vc)
2531            .arg(t0_dev)
2532            .arg(&kdk)
2533            .arg(&kdv)
2534            .arg(&ktb)
2535            .arg(&vtb);
2536        unsafe {
2537            b.launch(cfg)?;
2538        }
2539        Ok(())
2540    }
2541
2542    /// ROUND-STREAM: draft-chain pack + in-graph d2t remap (see kernels.cu headers).
2543    pub fn pack_tok_p(
2544        &self,
2545        tok: &CudaSlice<u32>,
2546        p: &CudaSlice<f32>,
2547        out: &mut CudaSlice<u32>,
2548        slot: usize,
2549    ) -> Result<(), Box<dyn std::error::Error>> {
2550        let f = self.func("pack_tok_p");
2551        let sl = slot as i32;
2552        let cfg = LaunchConfig {
2553            grid_dim: (1, 1, 1),
2554            block_dim: (32, 1, 1),
2555            shared_mem_bytes: 0,
2556        };
2557        let __s_b = self.gpu.stream();
2558        let mut b = __s_b.launch_builder(&f);
2559        b.arg(tok).arg(p).arg(out).arg(&sl);
2560        unsafe {
2561            b.launch(cfg)?;
2562        }
2563        Ok(())
2564    }
2565    pub fn tok_map_u32(
2566        &self,
2567        tok: &mut CudaSlice<u32>,
2568        map: &CudaSlice<u32>,
2569    ) -> Result<(), Box<dyn std::error::Error>> {
2570        let f = self.func("tok_map_u32");
2571        let cfg = LaunchConfig {
2572            grid_dim: (1, 1, 1),
2573            block_dim: (32, 1, 1),
2574            shared_mem_bytes: 0,
2575        };
2576        let __s_b = self.gpu.stream();
2577        let mut b = __s_b.launch_builder(&f);
2578        b.arg(tok).arg(map);
2579        unsafe {
2580            b.launch(cfg)?;
2581        }
2582        Ok(())
2583    }
2584
2585    /// ROUND-STREAM stage (c) 1: device verify-token assembly + p-min break derivation.
2586    #[allow(clippy::too_many_arguments)]
2587    pub fn spec_assemble_verify(
2588        &self,
2589        tokp: &CudaSlice<u32>,
2590        pend: &CudaSlice<u32>,
2591        d2t: Option<&CudaSlice<u32>>,
2592        vtok: &mut CudaSlice<u32>,
2593        brk: &mut CudaSlice<u32>,
2594        p_min: f32,
2595        k: usize,
2596        pmin0: bool,
2597    ) -> Result<(), Box<dyn std::error::Error>> {
2598        let f = self.func("spec_assemble_verify");
2599        let (ki, pm) = (k as i32, if pmin0 { 1i32 } else { 0i32 });
2600        let cfg = LaunchConfig {
2601            grid_dim: (1, 1, 1),
2602            block_dim: (32, 1, 1),
2603            shared_mem_bytes: 0,
2604        };
2605        let __s_b = self.gpu.stream();
2606        let mut b = __s_b.launch_builder(&f);
2607        match d2t {
2608            Some(m) => {
2609                b.arg(tokp)
2610                    .arg(pend)
2611                    .arg(m)
2612                    .arg(vtok)
2613                    .arg(brk)
2614                    .arg(&p_min)
2615                    .arg(&ki)
2616                    .arg(&pm);
2617                unsafe {
2618                    b.launch(cfg)?;
2619                }
2620            }
2621            None => {
2622                let null: u64 = 0;
2623                b.arg(tokp)
2624                    .arg(pend)
2625                    .arg(&null)
2626                    .arg(vtok)
2627                    .arg(brk)
2628                    .arg(&p_min)
2629                    .arg(&ki)
2630                    .arg(&pm);
2631                unsafe {
2632                    b.launch(cfg)?;
2633                }
2634            }
2635        }
2636        Ok(())
2637    }
2638
2639    /// ROUND-STREAM stage (b) 3b: recur-restore twins with device-j (see hybrid.cu headers).
2640    #[allow(clippy::too_many_arguments)]
2641    pub fn ssm_conv_ring_rebuild_dc(
2642        &self,
2643        qkv_tm: &CudaSlice<f32>,
2644        ring_old: &CudaSlice<f32>,
2645        conv_state: &mut CudaSlice<f32>,
2646        conv_dim: usize,
2647        acc: &CudaSlice<u32>,
2648        base: usize,
2649        t_v: usize,
2650        d_conv: usize,
2651    ) -> Result<(), Box<dyn std::error::Error>> {
2652        let f = self.func("ssm_conv_ring_rebuild_f32_dc");
2653        let n = conv_dim * (d_conv - 1);
2654        let cfg = LaunchConfig::for_num_elems(n as u32);
2655        let (cd, b0, tv, dc) = (conv_dim as i32, base as i32, t_v as i32, d_conv as i32);
2656        let __s_b = self.gpu.stream();
2657        let mut b = __s_b.launch_builder(&f);
2658        b.arg(qkv_tm)
2659            .arg(ring_old)
2660            .arg(conv_state)
2661            .arg(&cd)
2662            .arg(acc)
2663            .arg(&b0)
2664            .arg(&tv)
2665            .arg(&dc);
2666        unsafe {
2667            b.launch(cfg)?;
2668        }
2669        Ok(())
2670    }
2671    #[allow(clippy::too_many_arguments)]
2672    pub fn gdn_scan_s128_dc(
2673        &self,
2674        q: &CudaSlice<f32>,
2675        k: &CudaSlice<f32>,
2676        v: &CudaSlice<f32>,
2677        g: &CudaSlice<f32>,
2678        beta: &CudaSlice<f32>,
2679        state_in: &CudaSlice<f32>,
2680        state_out: &mut CudaSlice<f32>,
2681        o: &mut CudaSlice<f32>,
2682        n_head: usize,
2683        acc: &CudaSlice<u32>,
2684        base: usize,
2685        t_v: usize,
2686        scale: f32,
2687    ) -> Result<(), Box<dyn std::error::Error>> {
2688        let f = self.func("gdn_scan_s128_dc");
2689        const S_V: u32 = 128;
2690        const WARP: u32 = 32;
2691        const COLS_PER_BLOCK: u32 = 4;
2692        let cfg = LaunchConfig {
2693            grid_dim: (n_head as u32, 1, S_V / COLS_PER_BLOCK),
2694            block_dim: (WARP, COLS_PER_BLOCK, 1),
2695            shared_mem_bytes: 0,
2696        };
2697        let (h, b0, tv) = (n_head as i32, base as i32, t_v as i32);
2698        let __s_b = self.gpu.stream();
2699        let mut b = __s_b.launch_builder(&f);
2700        b.arg(q)
2701            .arg(k)
2702            .arg(v)
2703            .arg(g)
2704            .arg(beta)
2705            .arg(state_in)
2706            .arg(state_out)
2707            .arg(o)
2708            .arg(&h)
2709            .arg(acc)
2710            .arg(&b0)
2711            .arg(&tv)
2712            .arg(&scale);
2713        unsafe {
2714            b.launch(cfg)?;
2715        }
2716        Ok(())
2717    }
2718
2719    /// ROUND-STREAM stage (b) 3a: device per-layer KV-len rollback (see spec_rollback_kv).
2720    pub fn spec_rollback_kv(
2721        &self,
2722        len_ptrs: &CudaSlice<u64>,
2723        saved: &CudaSlice<i32>,
2724        acc: &CudaSlice<u32>,
2725        base: usize,
2726        n_layer: usize,
2727    ) -> Result<(), Box<dyn std::error::Error>> {
2728        let f = self.func("spec_rollback_kv");
2729        let (b, nl) = (base as i32, n_layer as i32);
2730        let cfg = LaunchConfig {
2731            grid_dim: (n_layer.div_ceil(64) as u32, 1, 1),
2732            block_dim: (64, 1, 1),
2733            shared_mem_bytes: 0,
2734        };
2735        let __s_bl = self.gpu.stream();
2736        let mut bl = __s_bl.launch_builder(&f);
2737        bl.arg(len_ptrs).arg(saved).arg(acc).arg(&b).arg(&nl);
2738        unsafe {
2739            bl.launch(cfg)?;
2740        }
2741        Ok(())
2742    }
2743
2744    /// OPTIPIPE increment 1: derive the K=1 successor-valid bit on device.
2745    pub fn spec_fork_valid(
2746        &self,
2747        acc: &CudaSlice<u32>,
2748        optimistic_pending: u32,
2749        valid: &mut CudaSlice<u32>,
2750    ) -> Result<(), Box<dyn std::error::Error>> {
2751        let f = self.func("spec_fork_valid");
2752        let cfg = LaunchConfig {
2753            grid_dim: (1, 1, 1),
2754            block_dim: (1, 1, 1),
2755            shared_mem_bytes: 0,
2756        };
2757        let __s_bl = self.gpu.stream();
2758        let mut bl = __s_bl.launch_builder(&f);
2759        bl.arg(acc).arg(&optimistic_pending).arg(valid);
2760        unsafe {
2761            bl.launch(cfg)?;
2762        }
2763        Ok(())
2764    }
2765
2766    /// OPTIPIPE increment 1: leave stage-local KV lengths on hit, restore them on miss.
2767    pub fn spec_fork_reconcile_kv(
2768        &self,
2769        len_ptrs: &CudaSlice<u64>,
2770        saved: &CudaSlice<i32>,
2771        acc: &CudaSlice<u32>,
2772        valid: &CudaSlice<u32>,
2773        base: usize,
2774        n_layer: usize,
2775    ) -> Result<(), Box<dyn std::error::Error>> {
2776        let f = self.func("spec_fork_reconcile_kv");
2777        let (b, nl) = (base as i32, n_layer as i32);
2778        let cfg = LaunchConfig {
2779            grid_dim: (n_layer.div_ceil(64) as u32, 1, 1),
2780            block_dim: (64, 1, 1),
2781            shared_mem_bytes: 0,
2782        };
2783        let __s_bl = self.gpu.stream();
2784        let mut bl = __s_bl.launch_builder(&f);
2785        bl.arg(len_ptrs)
2786            .arg(saved)
2787            .arg(acc)
2788            .arg(valid)
2789            .arg(&b)
2790            .arg(&nl);
2791        unsafe {
2792            bl.launch(cfg)?;
2793        }
2794        Ok(())
2795    }
2796
2797    /// OPTIPIPE increment 1: conditionally restore one stage-owned recurrent-state buffer.
2798    pub fn spec_fork_restore_f32(
2799        &self,
2800        snapshot: &CudaSlice<f32>,
2801        state: &mut CudaSlice<f32>,
2802        valid: &CudaSlice<u32>,
2803    ) -> Result<(), Box<dyn std::error::Error>> {
2804        assert_eq!(
2805            snapshot.len(),
2806            state.len(),
2807            "fork recurrent snapshot shape mismatch"
2808        );
2809        let f = self.func("spec_fork_restore_f32");
2810        let n = state.len() as i32;
2811        let blocks = state.len().div_ceil(256).min(65535).max(1) as u32;
2812        let cfg = LaunchConfig {
2813            grid_dim: (blocks, 1, 1),
2814            block_dim: (256, 1, 1),
2815            shared_mem_bytes: 0,
2816        };
2817        let __s_bl = self.gpu.stream();
2818        let mut bl = __s_bl.launch_builder(&f);
2819        bl.arg(snapshot).arg(state).arg(valid).arg(&n);
2820        unsafe {
2821            bl.launch(cfg)?;
2822        }
2823        Ok(())
2824    }
2825
2826    /// ROUND-STREAM stage (b): device next-round seed gather (see spec_seed_gather header).
2827    /// Caller D2Ds h_seed into fill_prev after (both slots carry the same value in every arm).
2828    pub fn spec_seed_gather(
2829        &self,
2830        vx: &CudaSlice<f32>,
2831        fill_prev: &CudaSlice<f32>,
2832        acc: &CudaSlice<u32>,
2833        h_seed: &mut CudaSlice<f32>,
2834        base: usize,
2835        n_embd: usize,
2836    ) -> Result<(), Box<dyn std::error::Error>> {
2837        let f = self.func("spec_seed_gather");
2838        let (b, ne) = (base as i32, n_embd as i32);
2839        let cfg = LaunchConfig {
2840            grid_dim: (n_embd.div_ceil(256) as u32, 1, 1),
2841            block_dim: (256, 1, 1),
2842            shared_mem_bytes: 0,
2843        };
2844        let __s_bl = self.gpu.stream();
2845        let mut bl = __s_bl.launch_builder(&f);
2846        bl.arg(vx)
2847            .arg(fill_prev)
2848            .arg(acc)
2849            .arg(h_seed)
2850            .arg(&b)
2851            .arg(&ne);
2852        unsafe {
2853            bl.launch(cfg)?;
2854        }
2855        Ok(())
2856    }
2857
2858    /// ROUND-STREAM stage (a): device greedy accept walk (see spec_accept_greedy header).
2859    pub fn spec_accept_greedy(
2860        &self,
2861        preds: &CudaSlice<u32>,
2862        draft: &CudaSlice<u32>,
2863        last_pred: u32,
2864        base: usize,
2865        k_round: usize,
2866        out: &mut CudaSlice<u32>,
2867    ) -> Result<(), Box<dyn std::error::Error>> {
2868        let f = self.func("spec_accept_greedy");
2869        let (b, k) = (base as i32, k_round as i32);
2870        let cfg = LaunchConfig {
2871            grid_dim: (1, 1, 1),
2872            block_dim: (32, 1, 1),
2873            shared_mem_bytes: 0,
2874        };
2875        let __s_bl = self.gpu.stream();
2876        let mut bl = __s_bl.launch_builder(&f);
2877        bl.arg(preds)
2878            .arg(draft)
2879            .arg(&last_pred)
2880            .arg(&b)
2881            .arg(&k)
2882            .arg(out);
2883        unsafe {
2884            bl.launch(cfg)?;
2885        }
2886        Ok(())
2887    }
2888
2889    // ================= SAMPLED-SPEC PRIMITIVES (spec_sample.cu, piece A) =================
2890    // Counter-based randomness: every call takes (seed, stream_pos) — the caller owns the
2891    // event counter (one per sampled token). temp <= 0 arms are exact greedy limits.
2892
2893    /// y = x/temp + Gumbel(Philox(seed, stream_pos)) over n logits (then run device argmax on y
2894    /// = one categorical sample at temperature `temp`). temp<=0: y = x (pure copy).
2895    pub fn gumbel_perturb(
2896        &self,
2897        x: &CudaSlice<f32>,
2898        y: &mut CudaSlice<f32>,
2899        n: usize,
2900        seed: u64,
2901        stream_pos: u32,
2902        temp: f32,
2903    ) -> Result<(), Box<dyn std::error::Error>> {
2904        let f = self.func("gumbel_perturb_f32");
2905        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
2906        let cfg = LaunchConfig {
2907            grid_dim: (n.div_ceil(256) as u32, 1, 1),
2908            block_dim: (256, 1, 1),
2909            shared_mem_bytes: 0,
2910        };
2911        let __s_b = self.gpu.stream();
2912        let mut b = __s_b.launch_builder(&f);
2913        b.arg(x)
2914            .arg(&mut *y)
2915            .arg(&ni)
2916            .arg(&slo)
2917            .arg(&shi)
2918            .arg(&stream_pos)
2919            .arg(&temp);
2920        unsafe {
2921            b.launch(cfg)?;
2922        }
2923        Ok(())
2924    }
2925
2926    /// GRAMMAR TOKEN MASK (constrained decoding, lane/constrained-full): ban every vocab id
2927    /// whose bit is unset in the packed llguidance bitset, IN PLACE on row `col` of a stacked
2928    /// [B, n_vocab] logits buffer. `mask` = the SimpleVob u32 words H2D'd verbatim
2929    /// (~n_vocab/8 bytes/step — trivial on PCIe); ids >= 32*mask_words (padded lm_head tail)
2930    /// are banned too, the device twin of constrained::apply_mask. Banned value -FLT_MAX ==
2931    /// the argmax/gumbel kernels' init sentinel, so a fully-banned tail can never win and
2932    /// ordering matches the host -inf mask bit-for-bit for every finite logit.
2933    pub fn mask_logits_col(
2934        &self,
2935        logits: &mut CudaSlice<f32>,
2936        mask: &CudaSlice<u32>,
2937        col: usize,
2938        n: usize,
2939        mask_words: usize,
2940    ) -> Result<(), Box<dyn std::error::Error>> {
2941        let f = self.func("mask_logits_f32");
2942        let (ci, ni, mw) = (col as i32, n as i32, mask_words as i32);
2943        let cfg = LaunchConfig {
2944            grid_dim: (n.div_ceil(256).min(1024) as u32, 1, 1),
2945            block_dim: (256, 1, 1),
2946            shared_mem_bytes: 0,
2947        };
2948        let __s_b = self.gpu.stream();
2949        let mut b = __s_b.launch_builder(&f);
2950        b.arg(&mut *logits).arg(mask).arg(&ci).arg(&ni).arg(&mw);
2951        unsafe {
2952            b.launch(cfg)?;
2953        }
2954        Ok(())
2955    }
2956
2957    /// Column-`col` twin of `gumbel_perturb` over stacked logits [B, n_vocab] (the batched
2958    /// serving tick's device sampler): y = x[col]/temp + gumbel(seed, stream_pos, lane).
2959    /// SAME kernel/Philox mapping as `gumbel_perturb` — bit-identical perturbation for the
2960    /// same (seed, stream_pos, temp) regardless of which batch column the row sits in
2961    /// (the lane index is the in-row position; `col` only moves the input pointer). That
2962    /// pointer-invariance IS the serving isolation contract for sampled rows.
2963    pub fn gumbel_perturb_col(
2964        &self,
2965        x: &CudaSlice<f32>,
2966        col: usize,
2967        y: &mut CudaSlice<f32>,
2968        n: usize,
2969        seed: u64,
2970        stream_pos: u32,
2971        temp: f32,
2972    ) -> Result<(), Box<dyn std::error::Error>> {
2973        let f = self.func("gumbel_perturb_f32");
2974        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
2975        let col_view = x.slice(col * n..(col + 1) * n);
2976        let cfg = LaunchConfig {
2977            grid_dim: (n.div_ceil(256) as u32, 1, 1),
2978            block_dim: (256, 1, 1),
2979            shared_mem_bytes: 0,
2980        };
2981        let __s_b = self.gpu.stream();
2982        let mut b = __s_b.launch_builder(&f);
2983        b.arg(&col_view)
2984            .arg(&mut *y)
2985            .arg(&ni)
2986            .arg(&slo)
2987            .arg(&shi)
2988            .arg(&stream_pos)
2989            .arg(&temp);
2990        unsafe {
2991            b.launch(cfg)?;
2992        }
2993        Ok(())
2994    }
2995
2996    /// Filtered twin of `gumbel_perturb_col`: the per-row (row_max, th) floor comes from
2997    /// DEVICE buffers (`filter_stats` output slots at `stat_idx`) — one filtered draw from
2998    /// the top-k/top-p/min-p-truncated softmax with no stat D2H and no row copy. Same
2999    /// Philox mapping as every gumbel kernel (pointer-invariant across batch columns —
3000    /// the serving isolation contract for sampled rows).
3001    #[allow(clippy::too_many_arguments)]
3002    pub fn gumbel_perturb_filtered_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        stat_max: &CudaSlice<f32>,
3012        stat_th: &CudaSlice<f32>,
3013        stat_idx: usize,
3014    ) -> Result<(), Box<dyn std::error::Error>> {
3015        let f = self.func("gumbel_perturb_filtered_col_f32");
3016        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3017        let (ci, si) = (col as i32, stat_idx as i32);
3018        let cfg = LaunchConfig {
3019            grid_dim: (n.div_ceil(256) as u32, 1, 1),
3020            block_dim: (256, 1, 1),
3021            shared_mem_bytes: 0,
3022        };
3023        let __s_b = self.gpu.stream();
3024        let mut b = __s_b.launch_builder(&f);
3025        b.arg(x)
3026            .arg(&ci)
3027            .arg(&mut *y)
3028            .arg(&ni)
3029            .arg(&slo)
3030            .arg(&shi)
3031            .arg(&stream_pos)
3032            .arg(&temp)
3033            .arg(stat_max)
3034            .arg(stat_th)
3035            .arg(&si);
3036        unsafe {
3037            b.launch(cfg)?;
3038        }
3039        Ok(())
3040    }
3041
3042    /// In-graph sampling-event counter bump (spec_sample.cu kernel 5): ctr[0] += 1. The sampled
3043    /// graph-draft chain replays with FIXED kernel args, so the Philox event counter must be
3044    /// DEVICE data — the host seeds it once per round; every replay bumps it before the perturb
3045    /// reads it (counter is data, not state — graph-replay-safe).
3046    pub fn sctr_inc(&self, ctr: &mut CudaSlice<u32>) -> Result<(), Box<dyn std::error::Error>> {
3047        let f = self.func("memra_sctr_inc");
3048        let cfg = LaunchConfig {
3049            grid_dim: (1, 1, 1),
3050            block_dim: (1, 1, 1),
3051            shared_mem_bytes: 0,
3052        };
3053        let __s_b = self.gpu.stream();
3054        let mut b = __s_b.launch_builder(&f);
3055        b.arg(&mut *ctr);
3056        unsafe {
3057            b.launch(cfg)?;
3058        }
3059        Ok(())
3060    }
3061
3062    /// Graph-capturable `gumbel_perturb`: the sampling-event counter comes from DEVICE memory
3063    /// (`ctr[0]`) instead of a host scalar. Identical math to `gumbel_perturb` at
3064    /// stream_pos == ctr[0] (same Philox call, same lane mapping) — the eager and graph sampled
3065    /// chains produce bit-identical perturbations for the same (seed, counter, temp).
3066    pub fn gumbel_perturb_ctr(
3067        &self,
3068        x: &CudaSlice<f32>,
3069        y: &mut CudaSlice<f32>,
3070        n: usize,
3071        seed: u64,
3072        ctr: &CudaSlice<u32>,
3073        temp: f32,
3074    ) -> Result<(), Box<dyn std::error::Error>> {
3075        let f = self.func("gumbel_perturb_ctr_f32");
3076        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3077        let cfg = LaunchConfig {
3078            grid_dim: (n.div_ceil(256) as u32, 1, 1),
3079            block_dim: (256, 1, 1),
3080            shared_mem_bytes: 0,
3081        };
3082        let __s_b = self.gpu.stream();
3083        let mut b = __s_b.launch_builder(&f);
3084        b.arg(x)
3085            .arg(&mut *y)
3086            .arg(&ni)
3087            .arg(&slo)
3088            .arg(&shi)
3089            .arg(ctr)
3090            .arg(&temp);
3091        unsafe {
3092            b.launch(cfg)?;
3093        }
3094        Ok(())
3095    }
3096
3097    /// out[pair] = softmax_temp(x[rows[pair]])[ids[pair]] for npair (row, id) pairs; rows index
3098    /// into x with `row_stride` f32s per row. temp<=0: out = 1.0 iff id is the row argmax
3099    /// (smallest-index tie-break — matches the argmax-gate contract).
3100    pub fn softmax_gather(
3101        &self,
3102        x: &CudaSlice<f32>,
3103        row_stride: usize,
3104        ids: &CudaSlice<u32>,
3105        rows: &CudaSlice<i32>,
3106        out: &mut CudaSlice<f32>,
3107        n: usize,
3108        npair: usize,
3109        temp: f32,
3110    ) -> Result<(), Box<dyn std::error::Error>> {
3111        let f = self.func("softmax_gather_f32");
3112        let (ni, rs) = (n as i32, row_stride as i64);
3113        let np = npair as i32;
3114        let cfg = LaunchConfig {
3115            grid_dim: (npair as u32, 1, 1),
3116            block_dim: (256, 1, 1),
3117            shared_mem_bytes: 0,
3118        };
3119        let __s_b = self.gpu.stream();
3120        let mut b = __s_b.launch_builder(&f);
3121        b.arg(x)
3122            .arg(&rs)
3123            .arg(ids)
3124            .arg(rows)
3125            .arg(&mut *out)
3126            .arg(&ni)
3127            .arg(&np)
3128            .arg(&temp);
3129        unsafe {
3130            b.launch(cfg)?;
3131        }
3132        Ok(())
3133    }
3134
3135    /// Sample token from norm(max(0, softmax_temp(p) - softmax_temp(q))) (q = None -> plain
3136    /// categorical from softmax_temp(p)). Row stats (max, sumexp at temp) must be precomputed
3137    /// (softmax_gather's pass-1 values; see spec.rs caller). Deterministic fixed-order CDF walk.
3138    pub fn residual_sample(
3139        &self,
3140        p: &CudaSlice<f32>,
3141        q: Option<&CudaSlice<f32>>,
3142        n: usize,
3143        temp: f32,
3144        seed: u64,
3145        stream_pos: u32,
3146        out_tok: &mut CudaSlice<u32>,
3147    ) -> Result<(), Box<dyn std::error::Error>> {
3148        let f = self.func("residual_sample_f32");
3149        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3150        let nth = 1024u32;
3151        let cfg = LaunchConfig {
3152            grid_dim: (1, 1, 1),
3153            block_dim: (nth, 1, 1),
3154            shared_mem_bytes: 0,
3155        };
3156        let has_q: i32 = q.is_some() as i32;
3157        let qbuf = q.unwrap_or(p); // dummy when absent; kernel gates on has_q
3158        let __s_b = self.gpu.stream();
3159        let mut b = __s_b.launch_builder(&f);
3160        b.arg(p)
3161            .arg(qbuf)
3162            .arg(&has_q)
3163            .arg(&ni)
3164            .arg(&temp)
3165            .arg(&slo)
3166            .arg(&shi)
3167            .arg(&stream_pos)
3168            .arg(&mut *out_tok);
3169        unsafe {
3170            b.launch(cfg)?;
3171        }
3172        Ok(())
3173    }
3174
3175    /// Access the shared MoE residency cache (EDGE-1 §B), building it on first use under
3176    /// MEMRA_MOE_CACHE. The closure runs while the lock is held — keep it to lookup/admit/issue, not
3177    /// the GEMM. `max_block_bytes` sizes the slots (largest of gate/up/down). Returns the closure's
3178    /// result. If MEMRA_MOE_CACHE is unset this is never called (the caller checks the env first).
3179    pub fn with_moe_cache<R>(
3180        &self,
3181        max_block_bytes: usize,
3182        f: impl FnOnce(
3183            &mut crate::moe_cache::MoeSlotCache,
3184            &Engine,
3185        ) -> Result<R, Box<dyn std::error::Error>>,
3186    ) -> Result<R, Box<dyn std::error::Error>> {
3187        let mut guard = self.moe_cache.lock().unwrap();
3188        if guard.is_none() {
3189            *guard = Some(crate::moe_cache::MoeSlotCache::new(self, max_block_bytes)?);
3190        }
3191        let cache = guard.as_mut().unwrap();
3192        f(cache, self)
3193    }
3194
3195    /// Freeze the already-built MoE residency set. This never constructs a cache: callers use it
3196    /// only after a real prefill has populated the machine-specific CPU/GPU working set.
3197    pub fn freeze_moe_cache(&self) {
3198        if let Some(cache) = self.moe_cache.lock().unwrap().as_mut() {
3199            cache.freeze();
3200        }
3201    }
3202
3203    /// The current residency set as (layer, proj, ex) triples, or None if no cache was built.
3204    /// Never constructs a cache.
3205    pub fn export_moe_residency(&self) -> Option<Vec<(u16, u8, u16)>> {
3206        self.moe_cache
3207            .lock()
3208            .unwrap()
3209            .as_ref()
3210            .map(crate::moe_cache::MoeSlotCache::export_residency)
3211    }
3212
3213    pub(crate) fn moe_cache_frozen(&self) -> bool {
3214        self.moe_cache
3215            .lock()
3216            .unwrap()
3217            .as_ref()
3218            .is_some_and(crate::moe_cache::MoeSlotCache::is_frozen)
3219    }
3220
3221    /// A frozen heterogeneous CPU/GPU expert split cannot use Hy3's ordinary batched prefill
3222    /// efficiently: T>=PRIME_MIN_T bypasses the CPU backend and transiently rereads every missing
3223    /// expert through the GPU spill path. Replay the short prompt through decode after freezing,
3224    /// while leaving the profiling warmup's established batched behavior untouched.
3225    /// (`pub`: run-gen's #46 batched-prime gate skips itself when generation will take the
3226    /// tokenwise arm anyway.)
3227    pub fn frozen_cpu_experts_prefer_tokenwise_prime(&self) -> bool {
3228        crate::cpu_experts::configured()
3229            && self.moe_cache_frozen()
3230            && std::env::var("MEMRA_CPU_EXPERT_BATCHED_PRIME").as_deref() != Ok("1")
3231    }
3232
3233    /// Install the loaded model's exact retained expert-block inventory before lazy cache build.
3234    pub(crate) fn configure_moe_cache_layout(&self, block_bytes: Vec<usize>) {
3235        assert!(
3236            self.moe_cache.lock().unwrap().is_none(),
3237            "MoE cache layout configured after cache construction"
3238        );
3239        *self.moe_cache_layout.lock().unwrap() = Some(block_bytes);
3240    }
3241
3242    pub(crate) fn moe_cache_layout(&self) -> Option<Vec<usize>> {
3243        self.moe_cache_layout.lock().unwrap().clone()
3244    }
3245
3246    /// True if the MoE residency cache is enabled (MEMRA_MOE_CACHE set).
3247    pub fn moe_cache_enabled() -> bool {
3248        std::env::var("MEMRA_MOE_CACHE").as_deref() != Ok("0")
3249    }
3250
3251    /// Snapshot the MoE cache counters (hits, misses, staged_bytes, n_slots) for the §D.4 PCIe gate.
3252    /// Returns None if the cache was never built (disabled or no MoE forward ran).
3253    pub fn moe_cache_stats(&self) -> Option<(u64, u64, u64, usize)> {
3254        let guard = self.moe_cache.lock().unwrap();
3255        guard
3256            .as_ref()
3257            .map(|c| (c.hits, c.misses, c.staged_bytes, c.n_slots()))
3258    }
3259
3260    /// Experimental CPU expert backend counters: completed layer calls, experts served, and the
3261    /// sum of backend wall nanoseconds. The timer includes explicit disk->RAM fills on cache misses;
3262    /// callers compare a before/after snapshot around a decode window.
3263    pub fn cpu_expert_stats(
3264        &self,
3265    ) -> Option<(u64, u64, u64, u64, u64, u64, u64, u64, u64, u64, u64)> {
3266        crate::cpu_experts::configured().then(crate::cpu_experts::stats)
3267    }
3268
3269    /// Caller-blocked nanoseconds at CPU expert joins. Compare before/after snapshots to measure
3270    /// the backend tail that resident-GPU expert work did not hide.
3271    pub fn cpu_expert_predictor_stats(&self) -> (u64, u64) {
3272        crate::cpu_experts::predictor_stats()
3273    }
3274
3275    pub fn cpu_expert_exposed_wait_ns(&self) -> Option<u64> {
3276        crate::cpu_experts::configured().then(crate::cpu_experts::exposed_wait_ns)
3277    }
3278
3279    /// CPU-routed expert selections grouped by how many of their three projections were already
3280    /// resident in HBM. This makes otherwise-stranded partial residency visible to tuning runs.
3281    pub fn cpu_expert_gpu_residency_stats(&self) -> Option<(u64, u64, u64)> {
3282        crate::cpu_experts::configured().then(crate::cpu_experts::incomplete_gpu_residency_stats)
3283    }
3284
3285    /// Positioned-read proof-backend counters:
3286    /// `(reads, bytes, read_errors, short_reads, mmap_fallbacks, buffer_waits, ring_full)`.
3287    pub fn moe_pread_stats(&self) -> Option<(u64, u64, u64, u64, u64, u64, u64)> {
3288        let guard = self.moe_cache.lock().unwrap();
3289        guard
3290            .as_ref()
3291            .and_then(|cache| cache.pread_stats())
3292            .map(|stats| {
3293                (
3294                    stats.reads,
3295                    stats.bytes,
3296                    stats.read_errors,
3297                    stats.short_reads,
3298                    stats.fallbacks,
3299                    stats.buffer_waits,
3300                    stats.ring_full,
3301                )
3302            })
3303    }
3304
3305    /// Spill configuration values that warned and substituted their documented defaults.
3306    pub fn spill_config_fallbacks(&self) -> u64 {
3307        crate::spill_pread::config_fallbacks()
3308    }
3309
3310    /// Reset the MoE cache perf counters (to separate warmup from steady-state windows).
3311    pub fn moe_cache_reset_counters(&self) {
3312        if let Some(c) = self.moe_cache.lock().unwrap().as_mut() {
3313            c.reset_counters();
3314        }
3315    }
3316
3317    pub fn htod_bytes(&self, v: &[u8]) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
3318        Ok(self.gpu.stream().clone_htod(v)?)
3319    }
3320
3321    /// `htod_bytes` with a mapped (uninit) tail pad: the wide-load expert dots read up to 6B
3322    /// past the final q4_0 block through their aligned window — the bytes never reach a
3323    /// result (funnelshift discards them) but must be mapped memory.
3324    pub fn htod_bytes_padded(
3325        &self,
3326        v: &[u8],
3327        pad: usize,
3328    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
3329        let mut d = self.alloc_u8_uninit(v.len() + pad)?;
3330        {
3331            let mut view = d.slice_mut(0..v.len());
3332            self.gpu.stream().memcpy_htod(v, &mut view)?;
3333        }
3334        Ok(d)
3335    }
3336
3337    /// Device-to-device copy of `src` into `dst[off..off+len]` (f32). For in-place KV append.
3338    pub fn copy_into(
3339        &self,
3340        dst: &mut CudaSlice<f32>,
3341        off: usize,
3342        src: &CudaSlice<f32>,
3343        len: usize,
3344    ) -> Result<(), Box<dyn std::error::Error>> {
3345        let mut view = dst.slice_mut(off..off + len);
3346        self.gpu
3347            .stream()
3348            .memcpy_dtod(&src.slice(0..len), &mut view)?;
3349        Ok(())
3350    }
3351
3352    /// View a sub-range of a device buffer (for attending over [0..len) of a KV cache).
3353    /// u8 twin of copy_into (D2D byte-range copy at an offset).
3354    pub fn copy_u8_into(
3355        &self,
3356        dst: &mut CudaSlice<u8>,
3357        off: usize,
3358        src: &CudaSlice<u8>,
3359        len: usize,
3360    ) -> Result<(), Box<dyn std::error::Error>> {
3361        let mut view = dst.slice_mut(off..off + len);
3362        self.gpu
3363            .stream()
3364            .memcpy_dtod(&src.slice(0..len), &mut view)?;
3365        Ok(())
3366    }
3367
3368    /// D2D byte-range copy with explicit source and destination offsets.
3369    pub fn copy_u8_range_into(
3370        &self,
3371        dst: &mut CudaSlice<u8>,
3372        dst_off: usize,
3373        src: &CudaSlice<u8>,
3374        src_off: usize,
3375        len: usize,
3376    ) -> Result<(), Box<dyn std::error::Error>> {
3377        let mut dst_view = dst.slice_mut(dst_off..dst_off + len);
3378        self.gpu
3379            .stream()
3380            .memcpy_dtod(&src.slice(src_off..src_off + len), &mut dst_view)?;
3381        Ok(())
3382    }
3383
3384    /// Resolve an absolute append slot to the Step35 SWA layer's physical rows. At wrap, copy
3385    /// only the aligned live prefix through temporary device storage and rebase it at row zero,
3386    /// keeping the audited attention range contiguous without changing its absolute start.
3387    pub fn prepare_kv_append(
3388        &self,
3389        kv: &mut crate::cache::KvLayer,
3390        retain_from: usize,
3391        append_rows: usize,
3392    ) -> Result<usize, Box<dyn std::error::Error>> {
3393        let Some(plan) = kv
3394            .ring
3395            .as_ref()
3396            .map(|ring| ring.append_plan(kv.len, retain_from, append_rows))
3397            .transpose()?
3398        else {
3399            return Ok(kv.len);
3400        };
3401        match plan {
3402            crate::cache::KvRingAppend::Contiguous { write_row } => Ok(write_row),
3403            crate::cache::KvRingAppend::Rebase {
3404                src_row,
3405                keep_rows,
3406                new_base,
3407                write_row,
3408            } => {
3409                if keep_rows > 0 {
3410                    let k_len = keep_rows * kv.k_tok_bytes;
3411                    let v_len = keep_rows * kv.v_tok_bytes;
3412                    let mut k_tmp = self.alloc_u8_uninit(k_len)?;
3413                    let mut v_tmp = self.alloc_u8_uninit(v_len)?;
3414                    self.copy_u8_range_into(&mut k_tmp, 0, &kv.k, src_row * kv.k_tok_bytes, k_len)?;
3415                    self.copy_u8_range_into(&mut v_tmp, 0, &kv.v, src_row * kv.v_tok_bytes, v_len)?;
3416                    self.copy_u8_into(&mut kv.k, 0, &k_tmp, k_len)?;
3417                    self.copy_u8_into(&mut kv.v, 0, &v_tmp, v_len)?;
3418                }
3419                kv.ring.as_mut().unwrap().apply_rebase(new_base);
3420                Ok(write_row)
3421            }
3422        }
3423    }
3424
3425    /// H2D write of `src` into `dst[off..off+src.len()]` (u8). In-place row updates for the
3426    /// adaptive trim head: no realloc, so captured graphs keep their baked addresses.
3427    pub fn htod_u8_into(
3428        &self,
3429        dst: &mut CudaSlice<u8>,
3430        off: usize,
3431        src: &[u8],
3432    ) -> Result<(), Box<dyn std::error::Error>> {
3433        let mut view = dst.slice_mut(off..off + src.len());
3434        self.gpu.stream().memcpy_htod(src, &mut view)?;
3435        Ok(())
3436    }
3437
3438    pub fn view<'a>(&self, b: &'a CudaSlice<f32>, len: usize) -> cudarc::driver::CudaView<'a, f32> {
3439        b.slice(0..len)
3440    }
3441
3442    /// View the first `len` BYTES of a u8 device buffer (quantized KV cache: [0..t_kv*tok_bytes)).
3443    /// Byte-range view (gemma4 R6 window offset into the quantized KV stream).
3444    pub fn view_u8_range<'a>(
3445        &self,
3446        b: &'a CudaSlice<u8>,
3447        start: usize,
3448        end: usize,
3449    ) -> cudarc::driver::CudaView<'a, u8> {
3450        b.slice(start..end)
3451    }
3452    pub fn view_u8<'a>(
3453        &self,
3454        b: &'a CudaSlice<u8>,
3455        len: usize,
3456    ) -> cudarc::driver::CudaView<'a, u8> {
3457        b.slice(0..len)
3458    }
3459
3460    /// Append-quantize ONE token's post-RoPE K (q8_0) and V (q5_1) into the resident byte caches at
3461    /// token index `t` (KVQUANT-PLAN §C). One CTA (one warp) per 32-element block; the kernel writes
3462    /// the f16 scale(s) + packed quants for K and V. k_row/v_row are f32 [kv_dim_k]/[kv_dim_v].
3463    pub fn append_kv_quantized(
3464        &self,
3465        k_row: &CudaSlice<f32>,
3466        v_row: &CudaSlice<f32>,
3467        kc: &mut CudaSlice<u8>,
3468        vc: &mut CudaSlice<u8>,
3469        t: usize,
3470        kv_dim_k: usize,
3471        kv_dim_v: usize,
3472        k_tok_bytes: usize,
3473        v_tok_bytes: usize,
3474        g: bool,
3475    ) -> Result<(), Box<dyn std::error::Error>> {
3476        let f = if g {
3477            self.func_g("append_quantize_kv_q8_0_q5_1")
3478        } else {
3479            self.func("append_quantize_kv_q8_0_q5_1")
3480        };
3481        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
3482        let cfg = LaunchConfig {
3483            grid_dim: (nblk, 1, 1),
3484            block_dim: (32, 1, 1),
3485            shared_mem_bytes: 0,
3486        };
3487        let (ti, kdk, kdv) = (t as i32, kv_dim_k as i32, kv_dim_v as i32);
3488        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3489        let __s_b = self.gpu.stream();
3490        let mut b = __s_b.launch_builder(&f);
3491        b.arg(k_row)
3492            .arg(v_row)
3493            .arg(kc)
3494            .arg(vc)
3495            .arg(&ti)
3496            .arg(&kdk)
3497            .arg(&kdv)
3498            .arg(&ktb)
3499            .arg(&vtb);
3500        unsafe {
3501            b.launch(cfg)?;
3502        }
3503        Ok(())
3504    }
3505
3506    /// Device-counter variant of `append_kv_quantized` (CUDA-GRAPH-PLAN Phase 2): the write slot
3507    /// `t` is read from `t_dev[0]` (a resident device i32[1]) instead of a host int arg, so the
3508    /// launch args are FIXED across decode steps (graph-capturable). Identical quant math.
3509    pub fn append_kv_quantized_dc(
3510        &self,
3511        k_row: &CudaSlice<f32>,
3512        v_row: &CudaSlice<f32>,
3513        kc: &mut CudaSlice<u8>,
3514        vc: &mut CudaSlice<u8>,
3515        t_dev: &CudaSlice<i32>,
3516        kv_dim_k: usize,
3517        kv_dim_v: usize,
3518        k_tok_bytes: usize,
3519        v_tok_bytes: usize,
3520        g: bool,
3521    ) -> Result<(), Box<dyn std::error::Error>> {
3522        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
3523        let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
3524        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3525        // PDL wave-B2: flash-module flavor mirrors the builder path's g flag exactly.
3526        if Self::pdl_on() && Self::pdl_wb_on() {
3527            use cudarc::driver::{DevicePtr, DevicePtrMut};
3528            let s = &self.gpu.stream();
3529            let (pk, _g0) = k_row.device_ptr(s);
3530            let (pv, _g1) = v_row.device_ptr(s);
3531            let (pkc, _g2) = kc.device_ptr_mut(s);
3532            let (pvc, _g3) = vc.device_ptr_mut(s);
3533            let (pt, _g4) = t_dev.device_ptr(s);
3534            let mut ps = [
3535                &pk as *const _ as *mut std::ffi::c_void,
3536                &pv as *const _ as *mut _,
3537                &pkc as *const _ as *mut _,
3538                &pvc as *const _ as *mut _,
3539                &pt as *const _ as *mut _,
3540                &kdk as *const _ as *mut _,
3541                &kdv as *const _ as *mut _,
3542                &ktb as *const _ as *mut _,
3543                &vtb as *const _ as *mut _,
3544            ];
3545            unsafe {
3546                self.launch_pdl_flash(
3547                    g,
3548                    "append_quantize_kv_q8_0_q5_1_dc",
3549                    (nblk, 1, 1),
3550                    (32, 1, 1),
3551                    0,
3552                    &mut ps,
3553                )?;
3554            }
3555            return Ok(());
3556        }
3557        let f = if g {
3558            self.func_g("append_quantize_kv_q8_0_q5_1_dc")
3559        } else {
3560            self.func("append_quantize_kv_q8_0_q5_1_dc")
3561        };
3562        let cfg = LaunchConfig {
3563            grid_dim: (nblk, 1, 1),
3564            block_dim: (32, 1, 1),
3565            shared_mem_bytes: 0,
3566        };
3567        let __s_b = self.gpu.stream();
3568        let mut b = __s_b.launch_builder(&f);
3569        b.arg(k_row)
3570            .arg(v_row)
3571            .arg(kc)
3572            .arg(vc)
3573            .arg(t_dev)
3574            .arg(&kdk)
3575            .arg(&kdv)
3576            .arg(&ktb)
3577            .arg(&vtb);
3578        unsafe {
3579            b.launch(cfg)?;
3580        }
3581        Ok(())
3582    }
3583
3584    /// Append-quantize T token rows in one shot (BATCHED PROMPT PRIME). k_rows/v_rows are
3585    /// token-major [T, kv_dim] post-RoPE f32; rows land at cache slots t0..t0+T. Default = the
3586    /// batched `_rows` kernel: one (nblk, T) launch whose per-(block,token) warp program is the
3587    /// per-token append kernel verbatim -> every written row is BIT-IDENTICAL to T sequential
3588    /// `append_kv_quantized_view` calls (kernel_check pins the bytes). MEMRA_PRIME_APPEND_LOOP=1
3589    /// forces the T-launch per-row loop (the A/B seam that measured the launch overhead).
3590    #[allow(clippy::too_many_arguments)]
3591    pub fn append_kv_quantized_rows(
3592        &self,
3593        k_rows: &CudaSlice<f32>,
3594        v_rows: &CudaSlice<f32>,
3595        kc: &mut CudaSlice<u8>,
3596        vc: &mut CudaSlice<u8>,
3597        t0: usize,
3598        t: usize,
3599        kv_dim_k: usize,
3600        kv_dim_v: usize,
3601        k_tok_bytes: usize,
3602        v_tok_bytes: usize,
3603        g: bool,
3604    ) -> Result<(), Box<dyn std::error::Error>> {
3605        if std::env::var("MEMRA_PRIME_APPEND_LOOP").is_ok() {
3606            for i in 0..t {
3607                let k_row = k_rows.slice(i * kv_dim_k..(i + 1) * kv_dim_k);
3608                let v_row = v_rows.slice(i * kv_dim_v..(i + 1) * kv_dim_v);
3609                self.append_kv_quantized_view(
3610                    &k_row,
3611                    &v_row,
3612                    kc,
3613                    vc,
3614                    t0 + i,
3615                    kv_dim_k,
3616                    kv_dim_v,
3617                    k_tok_bytes,
3618                    v_tok_bytes,
3619                    g,
3620                )?;
3621            }
3622            return Ok(());
3623        }
3624        let f = if g {
3625            self.func_g("append_quantize_kv_q8_0_q5_1_rows")
3626        } else {
3627            self.func("append_quantize_kv_q8_0_q5_1_rows")
3628        };
3629        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
3630        let cfg = LaunchConfig {
3631            grid_dim: (nblk, t as u32, 1),
3632            block_dim: (32, 1, 1),
3633            shared_mem_bytes: 0,
3634        };
3635        let (t0i, kdk, kdv) = (t0 as i32, kv_dim_k as i32, kv_dim_v as i32);
3636        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3637        let __s_b = self.gpu.stream();
3638        let mut b = __s_b.launch_builder(&f);
3639        b.arg(k_rows)
3640            .arg(v_rows)
3641            .arg(kc)
3642            .arg(vc)
3643            .arg(&t0i)
3644            .arg(&kdk)
3645            .arg(&kdv)
3646            .arg(&ktb)
3647            .arg(&vtb);
3648        unsafe {
3649            b.launch(cfg)?;
3650        }
3651        Ok(())
3652    }
3653
3654    /// Increment a device i32[1] counter in place (p[0] += 1) via the resident `inc_i32` kernel.
3655    /// Used to advance the device-resident seqlen/pos counters inside the decode-dc path (and,
3656    /// later, inside a captured graph) without a host round-trip.
3657    pub fn inc_seqlen(&self, p: &mut CudaSlice<i32>) -> Result<(), Box<dyn std::error::Error>> {
3658        let f = self.func("inc_i32");
3659        let cfg = LaunchConfig {
3660            grid_dim: (1, 1, 1),
3661            block_dim: (1, 1, 1),
3662            shared_mem_bytes: 0,
3663        };
3664        let __s_b = self.gpu.stream();
3665        let mut b = __s_b.launch_builder(&f);
3666        b.arg(p);
3667        unsafe {
3668            b.launch(cfg)?;
3669        }
3670        Ok(())
3671    }
3672
3673    /// Like `append_kv_quantized` but k_row/v_row are CudaViews (one token's row sliced out of a
3674    /// token-major [T, kv_dim] activation buffer — the MTP verify path appends T tokens).
3675    pub fn append_kv_quantized_view(
3676        &self,
3677        k_row: &cudarc::driver::CudaView<f32>,
3678        v_row: &cudarc::driver::CudaView<f32>,
3679        kc: &mut CudaSlice<u8>,
3680        vc: &mut CudaSlice<u8>,
3681        t: usize,
3682        kv_dim_k: usize,
3683        kv_dim_v: usize,
3684        k_tok_bytes: usize,
3685        v_tok_bytes: usize,
3686        g: bool,
3687    ) -> Result<(), Box<dyn std::error::Error>> {
3688        let f = if g {
3689            self.func_g("append_quantize_kv_q8_0_q5_1")
3690        } else {
3691            self.func("append_quantize_kv_q8_0_q5_1")
3692        };
3693        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
3694        let cfg = LaunchConfig {
3695            grid_dim: (nblk, 1, 1),
3696            block_dim: (32, 1, 1),
3697            shared_mem_bytes: 0,
3698        };
3699        let (ti, kdk, kdv) = (t as i32, kv_dim_k as i32, kv_dim_v as i32);
3700        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3701        let __s_b = self.gpu.stream();
3702        let mut b = __s_b.launch_builder(&f);
3703        b.arg(k_row)
3704            .arg(v_row)
3705            .arg(kc)
3706            .arg(vc)
3707            .arg(&ti)
3708            .arg(&kdk)
3709            .arg(&kdv)
3710            .arg(&ktb)
3711            .arg(&vtb);
3712        unsafe {
3713            b.launch(cfg)?;
3714        }
3715        Ok(())
3716    }
3717
3718    /// Device-to-device copy of a CudaView `src` into `dst[off..off+len]` (f32). Like `copy_into`
3719    /// but the source is a sub-view (e.g. one column of a token-major activation buffer).
3720    pub fn copy_view_into(
3721        &self,
3722        dst: &mut CudaSlice<f32>,
3723        off: usize,
3724        src: &cudarc::driver::CudaView<f32>,
3725        len: usize,
3726    ) -> Result<(), Box<dyn std::error::Error>> {
3727        let mut view = dst.slice_mut(off..off + len);
3728        self.gpu
3729            .stream()
3730            .memcpy_dtod(&src.slice(0..len), &mut view)?;
3731        Ok(())
3732    }
3733
3734    /// Real device-to-device COPY of `src` into a freshly allocated buffer (NOT an Arc clone).
3735    /// Used for cache snapshots (MTP-PLAN §D.4): `CudaSlice::clone()` only bumps a refcount and
3736    /// would alias the live buffer; this allocs new device memory and memcpy_dtod's the contents.
3737    pub fn clone_dtod(
3738        &self,
3739        src: &CudaSlice<f32>,
3740    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3741        let mut dst = self.gpu.stream().alloc_zeros::<f32>(src.len())?;
3742        self.gpu.stream().memcpy_dtod(src, &mut dst)?;
3743        Ok(dst)
3744    }
3745
3746    /// D2D row extraction: copy a view (e.g. one row of a [B, n] batch buffer) into `dst`.
3747    /// Stream-ordered, async — decode_batch's per-sequence row plumbing.
3748    pub fn dtod_copy_view(
3749        &self,
3750        src: &cudarc::driver::CudaView<f32>,
3751        dst: &mut CudaSlice<f32>,
3752    ) -> Result<(), Box<dyn std::error::Error>> {
3753        self.gpu.stream().memcpy_dtod(src, dst)?;
3754        Ok(())
3755    }
3756
3757    /// D2D i8 twin of `dtod_copy_view` (q8_1 activation rows).
3758    pub fn dtod_copy_view_i8(
3759        &self,
3760        src: &cudarc::driver::CudaView<i8>,
3761        dst: &mut CudaSlice<i8>,
3762    ) -> Result<(), Box<dyn std::error::Error>> {
3763        self.gpu.stream().memcpy_dtod(src, dst)?;
3764        Ok(())
3765    }
3766
3767    /// D2D row placement: copy `src` into `dst[offset .. offset+src.len()]`.
3768    pub fn dtod_copy_into(
3769        &self,
3770        src: &CudaSlice<f32>,
3771        dst: &mut CudaSlice<f32>,
3772        offset: usize,
3773    ) -> Result<(), Box<dyn std::error::Error>> {
3774        let n = src.len();
3775        let mut dv = dst.slice_mut(offset..offset + n);
3776        self.gpu.stream().memcpy_dtod(src, &mut dv)?;
3777        Ok(())
3778    }
3779
3780    /// Batched uniform-size D2D copy (engine-bundle slice 1, DSF-ROUNDCOST-20260820 §1.1):
3781    /// `n` disjoint regions of `words` f32 each; `table` = [src_0..src_{n-1}, dst_0..dst_{n-1}]
3782    /// raw device pointers. ONE kernel launch replaces `n` memcpy_dtod dispatches — the dspark
3783    /// round's snap/commit copy dribble (~0.9 ms/round of dispatch serialization measured).
3784    /// Bytes and stream order are identical to the memcpy sequence it replaces.
3785    pub fn copy_batch_uniform_f32(
3786        &self,
3787        table: &CudaSlice<u64>,
3788        n: usize,
3789        words: usize,
3790    ) -> Result<(), Box<dyn std::error::Error>> {
3791        if n == 0 || words == 0 {
3792            return Ok(());
3793        }
3794        debug_assert!(
3795            table.len() >= 2 * n,
3796            "pointer table must hold n srcs + n dsts"
3797        );
3798        let f = self.func("copy_batch_uniform_f32");
3799        // Enough blocks to stream a multi-MB region, few enough that (chunks x n) stays a
3800        // sane grid: 512K-word ssm regions get 48 grid-striding blocks each.
3801        let chunks = (words / 4).max(1).div_ceil(256).min(48) as u32;
3802        let (ni, wi) = (n as i32, words as i32);
3803        let cfg = LaunchConfig {
3804            grid_dim: (chunks, n as u32, 1),
3805            block_dim: (256, 1, 1),
3806            shared_mem_bytes: 0,
3807        };
3808        let __s = self.gpu.stream();
3809        let mut b = __s.launch_builder(&f);
3810        b.arg(table).arg(&ni).arg(&wi);
3811        unsafe {
3812            b.launch(cfg)?;
3813        }
3814        Ok(())
3815    }
3816
3817    /// H2D refresh of an EXISTING u64 pointer table IN PLACE (stable pointer — the batched
3818    /// state-copy tables are refreshed per round because the GDN ssm handles ping-pong).
3819    pub fn htod_u64_into(
3820        &self,
3821        v: &[u64],
3822        dst: &mut CudaSlice<u64>,
3823    ) -> Result<(), Box<dyn std::error::Error>> {
3824        let mut view = dst.slice_mut(0..v.len());
3825        self.gpu.stream().memcpy_htod(v, &mut view)?;
3826        Ok(())
3827    }
3828
3829    /// Indirect-source copy (engine-bundle slice 3): the src ADDRESS is loaded from a
3830    /// device pointer-table entry at run time, so a captured graph follows the gdn
3831    /// ping-pong through the same table its scan kernels read — a baked memcpy node
3832    /// would keep the capture-time physical buffer. `dst_off`/`words` in f32 elements.
3833    pub fn copy_indirect_src_f32(
3834        &self,
3835        src_entry: &cudarc::driver::CudaView<u64>,
3836        dst: &mut CudaSlice<f32>,
3837        dst_off: usize,
3838        words: usize,
3839    ) -> Result<(), Box<dyn std::error::Error>> {
3840        let f = self.func("copy_indirect_src_f32");
3841        let chunks = (words / 4).max(1).div_ceil(256).min(48) as u32;
3842        let wi = words as i32;
3843        let cfg = LaunchConfig {
3844            grid_dim: (chunks, 1, 1),
3845            block_dim: (256, 1, 1),
3846            shared_mem_bytes: 0,
3847        };
3848        let mut dv = dst.slice_mut(dst_off..dst_off + words);
3849        let __s = self.gpu.stream();
3850        let mut b = __s.launch_builder(&f);
3851        b.arg(src_entry).arg(&mut dv).arg(&wi);
3852        unsafe {
3853            b.launch(cfg)?;
3854        }
3855        Ok(())
3856    }
3857
3858    /// Uninitialized i8 device buffer (decode_batch q8_1 row scratch).
3859    pub fn uninit_i8(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
3860        self.alloc_uninit::<i8>(n)
3861    }
3862
3863    /// Resident-quantized linear (Stage-A: f32 dequant-in-kernel). y[m,out]=x[m,in]@W[out,in]^T.
3864    pub fn qmatvec(
3865        &self,
3866        w: &CudaSlice<u8>,
3867        x: &CudaSlice<f32>,
3868        m: usize,
3869        in_f: usize,
3870        out_f: usize,
3871        qtype: i32,
3872        row_bytes: usize,
3873    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3874        let f = self.func("qmatvec_f32");
3875        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
3876        let cfg = LaunchConfig {
3877            grid_dim: (out_f as u32, m as u32, 1),
3878            block_dim: (256, 1, 1),
3879            shared_mem_bytes: 0,
3880        };
3881        let (inf, outf, mi, qt, rb) =
3882            (in_f as i32, out_f as i32, m as i32, qtype, row_bytes as i64);
3883        let __s_b = self.gpu.stream();
3884        let mut b = __s_b.launch_builder(&f);
3885        b.arg(w)
3886            .arg(x)
3887            .arg(&mut y)
3888            .arg(&inf)
3889            .arg(&outf)
3890            .arg(&mi)
3891            .arg(&qt)
3892            .arg(&rb);
3893        unsafe {
3894            b.launch(cfg)?;
3895        }
3896        Ok(y)
3897    }
3898
3899    /// Allocate a reusable u8 GPU scratch buffer (for staged expert weights).
3900    pub fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
3901        let s = self.gpu.stream().alloc_zeros::<u8>(n)?;
3902        self.keep_if_capturing(&s);
3903        Ok(s)
3904    }
3905
3906    /// Uninitialized u8 scratch — skips alloc_zeros' memset. ONLY for staging buffers whose read
3907    /// range is fully overwritten by a stage_expert H2D before any kernel reads it (LAUNCH-STRUCTURE
3908    /// STAGE 2: the per-layer MoE scratch trio was 3 dead ~1MB memsets per layer per decode token).
3909    pub fn alloc_u8_uninit(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
3910        let s = unsafe { self.gpu.stream().alloc::<u8>(n)? };
3911        self.keep_if_capturing(&s);
3912        Ok(s)
3913    }
3914
3915    /// Zero a SUB-RANGE of an f32 buffer (CudaViewMut) — the row-sized memset the moe_out
3916    /// memset-elision uses for tokens that fall off the gdec fast path (LAUNCH-STRUCTURE STAGE 2).
3917    pub fn memset_zeros_view(
3918        &self,
3919        dst: &mut cudarc::driver::CudaViewMut<f32>,
3920    ) -> Result<(), Box<dyn std::error::Error>> {
3921        self.gpu.stream().memset_zeros(dst)?;
3922        Ok(())
3923    }
3924
3925    /// EDGE-1 staging: copy `host_bytes` (a sub-slice of a HostExps buffer) into `scratch`
3926    /// at byte offset `off` (async H2D on the default stream). Length is host_bytes.len().
3927    /// The qmatvec_view that reads `scratch[off..]` is enqueued on the SAME stream after this,
3928    /// so ordering is guaranteed without an explicit sync (Stage-1; Stage-2 prefetch on a 2nd
3929    /// stream would require an event).
3930    pub fn stage_expert(
3931        &self,
3932        host_bytes: &[u8],
3933        scratch: &mut CudaSlice<u8>,
3934        off: usize,
3935    ) -> Result<(), Box<dyn std::error::Error>> {
3936        let mut dst = scratch.slice_mut(off..off + host_bytes.len()); // CudaViewMut<u8>
3937        self.gpu.stream().memcpy_htod(host_bytes, &mut dst)?; // accepts &[u8] HostSlice src
3938        Ok(())
3939    }
3940
3941    /// EDGE-1 §A: fused MoE router. `logits` is the router output [t, n_expert] (device, f32, the
3942    /// `gate_inp @ z` result). Returns (sel_idx [t, n_used] i32, sel_w [t, n_used] f32): the top-k
3943    /// expert ids (DESC by prob, ascending-index tiebreak) and renormalized weights. Replaces the
3944    /// host dtoh + softmax-256 + stable DESC top-8 sort + renorm (hybrid_forward.rs ~281-298).
3945    /// One CTA per token row, 256 threads (one per expert).
3946    pub fn moe_router_topk(
3947        &self,
3948        logits: &CudaSlice<f32>,
3949        t: usize,
3950        n_expert: usize,
3951        n_used: usize,
3952    ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3953        let f = self.func("moe_router_topk_f32");
3954        let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?; // kernel fully overwrites
3955        let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?; // kernel fully overwrites
3956        let cfg = LaunchConfig {
3957            grid_dim: (t as u32, 1, 1),
3958            block_dim: (n_expert as u32, 1, 1),
3959            shared_mem_bytes: 0,
3960        };
3961        let (ne, nu) = (n_expert as i32, n_used as i32);
3962        let __s_b = self.gpu.stream();
3963        let mut b = __s_b.launch_builder(&f);
3964        b.arg(logits)
3965            .arg(&mut sel_idx)
3966            .arg(&mut sel_w)
3967            .arg(&ne)
3968            .arg(&nu);
3969        unsafe {
3970            b.launch(cfg)?;
3971        }
3972        Ok((sel_idx, sel_w))
3973    }
3974
3975    /// gemma4 twin: per-expert output scale folded into the topk renorm write (replaces the
3976    /// separate moe_w_exscale launch; value chain identical: (w/ws) * s[sel]).
3977    pub fn moe_router_topk_scaled(
3978        &self,
3979        logits: &CudaSlice<f32>,
3980        t: usize,
3981        n_expert: usize,
3982        n_used: usize,
3983        ex_scale: &CudaSlice<f32>,
3984    ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3985        // barrier-lean v2 twin (per-warp top-k + one-warp merge) FALSIFIED 2026-07-14:
3986        // bit-identical streams but −1.4% (26B plain N=3 interleaved) — at t=1 the grid is
3987        // ONE block, so the 6.6us is launch/dependency overhead, not the barrier chain;
3988        // fewer barriers bought nothing and the merge structure cost. jsonl is the record.
3989        let f = self.func("moe_router_topk_scaled_f32");
3990        let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?;
3991        let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?;
3992        let cfg = LaunchConfig {
3993            grid_dim: (t as u32, 1, 1),
3994            block_dim: (n_expert as u32, 1, 1),
3995            shared_mem_bytes: 0,
3996        };
3997        let (ne, nu) = (n_expert as i32, n_used as i32);
3998        let __s_b = self.gpu.stream();
3999        let mut b = __s_b.launch_builder(&f);
4000        b.arg(logits)
4001            .arg(&mut sel_idx)
4002            .arg(&mut sel_w)
4003            .arg(&ne)
4004            .arg(&nu)
4005            .arg(ex_scale);
4006        unsafe {
4007            b.launch(cfg)?;
4008        }
4009        Ok((sel_idx, sel_w))
4010    }
4011
4012    /// LAUNCH-STRUCTURE STAGE 1 (2026-07-05): fused router + SINGLE-SYNC host readback. The old
4013    /// MEMRA_FUSED_ROUTER path lost 2% at t=1 because it paid TWO full stream syncs (dtoh_i32 then
4014    /// dtoh, each = clone_dtoh + synchronize) + two alloc_zeros memsets per MoE layer, where the
4015    /// host route pays ONE sync on the 1KB logits dtoh. This variant: uninit outputs (kernel fully
4016    /// overwrites), both DtoH copies issued ASYNC into a persistent PINNED host staging buffer
4017    /// (flags=0 — cacheable, NOT cudarc's WRITECOMBINED default, so the host-side reads of sel/w
4018    /// stay cached), then ONE synchronize. Numerics identical to `moe_router_topk` (same kernel).
4019    pub fn moe_router_topk_host(
4020        &self,
4021        logits: &CudaSlice<f32>,
4022        t: usize,
4023        n_expert: usize,
4024        n_used: usize,
4025    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
4026        let f = self.func("moe_router_topk_f32");
4027        let n = t * n_used;
4028        let mut sel_idx = self.alloc_uninit::<i32>(n)?;
4029        let mut sel_w = self.alloc_uninit::<f32>(n)?;
4030        let cfg = LaunchConfig {
4031            grid_dim: (t as u32, 1, 1),
4032            block_dim: (n_expert as u32, 1, 1),
4033            shared_mem_bytes: 0,
4034        };
4035        let (ne, nu) = (n_expert as i32, n_used as i32);
4036        let __s_b = self.gpu.stream();
4037        let mut b = __s_b.launch_builder(&f);
4038        b.arg(logits)
4039            .arg(&mut sel_idx)
4040            .arg(&mut sel_w)
4041            .arg(&ne)
4042            .arg(&nu);
4043        unsafe {
4044            b.launch(cfg)?;
4045        }
4046        // single-sync readback: sel (i32) at offset 0, w (f32) at offset n*4 of the pinned stage.
4047        let bytes = n * 8;
4048        let mut guard = self.router_stage.lock().unwrap();
4049        if guard.as_ref().map(|p| p.cap < bytes).unwrap_or(true) {
4050            *guard = Some(PinnedStage::new(bytes.max(4096))?);
4051        }
4052        let stage = guard.as_mut().unwrap();
4053        let (si, sw) = unsafe {
4054            (
4055                std::slice::from_raw_parts_mut(stage.ptr as *mut i32, n),
4056                std::slice::from_raw_parts_mut(stage.ptr.add(n * 4) as *mut f32, n),
4057            )
4058        };
4059        self.gpu.stream().memcpy_dtoh(&sel_idx, si)?; // async (pinned dst)
4060        self.gpu.stream().memcpy_dtoh(&sel_w, sw)?; // async (pinned dst)
4061        self.gpu.stream().synchronize()?; // ONE sync for both
4062        Ok((si.iter().map(|&i| i as u32).collect(), sw.to_vec()))
4063    }
4064
4065    /// Device sigmoid router for Step-3.7 / DeepSeek-V3-class MoEs. `correction_bias` is added
4066    /// only to the top-k key; returned weights use the un-biased sigmoid score. `active` masks
4067    /// original expert ids before top-k. Exact key ties choose the smaller original id.
4068    #[allow(clippy::too_many_arguments)]
4069    pub fn moe_router_sigmoid_topk(
4070        &self,
4071        logits: &CudaSlice<f32>,
4072        t: usize,
4073        n_expert: usize,
4074        n_used: usize,
4075        active_count: usize,
4076        correction_bias: &CudaSlice<f32>,
4077        active: &CudaSlice<u8>,
4078        scaling_factor: f32,
4079        route_norm: bool,
4080    ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4081        crate::sigrouter_contract::validate_active_count(n_used, active_count)?;
4082        if n_expert == 0 || n_expert > 1024 || n_used == 0 || n_used > n_expert {
4083            return Err(format!(
4084                "sigmoid router shape unsupported: n_expert={n_expert}, n_used={n_used}",
4085            )
4086            .into());
4087        }
4088        if logits.len() < t * n_expert
4089            || correction_bias.len() != n_expert
4090            || active.len() != n_expert
4091        {
4092            return Err(format!(
4093                "sigmoid router buffer mismatch: logits={} bias={} active={} expected logits>={} row={}",
4094                logits.len(), correction_bias.len(), active.len(), t * n_expert, n_expert,
4095            ).into());
4096        }
4097        let f = self.func("moe_router_sigmoid_topk_f32");
4098        let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?;
4099        let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?;
4100        let threads = n_expert.div_ceil(32) * 32;
4101        let cfg = LaunchConfig {
4102            grid_dim: (t as u32, 1, 1),
4103            block_dim: (threads as u32, 1, 1),
4104            shared_mem_bytes: 0,
4105        };
4106        let (ne, nu, rn) = (n_expert as i32, n_used as i32, i32::from(route_norm));
4107        let __s_b = self.gpu.stream();
4108        let mut b = __s_b.launch_builder(&f);
4109        b.arg(logits)
4110            .arg(correction_bias)
4111            .arg(active)
4112            .arg(&mut sel_idx)
4113            .arg(&mut sel_w)
4114            .arg(&ne)
4115            .arg(&nu)
4116            .arg(&scaling_factor)
4117            .arg(&rn);
4118        unsafe {
4119            b.launch(cfg)?;
4120        }
4121        Ok((sel_idx, sel_w))
4122    }
4123
4124    /// Single-sync pinned readback twin of `moe_router_sigmoid_topk`. This preserves the existing
4125    /// grouped/staged dispatch contract while replacing the full-logit DtoH plus host sigmoid/sort.
4126    #[allow(clippy::too_many_arguments)]
4127    pub fn moe_router_sigmoid_topk_host(
4128        &self,
4129        logits: &CudaSlice<f32>,
4130        t: usize,
4131        n_expert: usize,
4132        n_used: usize,
4133        active_count: usize,
4134        correction_bias: &CudaSlice<f32>,
4135        active: &CudaSlice<u8>,
4136        scaling_factor: f32,
4137        route_norm: bool,
4138    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
4139        let (sel_idx, sel_w) = self.moe_router_sigmoid_topk(
4140            logits,
4141            t,
4142            n_expert,
4143            n_used,
4144            active_count,
4145            correction_bias,
4146            active,
4147            scaling_factor,
4148            route_norm,
4149        )?;
4150        let n = t * n_used;
4151        let bytes = n * 8;
4152        let mut guard = self.router_stage.lock().unwrap();
4153        if guard.as_ref().map(|p| p.cap < bytes).unwrap_or(true) {
4154            *guard = Some(PinnedStage::new(bytes.max(4096))?);
4155        }
4156        let stage = guard.as_mut().unwrap();
4157        let (si, sw) = unsafe {
4158            (
4159                std::slice::from_raw_parts_mut(stage.ptr as *mut i32, n),
4160                std::slice::from_raw_parts_mut(stage.ptr.add(n * 4) as *mut f32, n),
4161            )
4162        };
4163        self.gpu.stream().memcpy_dtoh(&sel_idx, si)?;
4164        self.gpu.stream().memcpy_dtoh(&sel_w, sw)?;
4165        self.gpu.stream().synchronize()?;
4166        Ok((si.iter().map(|&i| i as u32).collect(), sw.to_vec()))
4167    }
4168
4169    /// EDGE-1 §C.2: async H2D of `host_bytes` into `scratch[off..]` on the COPY stream, returning a
4170    /// recorded event the compute stream can `wait` on before the dependent GEMM. Used for in-token
4171    /// expert prefetch (pipeline by one). `host_bytes` should be pinned for a true DMA (§C.1).
4172    pub fn stage_expert_async(
4173        &self,
4174        host_bytes: &[u8],
4175        scratch: &mut CudaSlice<u8>,
4176        off: usize,
4177    ) -> Result<cudarc::driver::CudaEvent, Box<dyn std::error::Error>> {
4178        let mut dst = scratch.slice_mut(off..off + host_bytes.len());
4179        self.copy_stream.memcpy_htod(host_bytes, &mut dst)?;
4180        Ok(self.copy_stream.record_event(None)?)
4181    }
4182
4183    /// Make the compute stream wait for an async copy event (the consumer side of `stage_expert_async`).
4184    pub fn compute_wait(
4185        &self,
4186        ev: &cudarc::driver::CudaEvent,
4187    ) -> Result<(), Box<dyn std::error::Error>> {
4188        self.gpu.stream().wait(ev)?;
4189        Ok(())
4190    }
4191
4192    /// qmatvec over a byte sub-range of a (resident/scratch) CudaSlice<u8> holding ONE expert
4193    /// matrix. x is a CudaView<f32> (a sliced row of z, or a sliced activation). Reuses the
4194    /// validated qmatvec_f32 dequant path (NOT a fast path — the correctness gate). The
4195    /// CudaView base+offset pointer is honored by the launch arg.
4196    pub fn qmatvec_view(
4197        &self,
4198        w: &CudaSlice<u8>,
4199        range: std::ops::Range<usize>,
4200        x: &cudarc::driver::CudaView<f32>,
4201        m: usize,
4202        in_f: usize,
4203        out_f: usize,
4204        qtype: i32,
4205        row_bytes: usize,
4206    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4207        let f = self.func("qmatvec_f32");
4208        let wv = w.slice(range); // CudaView<u8>, offset honored
4209        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
4210        let cfg = LaunchConfig {
4211            grid_dim: (out_f as u32, m as u32, 1),
4212            block_dim: (256, 1, 1),
4213            shared_mem_bytes: 0,
4214        };
4215        let (inf, outf, mi, qt, rb) =
4216            (in_f as i32, out_f as i32, m as i32, qtype, row_bytes as i64);
4217        let __s_b = self.gpu.stream();
4218        let mut b = __s_b.launch_builder(&f);
4219        b.arg(&wv)
4220            .arg(x)
4221            .arg(&mut y)
4222            .arg(&inf)
4223            .arg(&outf)
4224            .arg(&mi)
4225            .arg(&qt)
4226            .arg(&rb);
4227        unsafe {
4228            b.launch(cfg)?;
4229        }
4230        Ok(y)
4231    }
4232
4233    /// STAGE-2 GROUPED DECODE (2026-07-04): one MoE layer's gate+up+SiLU for all `n_used` routed
4234    /// experts of ONE token in ONE launch (replaces 8x qmatvec(gate) + 8x qmatvec(up) + 8x
4235    /// silu_mul = 24 launches). `gp`/`up` are the 8 expert weight-block device pointers (SLRU
4236    /// cache slots — fixed-address, stable for the launch). Returns act [n_used, n_ff].
4237    /// BIT-IDENTICAL to the sequential chain: each dot reproduces qmatvec_f32's exact 256-thread
4238    /// reduction; the SiLU epilogue is silu_mul_f32's exact expression (see kernel header).
4239    #[allow(clippy::too_many_arguments)]
4240    /// dp4a q8 twins (MoE expert dp4a arc, 2026-07-06): same contract as the _f32 versions but
4241    /// consume a PRE-QUANTIZED q8_1 activation. FP-order differs from _f32 (int dot + warp tree)
4242    /// — the argmax/stream-identity battery arbitrates; MEMRA_MOE_Q8=0 restores f32.
4243    pub fn moe_gate_up_silu8_q8(
4244        &self,
4245        gp: WPtr8,
4246        up: WPtr8,
4247        aq: &CudaSlice<i8>,
4248        ad: &CudaSlice<f32>,
4249        in_f: usize,
4250        n_ff: usize,
4251        n_used: usize,
4252        qt_g: i32,
4253        qt_u: i32,
4254        rb_g: usize,
4255        rb_u: usize,
4256    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4257        let f = self.func("moe_gate_up_silu8_q8");
4258        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
4259        let cfg = LaunchConfig {
4260            grid_dim: (n_ff as u32, n_used as u32, 1),
4261            block_dim: (32, 1, 1),
4262            shared_mem_bytes: 0,
4263        };
4264        let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
4265        let __s_b = self.gpu.stream();
4266        let mut b = __s_b.launch_builder(&f);
4267        b.arg(&gp)
4268            .arg(&up)
4269            .arg(aq)
4270            .arg(ad)
4271            .arg(&mut act)
4272            .arg(&inf)
4273            .arg(&nff)
4274            .arg(&qt_g)
4275            .arg(&qt_u)
4276            .arg(&rbg)
4277            .arg(&rbu);
4278        unsafe {
4279            b.launch(cfg)?;
4280        }
4281        Ok(act)
4282    }
4283
4284    #[allow(clippy::too_many_arguments)]
4285    pub fn moe_down8_fma_q8(
4286        &self,
4287        dp: WPtr8,
4288        w: F32x8,
4289        aq2: &CudaSlice<i8>,
4290        ad2: &CudaSlice<f32>,
4291        dst: &mut cudarc::driver::CudaViewMut<f32>,
4292        in_f: usize,
4293        out_f: usize,
4294        n_used: usize,
4295        qt: i32,
4296        rb: usize,
4297    ) -> Result<(), Box<dyn std::error::Error>> {
4298        let f = self.func("moe_down8_fma_q8");
4299        let cfg = LaunchConfig {
4300            grid_dim: (out_f as u32, 1, 1),
4301            block_dim: (32, 1, 1),
4302            shared_mem_bytes: 0,
4303        };
4304        let (inf, outf, nu, rbi) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
4305        let __s_b = self.gpu.stream();
4306        let mut b = __s_b.launch_builder(&f);
4307        b.arg(&dp)
4308            .arg(&w)
4309            .arg(aq2)
4310            .arg(ad2)
4311            .arg(dst)
4312            .arg(&inf)
4313            .arg(&outf)
4314            .arg(&nu)
4315            .arg(&qt)
4316            .arg(&rbi);
4317        unsafe {
4318            b.launch(cfg)?;
4319        }
4320        Ok(())
4321    }
4322
4323    /// q8 sequential expert matvec (staged path twin of qmatvec_view for IQ3_S/IQ4_XS).
4324    pub fn qmatvec_expert_q8(
4325        &self,
4326        w: &CudaSlice<u8>,
4327        range: std::ops::Range<usize>,
4328        aq: &CudaSlice<i8>,
4329        ad: &CudaSlice<f32>,
4330        m: usize,
4331        in_f: usize,
4332        out_f: usize,
4333        qtype: i32,
4334        row_bytes: usize,
4335    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4336        let f = self.func("qmatvec_expert_q8");
4337        let wv = w.slice(range);
4338        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
4339        const ROWS: u32 = 4; // MEMRA_MMVQ_ROWS
4340        let cfg = LaunchConfig {
4341            grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, m as u32, 1),
4342            block_dim: (32, ROWS, 1),
4343            shared_mem_bytes: 0,
4344        };
4345        let (inf, outf, mi, rbi) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
4346        let __s_b = self.gpu.stream();
4347        let mut b = __s_b.launch_builder(&f);
4348        b.arg(&wv)
4349            .arg(aq)
4350            .arg(ad)
4351            .arg(&mut y)
4352            .arg(&inf)
4353            .arg(&outf)
4354            .arg(&mi)
4355            .arg(&qtype)
4356            .arg(&rbi);
4357        unsafe {
4358            b.launch(cfg)?;
4359        }
4360        Ok(y)
4361    }
4362
4363    pub fn moe_gate_up_silu8(
4364        &self,
4365        gp: WPtr8,
4366        up: WPtr8,
4367        x: &cudarc::driver::CudaView<f32>,
4368        in_f: usize,
4369        n_ff: usize,
4370        n_used: usize,
4371        qt_g: i32,
4372        qt_u: i32,
4373        rb_g: usize,
4374        rb_u: usize,
4375    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4376        let f = self.func("moe_gate_up_silu8_f32");
4377        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?; // fully overwritten
4378        let cfg = LaunchConfig {
4379            grid_dim: (n_ff as u32, n_used as u32, 1),
4380            block_dim: (256, 1, 1),
4381            shared_mem_bytes: 0,
4382        };
4383        let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
4384        let __s_b = self.gpu.stream();
4385        let mut b = __s_b.launch_builder(&f);
4386        b.arg(&gp)
4387            .arg(&up)
4388            .arg(x)
4389            .arg(&mut act)
4390            .arg(&inf)
4391            .arg(&nff)
4392            .arg(&qt_g)
4393            .arg(&qt_u)
4394            .arg(&rbg)
4395            .arg(&rbu);
4396        unsafe {
4397            b.launch(cfg)?;
4398        }
4399        Ok(act)
4400    }
4401
4402    /// STAGE-2 GROUPED DECODE: one MoE layer's down-proj + weighted accumulation for all `n_used`
4403    /// routed experts in ONE launch (replaces 8x qmatvec(down) + 8x axpy = 16 launches), writing
4404    /// the token's moe_out row DIRECTLY (`dst` is the zeroed row; the in-kernel slot-ordered
4405    /// __fmaf_rn chain starting at 0.0f reproduces the sequential axpy_f32 accumulation into the
4406    /// zeroed row bit-for-bit — the A2 byte-identity scheme at m=1).
4407    #[allow(clippy::too_many_arguments)]
4408    pub fn moe_down8_fma_into(
4409        &self,
4410        dp: WPtr8,
4411        w: F32x8,
4412        act: &CudaSlice<f32>,
4413        dst: &mut cudarc::driver::CudaViewMut<f32>,
4414        in_f: usize,
4415        out_f: usize,
4416        n_used: usize,
4417        qt: i32,
4418        rb: usize,
4419    ) -> Result<(), Box<dyn std::error::Error>> {
4420        let f = self.func("moe_down8_fma_f32");
4421        let cfg = LaunchConfig {
4422            grid_dim: (out_f as u32, 1, 1),
4423            block_dim: (256, 1, 1),
4424            shared_mem_bytes: 0,
4425        };
4426        let (inf, outf, nu, rbv) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
4427        let __s_b = self.gpu.stream();
4428        let mut b = __s_b.launch_builder(&f);
4429        b.arg(&dp)
4430            .arg(&w)
4431            .arg(act)
4432            .arg(dst)
4433            .arg(&inf)
4434            .arg(&outf)
4435            .arg(&nu)
4436            .arg(&qt)
4437            .arg(&rbv);
4438        unsafe {
4439            b.launch(cfg)?;
4440        }
4441        Ok(())
4442    }
4443
4444    /// LAUNCH-STRUCTURE STAGE 3: device-dispatch twin of `moe_gate_up_silu8` for FULLY-RESIDENT
4445    /// layers. The expert ids come from the router kernel's DEVICE `sel` output (no DtoH) and the
4446    /// weight pointers from the per-layer device table `[3, n_expert]` of slot base addresses.
4447    /// BIT-IDENTICAL math (same grid/block/reduction; only the pointer/id source differs).
4448    #[allow(clippy::too_many_arguments)]
4449    /// dp4a q8 twin of the _dev pair (resident-experts arc).
4450    ///
4451    /// GEOMETRY VARIANTS (multirow/occupancy arc 2026-07-05): all outputs are BIT-IDENTICAL to
4452    /// the base one-warp-per-(row,slot) kernel (same expert_dot_g g-order + warp tree per row;
4453    /// down's FMA chain stays slot-ordered serial). Seams:
4454    ///   MEMRA_MOE_DEVQ8_GU   = 0(base) | 1 | 2 | 4 -> _r{1,2,4} multirow twin (RPW rows/warp)
4455    ///                       | s2 (gate/up warp split) | s2z (s2 + WPB rows packed per block)
4456    ///                       | gs4 (gate/up x low/high-group 4-warp split, nsb==64 only)
4457    ///                       | u64 (nsb==64 unrolled ILP twin, geometry unchanged)
4458    ///   MEMRA_MOE_DEVQ8_WPB  = warps per block for _r twins / z-rows for s2z (default 4)
4459    ///   MEMRA_MOE_DEVQ8_DOWN = auto(default: w8h2 when in_f==512 & n_used<=8 — measured +3.8%
4460    ///                       decode on 35B/G7e) | 0 (base one-warp serial-slot) | 1 | 2 | 4 ->
4461    ///                       _w8r{1,2,4} slot-parallel twin | h2 (half-warp dual-row, nsb==16
4462    ///                       only) | w8h2 (h2 x slot-parallel)
4463    #[allow(clippy::too_many_arguments)]
4464    /// MoE PREFILL pair-batch matvec: one launch covers all (token,expert) pairs for one proj.
4465    #[allow(clippy::too_many_arguments)]
4466    pub fn moe_pairs_matvec_q8(
4467        &self,
4468        table: &CudaSlice<u64>,
4469        proj: i32,
4470        pair_tok: &CudaSlice<i32>,
4471        pair_ex: &CudaSlice<i32>,
4472        aq: &CudaSlice<i8>,
4473        ad: &CudaSlice<f32>,
4474        in_f: usize,
4475        out_f: usize,
4476        n_expert: usize,
4477        n_pairs: usize,
4478        qtype: i32,
4479        row_bytes: usize,
4480    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4481        let f = self.func("moe_pairs_matvec_q8");
4482        let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
4483        const ROWS: u32 = 4;
4484        let cfg = LaunchConfig {
4485            grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_pairs as u32, 1),
4486            block_dim: (32, ROWS, 1),
4487            shared_mem_bytes: 0,
4488        };
4489        let (inf, outf, ne, np, rbi) = (
4490            in_f as i32,
4491            out_f as i32,
4492            n_expert as i32,
4493            n_pairs as i32,
4494            row_bytes as i64,
4495        );
4496        let __s_b = self.gpu.stream();
4497        let mut b = __s_b.launch_builder(&f);
4498        b.arg(table)
4499            .arg(&proj)
4500            .arg(pair_tok)
4501            .arg(pair_ex)
4502            .arg(aq)
4503            .arg(ad)
4504            .arg(&mut y)
4505            .arg(&inf)
4506            .arg(&outf)
4507            .arg(&ne)
4508            .arg(&np)
4509            .arg(&qtype)
4510            .arg(&rbi);
4511        unsafe {
4512            b.launch(cfg)?;
4513        }
4514        Ok(y)
4515    }
4516
4517    /// Expert-major pair matvec (weight-reuse across each expert's token group).
4518    #[allow(clippy::too_many_arguments)]
4519    pub fn moe_pairs_matvec_q8_em(
4520        &self,
4521        table: &CudaSlice<u64>,
4522        proj: i32,
4523        ex_ids: &CudaSlice<i32>,
4524        ex_off: &CudaSlice<i32>,
4525        ex_pairs: &CudaSlice<i32>,
4526        pair_tok: &CudaSlice<i32>,
4527        aq: &CudaSlice<i8>,
4528        ad: &CudaSlice<f32>,
4529        in_f: usize,
4530        out_f: usize,
4531        n_expert: usize,
4532        n_active: usize,
4533        n_pairs: usize,
4534        qtype: i32,
4535        row_bytes: usize,
4536    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4537        let f = self.func("moe_pairs_matvec_q8_em");
4538        let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
4539        const ROWS: u32 = 4;
4540        let cfg = LaunchConfig {
4541            grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_active as u32, 1),
4542            block_dim: (32, ROWS, 1),
4543            shared_mem_bytes: 0,
4544        };
4545        let (inf, outf, ne, na, rbi) = (
4546            in_f as i32,
4547            out_f as i32,
4548            n_expert as i32,
4549            n_active as i32,
4550            row_bytes as i64,
4551        );
4552        let __s_b = self.gpu.stream();
4553        let mut b = __s_b.launch_builder(&f);
4554        b.arg(table)
4555            .arg(&proj)
4556            .arg(ex_ids)
4557            .arg(ex_off)
4558            .arg(ex_pairs)
4559            .arg(pair_tok)
4560            .arg(aq)
4561            .arg(ad)
4562            .arg(&mut y)
4563            .arg(&inf)
4564            .arg(&outf)
4565            .arg(&ne)
4566            .arg(&na)
4567            .arg(&qtype)
4568            .arg(&rbi);
4569        unsafe {
4570            b.launch(cfg)?;
4571        }
4572        Ok(y)
4573    }
4574
4575    // Decode-once expert-major MMQ (rung 3). Same CSR inputs/geometry as _em; kernel dequants each
4576    // weight group once per (row,group) then dp4a's across the expert's token group.
4577    #[allow(clippy::too_many_arguments)]
4578    pub fn moe_pairs_matvec_q8_dec(
4579        &self,
4580        table: &CudaSlice<u64>,
4581        proj: i32,
4582        ex_ids: &CudaSlice<i32>,
4583        ex_off: &CudaSlice<i32>,
4584        ex_pairs: &CudaSlice<i32>,
4585        pair_tok: &CudaSlice<i32>,
4586        aq: &CudaSlice<i8>,
4587        ad: &CudaSlice<f32>,
4588        in_f: usize,
4589        out_f: usize,
4590        n_expert: usize,
4591        n_active: usize,
4592        n_pairs: usize,
4593        qtype: i32,
4594        row_bytes: usize,
4595    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4596        let f = self.func("moe_pairs_matvec_q8_dec");
4597        let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
4598        const ROWS: u32 = 4;
4599        let cfg = LaunchConfig {
4600            grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_active as u32, 1),
4601            block_dim: (32, ROWS, 1),
4602            shared_mem_bytes: 0,
4603        };
4604        let (inf, outf, ne, na, rbi) = (
4605            in_f as i32,
4606            out_f as i32,
4607            n_expert as i32,
4608            n_active as i32,
4609            row_bytes as i64,
4610        );
4611        let __s_b = self.gpu.stream();
4612        let mut b = __s_b.launch_builder(&f);
4613        b.arg(table)
4614            .arg(&proj)
4615            .arg(ex_ids)
4616            .arg(ex_off)
4617            .arg(ex_pairs)
4618            .arg(pair_tok)
4619            .arg(aq)
4620            .arg(ad)
4621            .arg(&mut y)
4622            .arg(&inf)
4623            .arg(&outf)
4624            .arg(&ne)
4625            .arg(&na)
4626            .arg(&qtype)
4627            .arg(&rbi);
4628        unsafe {
4629            b.launch(cfg)?;
4630        }
4631        Ok(y)
4632    }
4633
4634    pub fn moe_pairs_gelu_mul(
4635        &self,
4636        gate: &CudaSlice<f32>,
4637        up: &CudaSlice<f32>,
4638        n: usize,
4639    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4640        let f = self.func("moe_pairs_gelu_mul");
4641        let mut act = self.alloc_uninit::<f32>(n)?;
4642        let cfg = LaunchConfig::for_num_elems(n as u32);
4643        let nl = n as i64;
4644        let __s_b = self.gpu.stream();
4645        let mut b = __s_b.launch_builder(&f);
4646        b.arg(gate).arg(up).arg(&mut act).arg(&nl);
4647        unsafe {
4648            b.launch(cfg)?;
4649        }
4650        Ok(act)
4651    }
4652
4653    pub fn moe_pairs_silu_mul(
4654        &self,
4655        gate: &CudaSlice<f32>,
4656        up: &CudaSlice<f32>,
4657        n: usize,
4658    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4659        let f = self.func("moe_pairs_silu_mul");
4660        let mut act = self.alloc_uninit::<f32>(n)?;
4661        let cfg = LaunchConfig::for_num_elems(n as u32);
4662        let nl = n as i64;
4663        let __s_b = self.gpu.stream();
4664        let mut b = __s_b.launch_builder(&f);
4665        b.arg(gate).arg(up).arg(&mut act).arg(&nl);
4666        unsafe {
4667            b.launch(cfg)?;
4668        }
4669        Ok(act)
4670    }
4671
4672    #[allow(clippy::too_many_arguments)]
4673    pub fn moe_pairs_scatter(
4674        &self,
4675        y_down: &CudaSlice<f32>,
4676        pair_w: &CudaSlice<f32>,
4677        tok_pair_off: &CudaSlice<i32>,
4678        tok_pair_ids: &CudaSlice<i32>,
4679        moe_out: &mut CudaSlice<f32>,
4680        t: usize,
4681        n_embd: usize,
4682    ) -> Result<(), Box<dyn std::error::Error>> {
4683        let f = self.func("moe_pairs_scatter");
4684        let cfg = LaunchConfig {
4685            grid_dim: (((n_embd + 255) / 256) as u32, t as u32, 1),
4686            block_dim: (256, 1, 1),
4687            shared_mem_bytes: 0,
4688        };
4689        let ne = n_embd as i32;
4690        let __s_b = self.gpu.stream();
4691        let mut b = __s_b.launch_builder(&f);
4692        b.arg(y_down)
4693            .arg(pair_w)
4694            .arg(tok_pair_off)
4695            .arg(tok_pair_ids)
4696            .arg(moe_out)
4697            .arg(&ne);
4698        unsafe {
4699            b.launch(cfg)?;
4700        }
4701        Ok(())
4702    }
4703
4704    /// gemma4 GELU twin of moe_gate_up_silu8_dev_q8 (base geometry — slot-packed j8/j8r2
4705    /// twins probed 2026-08-01 g26 decode dig: bit-identical rows, -2.5%/-2.9% whole-model
4706    /// decode x3 interleaved -> refuted and killed; research/g26-decode-20260801/receipts.md).
4707    #[allow(clippy::too_many_arguments)]
4708    pub fn moe_gate_up_gelu8_dev_q8(
4709        &self,
4710        table: &CudaSlice<u64>,
4711        sel: &cudarc::driver::CudaView<i32>,
4712        aq: &CudaSlice<i8>,
4713        ad: &CudaSlice<f32>,
4714        in_f: usize,
4715        n_ff: usize,
4716        n_used: usize,
4717        n_expert: usize,
4718        qt_g: i32,
4719        qt_u: i32,
4720        rb_g: usize,
4721        rb_u: usize,
4722    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4723        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
4724        let (inf, nff, ne, rbg, rbu) = (
4725            in_f as i32,
4726            n_ff as i32,
4727            n_expert as i32,
4728            rb_g as i64,
4729            rb_u as i64,
4730        );
4731        let f = self.func("moe_gate_up_gelu8_dev_q8");
4732        let cfg = LaunchConfig {
4733            grid_dim: (n_ff as u32, n_used as u32, 1),
4734            block_dim: (32, 1, 1),
4735            shared_mem_bytes: 0,
4736        };
4737        let __s_b = self.gpu.stream();
4738        let mut b = __s_b.launch_builder(&f);
4739        b.arg(table)
4740            .arg(sel)
4741            .arg(aq)
4742            .arg(ad)
4743            .arg(&mut act)
4744            .arg(&inf)
4745            .arg(&nff)
4746            .arg(&ne)
4747            .arg(&qt_g)
4748            .arg(&qt_u)
4749            .arg(&rbg)
4750            .arg(&rbu);
4751        unsafe {
4752            b.launch(cfg)?;
4753        }
4754        Ok(act)
4755    }
4756
4757    /// gemma4 GELU rows twin (verify): one launch over (n_ff, n_used, t).
4758    #[allow(clippy::too_many_arguments)]
4759    pub fn moe_gate_up_gelu8_dev_q8_rows(
4760        &self,
4761        table: &CudaSlice<u64>,
4762        sel: &CudaSlice<i32>,
4763        aq: &CudaSlice<i8>,
4764        ad: &CudaSlice<f32>,
4765        t: usize,
4766        in_f: usize,
4767        n_ff: usize,
4768        n_used: usize,
4769        n_expert: usize,
4770        qt_g: i32,
4771        qt_u: i32,
4772        rb_g: usize,
4773        rb_u: usize,
4774    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4775        let mut act = self.alloc_uninit::<f32>(t * n_used * n_ff)?;
4776        let (inf, nff, ne, rbg, rbu, nu) = (
4777            in_f as i32,
4778            n_ff as i32,
4779            n_expert as i32,
4780            rb_g as i64,
4781            rb_u as i64,
4782            n_used as i32,
4783        );
4784        let f = self.func("moe_gate_up_gelu8_dev_q8_rows");
4785        let cfg = LaunchConfig {
4786            grid_dim: (n_ff as u32, n_used as u32, t as u32),
4787            block_dim: (32, 1, 1),
4788            shared_mem_bytes: 0,
4789        };
4790        let __s_b = self.gpu.stream();
4791        let mut b = __s_b.launch_builder(&f);
4792        b.arg(table)
4793            .arg(sel)
4794            .arg(aq)
4795            .arg(ad)
4796            .arg(&mut act)
4797            .arg(&inf)
4798            .arg(&nff)
4799            .arg(&ne)
4800            .arg(&qt_g)
4801            .arg(&qt_u)
4802            .arg(&rbg)
4803            .arg(&rbu)
4804            .arg(&nu);
4805        unsafe {
4806            b.launch(cfg)?;
4807        }
4808        Ok(act)
4809    }
4810
4811    /// gemma4 GELU CSR twin (verify dedup: owner block serves every pair of its expert).
4812    #[allow(clippy::too_many_arguments)]
4813    pub fn moe_gate_up_gelu8_dev_q8_csr(
4814        &self,
4815        table: &CudaSlice<u64>,
4816        sel: &CudaSlice<i32>,
4817        aq: &CudaSlice<i8>,
4818        ad: &CudaSlice<f32>,
4819        n_pairs: usize,
4820        in_f: usize,
4821        n_ff: usize,
4822        n_used: usize,
4823        n_expert: usize,
4824        qt_g: i32,
4825        qt_u: i32,
4826        rb_g: usize,
4827        rb_u: usize,
4828    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4829        let mut act = self.alloc_uninit::<f32>(n_pairs * n_ff)?;
4830        let (inf, nff, ne, rbg, rbu, nu, npi) = (
4831            in_f as i32,
4832            n_ff as i32,
4833            n_expert as i32,
4834            rb_g as i64,
4835            rb_u as i64,
4836            n_used as i32,
4837            n_pairs as i32,
4838        );
4839        let f = self.func("moe_gate_up_gelu8_dev_q8_csr");
4840        let cfg = LaunchConfig {
4841            grid_dim: (n_ff as u32, n_pairs as u32, 1),
4842            block_dim: (32, 1, 1),
4843            shared_mem_bytes: 0,
4844        };
4845        let __s_b = self.gpu.stream();
4846        let mut b = __s_b.launch_builder(&f);
4847        b.arg(table)
4848            .arg(sel)
4849            .arg(aq)
4850            .arg(ad)
4851            .arg(&mut act)
4852            .arg(&inf)
4853            .arg(&nff)
4854            .arg(&ne)
4855            .arg(&qt_g)
4856            .arg(&qt_u)
4857            .arg(&rbg)
4858            .arg(&rbu)
4859            .arg(&nu)
4860            .arg(&npi);
4861        unsafe {
4862            b.launch(cfg)?;
4863        }
4864        Ok(act)
4865    }
4866
4867    /// gemma4 generic down rows twin (verify): one launch over (out_f, 1, t).
4868    #[allow(clippy::too_many_arguments)]
4869    pub fn moe_down8_fma_dev_q8_rows_g(
4870        &self,
4871        table: &CudaSlice<u64>,
4872        sel: &CudaSlice<i32>,
4873        w: &CudaSlice<f32>,
4874        aq2: &CudaSlice<i8>,
4875        ad2: &CudaSlice<f32>,
4876        dst: &mut CudaSlice<f32>,
4877        t: usize,
4878        in_f: usize,
4879        out_f: usize,
4880        n_used: usize,
4881        n_expert: usize,
4882        qt: i32,
4883        rb: usize,
4884    ) -> Result<(), Box<dyn std::error::Error>> {
4885        let (inf, outf, nu, ne, rbi) = (
4886            in_f as i32,
4887            out_f as i32,
4888            n_used as i32,
4889            n_expert as i32,
4890            rb as i64,
4891        );
4892        // Exact Step-3.7 B=1 shape: expose the eight independent slot dots as
4893        // eight warps, then replay the original slot-ordered FMA chain. Every
4894        // other shape retains the generic one-warp rows kernel.
4895        let step_b1_w8 = t == 1 && in_f == 1280 && out_f == 4096 && n_used == 8 && qt == QT_IQ4_XS;
4896        let f = self.func(if step_b1_w8 {
4897            "moe_down8_fma_dev_q8_rows_w8"
4898        } else {
4899            "moe_down8_fma_dev_q8_rows_g"
4900        });
4901        let cfg = LaunchConfig {
4902            grid_dim: (out_f as u32, 1, t as u32),
4903            block_dim: (32, if step_b1_w8 { 8 } else { 1 }, 1),
4904            shared_mem_bytes: 0,
4905        };
4906        let __s_b = self.gpu.stream();
4907        let mut b = __s_b.launch_builder(&f);
4908        b.arg(table)
4909            .arg(sel)
4910            .arg(w)
4911            .arg(aq2)
4912            .arg(ad2)
4913            .arg(dst)
4914            .arg(&inf)
4915            .arg(&outf)
4916            .arg(&nu)
4917            .arg(&ne)
4918            .arg(&qt)
4919            .arg(&rbi);
4920        unsafe {
4921            b.launch(cfg)?;
4922        }
4923        Ok(())
4924    }
4925
4926    /// rp_q4 microprobe (2026-07-10 verify-trunk lever): b4 GGUF-block layout vs the Q4_0
4927    /// split-plane twin on the wq-class shape. Returns (blk_us, rp_us) after asserting bitwise
4928    /// identity. Bench-only surface (rp_q4_probe bin); no production dispatch reads this.
4929    pub fn rp_probe_q4(&self, m: usize) -> Result<(f64, f64), Box<dyn std::error::Error>> {
4930        let (out_f, in_f) = (2048usize, 2816usize);
4931        let nblk = in_f / 32;
4932        let mut seed = 0x9E3779B97F4A7C15u64;
4933        let mut rng = move || {
4934            seed = seed
4935                .wrapping_mul(6364136223846793005)
4936                .wrapping_add(1442695040888963407);
4937            (seed >> 33) as u8
4938        };
4939        let mut w = vec![0u8; out_f * nblk * 18];
4940        for b in w.iter_mut() {
4941            *b = rng();
4942        }
4943        for r in 0..out_f {
4944            for g in 0..nblk {
4945                let off = (r * nblk + g) * 18;
4946                w[off] = 0x00;
4947                w[off + 1] = 0x2C; // sane half d
4948            }
4949        }
4950        let qplane = out_f * nblk * 16;
4951        let mut wrp = vec![0u8; w.len()];
4952        for r in 0..out_f {
4953            for g in 0..nblk {
4954                let src = &w[(r * nblk + g) * 18..(r * nblk + g) * 18 + 18];
4955                wrp[qplane + (r * nblk + g) * 2..qplane + (r * nblk + g) * 2 + 2]
4956                    .copy_from_slice(&src[0..2]);
4957                wrp[(r * nblk + g) * 16..(r * nblk + g) * 16 + 16].copy_from_slice(&src[2..18]);
4958            }
4959        }
4960        let w_d = self.htod_bytes(&w)?;
4961        let wrp_d = self.htod_bytes(&wrp)?;
4962        let mut aq = vec![0i8; m * in_f];
4963        for v in aq.iter_mut() {
4964            *v = rng() as i8;
4965        }
4966        let aq_d = self.htod_i8(&aq)?;
4967        let ad_d = self.htod(&vec![0.03125f32; m * nblk])?;
4968        let mut y0 = self.alloc_uninit::<f32>(m * out_f)?;
4969        let mut y1 = self.alloc_uninit::<f32>(m * out_f)?;
4970        const RPB: u32 = 4;
4971        let cfg = LaunchConfig {
4972            grid_dim: ((out_f as u32).div_ceil(RPB), 1, 1),
4973            block_dim: (32, RPB, 1),
4974            shared_mem_bytes: 0,
4975        };
4976        let (inf, outf, mi) = (in_f as i32, out_f as i32, m as i32);
4977        let (rb, qp) = ((nblk * 18) as i64, qplane as i64);
4978        let fb = self.func("qmatvec_q4_0_mmvq_b4");
4979        let fr = self.func("qmatvec_q4_0_mmvq_b4_rp");
4980        {
4981            let __s_b = self.gpu.stream();
4982            let mut b = __s_b.launch_builder(&fb);
4983            b.arg(&w_d)
4984                .arg(&aq_d)
4985                .arg(&ad_d)
4986                .arg(&mut y0)
4987                .arg(&inf)
4988                .arg(&outf)
4989                .arg(&mi)
4990                .arg(&rb);
4991            unsafe {
4992                b.launch(cfg)?;
4993            }
4994            let __s_b = self.gpu.stream();
4995            let mut b = __s_b.launch_builder(&fr);
4996            b.arg(&wrp_d)
4997                .arg(&aq_d)
4998                .arg(&ad_d)
4999                .arg(&mut y1)
5000                .arg(&inf)
5001                .arg(&outf)
5002                .arg(&mi)
5003                .arg(&qp);
5004            unsafe {
5005                b.launch(cfg)?;
5006            }
5007        }
5008        self.gpu.stream().synchronize()?;
5009        let (h0, h1) = (self.dtoh(&y0)?, self.dtoh(&y1)?);
5010        let nd = h0
5011            .iter()
5012            .zip(&h1)
5013            .filter(|(a, b)| a.to_bits() != b.to_bits())
5014            .count();
5015        if nd != 0 {
5016            return Err(format!("rp twin not bitwise: {nd}/{} diffs", h0.len()).into());
5017        }
5018        let mut time = |rp: bool| -> Result<f64, Box<dyn std::error::Error>> {
5019            self.gpu.stream().synchronize()?;
5020            let t0 = std::time::Instant::now();
5021            for _ in 0..500 {
5022                if rp {
5023                    let __s_b = self.gpu.stream();
5024                    let mut b = __s_b.launch_builder(&fr);
5025                    b.arg(&wrp_d)
5026                        .arg(&aq_d)
5027                        .arg(&ad_d)
5028                        .arg(&mut y1)
5029                        .arg(&inf)
5030                        .arg(&outf)
5031                        .arg(&mi)
5032                        .arg(&qp);
5033                    unsafe {
5034                        b.launch(cfg)?;
5035                    }
5036                } else {
5037                    let __s_b = self.gpu.stream();
5038                    let mut b = __s_b.launch_builder(&fb);
5039                    b.arg(&w_d)
5040                        .arg(&aq_d)
5041                        .arg(&ad_d)
5042                        .arg(&mut y0)
5043                        .arg(&inf)
5044                        .arg(&outf)
5045                        .arg(&mi)
5046                        .arg(&rb);
5047                    unsafe {
5048                        b.launch(cfg)?;
5049                    }
5050                }
5051            }
5052            self.gpu.stream().synchronize()?;
5053            Ok(t0.elapsed().as_secs_f64() * 1e6 / 500.0)
5054        };
5055        let _ = time(false)?;
5056        let _ = time(true)?; // warm
5057        Ok((time(false)?, time(true)?))
5058    }
5059
5060    /// Build the Q4_0 split-plane decode mirror for a 2D Quant tensor (device-side permutation,
5061    /// q4_0_split_rp_build). Raw bytes stay resident (prefill/gemm/Stage-A); the m<=8 decode
5062    /// dispatch prefers the mirror (_rp twins). No-op unless (Q4_0, 2D, mirror absent).
5063    /// VRAM cost == the tensor's weight size. MEMRA_Q4RP=0 disables at the call sites.
5064    pub fn build_q4_rp4(
5065        &self,
5066        t: &mut crate::model::GpuTensor,
5067    ) -> Result<(), Box<dyn std::error::Error>> {
5068        use crate::model::GpuTensor;
5069        let GpuTensor::Quant {
5070            bytes,
5071            qtype,
5072            row_bytes,
5073            ne,
5074            rp4,
5075            ..
5076        } = t
5077        else {
5078            return Ok(());
5079        };
5080        if *qtype != QT_Q4_0 || rp4.is_some() || ne.len() != 2 {
5081            return Ok(());
5082        }
5083        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5084        if in_f % 32 != 0 || *row_bytes != (in_f / 32) * 18 {
5085            return Ok(());
5086        }
5087        let nblk = in_f / 32;
5088        let mut dst = self.alloc_uninit::<u8>(out_f * nblk * 18)?;
5089        let f = self.func("q4_0_split_rp_build");
5090        let n = (out_f * nblk) as i32;
5091        let cfg = LaunchConfig {
5092            grid_dim: (((out_f * nblk) as u32).div_ceil(256), 1, 1),
5093            block_dim: (256, 1, 1),
5094            shared_mem_bytes: 0,
5095        };
5096        let (of, nb) = (out_f as i32, nblk as i32);
5097        let _ = n;
5098        let __s_b = self.gpu.stream();
5099        let mut b = __s_b.launch_builder(&f);
5100        b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
5101        unsafe {
5102            b.launch(cfg)?;
5103        }
5104        *rp4 = Some(dst);
5105        Ok(())
5106    }
5107
5108    /// Q8_0 twin of `build_q4_rp4` (H100 coalescing fix, 2026-07-26 ncu: GGUF 34B-stride
5109    /// weight loads hold Max Bandwidth at 41-46%; the split mirror makes them aligned 16B
5110    /// ldcs). Raw bytes stay resident (prefill GEMM/MMQ/fused m=1 launches read GGUF layout);
5111    /// the mmvq/batched decode arms prefer the mirror via `rp4`. Bit-identical outputs.
5112    pub fn build_q8_rp4(
5113        &self,
5114        t: &mut crate::model::GpuTensor,
5115    ) -> Result<(), Box<dyn std::error::Error>> {
5116        use crate::model::GpuTensor;
5117        let GpuTensor::Quant {
5118            bytes,
5119            qtype,
5120            row_bytes,
5121            ne,
5122            rp4,
5123            ..
5124        } = t
5125        else {
5126            return Ok(());
5127        };
5128        if *qtype != QT_Q8_0 || rp4.is_some() || ne.len() != 2 {
5129            return Ok(());
5130        }
5131        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5132        if in_f % 32 != 0 || *row_bytes != (in_f / 32) * 34 {
5133            return Ok(());
5134        }
5135        *rp4 = Some(self.build_q8_rp4_raw(bytes, in_f, out_f)?);
5136        Ok(())
5137    }
5138
5139    /// Raw rp-mirror build for gates/benches: split GGUF Q8_0 bytes into the qplane+dplane
5140    /// mirror without a GpuTensor (same kernel the loader path above uses).
5141    pub fn build_q8_rp4_raw(
5142        &self,
5143        bytes: &CudaSlice<u8>,
5144        in_f: usize,
5145        out_f: usize,
5146    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
5147        assert!(in_f % 32 == 0);
5148        let nblk = in_f / 32;
5149        let mut dst = self.alloc_uninit::<u8>(out_f * nblk * 34)?;
5150        let f = self.func("q8_0_split_rp_build");
5151        let cfg = LaunchConfig {
5152            grid_dim: (((out_f * nblk) as u32).div_ceil(256), 1, 1),
5153            block_dim: (256, 1, 1),
5154            shared_mem_bytes: 0,
5155        };
5156        let (of, nb) = (out_f as i32, nblk as i32);
5157        let __s_b = self.gpu.stream();
5158        let mut b = __s_b.launch_builder(&f);
5159        b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
5160        unsafe {
5161            b.launch(cfg)?;
5162        }
5163        Ok(dst)
5164    }
5165
5166    /// K-quant twins of `build_q8_rp4` (H100 K-quant coalescing fix, 2026-08-01 ncu on the
5167    /// q27 Q4_K_M decode: q4_K mmvq DRAM 41-54% with 65% excessive sectors, q6_K 40% with
5168    /// 78% — the 144B/210B superblock strides land every 4B weight load off-sector). The
5169    /// mirror re-packs each tensor into planes (q4_K: qs ++ 16B meta; q6_K: ql ++ qh ++
5170    /// scales ++ d — same total bytes) so every quant fetch is an aligned 16B ldcs. Raw
5171    /// bytes stay resident (prefill GEMM/dequant/Stage-A read GGUF layout); the mmvq/batched
5172    /// decode arms prefer the mirror via `rp4`. Bit-identical outputs.
5173    pub fn build_q4k_rp4(
5174        &self,
5175        t: &mut crate::model::GpuTensor,
5176    ) -> Result<(), Box<dyn std::error::Error>> {
5177        use crate::model::GpuTensor;
5178        let GpuTensor::Quant {
5179            bytes,
5180            qtype,
5181            row_bytes,
5182            ne,
5183            rp4,
5184            ..
5185        } = t
5186        else {
5187            return Ok(());
5188        };
5189        if *qtype != QT_Q4_K || rp4.is_some() || ne.len() != 2 {
5190            return Ok(());
5191        }
5192        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5193        if in_f % 256 != 0 || *row_bytes != (in_f / 256) * 144 {
5194            return Ok(());
5195        }
5196        *rp4 = Some(self.build_kq_rp4_raw(bytes, in_f, out_f, QT_Q4_K)?);
5197        Ok(())
5198    }
5199
5200    pub fn build_q6k_rp4(
5201        &self,
5202        t: &mut crate::model::GpuTensor,
5203    ) -> Result<(), Box<dyn std::error::Error>> {
5204        use crate::model::GpuTensor;
5205        let GpuTensor::Quant {
5206            bytes,
5207            qtype,
5208            row_bytes,
5209            ne,
5210            rp4,
5211            ..
5212        } = t
5213        else {
5214            return Ok(());
5215        };
5216        if *qtype != QT_Q6_K || rp4.is_some() || ne.len() != 2 {
5217            return Ok(());
5218        }
5219        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5220        if in_f % 256 != 0 || *row_bytes != (in_f / 256) * 210 {
5221            return Ok(());
5222        }
5223        *rp4 = Some(self.build_kq_rp4_raw(bytes, in_f, out_f, QT_Q6_K)?);
5224        Ok(())
5225    }
5226
5227    /// Raw K-quant rp-mirror build for gates/benches (same kernels the loader path uses).
5228    pub fn build_kq_rp4_raw(
5229        &self,
5230        bytes: &CudaSlice<u8>,
5231        in_f: usize,
5232        out_f: usize,
5233        qtype: i32,
5234    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
5235        assert!(in_f % 256 == 0);
5236        let nsbk = in_f / 256;
5237        let (sb_bytes, kname) = match qtype {
5238            QT_Q4_K => (144usize, "q4_K_split_rp_build"),
5239            QT_Q6_K => (210usize, "q6_K_split_rp_build"),
5240            _ => return Err(format!("build_kq_rp4_raw: qtype {qtype} has no rp mirror").into()),
5241        };
5242        let mut dst = self.alloc_uninit::<u8>(out_f * nsbk * sb_bytes)?;
5243        let f = self.func(kname);
5244        let cfg = LaunchConfig {
5245            grid_dim: (((out_f * nsbk) as u32).div_ceil(256), 1, 1),
5246            block_dim: (256, 1, 1),
5247            shared_mem_bytes: 0,
5248        };
5249        let (of, nb) = (out_f as i32, nsbk as i32);
5250        let __s_b = self.gpu.stream();
5251        let mut b = __s_b.launch_builder(&f);
5252        b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
5253        unsafe {
5254            b.launch(cfg)?;
5255        }
5256        Ok(dst)
5257    }
5258
5259    /// MEMRA_KQRP seam: the K-quant (q4_K/q6_K) split-plane decode mirrors at model load.
5260    /// Default follows the Q8RP convention — ON on the Hopper lane (80GB pays the mirror
5261    /// VRAM), OFF elsewhere (a 24GB card cannot hold model + mirror + KV for the big trunks).
5262    pub fn kqrp_enabled() -> bool {
5263        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5264        *ON.get_or_init(|| match std::env::var("MEMRA_KQRP").as_deref() {
5265            Ok("0") => false,
5266            Ok(_) => true,
5267            Err(_) => cfg!(memra_hopper_mma),
5268        })
5269    }
5270
5271    /// IN-PLACE split-plane swap (the 31B dense arc): build the split layout and REPLACE the
5272    /// GGUF bytes (zero extra steady-state VRAM — the transient peak is one tensor's size).
5273    /// The tensor's `rp` flag then routes every consumer (mmvq/batched `_rp` twins, the
5274    /// `qmatvec_gemm_q4_0_rp` prefill kernel). Callers gate on the fast path being active —
5275    /// the Stage-A f32 oracle (`MEMRA_FAST=0`) reads GGUF layout and must never see a swap.
5276    pub fn build_q4_rp_swap(
5277        &self,
5278        t: &mut crate::model::GpuTensor,
5279    ) -> Result<bool, Box<dyn std::error::Error>> {
5280        use crate::model::GpuTensor;
5281        // QTYPE GUARD IN THE SWAP ITSELF (gemma4 NVFP4mix prefill-NaN, 2026-08-17):
5282        // `rp4` is a SHARED field — the Q8RP walk parks Q8_0 MIRRORS there, and this
5283        // fn used to `take()` whatever it found. On a Q8_0-carrying gemma4-dense trunk
5284        // the swap hijacked those mirrors: `bytes` became split-plane in place, the
5285        // m<=16 `_rp` dispatch read them correctly (masking the corruption from every
5286        // decode pin), and every GGUF-layout prefill consumer (MMQ + GEMM) read the
5287        // fp16 d-plane as weights -> layer-0 NaN, <pad>-spam serving. Only a tensor
5288        // this fn's OWN builder serves may ever be swapped; everything else refuses
5289        // here, regardless of walk ordering.
5290        if !matches!(t, GpuTensor::Quant { qtype, .. } if *qtype == QT_Q4_0) {
5291            return Ok(false);
5292        }
5293        self.build_q4_rp4(t)?;
5294        self.gpu.stream().synchronize()?; // build kernel reads the GGUF bytes — drain BEFORE dropping them
5295        let GpuTensor::Quant { bytes, rp4, rp, .. } = t else {
5296            return Ok(false);
5297        };
5298        match rp4.take() {
5299            Some(split) => {
5300                *bytes = split; // the GGUF-layout buffer drops here
5301                *rp = true;
5302                Ok(true)
5303            }
5304            None => Ok(false),
5305        }
5306    }
5307
5308    /// MEMRA_Q4RP seam (default ON): the Q4_0 split-plane decode mirror at model load.
5309    pub fn q4rp_enabled() -> bool {
5310        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5311        *ON.get_or_init(|| {
5312            std::env::var("MEMRA_Q4RP")
5313                .map(|v| v != "0")
5314                .unwrap_or(true)
5315        })
5316    }
5317
5318    /// gemma4-E4B: dense [t][row_elems] gather of layer il's rows from the strided prologue
5319    /// buffer ([t][n_layer][n_epl]; off = il*n_epl, stride = n_layer*n_epl).
5320    pub fn copy_rows_strided(
5321        &self,
5322        src: &CudaSlice<f32>,
5323        dst: &mut CudaSlice<f32>,
5324        row_elems: usize,
5325        n_rows: usize,
5326        src_stride: usize,
5327        src_off: usize,
5328    ) -> Result<(), Box<dyn std::error::Error>> {
5329        let f = self.func("copy_rows_strided_f32");
5330        let cfg = LaunchConfig {
5331            grid_dim: (((row_elems as u32 + 255) / 256).max(1), n_rows as u32, 1),
5332            block_dim: (256, 1, 1),
5333            shared_mem_bytes: 0,
5334        };
5335        let (re, nr) = (row_elems as i32, n_rows as i32);
5336        let (st, off) = (src_stride as i64, src_off as i64);
5337        let __s_b = self.gpu.stream();
5338        let mut b = __s_b.launch_builder(&f);
5339        b.arg(src)
5340            .arg(&mut *dst)
5341            .arg(&re)
5342            .arg(&nr)
5343            .arg(&st)
5344            .arg(&off);
5345        unsafe {
5346            b.launch(cfg)?;
5347        }
5348        Ok(())
5349    }
5350
5351    /// Async device u32 store (value rides the kernel ARG — no host-memory transfer/sync).
5352    pub fn u32_set_k(
5353        &self,
5354        dst: &mut CudaSlice<u32>,
5355        v: u32,
5356        idx: usize,
5357    ) -> Result<(), Box<dyn std::error::Error>> {
5358        let f = self.func("u32_set_k");
5359        let cfg = LaunchConfig {
5360            grid_dim: (1, 1, 1),
5361            block_dim: (1, 1, 1),
5362            shared_mem_bytes: 0,
5363        };
5364        let ii = idx as i32;
5365        let __s_b = self.gpu.stream();
5366        let mut b = __s_b.launch_builder(&f);
5367        b.arg(dst).arg(&v).arg(&ii);
5368        unsafe {
5369            b.launch(cfg)?;
5370        }
5371        Ok(())
5372    }
5373
5374    /// counter += v (device-slot append advance; the +1 twin is `inc_seqlen`).
5375    pub fn i32_add_k(
5376        &self,
5377        d: &mut CudaSlice<i32>,
5378        v: i32,
5379    ) -> Result<(), Box<dyn std::error::Error>> {
5380        let f = self.func("i32_add_k");
5381        let cfg = LaunchConfig {
5382            grid_dim: (1, 1, 1),
5383            block_dim: (32, 1, 1),
5384            shared_mem_bytes: 0,
5385        };
5386        let __s_b = self.gpu.stream();
5387        let mut b = __s_b.launch_builder(&f);
5388        b.arg(d).arg(&v);
5389        unsafe {
5390            b.launch(cfg)?;
5391        }
5392        Ok(())
5393    }
5394
5395    /// pos rows from a device counter: dst[i] = ctr[0] + i (verify-stream rope positions).
5396    pub fn i32_iota_from(
5397        &self,
5398        ctr: &CudaSlice<i32>,
5399        dst: &mut CudaSlice<i32>,
5400        n: usize,
5401    ) -> Result<(), Box<dyn std::error::Error>> {
5402        let f = self.func("i32_iota_from");
5403        let cfg = LaunchConfig::for_num_elems(n as u32);
5404        let ni = n as i32;
5405        let __s_b = self.gpu.stream();
5406        let mut b = __s_b.launch_builder(&f);
5407        b.arg(ctr).arg(dst).arg(&ni);
5408        unsafe {
5409            b.launch(cfg)?;
5410        }
5411        Ok(())
5412    }
5413
5414    /// In-place trim-id translate: buf[idx] = map[buf[idx]] (FR-Spec d2t, async single-slot).
5415    pub fn u32_map_k(
5416        &self,
5417        buf: &mut CudaSlice<u32>,
5418        map: &CudaSlice<u32>,
5419        idx: usize,
5420    ) -> Result<(), Box<dyn std::error::Error>> {
5421        let f = self.func("u32_map_k");
5422        let cfg = LaunchConfig {
5423            grid_dim: (1, 1, 1),
5424            block_dim: (1, 1, 1),
5425            shared_mem_bytes: 0,
5426        };
5427        let ii = idx as i32;
5428        let __s_b = self.gpu.stream();
5429        let mut b = __s_b.launch_builder(&f);
5430        b.arg(buf).arg(map).arg(&ii);
5431        unsafe {
5432            b.launch(cfg)?;
5433        }
5434        Ok(())
5435    }
5436
5437    /// Pack a[off..off+n1] ++ b[0..n2] into one buffer (single dtoh follows).
5438    #[allow(clippy::too_many_arguments)]
5439    pub fn u32_pack2(
5440        &self,
5441        a: &CudaSlice<u32>,
5442        off_a: usize,
5443        n1: usize,
5444        b_in: &CudaSlice<u32>,
5445        n2: usize,
5446        out: &mut CudaSlice<u32>,
5447    ) -> Result<(), Box<dyn std::error::Error>> {
5448        let f = self.func("u32_pack2");
5449        let cfg = LaunchConfig::for_num_elems((n1 + n2) as u32);
5450        let (oa, i1, i2) = (off_a as i32, n1 as i32, n2 as i32);
5451        let __s_b = self.gpu.stream();
5452        let mut b = __s_b.launch_builder(&f);
5453        b.arg(a).arg(&oa).arg(&i1).arg(b_in).arg(&i2).arg(out);
5454        unsafe {
5455            b.launch(cfg)?;
5456        }
5457        Ok(())
5458    }
5459
5460    /// gemma4 R3 device fold: w[i] *= s[sel[i]] over the router's [n] (sel, w) pair.
5461    pub fn moe_w_exscale(
5462        &self,
5463        w: &mut CudaSlice<f32>,
5464        sel: &CudaSlice<i32>,
5465        s: &CudaSlice<f32>,
5466        n: usize,
5467    ) -> Result<(), Box<dyn std::error::Error>> {
5468        let f = self.func("moe_w_exscale");
5469        let cfg = LaunchConfig::for_num_elems(n as u32);
5470        let ni = n as i32;
5471        let __s_b = self.gpu.stream();
5472        let mut b = __s_b.launch_builder(&f);
5473        b.arg(w).arg(sel).arg(s).arg(&ni);
5474        unsafe {
5475            b.launch(cfg)?;
5476        }
5477        Ok(())
5478    }
5479
5480    /// Down-projection macro fold: w[i] *= macros[2*n_expert + sel[i]] on the device router
5481    /// weights (one launch per MoE layer, only for macro-carrying artifacts — see MoeWeights).
5482    pub fn moe_w_scale_by_expert(
5483        &self,
5484        w: &mut CudaSlice<f32>,
5485        sel: &CudaSlice<i32>,
5486        macros: &CudaSlice<f32>,
5487        n_expert: usize,
5488        n: usize,
5489    ) -> Result<(), Box<dyn std::error::Error>> {
5490        let f = self.func("moe_w_scale_by_expert");
5491        let cfg = LaunchConfig {
5492            grid_dim: (n.div_ceil(64) as u32, 1, 1),
5493            block_dim: (64, 1, 1),
5494            shared_mem_bytes: 0,
5495        };
5496        let (ne, nn) = (n_expert as i32, n as i32);
5497        let __s_b = self.gpu.stream();
5498        let mut b = __s_b.launch_builder(&f);
5499        b.arg(w).arg(sel).arg(macros).arg(&ne).arg(&nn);
5500        unsafe {
5501            b.launch(cfg)?;
5502        }
5503        Ok(())
5504    }
5505
5506    pub fn moe_gate_up_silu8_dev_q8(
5507        &self,
5508        table: &CudaSlice<u64>,
5509        sel: &cudarc::driver::CudaView<i32>,
5510        aq: &CudaSlice<i8>,
5511        ad: &CudaSlice<f32>,
5512        in_f: usize,
5513        n_ff: usize,
5514        n_used: usize,
5515        n_expert: usize,
5516        qt_g: i32,
5517        qt_u: i32,
5518        rb_g: usize,
5519        rb_u: usize,
5520        macros: &CudaSlice<f32>,
5521    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5522        static GU: std::sync::OnceLock<(String, u32)> = std::sync::OnceLock::new();
5523        let (mode, wpb) = GU.get_or_init(|| {
5524            let mode = std::env::var("MEMRA_MOE_DEVQ8_GU").unwrap_or_default();
5525            let wpb = std::env::var("MEMRA_MOE_DEVQ8_WPB")
5526                .ok()
5527                .and_then(|v| v.parse().ok())
5528                .unwrap_or(4u32)
5529                .clamp(1, 16);
5530            (mode, wpb)
5531        });
5532        let (mode, wpb) = (mode.as_str(), *wpb);
5533        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
5534        let (inf, nff, ne, rbg, rbu) = (
5535            in_f as i32,
5536            n_ff as i32,
5537            n_expert as i32,
5538            rb_g as i64,
5539            rb_u as i64,
5540        );
5541        let (f, cfg) = match mode {
5542            "1" | "2" | "4" => {
5543                let rpw: u32 = mode.parse().unwrap();
5544                let f = self.func(match rpw {
5545                    1 => "moe_gate_up_silu8_dev_q8_r1",
5546                    2 => "moe_gate_up_silu8_dev_q8_r2",
5547                    _ => "moe_gate_up_silu8_dev_q8_r4",
5548                });
5549                let rows_per_block = (rpw * wpb) as usize;
5550                let gx = n_ff.div_ceil(rows_per_block) as u32;
5551                (
5552                    f,
5553                    LaunchConfig {
5554                        grid_dim: (gx, n_used as u32, 1),
5555                        block_dim: (32, wpb, 1),
5556                        shared_mem_bytes: 0,
5557                    },
5558                )
5559            }
5560            "j8" if n_used <= 32 => (
5561                self.func("moe_gate_up_silu8_dev_q8_j8"),
5562                LaunchConfig {
5563                    grid_dim: (n_ff as u32, 1, 1),
5564                    block_dim: (32, n_used as u32, 1),
5565                    shared_mem_bytes: 0,
5566                },
5567            ),
5568            // SMEM-GRID twins (IQ3_S 2KB grid copied to shared, static smem — bit-identical dots)
5569            "vsm2" => {
5570                let f = self.func("moe_gate_up_silu8_dev_q8_vsm2");
5571                let sh = (rb_g + rb_u) as u32;
5572                use cudarc::driver::sys::CUfunction_attribute_enum as A;
5573                f.set_attribute(
5574                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
5575                    sh as i32,
5576                )?;
5577                (
5578                    f,
5579                    LaunchConfig {
5580                        grid_dim: (n_ff as u32, n_used as u32, 1),
5581                        block_dim: (32, 1, 1),
5582                        shared_mem_bytes: sh,
5583                    },
5584                )
5585            }
5586            "vsm" => {
5587                let f = self.func("moe_gate_up_silu8_dev_q8_vsm");
5588                let sh = (rb_g + rb_u) as u32;
5589                use cudarc::driver::sys::CUfunction_attribute_enum as A;
5590                f.set_attribute(
5591                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
5592                    sh as i32,
5593                )?;
5594                (
5595                    f,
5596                    LaunchConfig {
5597                        grid_dim: (n_ff as u32, n_used as u32, 1),
5598                        block_dim: (32, 1, 1),
5599                        shared_mem_bytes: sh,
5600                    },
5601                )
5602            }
5603            "sg" => (
5604                self.func("moe_gate_up_silu8_dev_q8_sg"),
5605                LaunchConfig {
5606                    grid_dim: (n_ff as u32, n_used as u32, 1),
5607                    block_dim: (32, 1, 1),
5608                    shared_mem_bytes: 0,
5609                },
5610            ),
5611            "j8sg" if n_used <= 32 => (
5612                self.func("moe_gate_up_silu8_dev_q8_j8sg"),
5613                LaunchConfig {
5614                    grid_dim: (n_ff as u32, 1, 1),
5615                    block_dim: (32, n_used as u32, 1),
5616                    shared_mem_bytes: 0,
5617                },
5618            ),
5619            "u64" if in_f == 2048 => (
5620                self.func("moe_gate_up_silu8_dev_q8_u64"),
5621                LaunchConfig {
5622                    grid_dim: (n_ff as u32, n_used as u32, 1),
5623                    block_dim: (32, 1, 1),
5624                    shared_mem_bytes: 0,
5625                },
5626            ),
5627            "gs4" if in_f == 2048 => (
5628                self.func("moe_gate_up_silu8_dev_q8_gs4"),
5629                LaunchConfig {
5630                    grid_dim: (n_ff as u32, n_used as u32, 1),
5631                    block_dim: (32, 4, 1),
5632                    shared_mem_bytes: 0,
5633                },
5634            ),
5635            // _v twin (down8 lane 2026-07-08): wide-load IQ4_XS dot, base geometry, bit-identical.
5636            "v" | "" => (
5637                self.func("moe_gate_up_silu8_dev_q8_v"),
5638                LaunchConfig {
5639                    grid_dim: (n_ff as u32, n_used as u32, 1),
5640                    block_dim: (32, 1, 1),
5641                    shared_mem_bytes: 0,
5642                },
5643            ),
5644            "s2" => (
5645                self.func("moe_gate_up_silu8_dev_q8_s2"),
5646                LaunchConfig {
5647                    grid_dim: (n_ff as u32, n_used as u32, 1),
5648                    block_dim: (32, 2, 1),
5649                    shared_mem_bytes: 0,
5650                },
5651            ),
5652            "s2z" => {
5653                let rz = wpb.min(16); // s2z smem tile is [16][2]
5654                (
5655                    self.func("moe_gate_up_silu8_dev_q8_s2z"),
5656                    LaunchConfig {
5657                        grid_dim: (n_ff.div_ceil(rz as usize) as u32, n_used as u32, 1),
5658                        block_dim: (32, 2, rz),
5659                        shared_mem_bytes: 0,
5660                    },
5661                )
5662            }
5663            _ => (
5664                self.func("moe_gate_up_silu8_dev_q8"),
5665                LaunchConfig {
5666                    grid_dim: (n_ff as u32, n_used as u32, 1),
5667                    block_dim: (32, 1, 1),
5668                    shared_mem_bytes: 0,
5669                },
5670            ),
5671        };
5672        let __s_b = self.gpu.stream();
5673        let mut b = __s_b.launch_builder(&f);
5674        b.arg(table)
5675            .arg(sel)
5676            .arg(aq)
5677            .arg(ad)
5678            .arg(&mut act)
5679            .arg(&inf)
5680            .arg(&nff)
5681            .arg(&ne)
5682            .arg(&qt_g)
5683            .arg(&qt_u)
5684            .arg(&rbg)
5685            .arg(&rbu)
5686            .arg(macros);
5687        unsafe {
5688            b.launch(cfg)?;
5689        }
5690        Ok(act)
5691    }
5692
5693    #[allow(clippy::too_many_arguments)]
5694    pub fn moe_down8_fma_dev_q8(
5695        &self,
5696        table: &CudaSlice<u64>,
5697        sel: &cudarc::driver::CudaView<i32>,
5698        w: &cudarc::driver::CudaView<f32>,
5699        aq2: &CudaSlice<i8>,
5700        ad2: &CudaSlice<f32>,
5701        dst: &mut cudarc::driver::CudaViewMut<f32>,
5702        in_f: usize,
5703        out_f: usize,
5704        n_used: usize,
5705        n_expert: usize,
5706        qt: i32,
5707        rb: usize,
5708    ) -> Result<(), Box<dyn std::error::Error>> {
5709        static DOWN: std::sync::OnceLock<String> = std::sync::OnceLock::new();
5710        let mode = DOWN.get_or_init(|| std::env::var("MEMRA_MOE_DEVQ8_DOWN").unwrap_or_default());
5711        let (inf, outf, nu, ne, rbi) = (
5712            in_f as i32,
5713            out_f as i32,
5714            n_used as i32,
5715            n_expert as i32,
5716            rb as i64,
5717        );
5718        // the w8 twins' smem tile is [RPW][8] — n_used must fit the 8-slot tile;
5719        // the h2 twins are nsb==16 (in_f==512) shape-gated.
5720        let (f, cfg) = match mode.as_str() {
5721            m @ ("1" | "2" | "4") if n_used <= 8 => {
5722                let rpw: usize = m.parse().unwrap();
5723                let f = self.func(match rpw {
5724                    1 => "moe_down8_fma_dev_q8_w8r1",
5725                    2 => "moe_down8_fma_dev_q8_w8r2",
5726                    _ => "moe_down8_fma_dev_q8_w8r4",
5727                });
5728                (
5729                    f,
5730                    LaunchConfig {
5731                        grid_dim: (out_f.div_ceil(rpw) as u32, 1, 1),
5732                        block_dim: (32, n_used as u32, 1),
5733                        shared_mem_bytes: 0,
5734                    },
5735                )
5736            }
5737            "h2" if in_f == 512 => (
5738                self.func("moe_down8_fma_dev_q8_h2"),
5739                LaunchConfig {
5740                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
5741                    block_dim: (32, 1, 1),
5742                    shared_mem_bytes: 0,
5743                },
5744            ),
5745            // "" = AUTO gemma shape (in_f==704): w8r2 measured +1 tok/s vs base (sweep
5746            // 1/2/4 -> 133.6/134.2/133.6, 2026-07-10); slot-ordered chain preserved.
5747            "" if in_f == 704 && n_used <= 8 => (
5748                self.func("moe_down8_fma_dev_q8_w8r2"),
5749                LaunchConfig {
5750                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
5751                    block_dim: (32, n_used as u32, 1),
5752                    shared_mem_bytes: 0,
5753                },
5754            ),
5755            // "" = AUTO: the measured winner for the 35B expert shape (arc 2026-07-05, +3.8%);
5756            // any shape the h2 kernels can't take (nsb!=16 / n_used>8) falls to base via `_`.
5757            // _v twins (down8 lane 2026-07-08): wide-load IQ4_XS dot, bit-identical outputs.
5758            "w8h2v" | "" if in_f == 512 && n_used <= 8 => (
5759                self.func("moe_down8_fma_dev_q8_w8h2v"),
5760                LaunchConfig {
5761                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
5762                    block_dim: (32, n_used as u32, 1),
5763                    shared_mem_bytes: 0,
5764                },
5765            ),
5766            "w8h2r2v" if in_f == 512 && n_used <= 8 => (
5767                self.func("moe_down8_fma_dev_q8_w8h2r2v"),
5768                LaunchConfig {
5769                    grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
5770                    block_dim: (32, n_used as u32, 1),
5771                    shared_mem_bytes: 0,
5772                },
5773            ),
5774            "w8h2r2" if in_f == 512 && n_used <= 8 => (
5775                self.func("moe_down8_fma_dev_q8_w8h2r2"),
5776                LaunchConfig {
5777                    grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
5778                    block_dim: (32, n_used as u32, 1),
5779                    shared_mem_bytes: 0,
5780                },
5781            ),
5782            "w8h2" if in_f == 512 && n_used <= 8 => (
5783                self.func("moe_down8_fma_dev_q8_w8h2"),
5784                LaunchConfig {
5785                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
5786                    block_dim: (32, n_used as u32, 1),
5787                    shared_mem_bytes: 0,
5788                },
5789            ),
5790            _ => (
5791                self.func("moe_down8_fma_dev_q8"),
5792                LaunchConfig {
5793                    grid_dim: (out_f as u32, 1, 1),
5794                    block_dim: (32, 1, 1),
5795                    shared_mem_bytes: 0,
5796                },
5797            ),
5798        };
5799        let __s_b = self.gpu.stream();
5800        let mut b = __s_b.launch_builder(&f);
5801        b.arg(table)
5802            .arg(sel)
5803            .arg(w)
5804            .arg(aq2)
5805            .arg(ad2)
5806            .arg(dst)
5807            .arg(&inf)
5808            .arg(&outf)
5809            .arg(&nu)
5810            .arg(&ne)
5811            .arg(&qt)
5812            .arg(&rbi);
5813        unsafe {
5814            b.launch(cfg)?;
5815        }
5816        Ok(())
5817    }
5818
5819    /// SMALL-M VERIFY rows twin (MEMRA_SPEC_M2, lane/spec-m2): ONE launch covers all `t` tokens
5820    /// of the spec verify's MoE dev gate/up (grid.z = token) — the _v geometry per token, with
5821    /// tok-offset sel/aq/ad/act pointers matching the serial loop's slices. BIT-IDENTICAL per
5822    /// token (see the kernel header). aq/ad are the BATCHED z-quantize ([t, in_f] rows —
5823    /// quantize_q8_1's per-32-block program is row-independent, so batched rows == the serial
5824    /// loop's per-token quantize_q8_1_view bytes). Returns act [t, n_used, n_ff].
5825    #[allow(clippy::too_many_arguments)]
5826    pub fn moe_gate_up_silu8_dev_q8_rows(
5827        &self,
5828        table: &CudaSlice<u64>,
5829        sel: &CudaSlice<i32>,
5830        aq: &CudaSlice<i8>,
5831        ad: &CudaSlice<f32>,
5832        t: usize,
5833        in_f: usize,
5834        n_ff: usize,
5835        n_used: usize,
5836        n_expert: usize,
5837        qt_g: i32,
5838        qt_u: i32,
5839        rb_g: usize,
5840        rb_u: usize,
5841        macros: &CudaSlice<f32>,
5842    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5843        let f = self.func("moe_gate_up_silu8_dev_q8_v_rows");
5844        let mut act = self.alloc_uninit::<f32>(t * n_used * n_ff)?;
5845        let cfg = LaunchConfig {
5846            grid_dim: (n_ff as u32, n_used as u32, t as u32),
5847            block_dim: (32, 1, 1),
5848            shared_mem_bytes: 0,
5849        };
5850        let (inf, nff, ne, nu, rbg, rbu) = (
5851            in_f as i32,
5852            n_ff as i32,
5853            n_expert as i32,
5854            n_used as i32,
5855            rb_g as i64,
5856            rb_u as i64,
5857        );
5858        let __s_b = self.gpu.stream();
5859        let mut b = __s_b.launch_builder(&f);
5860        b.arg(table)
5861            .arg(sel)
5862            .arg(aq)
5863            .arg(ad)
5864            .arg(&mut act)
5865            .arg(&inf)
5866            .arg(&nff)
5867            .arg(&ne)
5868            .arg(&qt_g)
5869            .arg(&qt_u)
5870            .arg(&rbg)
5871            .arg(&rbu)
5872            .arg(&nu)
5873            .arg(macros);
5874        unsafe {
5875            b.launch(cfg)?;
5876        }
5877        Ok(act)
5878    }
5879
5880    /// SMALL-M VERIFY rows twin of the down proj: w8h2v geometry per token on a grid.z token
5881    /// axis. Caller gates the w8h2v shape contract (in_f == 512, n_used <= 8) — same gate as
5882    /// the AUTO dispatch in `moe_down8_fma_dev_q8`. aq2/ad2 = batched act quantize
5883    /// ([t*n_used, in_f] rows). dst rows are FULLY overwritten per token.
5884    #[allow(clippy::too_many_arguments)]
5885    pub fn moe_down8_fma_dev_q8_rows(
5886        &self,
5887        table: &CudaSlice<u64>,
5888        sel: &CudaSlice<i32>,
5889        w: &CudaSlice<f32>,
5890        aq2: &CudaSlice<i8>,
5891        ad2: &CudaSlice<f32>,
5892        dst: &mut CudaSlice<f32>,
5893        t: usize,
5894        in_f: usize,
5895        out_f: usize,
5896        n_used: usize,
5897        n_expert: usize,
5898        qt: i32,
5899        rb: usize,
5900    ) -> Result<(), Box<dyn std::error::Error>> {
5901        assert!(
5902            in_f == 512 && n_used <= 8,
5903            "down rows twin is w8h2v shape-gated"
5904        );
5905        let f = self.func("moe_down8_fma_dev_q8_w8h2v_rows");
5906        let cfg = LaunchConfig {
5907            grid_dim: (out_f.div_ceil(2) as u32, 1, t as u32),
5908            block_dim: (32, n_used as u32, 1),
5909            shared_mem_bytes: 0,
5910        };
5911        let (inf, outf, nu, ne, rbi) = (
5912            in_f as i32,
5913            out_f as i32,
5914            n_used as i32,
5915            n_expert as i32,
5916            rb as i64,
5917        );
5918        let __s_b = self.gpu.stream();
5919        let mut b = __s_b.launch_builder(&f);
5920        b.arg(table)
5921            .arg(sel)
5922            .arg(w)
5923            .arg(aq2)
5924            .arg(ad2)
5925            .arg(dst)
5926            .arg(&inf)
5927            .arg(&outf)
5928            .arg(&nu)
5929            .arg(&ne)
5930            .arg(&qt)
5931            .arg(&rbi);
5932        unsafe {
5933            b.launch(cfg)?;
5934        }
5935        Ok(())
5936    }
5937
5938    /// CSR gate/up v3 (owner-scan dedup, no build kernel): qtypes {IQ4_XS, IQ3_S} (caller
5939    /// gates), grid.y = pair index; the first pair of each expert serves all its pairs.
5940    /// Bit-identical to moe_gate_up_silu8_dev_q8_v_rows (explicit-intrinsic accumulate).
5941    #[allow(clippy::too_many_arguments)]
5942    pub fn moe_gate_up_silu8_dev_q8_csr(
5943        &self,
5944        table: &CudaSlice<u64>,
5945        sel: &CudaSlice<i32>,
5946        aq: &CudaSlice<i8>,
5947        ad: &CudaSlice<f32>,
5948        n_pairs: usize,
5949        in_f: usize,
5950        n_ff: usize,
5951        n_used: usize,
5952        n_expert: usize,
5953        qt_g: i32,
5954        qt_u: i32,
5955        rb_g: usize,
5956        rb_u: usize,
5957    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5958        // NVFP4 experts take the NVFP4-specialized owner-scan twin (lane/moebatch-q35moe);
5959        // host gate guarantees qt_g == qt_u within a supported class.
5960        let f = if qt_g == crate::QT_NVFP4 {
5961            self.func("moe_gate_up_silu8_dev_q8_csr_nvfp4")
5962        } else {
5963            self.func("moe_gate_up_silu8_dev_q8_csr_iq4")
5964        };
5965        let mut act = self.alloc_uninit::<f32>(n_pairs * n_ff)?;
5966        let cfg = LaunchConfig {
5967            grid_dim: (n_ff as u32, n_pairs as u32, 1),
5968            block_dim: (32, 1, 1),
5969            shared_mem_bytes: 0,
5970        };
5971        let (inf, nff, ne, nu, npi, rbg, rbu) = (
5972            in_f as i32,
5973            n_ff as i32,
5974            n_expert as i32,
5975            n_used as i32,
5976            n_pairs as i32,
5977            rb_g as i64,
5978            rb_u as i64,
5979        );
5980        let __s_b = self.gpu.stream();
5981        let mut b = __s_b.launch_builder(&f);
5982        b.arg(table)
5983            .arg(sel)
5984            .arg(aq)
5985            .arg(ad)
5986            .arg(&mut act)
5987            .arg(&inf)
5988            .arg(&nff)
5989            .arg(&ne)
5990            .arg(&qt_g)
5991            .arg(&qt_u)
5992            .arg(&rbg)
5993            .arg(&rbu)
5994            .arg(&nu)
5995            .arg(&npi);
5996        unsafe {
5997            b.launch(cfg)?;
5998        }
5999        Ok(act)
6000    }
6001
6002    /// TEST SEAM (down8 lane 2026-07-08): launch a down dev_q8 variant BY NAME with its
6003    /// canonical geometry, bypassing the env-cached dispatch so moe-devq8-check can byte-
6004    /// compare variants in one process. Variants: "base", "w8h2", "w8h2r2", "w8h2v", "w8h2r2v".
6005    #[allow(clippy::too_many_arguments)]
6006    pub fn moe_down8_fma_dev_q8_variant(
6007        &self,
6008        variant: &str,
6009        table: &CudaSlice<u64>,
6010        sel: &cudarc::driver::CudaView<i32>,
6011        w: &cudarc::driver::CudaView<f32>,
6012        aq2: &CudaSlice<i8>,
6013        ad2: &CudaSlice<f32>,
6014        dst: &mut cudarc::driver::CudaViewMut<f32>,
6015        in_f: usize,
6016        out_f: usize,
6017        n_used: usize,
6018        n_expert: usize,
6019        qt: i32,
6020        rb: usize,
6021    ) -> Result<(), Box<dyn std::error::Error>> {
6022        let (inf, outf, nu, ne, rbi) = (
6023            in_f as i32,
6024            out_f as i32,
6025            n_used as i32,
6026            n_expert as i32,
6027            rb as i64,
6028        );
6029        let (f, cfg) = match variant {
6030            "w8h2" | "w8h2v" => (
6031                self.func(if variant == "w8h2" {
6032                    "moe_down8_fma_dev_q8_w8h2"
6033                } else {
6034                    "moe_down8_fma_dev_q8_w8h2v"
6035                }),
6036                LaunchConfig {
6037                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6038                    block_dim: (32, n_used as u32, 1),
6039                    shared_mem_bytes: 0,
6040                },
6041            ),
6042            "w8h2r2" | "w8h2r2v" => (
6043                self.func(if variant == "w8h2r2" {
6044                    "moe_down8_fma_dev_q8_w8h2r2"
6045                } else {
6046                    "moe_down8_fma_dev_q8_w8h2r2v"
6047                }),
6048                LaunchConfig {
6049                    grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
6050                    block_dim: (32, n_used as u32, 1),
6051                    shared_mem_bytes: 0,
6052                },
6053            ),
6054            _ => (
6055                self.func("moe_down8_fma_dev_q8"),
6056                LaunchConfig {
6057                    grid_dim: (out_f as u32, 1, 1),
6058                    block_dim: (32, 1, 1),
6059                    shared_mem_bytes: 0,
6060                },
6061            ),
6062        };
6063        let __s_b = self.gpu.stream();
6064        let mut b = __s_b.launch_builder(&f);
6065        b.arg(table)
6066            .arg(sel)
6067            .arg(w)
6068            .arg(aq2)
6069            .arg(ad2)
6070            .arg(dst)
6071            .arg(&inf)
6072            .arg(&outf)
6073            .arg(&nu)
6074            .arg(&ne)
6075            .arg(&qt)
6076            .arg(&rbi);
6077        unsafe {
6078            b.launch(cfg)?;
6079        }
6080        Ok(())
6081    }
6082
6083    /// TEST SEAM (down8 lane): gate_up twin of the above. Variants: "base", "v".
6084    #[allow(clippy::too_many_arguments)]
6085    pub fn moe_gate_up_silu8_dev_q8_variant(
6086        &self,
6087        variant: &str,
6088        table: &CudaSlice<u64>,
6089        sel: &cudarc::driver::CudaView<i32>,
6090        aq: &CudaSlice<i8>,
6091        ad: &CudaSlice<f32>,
6092        in_f: usize,
6093        n_ff: usize,
6094        n_used: usize,
6095        n_expert: usize,
6096        qt_g: i32,
6097        qt_u: i32,
6098        rb_g: usize,
6099        rb_u: usize,
6100    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6101        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
6102        let (inf, nff, ne, rbg, rbu) = (
6103            in_f as i32,
6104            n_ff as i32,
6105            n_expert as i32,
6106            rb_g as i64,
6107            rb_u as i64,
6108        );
6109        let f = self.func(if variant == "v" {
6110            "moe_gate_up_silu8_dev_q8_v"
6111        } else {
6112            "moe_gate_up_silu8_dev_q8"
6113        });
6114        let cfg = LaunchConfig {
6115            grid_dim: (n_ff as u32, n_used as u32, 1),
6116            block_dim: (32, 1, 1),
6117            shared_mem_bytes: 0,
6118        };
6119        let __s_b = self.gpu.stream();
6120        let mut b = __s_b.launch_builder(&f);
6121        b.arg(table)
6122            .arg(sel)
6123            .arg(aq)
6124            .arg(ad)
6125            .arg(&mut act)
6126            .arg(&inf)
6127            .arg(&nff)
6128            .arg(&ne)
6129            .arg(&qt_g)
6130            .arg(&qt_u)
6131            .arg(&rbg)
6132            .arg(&rbu);
6133        unsafe {
6134            b.launch(cfg)?;
6135        }
6136        Ok(act)
6137    }
6138
6139    pub fn moe_gate_up_silu8_dev(
6140        &self,
6141        table: &CudaSlice<u64>,
6142        sel: &cudarc::driver::CudaView<i32>,
6143        x: &cudarc::driver::CudaView<f32>,
6144        in_f: usize,
6145        n_ff: usize,
6146        n_used: usize,
6147        n_expert: usize,
6148        qt_g: i32,
6149        qt_u: i32,
6150        rb_g: usize,
6151        rb_u: usize,
6152        macros: &CudaSlice<f32>,
6153    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6154        let f = self.func("moe_gate_up_silu8_dev");
6155        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?; // fully overwritten
6156        let cfg = LaunchConfig {
6157            grid_dim: (n_ff as u32, n_used as u32, 1),
6158            block_dim: (256, 1, 1),
6159            shared_mem_bytes: 0,
6160        };
6161        let (inf, nff, ne, rbg, rbu) = (
6162            in_f as i32,
6163            n_ff as i32,
6164            n_expert as i32,
6165            rb_g as i64,
6166            rb_u as i64,
6167        );
6168        let __s_b = self.gpu.stream();
6169        let mut b = __s_b.launch_builder(&f);
6170        b.arg(table)
6171            .arg(sel)
6172            .arg(x)
6173            .arg(&mut act)
6174            .arg(&inf)
6175            .arg(&nff)
6176            .arg(&ne)
6177            .arg(&qt_g)
6178            .arg(&qt_u)
6179            .arg(&rbg)
6180            .arg(&rbu)
6181            .arg(macros);
6182        unsafe {
6183            b.launch(cfg)?;
6184        }
6185        Ok(act)
6186    }
6187
6188    /// LAUNCH-STRUCTURE STAGE 3: device-dispatch twin of `moe_down8_fma_into` — expert ids AND
6189    /// renormalized weights read from the router kernel's device output. BIT-IDENTICAL chain.
6190    #[allow(clippy::too_many_arguments)]
6191    pub fn moe_down8_fma_dev(
6192        &self,
6193        table: &CudaSlice<u64>,
6194        sel: &cudarc::driver::CudaView<i32>,
6195        w: &cudarc::driver::CudaView<f32>,
6196        act: &CudaSlice<f32>,
6197        dst: &mut cudarc::driver::CudaViewMut<f32>,
6198        in_f: usize,
6199        out_f: usize,
6200        n_used: usize,
6201        n_expert: usize,
6202        qt: i32,
6203        rb: usize,
6204    ) -> Result<(), Box<dyn std::error::Error>> {
6205        let f = self.func("moe_down8_fma_dev");
6206        let cfg = LaunchConfig {
6207            grid_dim: (out_f as u32, 1, 1),
6208            block_dim: (256, 1, 1),
6209            shared_mem_bytes: 0,
6210        };
6211        let (inf, outf, nu, ne, rbv) = (
6212            in_f as i32,
6213            out_f as i32,
6214            n_used as i32,
6215            n_expert as i32,
6216            rb as i64,
6217        );
6218        let __s_b = self.gpu.stream();
6219        let mut b = __s_b.launch_builder(&f);
6220        b.arg(table)
6221            .arg(sel)
6222            .arg(w)
6223            .arg(act)
6224            .arg(dst)
6225            .arg(&inf)
6226            .arg(&outf)
6227            .arg(&nu)
6228            .arg(&ne)
6229            .arg(&qt)
6230            .arg(&rbv);
6231        unsafe {
6232            b.launch(cfg)?;
6233        }
6234        Ok(())
6235    }
6236
6237    /// dst[i] += alpha * src[i], i in 0..n. dst is a CudaViewMut (a row of moe_out).
6238    pub fn axpy_into(
6239        &self,
6240        src: &CudaSlice<f32>,
6241        alpha: f32,
6242        dst: &mut cudarc::driver::CudaViewMut<f32>,
6243        n: usize,
6244    ) -> Result<(), Box<dyn std::error::Error>> {
6245        let f = self.func("axpy_f32");
6246        let cfg = LaunchConfig::for_num_elems(n as u32);
6247        let (a, ni) = (alpha, n as i32);
6248        let __s_b = self.gpu.stream();
6249        let mut b = __s_b.launch_builder(&f);
6250        b.arg(src).arg(dst).arg(&a).arg(&ni);
6251        unsafe {
6252            b.launch(cfg)?;
6253        }
6254        Ok(())
6255    }
6256
6257    /// dst[r*ncols + c] += src[r*ncols + c] * scale[r]. Per-row scalar accumulate (shared expert).
6258    pub fn add_scaled_rows(
6259        &self,
6260        src: &CudaSlice<f32>,
6261        scale: &CudaSlice<f32>,
6262        dst: &mut CudaSlice<f32>,
6263        ncols: usize,
6264        nrows: usize,
6265    ) -> Result<(), Box<dyn std::error::Error>> {
6266        let f = self.func("add_scaled_rows_f32");
6267        let cfg = LaunchConfig::for_num_elems((ncols * nrows) as u32);
6268        let (nc, nr) = (ncols as i32, nrows as i32);
6269        let __s_b = self.gpu.stream();
6270        let mut b = __s_b.launch_builder(&f);
6271        b.arg(src).arg(scale).arg(dst).arg(&nc).arg(&nr);
6272        unsafe {
6273            b.launch(cfg)?;
6274        }
6275        Ok(())
6276    }
6277
6278    // ======== A2 GROUPED MoE PREFILL KERNELS ========
6279
6280    /// Gather m_e rows from src[T, ncols] into dst[m_e, ncols] using index array idx[m_e].
6281    pub fn gather_rows(
6282        &self,
6283        src: &CudaSlice<f32>,
6284        idx: &CudaSlice<i32>,
6285        dst: &mut CudaSlice<f32>,
6286        ncols: usize,
6287        m_e: usize,
6288    ) -> Result<(), Box<dyn std::error::Error>> {
6289        let f = self.func("gather_rows_f32");
6290        let cfg = LaunchConfig::for_num_elems((m_e * ncols) as u32);
6291        let (nc, me) = (ncols as i32, m_e as i32);
6292        let __s_b = self.gpu.stream();
6293        let mut b = __s_b.launch_builder(&f);
6294        b.arg(src).arg(idx).arg(dst).arg(&nc).arg(&me);
6295        unsafe {
6296            b.launch(cfg)?;
6297        }
6298        Ok(())
6299    }
6300
6301    /// Scatter expert outputs into per-token slots: dst[tok_idx[r], slot_idx[r], :] = src[r, :] * weight[r].
6302    /// dst is [T, n_used, ncols], zero-initialized. Each (expert, token) pair maps to a unique slot.
6303    /// Scatter expert outputs into per-token slots (raw copy, no weight multiply).
6304    /// Weight stored into wbuf[tok*n_used + slot] for FMA in reduce step.
6305    pub fn scatter_slot(
6306        &self,
6307        src: &CudaSlice<f32>,
6308        tok_idx: &CudaSlice<i32>,
6309        slot_idx: &CudaSlice<i32>,
6310        weight: &CudaSlice<f32>,
6311        dst: &mut CudaSlice<f32>,
6312        wbuf: &mut CudaSlice<f32>,
6313        ncols: usize,
6314        n_used: usize,
6315        m_e: usize,
6316    ) -> Result<(), Box<dyn std::error::Error>> {
6317        let f = self.func("scatter_add_slot_f32");
6318        let cfg = LaunchConfig::for_num_elems((m_e * ncols) as u32);
6319        let (nc, nu, me) = (ncols as i32, n_used as i32, m_e as i32);
6320        let __s_b = self.gpu.stream();
6321        let mut b = __s_b.launch_builder(&f);
6322        b.arg(src)
6323            .arg(tok_idx)
6324            .arg(slot_idx)
6325            .arg(weight)
6326            .arg(dst)
6327            .arg(wbuf)
6328            .arg(&nc)
6329            .arg(&nu)
6330            .arg(&me);
6331        unsafe {
6332            b.launch(cfg)?;
6333        }
6334        Ok(())
6335    }
6336
6337    /// Reduce n_used slots per token: dst[t, col] = sum_s slots[t, s, col].
6338    /// Reduce n_used slots per token: dst[t, col] = sum_s FMA(wbuf[t,s], slots[t,s,col], acc).
6339    /// Uses FMA for bit-identity with the sequential axpy path.
6340    pub fn reduce_slots(
6341        &self,
6342        slots: &CudaSlice<f32>,
6343        wbuf: &CudaSlice<f32>,
6344        dst: &mut CudaSlice<f32>,
6345        ncols: usize,
6346        n_used: usize,
6347        t: usize,
6348    ) -> Result<(), Box<dyn std::error::Error>> {
6349        let f = self.func("reduce_slots_f32");
6350        let cfg = LaunchConfig::for_num_elems((t * ncols) as u32);
6351        let (nc, nu, ti) = (ncols as i32, n_used as i32, t as i32);
6352        let __s_b = self.gpu.stream();
6353        let mut b = __s_b.launch_builder(&f);
6354        b.arg(slots).arg(wbuf).arg(dst).arg(&nc).arg(&nu).arg(&ti);
6355        unsafe {
6356            b.launch(cfg)?;
6357        }
6358        Ok(())
6359    }
6360
6361    /// Stage-B: quantize activation [m,in] f32 -> q8_1 (int8 qs + per-block f32 scale).
6362    /// Quantize an activation [m, in_f] to q8_1 (int8 qs + per-32 f32 scale). Public so the
6363    /// forward can quantize a SHARED activation ONCE and feed it to several matmuls (gate+up
6364    /// share `z`; q/k/v and wqkv/gate/beta/alpha share `h`) — quantize_q8_1 was 13.5% of decode
6365    /// GPU time, ~half of it redundant re-quantization of the same row.
6366    /// quantize_q8_1 over a CudaView (a sliced z-row) — same kernel, offset-honoring arg.
6367    pub fn quantize_q8_1_view(
6368        &self,
6369        x: &cudarc::driver::CudaView<f32>,
6370        m: usize,
6371        in_f: usize,
6372    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6373        let f = self.func("quantize_q8_1");
6374        let nblk = in_f / 32;
6375        let mut q = self.alloc_uninit::<i8>(m * in_f)?;
6376        let mut d = self.alloc_uninit::<f32>(m * nblk)?;
6377        let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
6378        let (inf, mi) = (in_f as i32, m as i32);
6379        let __s_b = self.gpu.stream();
6380        let mut b = __s_b.launch_builder(&f);
6381        b.arg(x).arg(&mut q).arg(&mut d).arg(&inf).arg(&mi);
6382        unsafe {
6383            b.launch(cfg)?;
6384        }
6385        Ok((q, d))
6386    }
6387
6388    pub fn quantize_q8_1(
6389        &self,
6390        x: &CudaSlice<f32>,
6391        m: usize,
6392        in_f: usize,
6393    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6394        let nblk = in_f / 32;
6395        let mut q = self.alloc_uninit::<i8>(m * in_f)?; // full-overwrite output: skip memset
6396        let mut d = self.alloc_uninit::<f32>(m * nblk)?; // full-overwrite output: skip memset
6397        // WARP-PER-BLOCK kernel: one warp per 32-block -> m*in_f threads total.
6398        let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
6399        let (inf, mi) = (in_f as i32, m as i32);
6400        if Self::pdl_on() && Self::pdl_wb_on() {
6401            {
6402                use cudarc::driver::{DevicePtr, DevicePtrMut};
6403                let s = &self.gpu.stream();
6404                let (px, _g0) = x.device_ptr(s);
6405                let (pq, _g1) = q.device_ptr_mut(s);
6406                let (pd, _g2) = d.device_ptr_mut(s);
6407                let mut ps = [
6408                    &px as *const _ as *mut std::ffi::c_void,
6409                    &pq as *const _ as *mut _,
6410                    &pd as *const _ as *mut _,
6411                    &inf as *const _ as *mut _,
6412                    &mi as *const _ as *mut _,
6413                ];
6414                unsafe {
6415                    self.launch_pdl("quantize_q8_1", cfg.grid_dim, cfg.block_dim, &mut ps)?;
6416                }
6417            }
6418            return Ok((q, d));
6419        }
6420        let f = self.func("quantize_q8_1");
6421        let __s_b = self.gpu.stream();
6422        let mut b = __s_b.launch_builder(&f);
6423        b.arg(x).arg(&mut q).arg(&mut d).arg(&inf).arg(&mi);
6424        unsafe {
6425            b.launch(cfg)?;
6426        }
6427        Ok((q, d))
6428    }
6429
6430    /// Stage-C FP4: quantize activation [m,in] f32 -> e2m1 nibbles (aq4: u32 [m, in/8]) + per-16
6431    /// UE4M3 scale (ad4: u8 [m, in/16]), the layout the mxf4nvf4 block-scale GEMM B-operand wants.
6432    /// in_f must be a multiple of 64 (one NVFP4 K-block). One thread per (token, 16-block).
6433    pub fn quantize_fp4_act(
6434        &self,
6435        x: &CudaSlice<f32>,
6436        m: usize,
6437        in_f: usize,
6438    ) -> Result<(CudaSlice<u32>, CudaSlice<u8>), Box<dyn std::error::Error>> {
6439        let f = self.func("quantize_fp4_act");
6440        let nb16 = in_f / 16;
6441        let mut aq4 = self.alloc_uninit::<u32>(m * (in_f / 8))?; // full-overwrite output: skip memset
6442        let mut ad4 = self.alloc_uninit::<u8>(m * nb16)?; // full-overwrite output: skip memset
6443        let cfg = LaunchConfig::for_num_elems((m * nb16) as u32);
6444        let (inf, mi) = (in_f as i32, m as i32);
6445        let __s_b = self.gpu.stream();
6446        let mut b = __s_b.launch_builder(&f);
6447        b.arg(x).arg(&mut aq4).arg(&mut ad4).arg(&inf).arg(&mi);
6448        unsafe {
6449            b.launch(cfg)?;
6450        }
6451        Ok((aq4, ad4))
6452    }
6453
6454    /// Stage-C FP4 GEMM (NVFP4 weights): native mxf4nvf4 block-scale tensor-core matmul. Feeds raw
6455    /// e2m1 weight nibbles + raw UE4M3 micro-scales directly to mma.sync.m16n8k64 (762 TFLOP/s peak,
6456    /// 3.5x int8). Activation `x` is quantized to FP4 e2m1 here. NVFP4 per-tensor macro-scale applied
6457    /// post (scale==1.0 -> no-op). `bytes` = raw NVFP4 weight rows. Used by the MEMRA_FP4 prefill path.
6458    pub fn qmatvec_gemm_nvfp4_fp4(
6459        &self,
6460        bytes: &CudaSlice<u8>,
6461        x: &CudaSlice<f32>,
6462        m: usize,
6463        in_f: usize,
6464        out_f: usize,
6465        row_bytes: usize,
6466        scale: f32,
6467    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6468        assert!(
6469            in_f % 64 == 0,
6470            "FP4 GEMM requires in_f % 64 == 0, got {in_f}"
6471        );
6472        let (aq4, ad4) = self.quantize_fp4_act(x, m, in_f)?;
6473        let mut y = self.fp4_gemm_launch(bytes, &aq4, &ad4, m, in_f, out_f, row_bytes)?;
6474        if scale != 1.0 {
6475            self.scale_inplace(&mut y, scale, m * out_f)?;
6476        }
6477        Ok(y)
6478    }
6479
6480    /// Shared mxf4 GEMM launch (pre-quantized FP4 activation aq4/ad4). Same CTA tile as the int8 GEMM
6481    /// (BM=64 rows x BN=128 tokens, 4 warps). No macro-scale applied here.
6482    fn fp4_gemm_launch(
6483        &self,
6484        bytes: &CudaSlice<u8>,
6485        aq4: &CudaSlice<u32>,
6486        ad4: &CudaSlice<u8>,
6487        m: usize,
6488        in_f: usize,
6489        out_f: usize,
6490        row_bytes: usize,
6491    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6492        let f = self.func("qmatvec_gemm_nvfp4_fp4");
6493        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
6494        const BM: u32 = 64;
6495        const BN: u32 = 256;
6496        let cfg = LaunchConfig {
6497            grid_dim: ((out_f as u32 + BM - 1) / BM, (m as u32 + BN - 1) / BN, 1),
6498            block_dim: (32, 4, 1),
6499            shared_mem_bytes: 0,
6500        };
6501        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
6502        let __s_b = self.gpu.stream();
6503        let mut b = __s_b.launch_builder(&f);
6504        b.arg(bytes)
6505            .arg(aq4)
6506            .arg(ad4)
6507            .arg(&mut y)
6508            .arg(&inf)
6509            .arg(&outf)
6510            .arg(&mi)
6511            .arg(&rb);
6512        unsafe {
6513            b.launch(cfg)?;
6514        }
6515        Ok(y)
6516    }
6517
6518    /// Test entry (kernel_check): run the FP4 GEMM from raw bytes; NO macro-scale (caller compares bare).
6519    pub fn qmatvec_gemm_nvfp4_fp4_raw(
6520        &self,
6521        bytes: &CudaSlice<u8>,
6522        x: &CudaSlice<f32>,
6523        m: usize,
6524        in_f: usize,
6525        out_f: usize,
6526        row_bytes: usize,
6527    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6528        assert!(
6529            in_f % 64 == 0,
6530            "FP4 GEMM requires in_f % 64 == 0, got {in_f}"
6531        );
6532        let (aq4, ad4) = self.quantize_fp4_act(x, m, in_f)?;
6533        self.fp4_gemm_launch(bytes, &aq4, &ad4, m, in_f, out_f, row_bytes)
6534    }
6535
6536    /// Stage-B: Q8_0 weight x q8_1 activation int8 dp4a matmul. y[m,out]=x@W^T.
6537    pub fn qmatvec_q8_0_fast(
6538        &self,
6539        w: &CudaSlice<u8>,
6540        x: &CudaSlice<f32>,
6541        m: usize,
6542        in_f: usize,
6543        out_f: usize,
6544        row_bytes: usize,
6545    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6546        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
6547        let f = self.func("qmatvec_q8_0_dp4a");
6548        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
6549        let cfg = LaunchConfig {
6550            grid_dim: (out_f as u32, m as u32, 1),
6551            block_dim: (128, 1, 1),
6552            shared_mem_bytes: 0,
6553        };
6554        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
6555        let __s_b = self.gpu.stream();
6556        let mut b = __s_b.launch_builder(&f);
6557        b.arg(w)
6558            .arg(&aq)
6559            .arg(&ad)
6560            .arg(&mut y)
6561            .arg(&inf)
6562            .arg(&outf)
6563            .arg(&mi)
6564            .arg(&rb);
6565        unsafe {
6566            b.launch(cfg)?;
6567        }
6568        Ok(y)
6569    }
6570
6571    /// Stage-B: Q4_K weight x q8_1 activation int8 dp4a (decode). Min-offset via q8_1 sum term.
6572    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
6573    pub fn qmatvec_q4_K_fast(
6574        &self,
6575        w: &CudaSlice<u8>,
6576        x: &CudaSlice<f32>,
6577        m: usize,
6578        in_f: usize,
6579        out_f: usize,
6580        row_bytes: usize,
6581    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6582        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
6583        let f = self.func("qmatvec_q4_K_dp4a");
6584        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
6585        let cfg = LaunchConfig {
6586            grid_dim: (out_f as u32, m as u32, 1),
6587            block_dim: (128, 1, 1),
6588            shared_mem_bytes: 0,
6589        };
6590        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
6591        let __s_b = self.gpu.stream();
6592        let mut b = __s_b.launch_builder(&f);
6593        b.arg(w)
6594            .arg(&aq)
6595            .arg(&ad)
6596            .arg(&mut y)
6597            .arg(&inf)
6598            .arg(&outf)
6599            .arg(&mi)
6600            .arg(&rb);
6601        unsafe {
6602            b.launch(cfg)?;
6603        }
6604        Ok(y)
6605    }
6606
6607    /// Stage-B: Q6_K weight x q8_1 activation int8 dp4a (decode, symmetric).
6608    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
6609    pub fn qmatvec_q6_K_fast(
6610        &self,
6611        w: &CudaSlice<u8>,
6612        x: &CudaSlice<f32>,
6613        m: usize,
6614        in_f: usize,
6615        out_f: usize,
6616        row_bytes: usize,
6617    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6618        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
6619        let f = self.func("qmatvec_q6_K_dp4a");
6620        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
6621        let cfg = LaunchConfig {
6622            grid_dim: (out_f as u32, m as u32, 1),
6623            block_dim: (128, 1, 1),
6624            shared_mem_bytes: 0,
6625        };
6626        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
6627        let __s_b = self.gpu.stream();
6628        let mut b = __s_b.launch_builder(&f);
6629        b.arg(w)
6630            .arg(&aq)
6631            .arg(&ad)
6632            .arg(&mut y)
6633            .arg(&inf)
6634            .arg(&outf)
6635            .arg(&mi)
6636            .arg(&rb);
6637        unsafe {
6638            b.launch(cfg)?;
6639        }
6640        Ok(y)
6641    }
6642
6643    /// Stage-B: Q5_K weight x q8_1 activation int8 dp4a (decode). Min-offset via q8_1 sum term.
6644    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
6645    pub fn qmatvec_q5_K_fast(
6646        &self,
6647        w: &CudaSlice<u8>,
6648        x: &CudaSlice<f32>,
6649        m: usize,
6650        in_f: usize,
6651        out_f: usize,
6652        row_bytes: usize,
6653    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6654        self.qmatvec_dp4a_named("qmatvec_q5_K_dp4a", w, x, m, in_f, out_f, row_bytes)
6655    }
6656    /// Stage-B: Q3_K weight x q8_1 activation int8 dp4a (decode, symmetric).
6657    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
6658    pub fn qmatvec_q3_K_fast(
6659        &self,
6660        w: &CudaSlice<u8>,
6661        x: &CudaSlice<f32>,
6662        m: usize,
6663        in_f: usize,
6664        out_f: usize,
6665        row_bytes: usize,
6666    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6667        self.qmatvec_dp4a_named("qmatvec_q3_K_dp4a", w, x, m, in_f, out_f, row_bytes)
6668    }
6669    /// A6 split-plane twin of `qmatvec_nvfp4_fast` (weights repacked; used by the rp gates).
6670    pub fn qmatvec_nvfp4_fast_rp(
6671        &self,
6672        w: &CudaSlice<u8>,
6673        x: &CudaSlice<f32>,
6674        m: usize,
6675        in_f: usize,
6676        out_f: usize,
6677        row_bytes: usize,
6678    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6679        assert!(
6680            in_f % 64 == 0,
6681            "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
6682        );
6683        self.qmatvec_dp4a_named("qmatvec_nvfp4_dp4a_rp", w, x, m, in_f, out_f, row_bytes)
6684    }
6685    /// Stage-B: NVFP4 weight x q8_1 activation int8 dp4a (decode, symmetric, codebook lookup).
6686    pub fn qmatvec_nvfp4_fast(
6687        &self,
6688        w: &CudaSlice<u8>,
6689        x: &CudaSlice<f32>,
6690        m: usize,
6691        in_f: usize,
6692        out_f: usize,
6693        row_bytes: usize,
6694    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6695        // B1: the NVFP4 dp4a kernel maps two 32-elem q8_1 blocks onto one 64-elem block_nvfp4
6696        // (sblk = g >> 1). in_f must be a multiple of 64 or the last block reads a partial superblock.
6697        assert!(
6698            in_f % 64 == 0,
6699            "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
6700        );
6701        self.qmatvec_dp4a_named("qmatvec_nvfp4_dp4a", w, x, m, in_f, out_f, row_bytes)
6702    }
6703    /// Stage-B (optional perf): IQ4_XS codebook int8 dp4a.
6704    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
6705    pub fn qmatvec_iq4_XS_fast(
6706        &self,
6707        w: &CudaSlice<u8>,
6708        x: &CudaSlice<f32>,
6709        m: usize,
6710        in_f: usize,
6711        out_f: usize,
6712        row_bytes: usize,
6713    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6714        self.qmatvec_dp4a_named("qmatvec_iq4_XS_dp4a", w, x, m, in_f, out_f, row_bytes)
6715    }
6716
6717    /// Shared dp4a launcher: quantize_q8_1 then call the named kernel (grid (out,m), block 64).
6718    fn qmatvec_dp4a_named(
6719        &self,
6720        name: &str,
6721        w: &CudaSlice<u8>,
6722        x: &CudaSlice<f32>,
6723        m: usize,
6724        in_f: usize,
6725        out_f: usize,
6726        row_bytes: usize,
6727    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6728        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
6729        let f = self.func(name);
6730        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
6731        let cfg = LaunchConfig {
6732            grid_dim: (out_f as u32, m as u32, 1),
6733            block_dim: (128, 1, 1),
6734            shared_mem_bytes: 0,
6735        };
6736        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
6737        let __s_b = self.gpu.stream();
6738        let mut b = __s_b.launch_builder(&f);
6739        b.arg(w)
6740            .arg(&aq)
6741            .arg(&ad)
6742            .arg(&mut y)
6743            .arg(&inf)
6744            .arg(&outf)
6745            .arg(&mi)
6746            .arg(&rb);
6747        unsafe {
6748            b.launch(cfg)?;
6749        }
6750        Ok(y)
6751    }
6752
6753    pub fn htod(&self, v: &[f32]) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6754        Ok(self.gpu.stream().clone_htod(v)?)
6755    }
6756    pub fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
6757        Ok(self.gpu.stream().clone_htod(v)?)
6758    }
6759    /// i8 upload (moe-devq8-check: synthetic q8_1 activation bytes).
6760    pub fn htod_i8(&self, v: &[i8]) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
6761        Ok(self.gpu.stream().clone_htod(v)?)
6762    }
6763    pub fn htod_u64(&self, v: &[u64]) -> Result<CudaSlice<u64>, Box<dyn std::error::Error>> {
6764        Ok(self.gpu.stream().clone_htod(v)?)
6765    }
6766    /// View twin of `dtoh` (lean-logits component 3: D2H one row of a [B, n_vocab] stack).
6767    pub fn dtoh_view(
6768        &self,
6769        d: &cudarc::driver::CudaView<f32>,
6770    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
6771        let v = self.gpu.stream().clone_dtoh(d)?;
6772        self.gpu.stream().synchronize()?;
6773        Ok(v)
6774    }
6775    pub fn dtoh(&self, d: &CudaSlice<f32>) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
6776        let v = self.gpu.stream().clone_dtoh(d)?;
6777        self.gpu.stream().synchronize()?;
6778        Ok(v)
6779    }
6780    /// Queue two f32 device-to-host copies on the compute stream, then establish one host
6781    /// boundary for both. Hy3's CPU/GPU expert split needs the router logits and the MoE input;
6782    /// issuing them together avoids a second stream synchronization in every trunk layer.
6783    pub fn dtoh_pair(
6784        &self,
6785        a: &CudaSlice<f32>,
6786        b: &CudaSlice<f32>,
6787    ) -> Result<(Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
6788        let av = self.gpu.stream().clone_dtoh(a)?;
6789        let bv = self.gpu.stream().clone_dtoh(b)?;
6790        self.gpu.stream().synchronize()?;
6791        Ok((av, bv))
6792    }
6793    /// Device-to-host copy of an i32 buffer (fused-router sel_idx readback).
6794    pub fn dtoh_i32(&self, d: &CudaSlice<i32>) -> Result<Vec<i32>, Box<dyn std::error::Error>> {
6795        let v = self.gpu.stream().clone_dtoh(d)?;
6796        self.gpu.stream().synchronize()?;
6797        Ok(v)
6798    }
6799    /// Device-to-host copy of a u8 buffer (used to read back the quantized KV cache for validation).
6800    pub fn dtoh_u8(&self, d: &CudaSlice<u8>) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
6801        let v = self.gpu.stream().clone_dtoh(d)?;
6802        self.gpu.stream().synchronize()?;
6803        Ok(v)
6804    }
6805    pub fn dtoh_u8_view(
6806        &self,
6807        d: &cudarc::driver::CudaView<u8>,
6808    ) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
6809        let v = self.gpu.stream().clone_dtoh(d)?;
6810        self.gpu.stream().synchronize()?;
6811        Ok(v)
6812    }
6813    pub fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6814        let s = self.gpu.stream().alloc_zeros::<f32>(n)?;
6815        self.keep_if_capturing(&s);
6816        Ok(s)
6817    }
6818
6819    /// GPU-resident greedy argmax (CUDA-GRAPH-PLAN Phase 1): logits[n_vocab] -> token id in a
6820    /// resident device u32 [1]. PARALLEL 2-pass (RANK1 LEVER): the old single-CTA scan (one 256-thread
6821    /// block on one SM over 248K logits) was memory-starved at ~426us/token. Now pass 1 fans NB=256
6822    /// blocks across the SMs to saturate HBM, pass 2 reduces the NB partials. Bit-identical to host
6823    /// `argmax` (smallest index on tie). The whole point is NOT to dtoh logits — only a [1] u32 is read
6824    /// back (or kept resident for graph replay). Returns the device token buffer.
6825    /// Softmax probability of the (already-argmaxed) token `tok` under `logits` — the spec-decode
6826    /// p-min confidence signal. 2-pass like the parallel argmax; returns a device [1] f32.
6827    pub fn prob_of_token_device(
6828        &self,
6829        logits: &CudaSlice<f32>,
6830        tok: &CudaSlice<u32>,
6831        n_vocab: usize,
6832    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6833        let nb = ARGMAX_NB;
6834        let mut part = self.alloc_uninit::<f32>(nb)?;
6835        let mut p = self.alloc_uninit::<f32>(1)?;
6836        let f1 = self.func("prob_of_token_partial_f32");
6837        let cfg1 = LaunchConfig {
6838            grid_dim: (nb as u32, 1, 1),
6839            block_dim: (256, 1, 1),
6840            shared_mem_bytes: 0,
6841        };
6842        let nv = n_vocab as i32;
6843        let __s_b1 = self.gpu.stream();
6844        let mut b1 = __s_b1.launch_builder(&f1);
6845        b1.arg(logits).arg(tok).arg(&mut part).arg(&nv);
6846        unsafe {
6847            b1.launch(cfg1)?;
6848        }
6849        let f2 = self.func("prob_of_token_final_f32");
6850        let cfg2 = LaunchConfig {
6851            grid_dim: (1, 1, 1),
6852            block_dim: (256, 1, 1),
6853            shared_mem_bytes: 0,
6854        };
6855        let nbi = nb as i32;
6856        let __s_b2 = self.gpu.stream();
6857        let mut b2 = __s_b2.launch_builder(&f2);
6858        b2.arg(&part).arg(&mut p).arg(&nbi);
6859        unsafe {
6860            b2.launch(cfg2)?;
6861        }
6862        Ok(p)
6863    }
6864
6865    /// Like `prob_of_token_device` but writes into a PERSISTENT `p_out` buffer (stable pointer).
6866    /// Required for CUDA-graph capture of the draft chain: the captured prob kernels must write
6867    /// where the host reads the p-min confidence between replays. Same kernels, same math.
6868    /// Slot-addressed twin of `prob_of_token_device_into`: token read from `tok_all[tok_idx]`
6869    /// (a view at the slot), probability written to `p_out[p_idx]` — same two kernels, the
6870    /// pointers just land mid-buffer. Zero-sync (gemma confidence-adaptive draft depth).
6871    pub fn prob_of_token_device_col(
6872        &self,
6873        logits: &CudaSlice<f32>,
6874        tok_all: &CudaSlice<u32>,
6875        tok_idx: usize,
6876        p_out: &mut CudaSlice<f32>,
6877        p_idx: usize,
6878        n_vocab: usize,
6879    ) -> Result<(), Box<dyn std::error::Error>> {
6880        let tok_v = tok_all.slice(tok_idx..tok_idx + 1);
6881        let mut p_v = p_out.slice_mut(p_idx..p_idx + 1);
6882        let nb = ARGMAX_NB;
6883        let mut part = self.alloc_uninit::<f32>(nb)?;
6884        let f1 = self.func("prob_of_token_partial_f32");
6885        let cfg1 = LaunchConfig {
6886            grid_dim: (nb as u32, 1, 1),
6887            block_dim: (256, 1, 1),
6888            shared_mem_bytes: 0,
6889        };
6890        let nv = n_vocab as i32;
6891        let __s_b1 = self.gpu.stream();
6892        let mut b1 = __s_b1.launch_builder(&f1);
6893        b1.arg(logits).arg(&tok_v).arg(&mut part).arg(&nv);
6894        unsafe {
6895            b1.launch(cfg1)?;
6896        }
6897        let f2 = self.func("prob_of_token_final_f32");
6898        let cfg2 = LaunchConfig {
6899            grid_dim: (1, 1, 1),
6900            block_dim: (256, 1, 1),
6901            shared_mem_bytes: 0,
6902        };
6903        let nbi = nb as i32;
6904        let __s_b2 = self.gpu.stream();
6905        let mut b2 = __s_b2.launch_builder(&f2);
6906        b2.arg(&part).arg(&mut p_v).arg(&nbi);
6907        unsafe {
6908            b2.launch(cfg2)?;
6909        }
6910        Ok(())
6911    }
6912
6913    pub fn prob_of_token_device_into(
6914        &self,
6915        logits: &CudaSlice<f32>,
6916        tok: &CudaSlice<u32>,
6917        p_out: &mut CudaSlice<f32>,
6918        n_vocab: usize,
6919    ) -> Result<(), Box<dyn std::error::Error>> {
6920        let nb = ARGMAX_NB;
6921        let mut part = self.alloc_uninit::<f32>(nb)?;
6922        let f1 = self.func("prob_of_token_partial_f32");
6923        let cfg1 = LaunchConfig {
6924            grid_dim: (nb as u32, 1, 1),
6925            block_dim: (256, 1, 1),
6926            shared_mem_bytes: 0,
6927        };
6928        let nv = n_vocab as i32;
6929        let __s_b1 = self.gpu.stream();
6930        let mut b1 = __s_b1.launch_builder(&f1);
6931        b1.arg(logits).arg(tok).arg(&mut part).arg(&nv);
6932        unsafe {
6933            b1.launch(cfg1)?;
6934        }
6935        let f2 = self.func("prob_of_token_final_f32");
6936        let cfg2 = LaunchConfig {
6937            grid_dim: (1, 1, 1),
6938            block_dim: (256, 1, 1),
6939            shared_mem_bytes: 0,
6940        };
6941        let nbi = nb as i32;
6942        let __s_b2 = self.gpu.stream();
6943        let mut b2 = __s_b2.launch_builder(&f2);
6944        b2.arg(&part).arg(p_out).arg(&nbi);
6945        unsafe {
6946            b2.launch(cfg2)?;
6947        }
6948        Ok(())
6949    }
6950
6951    pub fn argmax_token_device(
6952        &self,
6953        logits: &CudaSlice<f32>,
6954        n_vocab: usize,
6955    ) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
6956        let mut tok = unsafe { self.gpu.stream().alloc::<u32>(1)? };
6957        self.argmax_token_device_into(logits, &mut tok, n_vocab)?;
6958        Ok(tok)
6959    }
6960    /// Like `argmax_token_device` but writes into a PERSISTENT `tok` buffer (stable pointer) instead
6961    /// of allocating a fresh one. Required for CUDA-graph capture: the captured argmax must write the
6962    /// next token into the SAME device buffer the next replay's embed_gather reads, so the buffer
6963    /// pointer is baked once and the token id never round-trips to host inside steady state. The
6964    /// pass-1 partials scratch (`argmax_partials`) is also a resident stable-pointer buffer so both
6965    /// captured passes bake fixed addresses.
6966    pub fn argmax_token_device_into(
6967        &self,
6968        logits: &CudaSlice<f32>,
6969        tok: &mut CudaSlice<u32>,
6970        n_vocab: usize,
6971    ) -> Result<(), Box<dyn std::error::Error>> {
6972        let nb = ARGMAX_NB;
6973        let f1 = self.func("argmax_partial_f32");
6974        let f2 = self.func("argmax_final_f32");
6975        let mut guard = self.argmax_partials.lock().unwrap();
6976        if guard.is_none() {
6977            // allocate ONCE; under generate_graph this runs in the tracking-off prime window so the
6978            // buffers carry no cudarc events (illegal inside capture).
6979            let pv = self.gpu.stream().alloc_zeros::<f32>(nb)?;
6980            let pi = self.gpu.stream().alloc_zeros::<i32>(nb)?;
6981            *guard = Some((pv, pi));
6982        }
6983        let (part_v, part_i) = guard.as_mut().unwrap();
6984        let nv = n_vocab as i32;
6985        let nbi = nb as i32;
6986        // pass 1: NB blocks x 256 threads grid-stride scan -> per-block (val, idx) partials.
6987        let cfg1 = LaunchConfig {
6988            grid_dim: (nb as u32, 1, 1),
6989            block_dim: (256, 1, 1),
6990            shared_mem_bytes: 0,
6991        };
6992        let __s_b1 = self.gpu.stream();
6993        let mut b1 = __s_b1.launch_builder(&f1);
6994        b1.arg(logits).arg(&mut *part_v).arg(&mut *part_i).arg(&nv);
6995        unsafe {
6996            b1.launch(cfg1)?;
6997        }
6998        // pass 2: one block reduces NB partials -> token_out[0].
6999        let cfg2 = LaunchConfig {
7000            grid_dim: (1, 1, 1),
7001            block_dim: (256, 1, 1),
7002            shared_mem_bytes: 0,
7003        };
7004        let __s_b2 = self.gpu.stream();
7005        let mut b2 = __s_b2.launch_builder(&f2);
7006        b2.arg(&*part_v).arg(&*part_i).arg(tok).arg(&nbi);
7007        unsafe {
7008            b2.launch(cfg2)?;
7009        }
7010        Ok(())
7011    }
7012    /// Column-`col` device argmax over a stacked verify-logits buffer [t, n_vocab] (spec accept
7013    /// walk): toks[out_idx] = argmax(logits[col*n_vocab .. (col+1)*n_vocab]). SAME 2-pass kernels
7014    /// and tie-break contract as `argmax_token_device_into` (bit-identical to host argmax,
7015    /// argmax_gate-validated) — only the input pointer (a column view) and the output slot differ.
7016    /// Lets the accept walk read ONE [t] u32 instead of dtoh'ing the full [t, n_vocab] logits.
7017    pub fn argmax_token_device_col(
7018        &self,
7019        logits: &CudaSlice<f32>,
7020        col: usize,
7021        n_vocab: usize,
7022        toks: &mut CudaSlice<u32>,
7023        out_idx: usize,
7024    ) -> Result<(), Box<dyn std::error::Error>> {
7025        let nb = ARGMAX_NB;
7026        let f1 = self.func("argmax_partial_f32");
7027        let f2 = self.func("argmax_final_f32");
7028        let mut guard = self.argmax_partials.lock().unwrap();
7029        if guard.is_none() {
7030            let pv = self.gpu.stream().alloc_zeros::<f32>(nb)?;
7031            let pi = self.gpu.stream().alloc_zeros::<i32>(nb)?;
7032            *guard = Some((pv, pi));
7033        }
7034        let (part_v, part_i) = guard.as_mut().unwrap();
7035        let col_view = logits.slice(col * n_vocab..(col + 1) * n_vocab);
7036        let nv = n_vocab as i32;
7037        let nbi = nb as i32;
7038        let cfg1 = LaunchConfig {
7039            grid_dim: (nb as u32, 1, 1),
7040            block_dim: (256, 1, 1),
7041            shared_mem_bytes: 0,
7042        };
7043        let __s_b1 = self.gpu.stream();
7044        let mut b1 = __s_b1.launch_builder(&f1);
7045        b1.arg(&col_view)
7046            .arg(&mut *part_v)
7047            .arg(&mut *part_i)
7048            .arg(&nv);
7049        unsafe {
7050            b1.launch(cfg1)?;
7051        }
7052        let mut tok_view = toks.slice_mut(out_idx..out_idx + 1);
7053        let cfg2 = LaunchConfig {
7054            grid_dim: (1, 1, 1),
7055            block_dim: (256, 1, 1),
7056            shared_mem_bytes: 0,
7057        };
7058        let __s_b2 = self.gpu.stream();
7059        let mut b2 = __s_b2.launch_builder(&f2);
7060        b2.arg(&*part_v).arg(&*part_i).arg(&mut tok_view).arg(&nbi);
7061        unsafe {
7062            b2.launch(cfg2)?;
7063        }
7064        Ok(())
7065    }
7066    /// Read back a device u32 buffer (the spec accept walk's [t] per-column argmax tokens).
7067    pub fn htod_u32_v(&self, v: &[u32]) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
7068        Ok(self.gpu.stream().clone_htod(v)?)
7069    }
7070    pub fn dtoh_u32(&self, d: &CudaSlice<u32>) -> Result<Vec<u32>, Box<dyn std::error::Error>> {
7071        let v = self.gpu.stream().clone_dtoh(d)?;
7072        self.gpu.stream().synchronize()?;
7073        Ok(v)
7074    }
7075    /// Allocate a zeroed device u32 buffer (persistent spec-loop prediction slots).
7076    /// H2D into an EXISTING u32 buffer (stable pointer — the per-step grammar-mask upload:
7077    /// contents change every step, the address must not, so a captured graph can read it).
7078    pub fn htod_u32_into(
7079        &self,
7080        dst: &mut CudaSlice<u32>,
7081        src: &[u32],
7082    ) -> Result<(), Box<dyn std::error::Error>> {
7083        let mut view = dst.slice_mut(0..src.len());
7084        self.gpu.stream().memcpy_htod(src, &mut view)?;
7085        Ok(())
7086    }
7087
7088    /// H2D into an existing i32 buffer. OPTIPIPE uses this to refresh a stage-local saved-len
7089    /// table without changing the device address its reconcile kernel consumes.
7090    pub fn htod_i32_into(
7091        &self,
7092        dst: &mut CudaSlice<i32>,
7093        src: &[i32],
7094    ) -> Result<(), Box<dyn std::error::Error>> {
7095        let mut view = dst.slice_mut(0..src.len());
7096        self.gpu.stream().memcpy_htod(src, &mut view)?;
7097        Ok(())
7098    }
7099
7100    pub fn alloc_u32_zeroed(&self, n: usize) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
7101        let s = self.gpu.stream().alloc_zeros::<u32>(n)?;
7102        self.keep_if_capturing(&s);
7103        Ok(s)
7104    }
7105    /// embed_gather into a PERSISTENT `x_out` buffer (stable pointer) for CUDA-graph capture (the
7106    /// embed output starts the per-step kernel chain and must be at a fixed address across replays).
7107    pub fn embed_gather_device_into(
7108        &self,
7109        embd: &CudaSlice<u8>,
7110        token_d: &CudaSlice<u32>,
7111        x_out: &mut CudaSlice<f32>,
7112        n_embd: usize,
7113        qtype: i32,
7114        row_bytes: usize,
7115    ) -> Result<(), Box<dyn std::error::Error>> {
7116        let f = self.func("embed_gather_u32");
7117        let cfg = LaunchConfig {
7118            grid_dim: (((n_embd as u32 + 255) / 256).max(1), 1, 1),
7119            block_dim: (256, 1, 1),
7120            shared_mem_bytes: 0,
7121        };
7122        let (ne, qt, rb) = (n_embd as i32, qtype, row_bytes as i64);
7123        let __s_b = self.gpu.stream();
7124        let mut b = __s_b.launch_builder(&f);
7125        b.arg(embd)
7126            .arg(token_d)
7127            .arg(x_out)
7128            .arg(&ne)
7129            .arg(&qt)
7130            .arg(&rb);
7131        unsafe {
7132            b.launch(cfg)?;
7133        }
7134        Ok(())
7135    }
7136    /// Read a [1] i32 device counter (pos / seqlen) back to host. Tiny D2H + sync.
7137    pub fn dtoh_i32_one(&self, d: &CudaSlice<i32>) -> Result<i32, Box<dyn std::error::Error>> {
7138        let v = self.gpu.stream().clone_dtoh(d)?;
7139        self.gpu.stream().synchronize()?;
7140        Ok(v[0])
7141    }
7142    /// Set a [1] i32 device counter IN PLACE (keeps the buffer pointer stable — required for the
7143    /// graph-resident pos/seqlen counters whose addresses are baked into captured graphs). Restores
7144    /// the counter value after the throwaway capture warmups corrupt it.
7145    /// ASYNC i32 single-slot store (value rides the kernel arg — no host-memory transfer/sync).
7146    /// The graph-arc device-len counters use this; set_i32_one below is the SYNCING pageable
7147    /// copy (fine at stream-idle boundaries, poison mid-round).
7148    pub fn i32_set_k(
7149        &self,
7150        dst: &mut CudaSlice<i32>,
7151        v: i32,
7152    ) -> Result<(), Box<dyn std::error::Error>> {
7153        let f = self.func("i32_set_k");
7154        let cfg = LaunchConfig {
7155            grid_dim: (1, 1, 1),
7156            block_dim: (1, 1, 1),
7157            shared_mem_bytes: 0,
7158        };
7159        let idx = 0i32;
7160        let __s_b = self.gpu.stream();
7161        let mut b = __s_b.launch_builder(&f);
7162        b.arg(dst).arg(&v).arg(&idx);
7163        unsafe {
7164            b.launch(cfg)?;
7165        }
7166        Ok(())
7167    }
7168
7169    pub fn set_i32_one(
7170        &self,
7171        d: &mut CudaSlice<i32>,
7172        v: i32,
7173    ) -> Result<(), Box<dyn std::error::Error>> {
7174        self.gpu.stream().memcpy_htod(&[v], d)?;
7175        Ok(())
7176    }
7177    /// Set a [1] u32 device buffer IN PLACE (stable pointer) — for the resident `token_d` counter
7178    /// during priming / capture-state restore.
7179    pub fn set_u32_one(
7180        &self,
7181        d: &mut CudaSlice<u32>,
7182        v: u32,
7183    ) -> Result<(), Box<dyn std::error::Error>> {
7184        self.gpu.stream().memcpy_htod(&[v], d)?;
7185        Ok(())
7186    }
7187    /// Read back a [1] u32 device buffer (the argmax token). One tiny D2H + sync.
7188    pub fn dtoh_u32_one(&self, d: &CudaSlice<u32>) -> Result<u32, Box<dyn std::error::Error>> {
7189        let v = self.gpu.stream().clone_dtoh(d)?;
7190        self.gpu.stream().synchronize()?;
7191        Ok(v[0])
7192    }
7193    /// Upload raw bytes to a resident device u8 buffer (e.g. the embed table for device gather).
7194    pub fn upload_u8(&self, bytes: &[u8]) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
7195        Ok(self.gpu.stream().clone_htod(bytes)?)
7196    }
7197    /// Embed-from-device (CUDA-GRAPH-PLAN Phase 1): gather+dequant the row for the token id in
7198    /// `token_d[0]` from the resident embed table `embd` -> x_out[n_embd]. Bit-identical to host
7199    /// EmbedHost::gather (same per-dtype `deq`). No host round-trip of the token id.
7200    pub fn embed_gather_device(
7201        &self,
7202        embd: &CudaSlice<u8>,
7203        token_d: &CudaSlice<u32>,
7204        n_embd: usize,
7205        qtype: i32,
7206        row_bytes: usize,
7207    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7208        let f = self.func("embed_gather_u32");
7209        let mut x = self.alloc_uninit::<f32>(n_embd)?;
7210        let cfg = LaunchConfig {
7211            grid_dim: (((n_embd as u32 + 255) / 256).max(1), 1, 1),
7212            block_dim: (256, 1, 1),
7213            shared_mem_bytes: 0,
7214        };
7215        let (ne, qt, rb) = (n_embd as i32, qtype, row_bytes as i64);
7216        let __s_b = self.gpu.stream();
7217        let mut b = __s_b.launch_builder(&f);
7218        b.arg(embd)
7219            .arg(token_d)
7220            .arg(&mut x)
7221            .arg(&ne)
7222            .arg(&qt)
7223            .arg(&rb);
7224        unsafe {
7225            b.launch(cfg)?;
7226        }
7227        Ok(x)
7228    }
7229
7230    /// T-token device embed gather (spec verify/replay): tokens uploaded as a tiny [T] u32 htod,
7231    /// rows dequanted on-device -> x[T, n_embd]. Replaces host per-row dequant + T*n_embd*4B htod
7232    /// (nsys: 84% of spec API time was HtoD). Bit-identical rows (same per-dtype deq).
7233    pub fn embed_gather_device_t(
7234        &self,
7235        embd: &CudaSlice<u8>,
7236        tokens: &[u32],
7237        n_embd: usize,
7238        qtype: i32,
7239        row_bytes: usize,
7240    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7241        let t = tokens.len();
7242        let tok_d = self.gpu.stream().clone_htod(tokens)?;
7243        let f = self.func("embed_gather_u32_t");
7244        let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
7245        let cfg = LaunchConfig {
7246            grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
7247            block_dim: (256, 1, 1),
7248            shared_mem_bytes: 0,
7249        };
7250        let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
7251        let __s_b = self.gpu.stream();
7252        let mut b = __s_b.launch_builder(&f);
7253        b.arg(embd)
7254            .arg(&tok_d)
7255            .arg(&mut x)
7256            .arg(&ne)
7257            .arg(&qt)
7258            .arg(&rb)
7259            .arg(&ti);
7260        unsafe {
7261            b.launch(cfg)?;
7262        }
7263        Ok(x)
7264    }
7265
7266    /// T-token embed gather from a DEVICE token buffer (round-stream stage c: the verify tokens
7267    /// are assembled on-device from the draft-chain pack slots; no host round trip). Same kernel
7268    /// as embed_gather_device_t — bit-identical rows.
7269    /// embed_gather over a token VIEW (spec round: tokens live in the round's batch buffer).
7270    pub fn embed_gather_device_tv(
7271        &self,
7272        embd: &CudaSlice<u8>,
7273        tok_v: &cudarc::driver::CudaView<u32>,
7274        t: usize,
7275        n_embd: usize,
7276        qtype: i32,
7277        row_bytes: usize,
7278    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7279        let f = self.func("embed_gather_u32_t");
7280        let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
7281        let cfg = LaunchConfig {
7282            grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
7283            block_dim: (256, 1, 1),
7284            shared_mem_bytes: 0,
7285        };
7286        let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
7287        let __s_b = self.gpu.stream();
7288        let mut b = __s_b.launch_builder(&f);
7289        b.arg(embd)
7290            .arg(tok_v)
7291            .arg(&mut x)
7292            .arg(&ne)
7293            .arg(&qt)
7294            .arg(&rb)
7295            .arg(&ti);
7296        unsafe {
7297            b.launch(cfg)?;
7298        }
7299        Ok(x)
7300    }
7301
7302    pub fn embed_gather_device_td(
7303        &self,
7304        embd: &CudaSlice<u8>,
7305        tok_d: &CudaSlice<u32>,
7306        t: usize,
7307        n_embd: usize,
7308        qtype: i32,
7309        row_bytes: usize,
7310    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7311        let f = self.func("embed_gather_u32_t");
7312        let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
7313        let cfg = LaunchConfig {
7314            grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
7315            block_dim: (256, 1, 1),
7316            shared_mem_bytes: 0,
7317        };
7318        let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
7319        let __s_b = self.gpu.stream();
7320        let mut b = __s_b.launch_builder(&f);
7321        b.arg(embd)
7322            .arg(tok_d)
7323            .arg(&mut x)
7324            .arg(&ne)
7325            .arg(&qt)
7326            .arg(&rb)
7327            .arg(&ti);
7328        unsafe {
7329            b.launch(cfg)?;
7330        }
7331        Ok(x)
7332    }
7333
7334    /// Uninitialized device buffer — SKIPS the memset that `alloc_zeros` always issues. Decode
7335    /// profile (nsys): ~1050 memsets/token = 6.5% of decode GPU time + ~half the launch count, the
7336    /// dominant contributor to the 19% inter-kernel idle gap and a blocker for clean CUDA-graph
7337    /// capture. Use ONLY for buffers a kernel FULLY overwrites (every element written, no `+=`).
7338    /// SAFETY: caller guarantees the producing kernel writes every element before any read.
7339    #[inline]
7340    /// Keep an allocation alive for the current capture (no-op when retain mode is off).
7341    fn keep_if_capturing<T: cudarc::driver::DeviceRepr + Send + 'static>(&self, s: &CudaSlice<T>) {
7342        if self
7343            .capture_keep_on
7344            .load(std::sync::atomic::Ordering::Relaxed)
7345        {
7346            self.capture_keep.lock().unwrap().push(Box::new(s.clone()));
7347        }
7348    }
7349
7350    fn alloc_uninit<T: cudarc::driver::DeviceRepr + Send + 'static>(
7351        &self,
7352        n: usize,
7353    ) -> Result<CudaSlice<T>, Box<dyn std::error::Error>> {
7354        let mut s = unsafe { self.gpu.stream().alloc::<T>(n)? };
7355        // MEMRA_DEBUG_ZERO_ALLOCS=1 (task #14 defect hunt): memset EVERY engine allocation —
7356        // the global uninit-read discriminator (the prime-fn-scoped zeroing experiment could
7357        // not cover engine-internal buffers). Debug-only: massive launch overhead.
7358        {
7359            static Z: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7360            if *Z.get_or_init(|| std::env::var("MEMRA_DEBUG_ZERO_ALLOCS").as_deref() == Ok("1")) {
7361                // raw D8 memset (T lacks ValidAsZeroBits in the generic bound)
7362                use cudarc::driver::DevicePtrMut;
7363                let n_bytes = s.len() * std::mem::size_of::<T>();
7364                let stream = self.gpu.stream();
7365                let (p_, _g) = s.device_ptr_mut(&stream);
7366                unsafe {
7367                    cudarc::driver::sys::cuMemsetD8Async(p_, 0, n_bytes, stream.cu_stream())
7368                        .result()?;
7369                }
7370            }
7371        }
7372        self.keep_if_capturing(&s);
7373        Ok(s)
7374    }
7375
7376    /// Public f32 uninitialized scratch (see `alloc_uninit`). For decode/forward scratch a kernel
7377    /// fully overwrites. SAFETY: producing kernel must write every element before any read.
7378    /// Uninitialized q8_1 activation pair (int8 + per-32 scales) — the fa combine q8-emit
7379    /// consumers alloc through this (m=1 decode arms).
7380    pub fn uninit_q8_pair(
7381        &self,
7382        n: usize,
7383    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7384        Ok((
7385            self.alloc_uninit::<i8>(n)?,
7386            self.alloc_uninit::<f32>(n / 32)?,
7387        ))
7388    }
7389
7390    pub fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7391        self.alloc_uninit::<f32>(n)
7392    }
7393
7394    /// i8 uninitialized scratch (same contract as `uninit`).
7395    pub fn alloc_i8_uninit(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
7396        self.alloc_uninit::<i8>(n)
7397    }
7398
7399    /// RMSNorm: x[ncols,nrows] row-major, weight[ncols] -> dst. One block/row, 256 threads.
7400    /// gemma4: 3 rms_norms of the SAME input in one launch (one reduction, three weights).
7401    /// Per-output bit-identical to three rms_norm calls (verbatim reduction/scale chain).
7402    #[allow(clippy::too_many_arguments)]
7403    pub fn rms_norm3(
7404        &self,
7405        x: &CudaSlice<f32>,
7406        w0: &CudaSlice<f32>,
7407        w1: &CudaSlice<f32>,
7408        w2: &CudaSlice<f32>,
7409        d0: &mut CudaSlice<f32>,
7410        d1: &mut CudaSlice<f32>,
7411        d2: &mut CudaSlice<f32>,
7412        ncols: usize,
7413        nrows: usize,
7414        eps: f32,
7415    ) -> Result<(), Box<dyn std::error::Error>> {
7416        let f = self.func("rms_norm3_f32");
7417        let cfg = LaunchConfig {
7418            grid_dim: (nrows as u32, 1, 1),
7419            block_dim: (rms_block(), 1, 1),
7420            shared_mem_bytes: 0,
7421        };
7422        let (nc, e) = (ncols as i32, eps);
7423        let __s_b = self.gpu.stream();
7424        let mut b = __s_b.launch_builder(&f);
7425        b.arg(x)
7426            .arg(w0)
7427            .arg(w1)
7428            .arg(w2)
7429            .arg(d0)
7430            .arg(d1)
7431            .arg(d2)
7432            .arg(&nc)
7433            .arg(&e);
7434        unsafe {
7435            b.launch(cfg)?;
7436        }
7437        Ok(())
7438    }
7439
7440    /// gemma4 fused q/k/v head norms (one launch, per-row rms_norm_f32-verbatim).
7441    #[allow(clippy::too_many_arguments)]
7442    /// True when the warp-per-row qkv norm would engage for (rows, ncols) — the emit lane
7443    /// piggybacks on the same conditions.
7444    pub fn qkvnorm_w_on_prefill(rows: usize, ncols: usize) -> bool {
7445        static WARP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7446        *WARP_ON.get_or_init(|| {
7447            std::env::var("MEMRA_QKVNORM_W")
7448                .map(|v| v != "0")
7449                .unwrap_or(true)
7450        }) && ncols % 4 == 0
7451            && rows >= 64
7452    }
7453
7454    /// w4 norm with bf16 V EMIT (31B glue lane): the v segment also writes its normed rows as
7455    /// bf16 (the FA V operand — bit-identical to a post-hoc f32_to_bf16). Prefill-depth only.
7456    #[allow(clippy::too_many_arguments)]
7457    pub fn rms_norm_qkv_w4b(
7458        &self,
7459        q: &CudaSlice<f32>,
7460        k: &CudaSlice<f32>,
7461        v: &CudaSlice<f32>,
7462        wq: &CudaSlice<f32>,
7463        wk: &CudaSlice<f32>,
7464        wv: &CudaSlice<f32>,
7465        dq: &mut CudaSlice<f32>,
7466        dk: &mut CudaSlice<f32>,
7467        dv: &mut CudaSlice<f32>,
7468        dvb: &mut CudaSlice<u8>,
7469        ncols: usize,
7470        rq: usize,
7471        rk: usize,
7472        eps: f32,
7473        vf16: bool,
7474    ) -> Result<(), Box<dyn std::error::Error>> {
7475        assert!(ncols % 4 == 0 && rq + 2 * rk >= 64);
7476        let f = self.func("rms_norm_qkv_w4b_f32");
7477        let rows = (rq + 2 * rk) as u32;
7478        let cfg = LaunchConfig {
7479            grid_dim: (rows.div_ceil(8), 1, 1),
7480            block_dim: (256, 1, 1),
7481            shared_mem_bytes: 0,
7482        };
7483        let (nc, rqi, rki, rvi, e) = (ncols as i32, rq as i32, rk as i32, rk as i32, eps);
7484        let vf = vf16 as i32;
7485        let __s_b = self.gpu.stream();
7486        let mut b = __s_b.launch_builder(&f);
7487        b.arg(q)
7488            .arg(k)
7489            .arg(v)
7490            .arg(wq)
7491            .arg(wk)
7492            .arg(wv)
7493            .arg(dq)
7494            .arg(dk)
7495            .arg(dv)
7496            .arg(&mut *dvb)
7497            .arg(&nc)
7498            .arg(&rqi)
7499            .arg(&rki)
7500            .arg(&rvi)
7501            .arg(&e)
7502            .arg(&vf);
7503        unsafe {
7504            b.launch(cfg)?;
7505        }
7506        Ok(())
7507    }
7508
7509    pub fn rms_norm_qkv(
7510        &self,
7511        q: &CudaSlice<f32>,
7512        k: &CudaSlice<f32>,
7513        v: &CudaSlice<f32>,
7514        wq: &CudaSlice<f32>,
7515        wk: &CudaSlice<f32>,
7516        wv: &CudaSlice<f32>,
7517        dq: &mut CudaSlice<f32>,
7518        dk: &mut CudaSlice<f32>,
7519        dv: &mut CudaSlice<f32>,
7520        ncols: usize,
7521        rq: usize,
7522        rk: usize,
7523        eps: f32,
7524    ) -> Result<(), Box<dyn std::error::Error>> {
7525        // Warp-per-row float4 twin (default; MEMRA_QKVNORM_W=0 reverts): the block-per-row form
7526        // spends 767us/launch on 17k+ 2KB rows at prefill depth (launch/reduce latency-bound,
7527        // ~92GB/s). Own numeric config (reduce order differs) — battery-gated.
7528        static WARP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7529        let warp_on = *WARP_ON.get_or_init(|| {
7530            std::env::var("MEMRA_QKVNORM_W")
7531                .map(|v| v != "0")
7532                .unwrap_or(true)
7533        });
7534        // rows >= 64 keeps decode (nh + 2*nkv rows) on the block-tree kernel — decode/verify/
7535        // replay numerics are untouched on every model; only prefill depth takes the new config.
7536        if warp_on && ncols % 4 == 0 && rq + 2 * rk >= 64 {
7537            let f = self.func("rms_norm_qkv_w4_f32");
7538            let rows = (rq + 2 * rk) as u32;
7539            let cfg = LaunchConfig {
7540                grid_dim: (rows.div_ceil(8), 1, 1),
7541                block_dim: (256, 1, 1),
7542                shared_mem_bytes: 0,
7543            };
7544            let (nc, rqi, rki, rvi, e) = (ncols as i32, rq as i32, rk as i32, rk as i32, eps);
7545            let __s_b = self.gpu.stream();
7546            let mut b = __s_b.launch_builder(&f);
7547            b.arg(q)
7548                .arg(k)
7549                .arg(v)
7550                .arg(wq)
7551                .arg(wk)
7552                .arg(wv)
7553                .arg(dq)
7554                .arg(dk)
7555                .arg(dv)
7556                .arg(&nc)
7557                .arg(&rqi)
7558                .arg(&rki)
7559                .arg(&rvi)
7560                .arg(&e);
7561            unsafe {
7562                b.launch(cfg)?;
7563            }
7564            return Ok(());
7565        }
7566        let f = self.func("rms_norm_qkv_f32");
7567        let grid = (rq + 2 * rk) as u32;
7568        let cfg = LaunchConfig {
7569            grid_dim: (grid, 1, 1),
7570            block_dim: (rms_block(), 1, 1),
7571            shared_mem_bytes: 0,
7572        };
7573        let (nc, rqi, rki, e) = (ncols as i32, rq as i32, rk as i32, eps);
7574        let __s_b = self.gpu.stream();
7575        let mut b = __s_b.launch_builder(&f);
7576        b.arg(q)
7577            .arg(k)
7578            .arg(v)
7579            .arg(wq)
7580            .arg(wk)
7581            .arg(wv)
7582            .arg(dq)
7583            .arg(dk)
7584            .arg(dv)
7585            .arg(&nc)
7586            .arg(&rqi)
7587            .arg(&rki)
7588            .arg(&e);
7589        unsafe {
7590            b.launch(cfg)?;
7591        }
7592        Ok(())
7593    }
7594
7595    /// gemma4 fused pair of rms_norms over two different inputs (same width).
7596    #[allow(clippy::too_many_arguments)]
7597    pub fn rms_norm2x(
7598        &self,
7599        a: &CudaSlice<f32>,
7600        bb: &CudaSlice<f32>,
7601        wa: &CudaSlice<f32>,
7602        wb: &CudaSlice<f32>,
7603        da: &mut CudaSlice<f32>,
7604        db: &mut CudaSlice<f32>,
7605        ncols: usize,
7606        nrows: usize,
7607        eps: f32,
7608    ) -> Result<(), Box<dyn std::error::Error>> {
7609        let f = self.func("rms_norm2x_f32");
7610        let cfg = LaunchConfig {
7611            grid_dim: (2 * nrows as u32, 1, 1),
7612            block_dim: (rms_block(), 1, 1),
7613            shared_mem_bytes: 0,
7614        };
7615        let (nc, nr, e) = (ncols as i32, nrows as i32, eps);
7616        let __s_b = self.gpu.stream();
7617        let mut b = __s_b.launch_builder(&f);
7618        b.arg(a)
7619            .arg(bb)
7620            .arg(wa)
7621            .arg(wb)
7622            .arg(da)
7623            .arg(db)
7624            .arg(&nc)
7625            .arg(&nr)
7626            .arg(&e);
7627        unsafe {
7628            b.launch(cfg)?;
7629        }
7630        Ok(())
7631    }
7632
7633    /// gemma4 R4: in-place final-logit softcap y = cap*tanh(y/cap).
7634    pub fn softcap(
7635        &self,
7636        y: &mut CudaSlice<f32>,
7637        cap: f32,
7638        n: usize,
7639    ) -> Result<(), Box<dyn std::error::Error>> {
7640        let f = self.func("softcap_f32");
7641        let cfg = LaunchConfig::for_num_elems(n as u32);
7642        let ni = n as i32;
7643        let __s_b = self.gpu.stream();
7644        let mut b = __s_b.launch_builder(&f);
7645        b.arg(y).arg(&cap).arg(&ni);
7646        unsafe {
7647            b.launch(cfg)?;
7648        }
7649        Ok(())
7650    }
7651
7652    /// gemma4 suppress-token mask: y[row][ids[j]] = -inf over t logits rows (fixed-arg launch —
7653    /// graph-capture safe; NOT monotonic like softcap, so it must run before any argmax).
7654    pub fn mask_ids_rows(
7655        &self,
7656        y: &mut CudaSlice<f32>,
7657        ids: &CudaSlice<i32>,
7658        n_ids: usize,
7659        n_vocab: usize,
7660        t: usize,
7661    ) -> Result<(), Box<dyn std::error::Error>> {
7662        let f = self.func("mask_ids_rows_f32");
7663        let cfg = LaunchConfig::for_num_elems((n_ids * t) as u32);
7664        let (ni, nv, ti) = (n_ids as i32, n_vocab as i32, t as i32);
7665        let __s_b = self.gpu.stream();
7666        let mut b = __s_b.launch_builder(&f);
7667        b.arg(y).arg(ids).arg(&ni).arg(&nv).arg(&ti);
7668        unsafe {
7669            b.launch(cfg)?;
7670        }
7671        Ok(())
7672    }
7673
7674    /// gemma4: res = (a+b)*c AND dst = rms_norm(res, w) in one launch.
7675    #[allow(clippy::too_many_arguments)]
7676    pub fn add_scale_rms_norm(
7677        &self,
7678        a: &CudaSlice<f32>,
7679        b_in: &CudaSlice<f32>,
7680        c: f32,
7681        w: &CudaSlice<f32>,
7682        res: &mut CudaSlice<f32>,
7683        dst: &mut CudaSlice<f32>,
7684        ncols: usize,
7685        nrows: usize,
7686        eps: f32,
7687    ) -> Result<(), Box<dyn std::error::Error>> {
7688        let f = self.func("add_scale_rms_norm_f32");
7689        let cfg = LaunchConfig {
7690            grid_dim: (nrows as u32, 1, 1),
7691            block_dim: (rms_block(), 1, 1),
7692            shared_mem_bytes: 0,
7693        };
7694        let (nc, e2) = (ncols as i32, eps);
7695        let __s_b = self.gpu.stream();
7696        let mut b = __s_b.launch_builder(&f);
7697        b.arg(a)
7698            .arg(b_in)
7699            .arg(&c)
7700            .arg(w)
7701            .arg(res)
7702            .arg(dst)
7703            .arg(&nc)
7704            .arg(&e2);
7705        unsafe {
7706            b.launch(cfg)?;
7707        }
7708        Ok(())
7709    }
7710
7711    /// gemma4: res = (a+b)*c AND the next layer's attn_norm EMITTED q8_1 in one launch.
7712    /// Quantize epilogue bit-identical to quantize_q8_1 (the rms_norm_q8_1 form).
7713    #[allow(clippy::too_many_arguments)]
7714    pub fn add_scale_rms_norm_q8_1(
7715        &self,
7716        a: &CudaSlice<f32>,
7717        b_in: &CudaSlice<f32>,
7718        c: f32,
7719        w: &CudaSlice<f32>,
7720        res: &mut CudaSlice<f32>,
7721        ncols: usize,
7722        nrows: usize,
7723        eps: f32,
7724    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7725        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
7726        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
7727        let (nc, e2) = (ncols as i32, eps);
7728        if Self::pdl_on() && Self::pdl_wb_on() {
7729            {
7730                use cudarc::driver::{DevicePtr, DevicePtrMut};
7731                let s = &self.gpu.stream();
7732                let (pa, _g0) = a.device_ptr(s);
7733                let (pb, _g1) = b_in.device_ptr(s);
7734                let (pw, _g2) = w.device_ptr(s);
7735                let (pr, _g3) = res.device_ptr_mut(s);
7736                let (pq, _g4) = out_q.device_ptr_mut(s);
7737                let (pd, _g5) = out_d.device_ptr_mut(s);
7738                let mut ps = [
7739                    &pa as *const _ as *mut std::ffi::c_void,
7740                    &pb as *const _ as *mut _,
7741                    &c as *const _ as *mut _,
7742                    &pw as *const _ as *mut _,
7743                    &pr as *const _ as *mut _,
7744                    &pq as *const _ as *mut _,
7745                    &pd as *const _ as *mut _,
7746                    &nc as *const _ as *mut _,
7747                    &e2 as *const _ as *mut _,
7748                ];
7749                unsafe {
7750                    self.launch_pdl(
7751                        "add_scale_rms_norm_q8_1",
7752                        (nrows as u32, 1, 1),
7753                        (rms_block(), 1, 1),
7754                        &mut ps,
7755                    )?;
7756                }
7757            }
7758            return Ok((out_q, out_d));
7759        }
7760        let f = self.func("add_scale_rms_norm_q8_1");
7761        let cfg = LaunchConfig {
7762            grid_dim: (nrows as u32, 1, 1),
7763            block_dim: (rms_block(), 1, 1),
7764            shared_mem_bytes: 0,
7765        };
7766        let __s_b = self.gpu.stream();
7767        let mut b = __s_b.launch_builder(&f);
7768        b.arg(a)
7769            .arg(b_in)
7770            .arg(&c)
7771            .arg(w)
7772            .arg(res)
7773            .arg(&mut out_q)
7774            .arg(&mut out_d)
7775            .arg(&nc)
7776            .arg(&e2);
7777        unsafe {
7778            b.launch(cfg)?;
7779        }
7780        Ok((out_q, out_d))
7781    }
7782
7783    /// Slot-fed add_scale_rms_norm_q8_1 twin (alloc-free capture lane).
7784    #[allow(clippy::too_many_arguments)]
7785    pub fn add_scale_rms_norm_q8_1_into(
7786        &self,
7787        a: &CudaSlice<f32>,
7788        b_in: &CudaSlice<f32>,
7789        c: f32,
7790        w: &CudaSlice<f32>,
7791        res: &mut CudaSlice<f32>,
7792        ncols: usize,
7793        nrows: usize,
7794        eps: f32,
7795        out_q: &mut CudaSlice<i8>,
7796        out_d: &mut CudaSlice<f32>,
7797    ) -> Result<(), Box<dyn std::error::Error>> {
7798        debug_assert!(out_q.len() >= nrows * ncols && out_d.len() >= nrows * (ncols / 32));
7799        let (nc, e2) = (ncols as i32, eps);
7800        if Self::pdl_on() && Self::pdl_wb_on() {
7801            use cudarc::driver::{DevicePtr, DevicePtrMut};
7802            let s = &self.gpu.stream();
7803            let (pa, _g0) = a.device_ptr(s);
7804            let (pb, _g1) = b_in.device_ptr(s);
7805            let (pw, _g2) = w.device_ptr(s);
7806            let (pr, _g3) = res.device_ptr_mut(s);
7807            let (pq, _g4) = out_q.device_ptr_mut(s);
7808            let (pd, _g5) = out_d.device_ptr_mut(s);
7809            let mut ps = [
7810                &pa as *const _ as *mut std::ffi::c_void,
7811                &pb as *const _ as *mut _,
7812                &c as *const _ as *mut _,
7813                &pw as *const _ as *mut _,
7814                &pr as *const _ as *mut _,
7815                &pq as *const _ as *mut _,
7816                &pd as *const _ as *mut _,
7817                &nc as *const _ as *mut _,
7818                &e2 as *const _ as *mut _,
7819            ];
7820            unsafe {
7821                self.launch_pdl(
7822                    "add_scale_rms_norm_q8_1",
7823                    (nrows as u32, 1, 1),
7824                    (rms_block(), 1, 1),
7825                    &mut ps,
7826                )?;
7827            }
7828            return Ok(());
7829        }
7830        let f = self.func("add_scale_rms_norm_q8_1");
7831        let cfg = LaunchConfig {
7832            grid_dim: (nrows as u32, 1, 1),
7833            block_dim: (rms_block(), 1, 1),
7834            shared_mem_bytes: 0,
7835        };
7836        let __s_b = self.gpu.stream();
7837        let mut b = __s_b.launch_builder(&f);
7838        b.arg(a)
7839            .arg(b_in)
7840            .arg(&c)
7841            .arg(w)
7842            .arg(res)
7843            .arg(&mut *out_q)
7844            .arg(&mut *out_d)
7845            .arg(&nc)
7846            .arg(&e2);
7847        unsafe {
7848            b.launch(cfg)?;
7849        }
7850        Ok(())
7851    }
7852
7853    /// E4B glue fusion: rms(a, wa) prologue + the add_scale_rms_norm_q8_1 program — one launch
7854    /// replaces the per-layer rms_norm_f32(y) + emit pair in the PLE tail.
7855    #[allow(clippy::too_many_arguments)]
7856    pub fn rms_pre_add_scale_rms_norm_q8_1(
7857        &self,
7858        a: &CudaSlice<f32>,
7859        wa: &CudaSlice<f32>,
7860        b_in: &CudaSlice<f32>,
7861        c: f32,
7862        w: &CudaSlice<f32>,
7863        res: &mut CudaSlice<f32>,
7864        ncols: usize,
7865        nrows: usize,
7866        eps: f32,
7867    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7868        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
7869        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
7870        let (nc, e2) = (ncols as i32, eps);
7871        if Self::pdl_on() {
7872            {
7873                use cudarc::driver::{DevicePtr, DevicePtrMut};
7874                let s = &self.gpu.stream();
7875                let (pa, _g0) = a.device_ptr(s);
7876                let (pwa, _g1) = wa.device_ptr(s);
7877                let (pb, _g2) = b_in.device_ptr(s);
7878                let (pw, _g3) = w.device_ptr(s);
7879                let (pr, _g4) = res.device_ptr_mut(s);
7880                let (pq, _g5) = out_q.device_ptr_mut(s);
7881                let (pd, _g6) = out_d.device_ptr_mut(s);
7882                let mut ps = [
7883                    &pa as *const _ as *mut std::ffi::c_void,
7884                    &pwa as *const _ as *mut _,
7885                    &pb as *const _ as *mut _,
7886                    &c as *const _ as *mut _,
7887                    &pw as *const _ as *mut _,
7888                    &pr as *const _ as *mut _,
7889                    &pq as *const _ as *mut _,
7890                    &pd as *const _ as *mut _,
7891                    &nc as *const _ as *mut _,
7892                    &e2 as *const _ as *mut _,
7893                ];
7894                unsafe {
7895                    self.launch_pdl(
7896                        "rms_pre_add_scale_rms_norm_q8_1",
7897                        (nrows as u32, 1, 1),
7898                        (rms_block(), 1, 1),
7899                        &mut ps,
7900                    )?;
7901                }
7902            }
7903            return Ok((out_q, out_d));
7904        }
7905        let f = self.func("rms_pre_add_scale_rms_norm_q8_1");
7906        let cfg = LaunchConfig {
7907            grid_dim: (nrows as u32, 1, 1),
7908            block_dim: (rms_block(), 1, 1),
7909            shared_mem_bytes: 0,
7910        };
7911        let __s_b = self.gpu.stream();
7912        let mut b = __s_b.launch_builder(&f);
7913        b.arg(a)
7914            .arg(wa)
7915            .arg(b_in)
7916            .arg(&c)
7917            .arg(w)
7918            .arg(res)
7919            .arg(&mut out_q)
7920            .arg(&mut out_d)
7921            .arg(&nc)
7922            .arg(&e2);
7923        unsafe {
7924            b.launch(cfg)?;
7925        }
7926        Ok((out_q, out_d))
7927    }
7928
7929    /// GELU(tanh)*up with the activation emitted q8_1 alongside f32 (glue-fusion lane): the
7930    /// consumer matmul rides matmul_pre, killing its standalone quantize_q8_1 launch.
7931    pub fn gelu_tanh_mul_q8_1(
7932        &self,
7933        gate: &CudaSlice<f32>,
7934        up: &cudarc::driver::CudaView<f32>,
7935        act: &mut CudaSlice<f32>,
7936        ncols: usize,
7937        nrows: usize,
7938    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7939        debug_assert!(ncols % 128 == 0);
7940        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
7941        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
7942        let nc = ncols as i32;
7943        if Self::pdl_on() {
7944            {
7945                use cudarc::driver::{DevicePtr, DevicePtrMut};
7946                let s = &self.gpu.stream();
7947                let (pg, _g0) = gate.device_ptr(s);
7948                let (pu, _g1) = up.device_ptr(s);
7949                let (pact, _g2) = act.device_ptr_mut(s);
7950                let (pq, _g3) = out_q.device_ptr_mut(s);
7951                let (pd, _g4) = out_d.device_ptr_mut(s);
7952                let mut ps = [
7953                    &pg as *const _ as *mut std::ffi::c_void,
7954                    &pu as *const _ as *mut _,
7955                    &pact as *const _ as *mut _,
7956                    &pq as *const _ as *mut _,
7957                    &pd as *const _ as *mut _,
7958                    &nc as *const _ as *mut _,
7959                ];
7960                unsafe {
7961                    self.launch_pdl(
7962                        "gelu_tanh_mul_q8_1",
7963                        (nrows as u32, 1, 1),
7964                        (rms_block(), 1, 1),
7965                        &mut ps,
7966                    )?;
7967                }
7968            }
7969            return Ok((out_q, out_d));
7970        }
7971        let f = self.func("gelu_tanh_mul_q8_1");
7972        let cfg = LaunchConfig {
7973            grid_dim: (nrows as u32, 1, 1),
7974            block_dim: (rms_block(), 1, 1),
7975            shared_mem_bytes: 0,
7976        };
7977        let __s_b = self.gpu.stream();
7978        let mut b = __s_b.launch_builder(&f);
7979        b.arg(gate)
7980            .arg(up)
7981            .arg(act)
7982            .arg(&mut out_q)
7983            .arg(&mut out_d)
7984            .arg(&nc);
7985        unsafe {
7986            b.launch(cfg)?;
7987        }
7988        Ok((out_q, out_d))
7989    }
7990
7991    /// Slot-fed gelu_tanh_mul_q8_1 twin (alloc-free capture lane; incl. the PDL arm).
7992    #[allow(clippy::too_many_arguments)]
7993    pub fn gelu_tanh_mul_q8_1_into(
7994        &self,
7995        gate: &CudaSlice<f32>,
7996        up: &cudarc::driver::CudaView<f32>,
7997        act: &mut CudaSlice<f32>,
7998        ncols: usize,
7999        nrows: usize,
8000        out_q: &mut CudaSlice<i8>,
8001        out_d: &mut CudaSlice<f32>,
8002    ) -> Result<(), Box<dyn std::error::Error>> {
8003        debug_assert!(ncols % 128 == 0);
8004        debug_assert!(out_q.len() >= nrows * ncols && out_d.len() >= nrows * (ncols / 32));
8005        let nc = ncols as i32;
8006        if Self::pdl_on() {
8007            use cudarc::driver::{DevicePtr, DevicePtrMut};
8008            let s = &self.gpu.stream();
8009            let (pg, _g0) = gate.device_ptr(s);
8010            let (pu, _g1) = up.device_ptr(s);
8011            let (pact, _g2) = act.device_ptr_mut(s);
8012            let (pq, _g3) = out_q.device_ptr_mut(s);
8013            let (pd, _g4) = out_d.device_ptr_mut(s);
8014            let mut ps = [
8015                &pg as *const _ as *mut std::ffi::c_void,
8016                &pu as *const _ as *mut _,
8017                &pact as *const _ as *mut _,
8018                &pq as *const _ as *mut _,
8019                &pd as *const _ as *mut _,
8020                &nc as *const _ as *mut _,
8021            ];
8022            unsafe {
8023                self.launch_pdl(
8024                    "gelu_tanh_mul_q8_1",
8025                    (nrows as u32, 1, 1),
8026                    (rms_block(), 1, 1),
8027                    &mut ps,
8028                )?;
8029            }
8030            return Ok(());
8031        }
8032        let f = self.func("gelu_tanh_mul_q8_1");
8033        let cfg = LaunchConfig {
8034            grid_dim: (nrows as u32, 1, 1),
8035            block_dim: (rms_block(), 1, 1),
8036            shared_mem_bytes: 0,
8037        };
8038        let __s_b = self.gpu.stream();
8039        let mut b = __s_b.launch_builder(&f);
8040        b.arg(gate)
8041            .arg(up)
8042            .arg(&mut *act)
8043            .arg(&mut *out_q)
8044            .arg(&mut *out_d)
8045            .arg(&nc);
8046        unsafe {
8047            b.launch(cfg)?;
8048        }
8049        Ok(())
8050    }
8051
8052    /// gemma4: add + rms_norm3 with outputs 0/2 emitted q8_1 (zsh + moe_in) and 1 f32 (router).
8053    #[allow(clippy::too_many_arguments)]
8054    pub fn add_rms_norm3_q8z(
8055        &self,
8056        a: &CudaSlice<f32>,
8057        b_in: &CudaSlice<f32>,
8058        w0: &CudaSlice<f32>,
8059        w1: &CudaSlice<f32>,
8060        w2: &CudaSlice<f32>,
8061        res: &mut CudaSlice<f32>,
8062        out1: &mut CudaSlice<f32>,
8063        ncols: usize,
8064        nrows: usize,
8065        eps: f32,
8066    ) -> Result<
8067        (
8068            (CudaSlice<i8>, CudaSlice<f32>),
8069            (CudaSlice<i8>, CudaSlice<f32>),
8070        ),
8071        Box<dyn std::error::Error>,
8072    > {
8073        let mut q0 = self.alloc_uninit::<i8>(nrows * ncols)?;
8074        let mut d0 = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
8075        let mut q2 = self.alloc_uninit::<i8>(nrows * ncols)?;
8076        let mut d2 = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
8077        let f = self.func("add_rms_norm3_q8z_f32");
8078        let cfg = LaunchConfig {
8079            grid_dim: (nrows as u32, 1, 1),
8080            block_dim: (rms_block(), 1, 1),
8081            shared_mem_bytes: 0,
8082        };
8083        let (nc, e2) = (ncols as i32, eps);
8084        let __s_b = self.gpu.stream();
8085        let mut b = __s_b.launch_builder(&f);
8086        b.arg(a)
8087            .arg(b_in)
8088            .arg(w0)
8089            .arg(w1)
8090            .arg(w2)
8091            .arg(res)
8092            .arg(&mut q0)
8093            .arg(&mut d0)
8094            .arg(out1)
8095            .arg(&mut q2)
8096            .arg(&mut d2)
8097            .arg(&nc)
8098            .arg(&e2);
8099        unsafe {
8100            b.launch(cfg)?;
8101        }
8102        Ok(((q0, d0), (q2, d2)))
8103    }
8104
8105    /// gemma4: res = a+b AND the three rms_norms of res in one launch.
8106    #[allow(clippy::too_many_arguments)]
8107    pub fn add_rms_norm3(
8108        &self,
8109        a: &CudaSlice<f32>,
8110        b_in: &CudaSlice<f32>,
8111        w0: &CudaSlice<f32>,
8112        w1: &CudaSlice<f32>,
8113        w2: &CudaSlice<f32>,
8114        res: &mut CudaSlice<f32>,
8115        d0: &mut CudaSlice<f32>,
8116        d1: &mut CudaSlice<f32>,
8117        d2: &mut CudaSlice<f32>,
8118        ncols: usize,
8119        nrows: usize,
8120        eps: f32,
8121    ) -> Result<(), Box<dyn std::error::Error>> {
8122        let f = self.func("add_rms_norm3_f32");
8123        let cfg = LaunchConfig {
8124            grid_dim: (nrows as u32, 1, 1),
8125            block_dim: (rms_block(), 1, 1),
8126            shared_mem_bytes: 0,
8127        };
8128        let (nc, e2) = (ncols as i32, eps);
8129        let __s_b = self.gpu.stream();
8130        let mut b = __s_b.launch_builder(&f);
8131        b.arg(a)
8132            .arg(b_in)
8133            .arg(w0)
8134            .arg(w1)
8135            .arg(w2)
8136            .arg(res)
8137            .arg(d0)
8138            .arg(d1)
8139            .arg(d2)
8140            .arg(&nc)
8141            .arg(&e2);
8142        unsafe {
8143            b.launch(cfg)?;
8144        }
8145        Ok(())
8146    }
8147
8148    /// dst = (a + b) * c (residual add + layer scale, one launch).
8149    pub fn add_scale(
8150        &self,
8151        a: &CudaSlice<f32>,
8152        b_in: &CudaSlice<f32>,
8153        c: f32,
8154        dst: &mut CudaSlice<f32>,
8155        n: usize,
8156    ) -> Result<(), Box<dyn std::error::Error>> {
8157        let f = self.func("add_scale_f32");
8158        let cfg = LaunchConfig::for_num_elems(n as u32);
8159        let ni = n as i32;
8160        let __s_b = self.gpu.stream();
8161        let mut b = __s_b.launch_builder(&f);
8162        b.arg(a).arg(b_in).arg(&c).arg(dst).arg(&ni);
8163        unsafe {
8164            b.launch(cfg)?;
8165        }
8166        Ok(())
8167    }
8168
8169    /// Vision-tower LayerNorm (with bias) over [nrows, ncols] — lane/vision.
8170    pub fn layer_norm_bias(
8171        &self,
8172        x: &CudaSlice<f32>,
8173        w: &CudaSlice<f32>,
8174        b: &CudaSlice<f32>,
8175        dst: &mut CudaSlice<f32>,
8176        ncols: usize,
8177        nrows: usize,
8178        eps: f32,
8179    ) -> Result<(), Box<dyn std::error::Error>> {
8180        let f = self.func("layer_norm_bias_f32");
8181        let (nc, e) = (ncols as i32, eps);
8182        let cfg = LaunchConfig {
8183            grid_dim: (nrows as u32, 1, 1),
8184            block_dim: (256, 1, 1),
8185            shared_mem_bytes: 0,
8186        };
8187        let __s_b = self.gpu.stream();
8188        let mut lb = __s_b.launch_builder(&f);
8189        lb.arg(x).arg(w).arg(b).arg(&mut *dst).arg(&nc).arg(&e);
8190        unsafe {
8191            lb.launch(cfg)?;
8192        }
8193        Ok(())
8194    }
8195
8196    /// gelu_pytorch_tanh elementwise (vision tower MLP activation).
8197    pub fn gelu_tanh(
8198        &self,
8199        x: &CudaSlice<f32>,
8200        dst: &mut CudaSlice<f32>,
8201        n: usize,
8202    ) -> Result<(), Box<dyn std::error::Error>> {
8203        let f = self.func("gelu_tanh_f32");
8204        let ni = n as i64;
8205        let cfg = LaunchConfig {
8206            grid_dim: (n.div_ceil(256) as u32, 1, 1),
8207            block_dim: (256, 1, 1),
8208            shared_mem_bytes: 0,
8209        };
8210        let __s_b = self.gpu.stream();
8211        let mut lb = __s_b.launch_builder(&f);
8212        lb.arg(x).arg(&mut *dst).arg(&ni);
8213        unsafe {
8214            lb.launch(cfg)?;
8215        }
8216        Ok(())
8217    }
8218
8219    /// In-place row softmax over [nrows, ncols] (bidirectional vision attention).
8220    pub fn row_softmax(
8221        &self,
8222        x: &mut CudaSlice<f32>,
8223        ncols: usize,
8224        nrows: usize,
8225    ) -> Result<(), Box<dyn std::error::Error>> {
8226        let f = self.func("row_softmax_f32");
8227        let nc = ncols as i32;
8228        let cfg = LaunchConfig {
8229            grid_dim: (nrows as u32, 1, 1),
8230            block_dim: (256, 1, 1),
8231            shared_mem_bytes: 0,
8232        };
8233        let __s_b = self.gpu.stream();
8234        let mut lb = __s_b.launch_builder(&f);
8235        lb.arg(&mut *x).arg(&nc);
8236        unsafe {
8237            lb.launch(cfg)?;
8238        }
8239        Ok(())
8240    }
8241
8242    pub fn rms_norm(
8243        &self,
8244        x: &CudaSlice<f32>,
8245        w: &CudaSlice<f32>,
8246        dst: &mut CudaSlice<f32>,
8247        ncols: usize,
8248        nrows: usize,
8249        eps: f32,
8250    ) -> Result<(), Box<dyn std::error::Error>> {
8251        let (nc, e) = (ncols as i32, eps);
8252        if Self::pdl_on() && Self::pdl_wb_on() {
8253            use cudarc::driver::{DevicePtr, DevicePtrMut};
8254            let s = &self.gpu.stream();
8255            let (px, _g0) = x.device_ptr(s);
8256            let (pw, _g1) = w.device_ptr(s);
8257            let (pd, _g2) = dst.device_ptr_mut(s);
8258            let mut ps = [
8259                &px as *const _ as *mut std::ffi::c_void,
8260                &pw as *const _ as *mut _,
8261                &pd as *const _ as *mut _,
8262                &nc as *const _ as *mut _,
8263                &e as *const _ as *mut _,
8264            ];
8265            unsafe {
8266                self.launch_pdl(
8267                    "rms_norm_f32",
8268                    (nrows as u32, 1, 1),
8269                    (rms_block(), 1, 1),
8270                    &mut ps,
8271                )?;
8272            }
8273            return Ok(());
8274        }
8275        let f = self.func("rms_norm_f32");
8276        let cfg = LaunchConfig {
8277            grid_dim: (nrows as u32, 1, 1),
8278            block_dim: (rms_block(), 1, 1),
8279            shared_mem_bytes: 0,
8280        };
8281        let __s_b = self.gpu.stream();
8282        let mut b = __s_b.launch_builder(&f);
8283        b.arg(x).arg(w).arg(dst).arg(&nc).arg(&e);
8284        unsafe {
8285            b.launch(cfg)?;
8286        }
8287        Ok(())
8288    }
8289
8290    /// RMS-norm with blockDim=1024 — BIT-IDENTICAL to the fused `rms_norm_q8_1` and
8291    /// `add_rms_norm_q8_1` kernels' sum-of-squares reduction. The spec verify path MUST use this
8292    /// to match decode's FP accumulation order: the standard `rms_norm` at blockDim=256 has a
8293    /// different per-thread stride (ncols/256 partials vs ncols/1024 partials) and therefore a
8294    /// different shfl-tree reduction that can shift `scale = rsqrt(sum/n + eps)` by ULPs, causing
8295    /// divergence through the GDN scan and argmax flips on the 9B text prompt. The underlying
8296    /// `rms_norm_f32` kernel supports any blockDim (generic reduce with shared[32]).
8297    pub fn rms_norm_decode(
8298        &self,
8299        x: &CudaSlice<f32>,
8300        w: &CudaSlice<f32>,
8301        dst: &mut CudaSlice<f32>,
8302        ncols: usize,
8303        nrows: usize,
8304        eps: f32,
8305    ) -> Result<(), Box<dyn std::error::Error>> {
8306        let f = self.func("rms_norm_f32");
8307        let cfg = LaunchConfig {
8308            grid_dim: (nrows as u32, 1, 1),
8309            block_dim: (1024, 1, 1),
8310            shared_mem_bytes: 0,
8311        };
8312        let (nc, e) = (ncols as i32, eps);
8313        let __s_b = self.gpu.stream();
8314        let mut b = __s_b.launch_builder(&f);
8315        b.arg(x).arg(w).arg(dst).arg(&nc).arg(&e);
8316        unsafe {
8317            b.launch(cfg)?;
8318        }
8319        Ok(())
8320    }
8321
8322    /// DECODE GLUE-FUSION LEVER: `z = rms_norm(x)*w` emitted DIRECTLY as q8_1 (no f32 `z` materialized,
8323    /// no standalone quantize_q8_1 launch). Returns (out_q [nrows*ncols i8], out_d [nrows*nblk f32])
8324    /// ready to feed matmul_pre. BIT-IDENTICAL to rms_norm + quantize_q8_1. ncols % 32 == 0.
8325    pub fn rms_norm_q8_1(
8326        &self,
8327        x: &CudaSlice<f32>,
8328        w: &CudaSlice<f32>,
8329        ncols: usize,
8330        nrows: usize,
8331        eps: f32,
8332    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8333        let nblk = ncols / 32;
8334        let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
8335        let mut d = self.alloc_uninit::<f32>(nrows * nblk)?;
8336        let (nc, e) = (ncols as i32, eps);
8337        if Self::pdl_on() {
8338            {
8339                use cudarc::driver::{DevicePtr, DevicePtrMut};
8340                let s = &self.gpu.stream();
8341                let (px, _g0) = x.device_ptr(s);
8342                let (pw, _g1) = w.device_ptr(s);
8343                let (pq, _g2) = q.device_ptr_mut(s);
8344                let (pd, _g3) = d.device_ptr_mut(s);
8345                let mut ps = [
8346                    &px as *const _ as *mut std::ffi::c_void,
8347                    &pw as *const _ as *mut _,
8348                    &pq as *const _ as *mut _,
8349                    &pd as *const _ as *mut _,
8350                    &nc as *const _ as *mut _,
8351                    &e as *const _ as *mut _,
8352                ];
8353                unsafe {
8354                    self.launch_pdl("rms_norm_q8_1", (nrows as u32, 1, 1), (1024, 1, 1), &mut ps)?;
8355                }
8356            }
8357            return Ok((q, d));
8358        }
8359        let f = self.func("rms_norm_q8_1");
8360        // 1024 threads: decode is nrows=1 -> ONE CTA; 32 warps hide the pass1->pass2 latency
8361        // (s[32] reduce already sized for 32 warps). Same shape math at any blockDim.
8362        let cfg = LaunchConfig {
8363            grid_dim: (nrows as u32, 1, 1),
8364            block_dim: (1024, 1, 1),
8365            shared_mem_bytes: 0,
8366        };
8367        let __s_b = self.gpu.stream();
8368        let mut b = __s_b.launch_builder(&f);
8369        b.arg(x).arg(w).arg(&mut q).arg(&mut d).arg(&nc).arg(&e);
8370        unsafe {
8371            b.launch(cfg)?;
8372        }
8373        Ok((q, d))
8374    }
8375
8376    /// Slot-fed rms_norm_q8_1 twin (alloc-free capture lane): identical launch (incl. the
8377    /// PDL arm), caller-owned outputs.
8378    pub fn rms_norm_q8_1_into(
8379        &self,
8380        x: &CudaSlice<f32>,
8381        w: &CudaSlice<f32>,
8382        ncols: usize,
8383        nrows: usize,
8384        eps: f32,
8385        q: &mut CudaSlice<i8>,
8386        d: &mut CudaSlice<f32>,
8387    ) -> Result<(), Box<dyn std::error::Error>> {
8388        let nblk = ncols / 32;
8389        debug_assert!(q.len() >= nrows * ncols && d.len() >= nrows * nblk);
8390        let (nc, e) = (ncols as i32, eps);
8391        if Self::pdl_on() {
8392            use cudarc::driver::{DevicePtr, DevicePtrMut};
8393            let s = &self.gpu.stream();
8394            let (px, _g0) = x.device_ptr(s);
8395            let (pw, _g1) = w.device_ptr(s);
8396            let (pq, _g2) = q.device_ptr_mut(s);
8397            let (pd, _g3) = d.device_ptr_mut(s);
8398            let mut ps = [
8399                &px as *const _ as *mut std::ffi::c_void,
8400                &pw as *const _ as *mut _,
8401                &pq as *const _ as *mut _,
8402                &pd as *const _ as *mut _,
8403                &nc as *const _ as *mut _,
8404                &e as *const _ as *mut _,
8405            ];
8406            unsafe {
8407                self.launch_pdl("rms_norm_q8_1", (nrows as u32, 1, 1), (1024, 1, 1), &mut ps)?;
8408            }
8409            return Ok(());
8410        }
8411        let f = self.func("rms_norm_q8_1");
8412        let cfg = LaunchConfig {
8413            grid_dim: (nrows as u32, 1, 1),
8414            block_dim: (1024, 1, 1),
8415            shared_mem_bytes: 0,
8416        };
8417        let __s_b = self.gpu.stream();
8418        let mut b = __s_b.launch_builder(&f);
8419        b.arg(x).arg(w).arg(&mut *q).arg(&mut *d).arg(&nc).arg(&e);
8420        unsafe {
8421            b.launch(cfg)?;
8422        }
8423        Ok(())
8424    }
8425
8426    /// Slot-fed quantize_q8_1 twin (alloc-free capture lane).
8427    pub fn quantize_q8_1_into(
8428        &self,
8429        x: &CudaSlice<f32>,
8430        m: usize,
8431        in_f: usize,
8432        q: &mut CudaSlice<i8>,
8433        d: &mut CudaSlice<f32>,
8434    ) -> Result<(), Box<dyn std::error::Error>> {
8435        let nblk = in_f / 32;
8436        debug_assert!(q.len() >= m * in_f && d.len() >= m * nblk);
8437        let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
8438        let (inf, mi) = (in_f as i32, m as i32);
8439        if Self::pdl_on() && Self::pdl_wb_on() {
8440            use cudarc::driver::{DevicePtr, DevicePtrMut};
8441            let s = &self.gpu.stream();
8442            let (px, _g0) = x.device_ptr(s);
8443            let (pq, _g1) = q.device_ptr_mut(s);
8444            let (pd, _g2) = d.device_ptr_mut(s);
8445            let mut ps = [
8446                &px as *const _ as *mut std::ffi::c_void,
8447                &pq as *const _ as *mut _,
8448                &pd as *const _ as *mut _,
8449                &inf as *const _ as *mut _,
8450                &mi as *const _ as *mut _,
8451            ];
8452            unsafe {
8453                self.launch_pdl("quantize_q8_1", cfg.grid_dim, cfg.block_dim, &mut ps)?;
8454            }
8455            return Ok(());
8456        }
8457        let f = self.func("quantize_q8_1");
8458        let __s_b = self.gpu.stream();
8459        let mut b = __s_b.launch_builder(&f);
8460        b.arg(x).arg(&mut *q).arg(&mut *d).arg(&inf).arg(&mi);
8461        unsafe {
8462            b.launch(cfg)?;
8463        }
8464        Ok(())
8465    }
8466
8467    /// DECODE GLUE-FUSION LEVER: `res = a+b; z = rms_norm(res)*w` with z emitted as q8_1. `res` is
8468    /// still written (the post-ffn residual add reads it). Fuses add_rms_norm + quantize_q8_1.
8469    /// Returns (out_q, out_d) for matmul_pre. BIT-IDENTICAL. ncols % 32 == 0.
8470    pub fn add_rms_norm_q8_1(
8471        &self,
8472        a: &CudaSlice<f32>,
8473        b_in: &CudaSlice<f32>,
8474        w: &CudaSlice<f32>,
8475        res: &mut CudaSlice<f32>,
8476        ncols: usize,
8477        nrows: usize,
8478        eps: f32,
8479    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8480        let nblk = ncols / 32;
8481        let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
8482        let mut d = self.alloc_uninit::<f32>(nrows * nblk)?;
8483        let f = self.func("add_rms_norm_q8_1");
8484        // 1024 threads: same single-CTA-at-decode reasoning as rms_norm_q8_1.
8485        let cfg = LaunchConfig {
8486            grid_dim: (nrows as u32, 1, 1),
8487            block_dim: (1024, 1, 1),
8488            shared_mem_bytes: 0,
8489        };
8490        let (nc, e) = (ncols as i32, eps);
8491        let __s_bld = self.gpu.stream();
8492        let mut bld = __s_bld.launch_builder(&f);
8493        bld.arg(a)
8494            .arg(b_in)
8495            .arg(w)
8496            .arg(res)
8497            .arg(&mut q)
8498            .arg(&mut d)
8499            .arg(&nc)
8500            .arg(&e);
8501        unsafe {
8502            bld.launch(cfg)?;
8503        }
8504        Ok((q, d))
8505    }
8506
8507    /// RANK3 LEVER (add+rmsnorm fuse): `res = a + b; dst = rms_norm(res) * w` in ONE launch. Fuses
8508    /// e.add(a,b,res) + e.rms_norm(res,w,dst), removing one launch + one HBM read of the residual per
8509    /// residual+norm pair. BIT-IDENTICAL to the two-kernel sequence (same IEEE add, same reduction).
8510    pub fn add_rms_norm(
8511        &self,
8512        a: &CudaSlice<f32>,
8513        b: &CudaSlice<f32>,
8514        w: &CudaSlice<f32>,
8515        res: &mut CudaSlice<f32>,
8516        dst: &mut CudaSlice<f32>,
8517        ncols: usize,
8518        nrows: usize,
8519        eps: f32,
8520    ) -> Result<(), Box<dyn std::error::Error>> {
8521        let (nc, e) = (ncols as i32, eps);
8522        if Self::pdl_on() && Self::pdl_wb_on() {
8523            use cudarc::driver::{DevicePtr, DevicePtrMut};
8524            let s = &self.gpu.stream();
8525            let (pa, _g0) = a.device_ptr(s);
8526            let (pb, _g1) = b.device_ptr(s);
8527            let (pw, _g2) = w.device_ptr(s);
8528            let (pr, _g3) = res.device_ptr_mut(s);
8529            let (pd, _g4) = dst.device_ptr_mut(s);
8530            let mut ps = [
8531                &pa as *const _ as *mut std::ffi::c_void,
8532                &pb as *const _ as *mut _,
8533                &pw as *const _ as *mut _,
8534                &pr as *const _ as *mut _,
8535                &pd as *const _ as *mut _,
8536                &nc as *const _ as *mut _,
8537                &e as *const _ as *mut _,
8538            ];
8539            unsafe {
8540                self.launch_pdl(
8541                    "add_rms_norm_f32",
8542                    (nrows as u32, 1, 1),
8543                    (rms_block(), 1, 1),
8544                    &mut ps,
8545                )?;
8546            }
8547            return Ok(());
8548        }
8549        let f = self.func("add_rms_norm_f32");
8550        let cfg = LaunchConfig {
8551            grid_dim: (nrows as u32, 1, 1),
8552            block_dim: (rms_block(), 1, 1),
8553            shared_mem_bytes: 0,
8554        };
8555        let __s_b2 = self.gpu.stream();
8556        let mut b2 = __s_b2.launch_builder(&f);
8557        b2.arg(a)
8558            .arg(b)
8559            .arg(w)
8560            .arg(&mut *res)
8561            .arg(&mut *dst)
8562            .arg(&nc)
8563            .arg(&e);
8564        unsafe {
8565            b2.launch(cfg)?;
8566        }
8567        Ok(())
8568    }
8569
8570    /// E4B glue fusion: rms(a, wa) prologue + add_rms_norm — folds the post-attn norm into
8571    /// the tail entry (res = rms(a)*wa + b; dst = rms(res)*w).
8572    #[allow(clippy::too_many_arguments)]
8573    pub fn rms_pre_add_rms_norm(
8574        &self,
8575        a: &CudaSlice<f32>,
8576        wa: &CudaSlice<f32>,
8577        b: &CudaSlice<f32>,
8578        w: &CudaSlice<f32>,
8579        res: &mut CudaSlice<f32>,
8580        dst: &mut CudaSlice<f32>,
8581        ncols: usize,
8582        nrows: usize,
8583        eps: f32,
8584    ) -> Result<(), Box<dyn std::error::Error>> {
8585        let f = self.func("rms_pre_add_rms_norm_f32");
8586        let cfg = LaunchConfig {
8587            grid_dim: (nrows as u32, 1, 1),
8588            block_dim: (rms_block(), 1, 1),
8589            shared_mem_bytes: 0,
8590        };
8591        let (nc, e) = (ncols as i32, eps);
8592        let __s_b2 = self.gpu.stream();
8593        let mut b2 = __s_b2.launch_builder(&f);
8594        b2.arg(a)
8595            .arg(wa)
8596            .arg(b)
8597            .arg(w)
8598            .arg(&mut *res)
8599            .arg(&mut *dst)
8600            .arg(&nc)
8601            .arg(&e);
8602        unsafe {
8603            b2.launch(cfg)?;
8604        }
8605        Ok(())
8606    }
8607
8608    /// wave-2 fold: rms(a,wa) + add + ffn-norm with zsh EMITTED q8_1 (fused2 consumes it).
8609    #[allow(clippy::too_many_arguments)]
8610    pub fn rms_pre_add_rms_norm_q8z(
8611        &self,
8612        a: &CudaSlice<f32>,
8613        wa: &CudaSlice<f32>,
8614        b: &CudaSlice<f32>,
8615        w: &CudaSlice<f32>,
8616        res: &mut CudaSlice<f32>,
8617        dst: &mut CudaSlice<f32>,
8618        ncols: usize,
8619        nrows: usize,
8620        eps: f32,
8621    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8622        debug_assert!(ncols % 128 == 0);
8623        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
8624        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
8625        let (nc, e) = (ncols as i32, eps);
8626        if Self::pdl_on() {
8627            {
8628                use cudarc::driver::{DevicePtr, DevicePtrMut};
8629                let s = &self.gpu.stream();
8630                let (pa, _g0) = a.device_ptr(s);
8631                let (pwa, _g1) = wa.device_ptr(s);
8632                let (pb, _g2) = b.device_ptr(s);
8633                let (pw, _g3) = w.device_ptr(s);
8634                let (pr, _g4) = res.device_ptr_mut(s);
8635                let (pdst, _g5) = dst.device_ptr_mut(s);
8636                let (pq, _g6) = out_q.device_ptr_mut(s);
8637                let (pd, _g7) = out_d.device_ptr_mut(s);
8638                let mut ps = [
8639                    &pa as *const _ as *mut std::ffi::c_void,
8640                    &pwa as *const _ as *mut _,
8641                    &pb as *const _ as *mut _,
8642                    &pw as *const _ as *mut _,
8643                    &pr as *const _ as *mut _,
8644                    &pdst as *const _ as *mut _,
8645                    &pq as *const _ as *mut _,
8646                    &pd as *const _ as *mut _,
8647                    &nc as *const _ as *mut _,
8648                    &e as *const _ as *mut _,
8649                ];
8650                unsafe {
8651                    self.launch_pdl(
8652                        "rms_pre_add_rms_norm_q8z_f32",
8653                        (nrows as u32, 1, 1),
8654                        (rms_block(), 1, 1),
8655                        &mut ps,
8656                    )?;
8657                }
8658            }
8659            return Ok((out_q, out_d));
8660        }
8661        let f = self.func("rms_pre_add_rms_norm_q8z_f32");
8662        let cfg = LaunchConfig {
8663            grid_dim: (nrows as u32, 1, 1),
8664            block_dim: (rms_block(), 1, 1),
8665            shared_mem_bytes: 0,
8666        };
8667        let __s_b2 = self.gpu.stream();
8668        let mut b2 = __s_b2.launch_builder(&f);
8669        b2.arg(a)
8670            .arg(wa)
8671            .arg(b)
8672            .arg(w)
8673            .arg(&mut *res)
8674            .arg(&mut *dst)
8675            .arg(&mut out_q)
8676            .arg(&mut out_d)
8677            .arg(&nc)
8678            .arg(&e);
8679        unsafe {
8680            b2.launch(cfg)?;
8681        }
8682        Ok((out_q, out_d))
8683    }
8684
8685    /// Slot-fed twin of `rms_pre_add_rms_norm_q8z` (gemma4 pn-fold, slotted/graph arm):
8686    /// identical kernel, caller-owned outputs, PLAIN launch only — the dc_slotted capture
8687    /// body must stay attribute-free (the fused2_into precedent).
8688    #[allow(clippy::too_many_arguments)]
8689    pub fn rms_pre_add_rms_norm_q8z_into(
8690        &self,
8691        a: &CudaSlice<f32>,
8692        wa: &CudaSlice<f32>,
8693        b: &CudaSlice<f32>,
8694        w: &CudaSlice<f32>,
8695        res: &mut CudaSlice<f32>,
8696        dst: &mut CudaSlice<f32>,
8697        ncols: usize,
8698        nrows: usize,
8699        eps: f32,
8700        out_q: &mut CudaSlice<i8>,
8701        out_d: &mut CudaSlice<f32>,
8702    ) -> Result<(), Box<dyn std::error::Error>> {
8703        debug_assert!(ncols % 128 == 0);
8704        let (nc, e) = (ncols as i32, eps);
8705        let f = self.func("rms_pre_add_rms_norm_q8z_f32");
8706        let cfg = LaunchConfig {
8707            grid_dim: (nrows as u32, 1, 1),
8708            block_dim: (rms_block(), 1, 1),
8709            shared_mem_bytes: 0,
8710        };
8711        let __s_b = self.gpu.stream();
8712        let mut b2 = __s_b.launch_builder(&f);
8713        b2.arg(a)
8714            .arg(wa)
8715            .arg(b)
8716            .arg(w)
8717            .arg(&mut *res)
8718            .arg(&mut *dst)
8719            .arg(&mut *out_q)
8720            .arg(&mut *out_d)
8721            .arg(&nc)
8722            .arg(&e);
8723        unsafe {
8724            b2.launch(cfg)?;
8725        }
8726        Ok(())
8727    }
8728
8729    /// Slot-fed twin of `rms_pre_add_scale_rms_norm_q8_1` (gemma4 pn-fold exit, slotted
8730    /// arm): identical kernel, caller-owned outputs, PLAIN launch (capture-safe).
8731    #[allow(clippy::too_many_arguments)]
8732    pub fn rms_pre_add_scale_rms_norm_q8_1_into(
8733        &self,
8734        a: &CudaSlice<f32>,
8735        wa: &CudaSlice<f32>,
8736        b_in: &CudaSlice<f32>,
8737        c: f32,
8738        w: &CudaSlice<f32>,
8739        res: &mut CudaSlice<f32>,
8740        ncols: usize,
8741        nrows: usize,
8742        eps: f32,
8743        out_q: &mut CudaSlice<i8>,
8744        out_d: &mut CudaSlice<f32>,
8745    ) -> Result<(), Box<dyn std::error::Error>> {
8746        debug_assert!(ncols % 128 == 0);
8747        let (nc, e2) = (ncols as i32, eps);
8748        let f = self.func("rms_pre_add_scale_rms_norm_q8_1");
8749        let cfg = LaunchConfig {
8750            grid_dim: (nrows as u32, 1, 1),
8751            block_dim: (rms_block(), 1, 1),
8752            shared_mem_bytes: 0,
8753        };
8754        let __s_b = self.gpu.stream();
8755        let mut b2 = __s_b.launch_builder(&f);
8756        b2.arg(a)
8757            .arg(wa)
8758            .arg(b_in)
8759            .arg(&c)
8760            .arg(w)
8761            .arg(&mut *res)
8762            .arg(&mut *out_q)
8763            .arg(&mut *out_d)
8764            .arg(&nc)
8765            .arg(&e2);
8766        unsafe {
8767            b2.launch(cfg)?;
8768        }
8769        Ok(())
8770    }
8771
8772    /// gemma4 pn-fold seam (GAP-DIAGNOSIS verdict 7, the E4B glue backport): the dense
8773    /// decode/verify/slotted trio folds post_attn_norm into the tail entry
8774    /// (rms_pre_add_rms_norm[_q8z]) and post_ffw_norm into the residual exit
8775    /// (rms_pre_add_scale_rms_norm_q8_1). BITS-CHANGING vs the two-launch chain (the
8776    /// single-phase reduction's expansion rounding — E4B receipts); every arm moves
8777    /// together so decode == verify == graph parity holds BY CONSTRUCTION within either
8778    /// seam value. MEMRA_G4_PNFOLD=0 restores the unfused chain everywhere.
8779    pub fn g4_pnfold_on() -> bool {
8780        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8781        *ON.get_or_init(|| {
8782            std::env::var("MEMRA_G4_PNFOLD")
8783                .map(|v| v != "0")
8784                .unwrap_or(true)
8785        })
8786    }
8787
8788    /// wave-4b: OUT-dim concat of three Q4_0 tensors (same in_features; rows are independent
8789    /// blocks, so the concat is a D2D byte concat of the GGUF-layout planes). Returns None
8790    /// off-class (non-Q4_0, mismatched widths, or any tensor already rp-swapped in place).
8791    pub fn build_q4_out_concat3(
8792        &self,
8793        w0: &crate::model::GpuTensor,
8794        w1: &crate::model::GpuTensor,
8795        w2: &crate::model::GpuTensor,
8796    ) -> Result<Option<crate::model::GpuTensor>, Box<dyn std::error::Error>> {
8797        use crate::model::GpuTensor;
8798        let part = |w: &GpuTensor| -> Option<(usize, usize)> {
8799            match w {
8800                GpuTensor::Quant {
8801                    qtype,
8802                    row_bytes,
8803                    rp,
8804                    ..
8805                } if *qtype == QT_Q4_0 && !*rp => Some((*row_bytes, w.out_features())),
8806                _ => None,
8807            }
8808        };
8809        let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (part(w0), part(w1), part(w2))
8810        else {
8811            return Ok(None);
8812        };
8813        if rb0 != rb1
8814            || rb0 != rb2
8815            || w0.in_features() != w1.in_features()
8816            || w0.in_features() != w2.in_features()
8817        {
8818            return Ok(None);
8819        }
8820        fn bytes_of(w: &crate::model::GpuTensor) -> &CudaSlice<u8> {
8821            match w {
8822                crate::model::GpuTensor::Quant { bytes, .. } => bytes,
8823                _ => unreachable!(),
8824            }
8825        }
8826        let (b0, b1, b2) = (bytes_of(w0), bytes_of(w1), bytes_of(w2));
8827        let total = rb0 * (o0 + o1 + o2);
8828        let mut cat = self.alloc_u8(total)?;
8829        self.copy_u8_into(&mut cat, 0, b0, rb0 * o0)?;
8830        self.copy_u8_into(&mut cat, rb0 * o0, b1, rb1 * o1)?;
8831        self.copy_u8_into(&mut cat, rb0 * (o0 + o1), b2, rb2 * o2)?;
8832        Ok(Some(GpuTensor::Quant {
8833            bytes: cat,
8834            qtype: QT_Q4_0,
8835            row_bytes: rb0,
8836            ne: vec![w0.in_features() as u64, (o0 + o1 + o2) as u64],
8837            scale: 1.0,
8838            rp: false,
8839            #[cfg(memra_cutlass)]
8840            cutlass: None,
8841            fp8: None,
8842            blk: None,
8843            rp4: None,
8844            f16: None,
8845        }))
8846    }
8847
8848    /// FULL-WIDTH-ROPE CONTRACT for the fused rms_norm+qkv+rope kernels
8849    /// (lane/graph-s-key-exactness-20260819, probe O-10 follow-up).
8850    ///
8851    /// `rms_norm_qkv_rope_f32` / `_cat_f32` (`cu/kernels.cu`) and
8852    /// `rms_norm_qkv_rope_append_body` (`cu/flash_attn.cu`) take NO `n_dims`/`n_rot` argument.
8853    /// They compute `int half = ncols / 2` and rotate the FULL head width by construction — the
8854    /// standalone `rope_neox*` kernels take `n_dims` and early-return above it, these do not.
8855    ///
8856    /// Every call site today is a gemma-4 arm, where `n_rot == head_dim`, so nothing is wrong
8857    /// now. But this is the n_rot bug ONE FUSION away: fuse a partial-rotary arch onto these
8858    /// kernels (qwen3.5 = 64 rotary dims of a 256-wide head; step35 full-attn = 64 of 128) and
8859    /// 192 dims that must pass through unrotated get rotated silently — no error, no NaN, just a
8860    /// wrong model. The n_rot lane already paid for that class once, in the config derivation.
8861    ///
8862    /// So the fusions now take the layer's DERIVED rope width and refuse anything but full
8863    /// width. A future partial-rotary caller fails at its first launch with the geometry named
8864    /// instead of serving quietly wrong logits.
8865    fn full_width_rope_only(
8866        kernel: &str,
8867        n_rot: usize,
8868        head_dim: usize,
8869    ) -> Result<(), Box<dyn std::error::Error>> {
8870        if n_rot == head_dim {
8871            return Ok(());
8872        }
8873        Err(format!(
8874            "{kernel}: PARTIAL ROTARY REFUSED — n_rot {n_rot} != head_dim {head_dim}. This fused \
8875             rms_norm+qkv+rope kernel carries no n_dims parameter and rotates the full head \
8876             width (half = ncols/2), so it would rotate dims {n_rot}..{head_dim} that must pass \
8877             through unrotated. Use the split path (rms_norm_qkv + rope_neox/rope_neox2 with \
8878             n_dims={n_rot}), or add an n_dims early-return to the kernel and widen this guard."
8879        )
8880        .into())
8881    }
8882
8883    /// wave-4b: the qkv-cat twin — one contiguous [rq+2*rk, hd] input from the concat matvec.
8884    /// `n_rot` is the layer's derived rotary width and MUST equal `head_dim`
8885    /// ([`Engine::full_width_rope_only`]).
8886    #[allow(clippy::too_many_arguments)]
8887    pub fn rms_norm_qkv_rope_cat(
8888        &self,
8889        qkv: &CudaSlice<f32>,
8890        wq: &CudaSlice<f32>,
8891        wk: &CudaSlice<f32>,
8892        wv: &CudaSlice<f32>,
8893        q: &mut CudaSlice<f32>,
8894        k: &mut CudaSlice<f32>,
8895        v: &mut CudaSlice<f32>,
8896        head_dim: usize,
8897        n_rot: usize,
8898        rq: usize,
8899        rk: usize,
8900        pos: &CudaSlice<i32>,
8901        nh_q: usize,
8902        nh_k: usize,
8903        base: f32,
8904        freq_scale: f32,
8905        ff: Option<&CudaSlice<f32>>,
8906        eps: f32,
8907    ) -> Result<(), Box<dyn std::error::Error>> {
8908        Self::full_width_rope_only("rms_norm_qkv_rope_cat", n_rot, head_dim)?;
8909        let rows = rq + rk + rk;
8910        let theta_scale = base.powf(-2.0 / head_dim as f32);
8911        let (nc, rqi, rki, nhq, nhk) = (
8912            head_dim as i32,
8913            rq as i32,
8914            rk as i32,
8915            nh_q as i32,
8916            nh_k as i32,
8917        );
8918        if Self::pdl_on() {
8919            use cudarc::driver::{DevicePtr, DevicePtrMut};
8920            let s = &self.gpu.stream();
8921            let (pqkv, _g0) = qkv.device_ptr(s);
8922            let (pwq, _g1) = wq.device_ptr(s);
8923            let (pwk, _g2) = wk.device_ptr(s);
8924            let (pwv, _g3) = wv.device_ptr(s);
8925            let (pq, _g4) = q.device_ptr_mut(s);
8926            let (pk, _g5) = k.device_ptr_mut(s);
8927            let (pv, _g6) = v.device_ptr_mut(s);
8928            let (ppos, _g7) = pos.device_ptr(s);
8929            let (pff, _g8) = match ff {
8930                Some(t) => {
8931                    let (p, g) = t.device_ptr(s);
8932                    (p, Some(g))
8933                }
8934                None => (0, None),
8935            };
8936            let mut ps = [
8937                &pqkv as *const _ as *mut std::ffi::c_void,
8938                &pwq as *const _ as *mut _,
8939                &pwk as *const _ as *mut _,
8940                &pwv as *const _ as *mut _,
8941                &pq as *const _ as *mut _,
8942                &pk as *const _ as *mut _,
8943                &pv as *const _ as *mut _,
8944                &nc as *const _ as *mut _,
8945                &rqi as *const _ as *mut _,
8946                &rki as *const _ as *mut _,
8947                &ppos as *const _ as *mut _,
8948                &nhq as *const _ as *mut _,
8949                &nhk as *const _ as *mut _,
8950                &theta_scale as *const _ as *mut _,
8951                &freq_scale as *const _ as *mut _,
8952                &pff as *const _ as *mut _,
8953                &eps as *const _ as *mut _,
8954            ];
8955            unsafe {
8956                self.launch_pdl(
8957                    "rms_norm_qkv_rope_cat_f32",
8958                    (rows as u32, 1, 1),
8959                    (rms_block(), 1, 1),
8960                    &mut ps,
8961                )?;
8962            }
8963            return Ok(());
8964        }
8965        let f = self.func("rms_norm_qkv_rope_cat_f32");
8966        let cfg = LaunchConfig {
8967            grid_dim: (rows as u32, 1, 1),
8968            block_dim: (rms_block(), 1, 1),
8969            shared_mem_bytes: 0,
8970        };
8971        let __s_b = self.gpu.stream();
8972        let mut b = __s_b.launch_builder(&f);
8973        match ff {
8974            Some(t) => {
8975                b.arg(qkv)
8976                    .arg(wq)
8977                    .arg(wk)
8978                    .arg(wv)
8979                    .arg(&mut *q)
8980                    .arg(&mut *k)
8981                    .arg(&mut *v)
8982                    .arg(&nc)
8983                    .arg(&rqi)
8984                    .arg(&rki)
8985                    .arg(pos)
8986                    .arg(&nhq)
8987                    .arg(&nhk)
8988                    .arg(&theta_scale)
8989                    .arg(&freq_scale)
8990                    .arg(t)
8991                    .arg(&eps);
8992                unsafe {
8993                    b.launch(cfg)?;
8994                }
8995            }
8996            None => {
8997                let null: u64 = 0;
8998                b.arg(qkv)
8999                    .arg(wq)
9000                    .arg(wk)
9001                    .arg(wv)
9002                    .arg(&mut *q)
9003                    .arg(&mut *k)
9004                    .arg(&mut *v)
9005                    .arg(&nc)
9006                    .arg(&rqi)
9007                    .arg(&rki)
9008                    .arg(pos)
9009                    .arg(&nhq)
9010                    .arg(&nhk)
9011                    .arg(&theta_scale)
9012                    .arg(&freq_scale)
9013                    .arg(&null)
9014                    .arg(&eps);
9015                unsafe {
9016                    b.launch(cfg)?;
9017                }
9018            }
9019        }
9020        Ok(())
9021    }
9022
9023    /// wave-3 fold: rms_norm_qkv + rope_neox2 in ONE launch (n_dims == head_dim; ff nullable).
9024    /// `n_rot` is the layer's derived rotary width and MUST equal `head_dim`
9025    /// ([`Engine::full_width_rope_only`]).
9026    #[allow(clippy::too_many_arguments)]
9027    pub fn rms_norm_qkv_rope(
9028        &self,
9029        q0: &CudaSlice<f32>,
9030        k0: &CudaSlice<f32>,
9031        v0: &CudaSlice<f32>,
9032        wq: &CudaSlice<f32>,
9033        wk: &CudaSlice<f32>,
9034        wv: &CudaSlice<f32>,
9035        q: &mut CudaSlice<f32>,
9036        k: &mut CudaSlice<f32>,
9037        v: &mut CudaSlice<f32>,
9038        head_dim: usize,
9039        n_rot: usize,
9040        rq: usize,
9041        rk: usize,
9042        pos: &CudaSlice<i32>,
9043        nh_q: usize,
9044        nh_k: usize,
9045        base: f32,
9046        freq_scale: f32,
9047        ff: Option<&CudaSlice<f32>>,
9048        eps: f32,
9049    ) -> Result<(), Box<dyn std::error::Error>> {
9050        Self::full_width_rope_only("rms_norm_qkv_rope", n_rot, head_dim)?;
9051        let f = self.func("rms_norm_qkv_rope_f32");
9052        let rows = rq + rk + rk; // q rows + k rows + v rows (rk == rv)
9053        let cfg = LaunchConfig {
9054            grid_dim: (rows as u32, 1, 1),
9055            block_dim: (rms_block(), 1, 1),
9056            shared_mem_bytes: 0,
9057        };
9058        let theta_scale = base.powf(-2.0 / head_dim as f32);
9059        let (nc, rqi, rki, nhq, nhk) = (
9060            head_dim as i32,
9061            rq as i32,
9062            rk as i32,
9063            nh_q as i32,
9064            nh_k as i32,
9065        );
9066        let __s_b = self.gpu.stream();
9067        let mut b = __s_b.launch_builder(&f);
9068        match ff {
9069            Some(t) => {
9070                b.arg(q0)
9071                    .arg(k0)
9072                    .arg(v0)
9073                    .arg(wq)
9074                    .arg(wk)
9075                    .arg(wv)
9076                    .arg(&mut *q)
9077                    .arg(&mut *k)
9078                    .arg(&mut *v)
9079                    .arg(&nc)
9080                    .arg(&rqi)
9081                    .arg(&rki)
9082                    .arg(pos)
9083                    .arg(&nhq)
9084                    .arg(&nhk)
9085                    .arg(&theta_scale)
9086                    .arg(&freq_scale)
9087                    .arg(t)
9088                    .arg(&eps);
9089                unsafe {
9090                    b.launch(cfg)?;
9091                }
9092            }
9093            None => {
9094                let null: u64 = 0;
9095                b.arg(q0)
9096                    .arg(k0)
9097                    .arg(v0)
9098                    .arg(wq)
9099                    .arg(wk)
9100                    .arg(wv)
9101                    .arg(&mut *q)
9102                    .arg(&mut *k)
9103                    .arg(&mut *v)
9104                    .arg(&nc)
9105                    .arg(&rqi)
9106                    .arg(&rki)
9107                    .arg(pos)
9108                    .arg(&nhq)
9109                    .arg(&nhk)
9110                    .arg(&theta_scale)
9111                    .arg(&freq_scale)
9112                    .arg(&null)
9113                    .arg(&eps);
9114                unsafe {
9115                    b.launch(cfg)?;
9116                }
9117            }
9118        }
9119        Ok(())
9120    }
9121
9122    /// FUSED norm+rope+APPEND (m=1 decode, 2026-07-23): one launch replaces the
9123    /// rms_norm_qkv_rope + append_kv_quantized_dc pair. Kernel lives in the flash fatbins
9124    /// (format-flavored quant tail) — `g` must mirror the append path's flavor exactly.
9125    /// `n_rot` is the layer's derived rotary width and MUST equal `head_dim`
9126    /// ([`Engine::full_width_rope_only`]).
9127    #[allow(clippy::too_many_arguments)]
9128    pub fn rms_norm_qkv_rope_append_dc(
9129        &self,
9130        q0: &CudaSlice<f32>,
9131        k0: &CudaSlice<f32>,
9132        v0: &CudaSlice<f32>,
9133        wq: &CudaSlice<f32>,
9134        wk: &CudaSlice<f32>,
9135        wv: &CudaSlice<f32>,
9136        q: &mut CudaSlice<f32>,
9137        k: &mut CudaSlice<f32>,
9138        v: &mut CudaSlice<f32>,
9139        head_dim: usize,
9140        n_rot: usize,
9141        rq: usize,
9142        rk: usize,
9143        pos: &CudaSlice<i32>,
9144        nh_q: usize,
9145        nh_k: usize,
9146        base: f32,
9147        freq_scale: f32,
9148        ff: Option<&CudaSlice<f32>>,
9149        eps: f32,
9150        kc: &mut CudaSlice<u8>,
9151        vc: &mut CudaSlice<u8>,
9152        t_dev: &CudaSlice<i32>,
9153        k_tok_bytes: usize,
9154        v_tok_bytes: usize,
9155        g: bool,
9156    ) -> Result<(), Box<dyn std::error::Error>> {
9157        Self::full_width_rope_only("rms_norm_qkv_rope_append_dc", n_rot, head_dim)?;
9158        let rows = rq + rk + rk;
9159        let theta_scale = base.powf(-2.0 / head_dim as f32);
9160        let (nc, rqi, rki, nhq, nhk) = (
9161            head_dim as i32,
9162            rq as i32,
9163            rk as i32,
9164            nh_q as i32,
9165            nh_k as i32,
9166        );
9167        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
9168        if Self::pdl_on() && Self::pdl_wb_on() {
9169            use cudarc::driver::{DevicePtr, DevicePtrMut};
9170            let s = &self.gpu.stream();
9171            let (p0, _a0) = q0.device_ptr(s);
9172            let (p1, _a1) = k0.device_ptr(s);
9173            let (p2, _a2) = v0.device_ptr(s);
9174            let (pwq, _a3) = wq.device_ptr(s);
9175            let (pwk, _a4) = wk.device_ptr(s);
9176            let (pwv, _a5) = wv.device_ptr(s);
9177            let (pq, _a6) = q.device_ptr_mut(s);
9178            let (pk, _a7) = k.device_ptr_mut(s);
9179            let (pv, _a8) = v.device_ptr_mut(s);
9180            let (pp, _a9) = pos.device_ptr(s);
9181            let pff: u64 = match ff {
9182                Some(t) => {
9183                    let (p, _gg) = t.device_ptr(s);
9184                    p as u64
9185                }
9186                None => 0,
9187            };
9188            let (pkc, _a10) = kc.device_ptr_mut(s);
9189            let (pvc, _a11) = vc.device_ptr_mut(s);
9190            let (pt, _a12) = t_dev.device_ptr(s);
9191            let mut ps = [
9192                &p0 as *const _ as *mut std::ffi::c_void,
9193                &p1 as *const _ as *mut _,
9194                &p2 as *const _ as *mut _,
9195                &pwq as *const _ as *mut _,
9196                &pwk as *const _ as *mut _,
9197                &pwv as *const _ as *mut _,
9198                &pq as *const _ as *mut _,
9199                &pk as *const _ as *mut _,
9200                &pv as *const _ as *mut _,
9201                &nc as *const _ as *mut _,
9202                &rqi as *const _ as *mut _,
9203                &rki as *const _ as *mut _,
9204                &pp as *const _ as *mut _,
9205                &nhq as *const _ as *mut _,
9206                &nhk as *const _ as *mut _,
9207                &theta_scale as *const _ as *mut _,
9208                &freq_scale as *const _ as *mut _,
9209                &pff as *const _ as *mut _,
9210                &eps as *const _ as *mut _,
9211                &pkc as *const _ as *mut _,
9212                &pvc as *const _ as *mut _,
9213                &pt as *const _ as *mut _,
9214                &ktb as *const _ as *mut _,
9215                &vtb as *const _ as *mut _,
9216            ];
9217            unsafe {
9218                self.launch_pdl_flash(
9219                    g,
9220                    "rms_norm_qkv_rope_append_dc_f32",
9221                    (rows as u32, 1, 1),
9222                    (rms_block(), 1, 1),
9223                    0,
9224                    &mut ps,
9225                )?;
9226            }
9227            return Ok(());
9228        }
9229        let f = if g {
9230            self.func_g("rms_norm_qkv_rope_append_dc_f32")
9231        } else {
9232            self.func("rms_norm_qkv_rope_append_dc_f32")
9233        };
9234        let cfg = LaunchConfig {
9235            grid_dim: (rows as u32, 1, 1),
9236            block_dim: (rms_block(), 1, 1),
9237            shared_mem_bytes: 0,
9238        };
9239        let __s_b = self.gpu.stream();
9240        let mut b = __s_b.launch_builder(&f);
9241        match ff {
9242            Some(t) => {
9243                b.arg(q0)
9244                    .arg(k0)
9245                    .arg(v0)
9246                    .arg(wq)
9247                    .arg(wk)
9248                    .arg(wv)
9249                    .arg(&mut *q)
9250                    .arg(&mut *k)
9251                    .arg(&mut *v)
9252                    .arg(&nc)
9253                    .arg(&rqi)
9254                    .arg(&rki)
9255                    .arg(pos)
9256                    .arg(&nhq)
9257                    .arg(&nhk)
9258                    .arg(&theta_scale)
9259                    .arg(&freq_scale)
9260                    .arg(t)
9261                    .arg(&eps)
9262                    .arg(&mut *kc)
9263                    .arg(&mut *vc)
9264                    .arg(t_dev)
9265                    .arg(&ktb)
9266                    .arg(&vtb);
9267                unsafe {
9268                    b.launch(cfg)?;
9269                }
9270            }
9271            None => {
9272                let null: u64 = 0;
9273                b.arg(q0)
9274                    .arg(k0)
9275                    .arg(v0)
9276                    .arg(wq)
9277                    .arg(wk)
9278                    .arg(wv)
9279                    .arg(&mut *q)
9280                    .arg(&mut *k)
9281                    .arg(&mut *v)
9282                    .arg(&nc)
9283                    .arg(&rqi)
9284                    .arg(&rki)
9285                    .arg(pos)
9286                    .arg(&nhq)
9287                    .arg(&nhk)
9288                    .arg(&theta_scale)
9289                    .arg(&freq_scale)
9290                    .arg(&null)
9291                    .arg(&eps)
9292                    .arg(&mut *kc)
9293                    .arg(&mut *vc)
9294                    .arg(t_dev)
9295                    .arg(&ktb)
9296                    .arg(&vtb);
9297                unsafe {
9298                    b.launch(cfg)?;
9299                }
9300            }
9301        }
9302        Ok(())
9303    }
9304
9305    /// wave-2 fold: a + b with the sum emitted q8_1 alongside f32.
9306    /// Host-len twin of `rms_norm_qkv_rope_append_dc` (zoo-fusion arc): the eager decode
9307    /// arm tracks kv length on host (`kvl.len`), so the append slot rides the launch arg
9308    /// instead of the device counter. Kernel body is SHARED with the _dc entry (one
9309    /// inlined body — bit-identical to the rms_norm_qkv_rope + append pair it replaces,
9310    /// same law as the dc fold). `n_rot` is the layer's derived rotary width and MUST equal
9311    /// `head_dim` ([`Engine::full_width_rope_only`]).
9312    #[allow(clippy::too_many_arguments)]
9313    pub fn rms_norm_qkv_rope_append(
9314        &self,
9315        q0: &CudaSlice<f32>,
9316        k0: &CudaSlice<f32>,
9317        v0: &CudaSlice<f32>,
9318        wq: &CudaSlice<f32>,
9319        wk: &CudaSlice<f32>,
9320        wv: &CudaSlice<f32>,
9321        q: &mut CudaSlice<f32>,
9322        k: &mut CudaSlice<f32>,
9323        v: &mut CudaSlice<f32>,
9324        head_dim: usize,
9325        n_rot: usize,
9326        rq: usize,
9327        rk: usize,
9328        pos: &CudaSlice<i32>,
9329        nh_q: usize,
9330        nh_k: usize,
9331        base: f32,
9332        freq_scale: f32,
9333        ff: Option<&CudaSlice<f32>>,
9334        eps: f32,
9335        kc: &mut CudaSlice<u8>,
9336        vc: &mut CudaSlice<u8>,
9337        t: usize,
9338        k_tok_bytes: usize,
9339        v_tok_bytes: usize,
9340        g: bool,
9341    ) -> Result<(), Box<dyn std::error::Error>> {
9342        Self::full_width_rope_only("rms_norm_qkv_rope_append", n_rot, head_dim)?;
9343        let rows = rq + rk + rk;
9344        let theta_scale = base.powf(-2.0 / head_dim as f32);
9345        let (nc, rqi, rki, nhq, nhk) = (
9346            head_dim as i32,
9347            rq as i32,
9348            rk as i32,
9349            nh_q as i32,
9350            nh_k as i32,
9351        );
9352        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
9353        let ti = t as i32;
9354        if Self::pdl_on() && Self::pdl_wb_on() {
9355            use cudarc::driver::{DevicePtr, DevicePtrMut};
9356            let s = &self.gpu.stream();
9357            let (p0, _a0) = q0.device_ptr(s);
9358            let (p1, _a1) = k0.device_ptr(s);
9359            let (p2, _a2) = v0.device_ptr(s);
9360            let (pwq, _a3) = wq.device_ptr(s);
9361            let (pwk, _a4) = wk.device_ptr(s);
9362            let (pwv, _a5) = wv.device_ptr(s);
9363            let (pq, _a6) = q.device_ptr_mut(s);
9364            let (pk, _a7) = k.device_ptr_mut(s);
9365            let (pv, _a8) = v.device_ptr_mut(s);
9366            let (pp, _a9) = pos.device_ptr(s);
9367            let pff: u64 = match ff {
9368                Some(t) => {
9369                    let (p, _gg) = t.device_ptr(s);
9370                    p as u64
9371                }
9372                None => 0,
9373            };
9374            let (pkc, _a10) = kc.device_ptr_mut(s);
9375            let (pvc, _a11) = vc.device_ptr_mut(s);
9376            let mut ps = [
9377                &p0 as *const _ as *mut std::ffi::c_void,
9378                &p1 as *const _ as *mut _,
9379                &p2 as *const _ as *mut _,
9380                &pwq as *const _ as *mut _,
9381                &pwk as *const _ as *mut _,
9382                &pwv as *const _ as *mut _,
9383                &pq as *const _ as *mut _,
9384                &pk as *const _ as *mut _,
9385                &pv as *const _ as *mut _,
9386                &nc as *const _ as *mut _,
9387                &rqi as *const _ as *mut _,
9388                &rki as *const _ as *mut _,
9389                &pp as *const _ as *mut _,
9390                &nhq as *const _ as *mut _,
9391                &nhk as *const _ as *mut _,
9392                &theta_scale as *const _ as *mut _,
9393                &freq_scale as *const _ as *mut _,
9394                &pff as *const _ as *mut _,
9395                &eps as *const _ as *mut _,
9396                &pkc as *const _ as *mut _,
9397                &pvc as *const _ as *mut _,
9398                &ti as *const _ as *mut _,
9399                &ktb as *const _ as *mut _,
9400                &vtb as *const _ as *mut _,
9401            ];
9402            unsafe {
9403                self.launch_pdl_flash(
9404                    g,
9405                    "rms_norm_qkv_rope_append_f32",
9406                    (rows as u32, 1, 1),
9407                    (rms_block(), 1, 1),
9408                    0,
9409                    &mut ps,
9410                )?;
9411            }
9412            return Ok(());
9413        }
9414        let f = if g {
9415            self.func_g("rms_norm_qkv_rope_append_f32")
9416        } else {
9417            self.func("rms_norm_qkv_rope_append_f32")
9418        };
9419        let cfg = LaunchConfig {
9420            grid_dim: (rows as u32, 1, 1),
9421            block_dim: (rms_block(), 1, 1),
9422            shared_mem_bytes: 0,
9423        };
9424        let __s_b = self.gpu.stream();
9425        let mut b = __s_b.launch_builder(&f);
9426        let null: u64 = 0;
9427        b.arg(q0)
9428            .arg(k0)
9429            .arg(v0)
9430            .arg(wq)
9431            .arg(wk)
9432            .arg(wv)
9433            .arg(&mut *q)
9434            .arg(&mut *k)
9435            .arg(&mut *v)
9436            .arg(&nc)
9437            .arg(&rqi)
9438            .arg(&rki)
9439            .arg(pos)
9440            .arg(&nhq)
9441            .arg(&nhk)
9442            .arg(&theta_scale)
9443            .arg(&freq_scale);
9444        match ff {
9445            Some(t) => {
9446                b.arg(t);
9447            }
9448            None => {
9449                b.arg(&null);
9450            }
9451        }
9452        b.arg(&eps)
9453            .arg(&mut *kc)
9454            .arg(&mut *vc)
9455            .arg(&ti)
9456            .arg(&ktb)
9457            .arg(&vtb);
9458        unsafe {
9459            b.launch(cfg)?;
9460        }
9461        Ok(())
9462    }
9463
9464    pub fn add_q8_1(
9465        &self,
9466        a: &CudaSlice<f32>,
9467        b: &CudaSlice<f32>,
9468        res: &mut CudaSlice<f32>,
9469        ncols: usize,
9470        nrows: usize,
9471    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9472        debug_assert!(ncols % 128 == 0);
9473        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
9474        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9475        let f = self.func("add_q8_1_f32");
9476        let cfg = LaunchConfig {
9477            grid_dim: (nrows as u32, 1, 1),
9478            block_dim: (rms_block(), 1, 1),
9479            shared_mem_bytes: 0,
9480        };
9481        let nc = ncols as i32;
9482        let __s_b2 = self.gpu.stream();
9483        let mut b2 = __s_b2.launch_builder(&f);
9484        b2.arg(a)
9485            .arg(b)
9486            .arg(&mut *res)
9487            .arg(&mut out_q)
9488            .arg(&mut out_d)
9489            .arg(&nc);
9490        unsafe {
9491            b2.launch(cfg)?;
9492        }
9493        Ok((out_q, out_d))
9494    }
9495
9496    /// E4B FFN-tail exit fusion (glue wave 5): resid = b + rms(a, wa) emitted f32 + q8_1 pair
9497    /// in ONE launch — replaces rms_norm(a,wa->sn) + add_q8_1(sn,b). Same rms_block() config
9498    /// as both parents (bit-identity: identical reduction + quad-walk quantize).
9499    pub fn rms_pre_add_q8_1(
9500        &self,
9501        a: &CudaSlice<f32>,
9502        wa: &CudaSlice<f32>,
9503        b: &CudaSlice<f32>,
9504        res: &mut CudaSlice<f32>,
9505        ncols: usize,
9506        nrows: usize,
9507        eps: f32,
9508    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9509        debug_assert!(ncols % 128 == 0);
9510        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
9511        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9512        let f = self.func("rms_pre_add_q8_1_f32");
9513        let cfg = LaunchConfig {
9514            grid_dim: (nrows as u32, 1, 1),
9515            block_dim: (rms_block(), 1, 1),
9516            shared_mem_bytes: 0,
9517        };
9518        let (nc, ep) = (ncols as i32, eps);
9519        let __s_b2 = self.gpu.stream();
9520        let mut b2 = __s_b2.launch_builder(&f);
9521        b2.arg(a)
9522            .arg(wa)
9523            .arg(b)
9524            .arg(&mut *res)
9525            .arg(&mut out_q)
9526            .arg(&mut out_d)
9527            .arg(&nc)
9528            .arg(&ep);
9529        unsafe {
9530            b2.launch(cfg)?;
9531        }
9532        Ok((out_q, out_d))
9533    }
9534
9535    /// L2 norm per row (head_dim), no weight.
9536    /// PREFILL l2 dispatch (round 27): the warp-per-row float4 v2 when the numeric-config
9537    /// seam allows (MEMRA_L2_V2, default ON, d_state==128 only); else the strided kernel.
9538    pub fn l2_v2_on(ncols: usize) -> bool {
9539        ncols == 128 && std::env::var("MEMRA_L2_V2").as_deref() != Ok("0")
9540    }
9541
9542    pub fn l2_norm_pp(
9543        &self,
9544        x: &CudaSlice<f32>,
9545        dst: &mut CudaSlice<f32>,
9546        dst16: Option<&mut CudaSlice<u8>>,
9547        ncols: usize,
9548        nrows: usize,
9549        eps: f32,
9550    ) -> Result<(), Box<dyn std::error::Error>> {
9551        if Self::l2_v2_on(ncols) {
9552            let f = self.func("l2_norm_pp_v2_f32");
9553            let rows_per_block = 8u32; // 256 threads = 8 warps = 8 rows
9554            let cfg = LaunchConfig {
9555                grid_dim: ((nrows as u32).div_ceil(rows_per_block), 1, 1),
9556                block_dim: (256, 1, 1),
9557                shared_mem_bytes: 0,
9558            };
9559            let (nc, nr, e) = (ncols as i32, nrows as i32, eps);
9560            // mirror-fold: bf16 twin address by value (0 = skip; matches the nullable param)
9561            let d16: u64 = match dst16 {
9562                Some(d) => self.addr_u8(d),
9563                None => 0,
9564            };
9565            let __s_b = self.gpu.stream();
9566            let mut b = __s_b.launch_builder(&f);
9567            b.arg(x).arg(dst).arg(&d16).arg(&nc).arg(&nr).arg(&e);
9568            unsafe {
9569                b.launch(cfg)?;
9570            }
9571            return Ok(());
9572        }
9573        self.l2_norm(x, dst, ncols, nrows, eps)
9574    }
9575
9576    pub fn l2_norm(
9577        &self,
9578        x: &CudaSlice<f32>,
9579        dst: &mut CudaSlice<f32>,
9580        ncols: usize,
9581        nrows: usize,
9582        eps: f32,
9583    ) -> Result<(), Box<dyn std::error::Error>> {
9584        let f = self.func("l2_norm_f32");
9585        let cfg = LaunchConfig {
9586            grid_dim: (nrows as u32, 1, 1),
9587            block_dim: (256, 1, 1),
9588            shared_mem_bytes: 0,
9589        };
9590        let (nc, e) = (ncols as i32, eps);
9591        let __s_b = self.gpu.stream();
9592        let mut b = __s_b.launch_builder(&f);
9593        b.arg(x).arg(dst).arg(&nc).arg(&e);
9594        unsafe {
9595            b.launch(cfg)?;
9596        }
9597        Ok(())
9598    }
9599
9600    /// L2-norm with blockDim=32 (warp-tree reduction) — BIT-IDENTICAL to gdn_prep_decode_f32's
9601    /// per-warp L2 norm. The verify path MUST use this to match decode's FP accumulation order:
9602    /// l2_norm at blockDim=256 produces a different shfl-tree reduction of the 128-element
9603    /// squared-sum (pairwise tree vs serial-4-then-warp-tree), causing ULP differences that
9604    /// propagate through gdn_scan and flip argmax on marginal logits.
9605    pub fn l2_norm_decode(
9606        &self,
9607        x: &CudaSlice<f32>,
9608        dst: &mut CudaSlice<f32>,
9609        ncols: usize,
9610        nrows: usize,
9611        eps: f32,
9612    ) -> Result<(), Box<dyn std::error::Error>> {
9613        let f = self.func("l2_norm_f32");
9614        let cfg = LaunchConfig {
9615            grid_dim: (nrows as u32, 1, 1),
9616            block_dim: (32, 1, 1),
9617            shared_mem_bytes: 0,
9618        };
9619        let (nc, e) = (ncols as i32, eps);
9620        let __s_b = self.gpu.stream();
9621        let mut b = __s_b.launch_builder(&f);
9622        b.arg(x).arg(dst).arg(&nc).arg(&e);
9623        unsafe {
9624            b.launch(cfg)?;
9625        }
9626        Ok(())
9627    }
9628
9629    /// RoPE NEOX in-place. x:[head_dim, n_heads, n_tokens], pos:[n_tokens].
9630    pub fn rope_neox(
9631        &self,
9632        x: &mut CudaSlice<f32>,
9633        pos: &CudaSlice<i32>,
9634        head_dim: usize,
9635        n_dims: usize,
9636        n_heads: usize,
9637        n_tokens: usize,
9638        freq_base: f32,
9639        freq_scale: f32,
9640    ) -> Result<(), Box<dyn std::error::Error>> {
9641        let f = self.func("rope_neox_f32");
9642        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
9643        let grid = (n_heads * n_tokens) as u32;
9644        let cfg = LaunchConfig {
9645            grid_dim: (grid, 1, 1),
9646            block_dim: ((head_dim / 2) as u32, 1, 1),
9647            shared_mem_bytes: 0,
9648        };
9649        let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
9650        let __s_b = self.gpu.stream();
9651        let mut b = __s_b.launch_builder(&f);
9652        b.arg(x)
9653            .arg(pos)
9654            .arg(&hd)
9655            .arg(&nd)
9656            .arg(&nh)
9657            .arg(&theta_scale)
9658            .arg(&freq_scale);
9659        unsafe {
9660            b.launch(cfg)?;
9661        }
9662        Ok(())
9663    }
9664
9665    /// RoPE NEOX with per-dim freq factors (gemma4 global layers, rope_freqs.weight [n_dims/2]).
9666    pub fn rope_neox_ff(
9667        &self,
9668        x: &mut CudaSlice<f32>,
9669        pos: &CudaSlice<i32>,
9670        head_dim: usize,
9671        n_dims: usize,
9672        n_heads: usize,
9673        n_tokens: usize,
9674        freq_base: f32,
9675        freq_scale: f32,
9676        ff: &CudaSlice<f32>,
9677    ) -> Result<(), Box<dyn std::error::Error>> {
9678        let f = self.func("rope_neox_ff_f32");
9679        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
9680        let grid = (n_heads * n_tokens) as u32;
9681        let cfg = LaunchConfig {
9682            grid_dim: (grid, 1, 1),
9683            block_dim: ((head_dim / 2) as u32, 1, 1),
9684            shared_mem_bytes: 0,
9685        };
9686        let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
9687        let __s_b = self.gpu.stream();
9688        let mut b = __s_b.launch_builder(&f);
9689        b.arg(x)
9690            .arg(pos)
9691            .arg(&hd)
9692            .arg(&nd)
9693            .arg(&nh)
9694            .arg(&theta_scale)
9695            .arg(&freq_scale)
9696            .arg(ff);
9697        unsafe {
9698            b.launch(cfg)?;
9699        }
9700        Ok(())
9701    }
9702
9703    /// gemma4: rope q and k in one launch (per-row chain = rope_neox / rope_neox_ff verbatim).
9704    #[allow(clippy::too_many_arguments)]
9705    pub fn rope_neox2(
9706        &self,
9707        q: &mut CudaSlice<f32>,
9708        k: &mut CudaSlice<f32>,
9709        pos: &CudaSlice<i32>,
9710        head_dim: usize,
9711        n_dims: usize,
9712        nh_q: usize,
9713        nh_k: usize,
9714        n_tokens: usize,
9715        freq_base: f32,
9716        freq_scale: f32,
9717        ff: Option<&CudaSlice<f32>>,
9718    ) -> Result<(), Box<dyn std::error::Error>> {
9719        let f = self.func("rope_neox2_f32");
9720        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
9721        let grid = ((nh_q + nh_k) * n_tokens) as u32;
9722        let cfg = LaunchConfig {
9723            grid_dim: (grid, 1, 1),
9724            block_dim: ((head_dim / 2) as u32, 1, 1),
9725            shared_mem_bytes: 0,
9726        };
9727        let (hd, nd, nq, nk, nt) = (
9728            head_dim as i32,
9729            n_dims as i32,
9730            nh_q as i32,
9731            nh_k as i32,
9732            n_tokens as i32,
9733        );
9734        let __s_b = self.gpu.stream();
9735        let mut b = __s_b.launch_builder(&f);
9736        b.arg(q)
9737            .arg(k)
9738            .arg(pos)
9739            .arg(&hd)
9740            .arg(&nd)
9741            .arg(&nq)
9742            .arg(&nk)
9743            .arg(&nt)
9744            .arg(&theta_scale)
9745            .arg(&freq_scale);
9746        match ff {
9747            Some(ffv) => {
9748                b.arg(ffv);
9749                unsafe {
9750                    b.launch(cfg)?;
9751                }
9752            }
9753            None => {
9754                let null: u64 = 0;
9755                b.arg(&null);
9756                unsafe {
9757                    b.launch(cfg)?;
9758                }
9759            }
9760        }
9761        Ok(())
9762    }
9763
9764    /// gemma4 R1: dst = GELU_tanh(gate) * up.
9765    pub fn gelu_tanh_mul(
9766        &self,
9767        gate: &CudaSlice<f32>,
9768        up: &CudaSlice<f32>,
9769        dst: &mut CudaSlice<f32>,
9770        n: usize,
9771    ) -> Result<(), Box<dyn std::error::Error>> {
9772        let f = self.func("gelu_tanh_mul_f32");
9773        let cfg = LaunchConfig::for_num_elems(n as u32);
9774        let ni = n as i32;
9775        let __s_b = self.gpu.stream();
9776        let mut b = __s_b.launch_builder(&f);
9777        b.arg(gate).arg(up).arg(dst).arg(&ni);
9778        unsafe {
9779            b.launch(cfg)?;
9780        }
9781        Ok(())
9782    }
9783
9784    pub fn silu_mul(
9785        &self,
9786        gate: &CudaSlice<f32>,
9787        up: &CudaSlice<f32>,
9788        dst: &mut CudaSlice<f32>,
9789        n: usize,
9790    ) -> Result<(), Box<dyn std::error::Error>> {
9791        let f = self.func("silu_mul_f32");
9792        // float4 kernel: one thread per 4 elements (tail handled in-kernel)
9793        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
9794        let ni = n as i32;
9795        let __s_b = self.gpu.stream();
9796        let mut b = __s_b.launch_builder(&f);
9797        b.arg(gate).arg(up).arg(dst).arg(&ni);
9798        unsafe {
9799            b.launch(cfg)?;
9800        }
9801        Ok(())
9802    }
9803
9804    /// f16out twin of `silu_mul` (task #17): the epilogue also emits the fp16 GEMM operand
9805    /// for the down projection — kills the standalone convert pass. Bit-identical class.
9806    pub fn silu_mul_f16out(
9807        &self,
9808        gate: &CudaSlice<f32>,
9809        up: &CudaSlice<f32>,
9810        dst: &mut CudaSlice<f32>,
9811        dst16: &mut CudaSlice<u8>,
9812        n: usize,
9813    ) -> Result<(), Box<dyn std::error::Error>> {
9814        let f = self.func("silu_mul_f16out_f32");
9815        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
9816        let ni = n as i32;
9817        let __s_b = self.gpu.stream();
9818        let mut b = __s_b.launch_builder(&f);
9819        b.arg(gate).arg(up).arg(dst).arg(dst16).arg(&ni);
9820        unsafe {
9821            b.launch(cfg)?;
9822        }
9823        Ok(())
9824    }
9825
9826    /// FFN SwiGLU epilogue fusion (RANK3 LEVER 2): `dst = silu(gate*gs) * (up*us)` in ONE launch,
9827    /// folding the per-tensor NVFP4 macro-scale (`gs`,`us`) that would otherwise be two separate
9828    /// `scale_inplace` launches on the gate/up matmul outputs. BIT-IDENTICAL to
9829    /// scale_inplace(gate,gs); scale_inplace(up,us); silu_mul(gate,up,dst) — identical float ops in
9830    /// identical order. For non-NVFP4 weights gs==us==1.0 -> identical to `silu_mul`. Net: -2
9831    /// launches per dense FFN layer (the gate+up post-matmul scales).
9832    pub fn silu_mul_scaled(
9833        &self,
9834        gate: &CudaSlice<f32>,
9835        up: &CudaSlice<f32>,
9836        gs: f32,
9837        us: f32,
9838        dst: &mut CudaSlice<f32>,
9839        n: usize,
9840    ) -> Result<(), Box<dyn std::error::Error>> {
9841        let f = self.func("silu_mul_scaled_f32");
9842        let cfg = LaunchConfig::for_num_elems(n as u32);
9843        let ni = n as i32;
9844        let (gsf, usf) = (gs, us);
9845        let __s_b = self.gpu.stream();
9846        let mut b = __s_b.launch_builder(&f);
9847        b.arg(gate).arg(up).arg(&gsf).arg(&usf).arg(dst).arg(&ni);
9848        unsafe {
9849            b.launch(cfg)?;
9850        }
9851        Ok(())
9852    }
9853
9854    /// swigluoai (MiniMax-M3 / GPT-OSS): clamped SwiGLU epilogue, math 1:1 vs llama.cpp
9855    /// ggml_cuda_op_swiglu_oai_single. `dst = swish_alpha(clamp(gate*gs)) * (1 + clamp(up*us))`.
9856    /// gs/us fold the NVFP4 macro-scales exactly like `silu_mul_scaled`.
9857    #[allow(clippy::too_many_arguments)]
9858    pub fn swigluoai_mul_scaled(
9859        &self,
9860        gate: &CudaSlice<f32>,
9861        up: &CudaSlice<f32>,
9862        gs: f32,
9863        us: f32,
9864        alpha: f32,
9865        limit: f32,
9866        dst: &mut CudaSlice<f32>,
9867        n: usize,
9868    ) -> Result<(), Box<dyn std::error::Error>> {
9869        let f = self.func("swigluoai_mul_scaled_f32");
9870        let cfg = LaunchConfig::for_num_elems(n as u32);
9871        let ni = n as i32;
9872        let __s_b = self.gpu.stream();
9873        let mut b = __s_b.launch_builder(&f);
9874        b.arg(gate)
9875            .arg(up)
9876            .arg(&gs)
9877            .arg(&us)
9878            .arg(&alpha)
9879            .arg(&limit)
9880            .arg(dst)
9881            .arg(&ni);
9882        unsafe {
9883            b.launch(cfg)?;
9884        }
9885        Ok(())
9886    }
9887
9888    /// RANK2 LEVER (q8_1 quant-fold): SwiGLU epilogue that EMITS the q8_1 quantization of `act`
9889    /// directly (aq int8 [n] + ad f32 [n/32]), so ffn_down's standalone `quantize_q8_1` launch is
9890    /// removed — the down-proj activation has one consumer, so the quant folds into the producer for
9891    /// free (no extra HBM read; no f32 `act` write). gs/us fold the gate/up NVFP4 macro-scales like
9892    /// `silu_mul_scaled`. BIT-IDENTICAL q8_1 to silu_mul_scaled(...) then quantize_q8_1(...). Only
9893    /// valid when ffn_down uses the q8_1 dp4a/mmvq path; the caller checks `uses_q8_1_fast(ffn_down)`.
9894    /// n must be a multiple of 32 (n_ff always is).
9895    pub fn silu_mul_scaled_q8_1(
9896        &self,
9897        gate: &CudaSlice<f32>,
9898        up: &CudaSlice<f32>,
9899        gs: f32,
9900        us: f32,
9901        n: usize,
9902    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9903        let f = self.func("silu_mul_scaled_q8_1");
9904        let nblk = n / 32;
9905        let mut aq = self.alloc_uninit::<i8>(n)?; // full-overwrite output
9906        let mut ad = self.alloc_uninit::<f32>(nblk)?; // full-overwrite output
9907        // WARP-PER-BLOCK kernel: one warp (32 lanes) per 32-block -> n threads total.
9908        let cfg = LaunchConfig::for_num_elems(n as u32);
9909        let (gsf, usf, ni) = (gs, us, n as i32);
9910        let __s_b = self.gpu.stream();
9911        let mut b = __s_b.launch_builder(&f);
9912        b.arg(gate)
9913            .arg(up)
9914            .arg(&gsf)
9915            .arg(&usf)
9916            .arg(&mut aq)
9917            .arg(&mut ad)
9918            .arg(&ni);
9919        unsafe {
9920            b.launch(cfg)?;
9921        }
9922        Ok((aq, ad))
9923    }
9924
9925    pub fn add(
9926        &self,
9927        a: &CudaSlice<f32>,
9928        b_in: &CudaSlice<f32>,
9929        dst: &mut CudaSlice<f32>,
9930        n: usize,
9931    ) -> Result<(), Box<dyn std::error::Error>> {
9932        let f = self.func("add_f32");
9933        // float4 kernel: one thread per 4 elements (tail handled in-kernel)
9934        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
9935        let ni = n as i32;
9936        let __s_bld = self.gpu.stream();
9937        let mut bld = __s_bld.launch_builder(&f);
9938        bld.arg(a).arg(b_in).arg(dst).arg(&ni);
9939        unsafe {
9940            bld.launch(cfg)?;
9941        }
9942        Ok(())
9943    }
9944
9945    pub fn mul(
9946        &self,
9947        a: &CudaSlice<f32>,
9948        b_in: &CudaSlice<f32>,
9949        dst: &mut CudaSlice<f32>,
9950        n: usize,
9951    ) -> Result<(), Box<dyn std::error::Error>> {
9952        let f = self.func("mul_f32");
9953        let cfg = LaunchConfig::for_num_elems(n as u32);
9954        let ni = n as i32;
9955        let __s_bld = self.gpu.stream();
9956        let mut bld = __s_bld.launch_builder(&f);
9957        bld.arg(a).arg(b_in).arg(dst).arg(&ni);
9958        unsafe {
9959            bld.launch(cfg)?;
9960        }
9961        Ok(())
9962    }
9963
9964    /// Unified weight-tensor matmul: dispatches quant tensors to qmatvec (weights packed) and
9965    /// float tensors to cuBLASLt. y[m,out] = x[m,in] @ W[out,in]^T.
9966    pub fn matmul(
9967        &self,
9968        w: &crate::model::GpuTensor,
9969        x: &CudaSlice<f32>,
9970        m: usize,
9971    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9972        use crate::model::GpuTensor;
9973        let in_f = w.in_features();
9974        let out_f = w.out_features();
9975        // PREFILL (T>1) ROOT FIX: batched tensor-core int8 GEMM. Decodes each weight tile to int8
9976        // in smem ONCE and reuses across all tokens via mma — vs the dp4a matvec's per-token weight
9977        // re-read. Only the 4 daily-hot dtypes; m=1 decode keeps dp4a (it's bandwidth-bound, mma
9978        // gives nothing). Quantize the activation once here then call the GEMM.
9979        // m cutoff FIXED at 16: the m=4 MMA-verify A/B (2026-07-06, was MEMRA_GEMM_M) measured
9980        // NEGATIVE — the MMA tile grid starves at m=4 (BN=256 -> grid.y=1) and its FP order
9981        // shifted verify argmax at tight margins. Do not lower without re-running that battery.
9982        #[allow(non_snake_case)]
9983        // VERIFY-EXACT scope pushes the GEMM crossover out of reach (usize::MAX) — the
9984        // t>=16 dflash verify must ride the decode-exact batched class (parity law).
9985        let GEMM_M_THRESHOLD = if self.verify_exact_on() {
9986            usize::MAX
9987        } else {
9988            16usize
9989        };
9990
9991        // PREFILL GEMM (m>=16). ACCURACY-FIRST dispatch (2026-06-28, prefill-gemm-beat-research wf
9992        // wllbyo6vc step 1): the int8 W4A8 GEMM (qmatvec_gemm, q8_1 activation, s32 accumulate) is
9993        // ACCURATE (prefill logit maxdiff 0.159, < dp4a 0.55) and the default. The FP4 W4A4 mxf4 path
9994        // (try_fp4_gemm) quantizes the ACTIVATION to e2m1 4-bit (8 magnitude levels) -> maxdiff 1.0
9995        // when combined — a real accuracy loss, NOT a math bug. So FP4-W4A4 is taken ONLY under the
9996        // explicit MEMRA_FP4 opt-in AND it must come SECOND (int8 W4A8 is the correct default for NVFP4).
9997        // The workflow plan rebuilds the FP4 path (kill per-K repack, widen K, deepen pipeline, TMA) to
9998        // be both fast AND accurate; until then NVFP4 prefill defaults to the accurate int8 GEMM.
9999        // TINY-OUT_F GUARD (2026-06-28, ncu trace): the tiling GEMM's grid is (ceil(out_f/BM=64),
10000        // ceil(m/BN=256)). For tiny out_f (ssm_beta/ssm_alpha out_f=num_v_heads~32), grid.x=1 -> only
10001        // ceil(m/256) CTAs (e.g. 2 for m=512) on 82 SMs = 0.39% SM throughput, 852us EACH (measured
10002        // worst offender). The dp4a path grids (out_f, m) = far more CTAs, filling the GPU. So route
10003        // out_f < 2*BM to dp4a (skip the tiling GEMM which structurally can't fill the SMs here).
10004        const GEMM_MIN_OUT_F: usize = 128; // 2*BM; below this the GEMM grid.x starves the 82 SMs
10005        // VENDORED llama MMQ prefill GEMMs. NVFP4 W4A8 is DEFAULT-ON (2026-07-05 flip: same int8
10006        // accuracy class as the int8 GEMM below at ~1.9x pp512, rp-loader coexists with the A6
10007        // repack; MEMRA_MMQ_W4A8=0 = escape hatch). W4A4 mxf4nvf4 + Q4_K/Q5_K stay behind MEMRA_MMQ=1.
10008        // The env policy lives in mmq_supports/qmatvec_mmq. Feeds raw f32 activation `x` (the
10009        // launcher quantizes internally). out_f>=MMQ_Y/2 keeps the tile grid from starving the SMs.
10010        // FP8-ACT PREFILL (MEMRA_PP_FP8=1, probe verdict 2026-07-08): F8-E4M3-origin projections
10011        // carry their raw e4m3 device bytes (the `fp8` operand stashed at load next to the Q8_0
10012        // re-encode) — cuBLASLt FP8 TN at 620-795 TF vs 47-72 TF for this class's int8 GEMM.
10013        // Weight side EXACT (checkpoint bytes); activation rides ONE per-batch e4m3 scale
10014        // (amax/448) folded with weight_scale in-GEMM. Prefill only; decode keeps Q8_0 untouched.
10015        if m >= GEMM_M_THRESHOLD {
10016            if let Some(y) = self.try_fp8_gemm(w, x, m)? {
10017                return Ok(y);
10018            }
10019            // PER-BLOCK FP8 MMQ (lane/fp8-mmq): the block-128 class try_fp8_gemm skips (cuBLASLt
10020            // takes no block grid on sm_120). Exact per block — the checkpoint's e4m3 bytes and its
10021            // f32 grid go into the tile unchanged. TWO SOURCES, TWO DEFAULTS: the load-time stash is
10022            // opt-in (MEMRA_FP8_MMQ=1), the native-resident QT_F8_E4M3_BLK grid is DEFAULT ON
10023            // (MEMRA_FP8_MMQ=0 reverts it to dequant-per-call) — see fp8_ffi.rs for why the same
10024            // tile defaults differently by operand source.
10025            if let Some(y) = self.try_fp8_blk_mmq(w, x, m)? {
10026                return Ok(y);
10027            }
10028            // FP16-mirror prefill (MEMRA_PP_F16=1, probe 2026-07-26: 3.2-3.7x the MMQ class).
10029            // Mirror presence IS the gate (only built under the env). Decode never reaches here.
10030            if let Some(y) = self.try_f16_gemm(w, x, m)? {
10031                return Ok(y);
10032            }
10033        }
10034        // F8-E4M3 BLOCK-128 (QT_F8_E4M3_BLK, lane/fp8-blk128-decode). TWO arms, split at the SAME
10035        // m threshold the rest of this method uses:
10036        //   * m >= threshold (prefill): dequant-per-call to the ARM B' Q8_0 slab and recurse, so
10037        //     prefill keeps the floor's kernels AND the floor's bits (try_e4m3_blk_prefill).
10038        //   * m <  threshold: the native per-block GEMV — m=1 decode and the m=2..15 verify tiers.
10039        //     grid.y=m runs the exact m=1 program per (token,row), so the decode-parity law holds
10040        //     across every tier by construction with no batched twin needed.
10041        //
10042        // NOT gated on `fast`: this dtype has no dp4a twin and no Stage-A f32-dequant oracle (the
10043        // generic `deq()` switch has no block-scale input), exactly as QT_F8_E4M3 has none, so
10044        // MEMRA_FAST=0 cannot route it anywhere else. Placed before every GEMM/MMQ arm below
10045        // because gemm_supports/mmq_supports/mmvq_supports all deliberately REFUSE this qtype —
10046        // reaching the generic tail would panic rather than produce wrong numbers, and this pair of
10047        // arms is what makes sure it never gets there.
10048        if let GpuTensor::Quant { qtype, .. } = w {
10049            if *qtype == QT_F8_E4M3_BLK {
10050                if m >= GEMM_M_THRESHOLD {
10051                    if let Some(y) = self.try_e4m3_blk_prefill(w, x, m)? {
10052                        return Ok(y);
10053                    }
10054                }
10055                let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10056                if let Some(y) = self.try_e4m3_blk_pre(w, &aq, &ad, m)? {
10057                    return Ok(y);
10058                }
10059            }
10060        }
10061        if m >= GEMM_M_THRESHOLD && out_f >= GEMM_MIN_OUT_F && self.mmq_supports(w) {
10062            return self.qmatvec_mmq(w, x, m);
10063        }
10064        if m >= GEMM_M_THRESHOLD && out_f >= GEMM_MIN_OUT_F && self.gemm_supports(w) {
10065            let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10066            return self.qmatvec_gemm(w, &aq, &ad, m);
10067        }
10068        // FP4 W4A4 only as an explicit speed/accuracy tradeoff opt-in, and only if the int8 GEMM
10069        // above didn't already handle this weight (e.g. NVFP4 with in_f%64!=0, or MEMRA_NO_GEMM set).
10070        if m >= GEMM_M_THRESHOLD {
10071            if let Some(y) = self.try_fp4_gemm(w, x, m, in_f, out_f)? {
10072                return Ok(y);
10073            }
10074        }
10075        // Stage-B fast int8 dp4a is the DEFAULT since 2026-07-08 (it has been the daily path
10076        // for weeks; the old opt-in flag was a silent-slow-path landmine). MEMRA_FAST=0 reverts
10077        // to Stage-A f32-dequant (the correctness oracle path).
10078        let fast = std::env::var("MEMRA_FAST").as_deref() != Ok("0");
10079        // PERF-3 decode-GEMV: m=1 warp-per-row MMVQ (MEMRA_MMVQ). The big decode matvecs reach
10080        // `matmul` directly (ffn_down, lm_head output, wo), so route them here too — not only the
10081        // matmul_pre siblings. qmatvec_mmvq_raw quantizes the activation internally (q8_1) like the
10082        // _fast paths; the NVFP4 macro-scale is applied by the `scale != 1.0` block below.
10083        if m == 1 && fast {
10084            if let GpuTensor::Quant {
10085                bytes,
10086                qtype,
10087                row_bytes,
10088                rp,
10089                rp4,
10090                scale,
10091                ..
10092            } = w
10093            {
10094                if self.mmvq_supports(*qtype) {
10095                    // NVFP4 macro-scale rides the kernel's fused epilogue arg (one launch total);
10096                    // non-NVFP4 has scale==1.0 so qmatvec_mmvq skips scale_inplace either way.
10097                    // Q4_0 split-plane mirror (rp4): the decode arm reads it via the _rp twins.
10098                    let (bytes, rp) = match rp4 {
10099                        Some(m4) => (m4, true),
10100                        None => (bytes, *rp),
10101                    };
10102                    let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10103                    return self.qmatvec_mmvq(
10104                        bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, *scale, rp,
10105                    );
10106                }
10107            }
10108        }
10109        // BATCHED weight-resident matvec for the m=2-4 band (the MTP/verify forward's ffn_down, wo, and
10110        // lm_head `output` reach `matmul` directly at m=T=2..4). Walks the weight ONCE, dp4a vs all m
10111        // activation columns -> 1 weight read for m tokens (vs grid.y=m re-reading m times below). Quant
10112        // the activation once here (q8_1) like the _fast paths; macro-scale applied via the scale!=1.0
10113        // block below. MEMRA_NO_BATCHED -> per-m path.
10114        //
10115        // DECODE-PARITY GATE (2026-07-07, the 9B synth K=3/4/6 spec FAIL root cause): the batched
10116        // kernels are bit-identical per (token,row) to MMVQ's 32-thread warp reduce, NOT to the
10117        // dp4a kernels' 128-thread two-level reduce. Without MEMRA_MMVQ the m=1 decode chain rides
10118        // dp4a, so a verify riding batched here has a DIFFERENT FP order than the decode it must
10119        // match bit-for-bit — greedy spec flips at tight-margin tokens (the old HANDOVER "ENV LAW:
10120        // FAST+MMVQ both required" footgun, closed here). Parity law: the m>1 kernel CLASS must be
10121        // a pure function of (dtype, env) equal to the m=1 class — batched iff MMVQ. Without MMVQ
10122        // the verify falls to the per-m grid.y=m dp4a path below (each column = the exact m=1
10123        // dp4a program). MEMRA_MMVQ=1 (the daily config) is dispatch-unchanged.
10124        if (2..=16).contains(&m)
10125            && fast
10126            && std::env::var("MEMRA_NO_BATCHED").is_err()
10127            && (m <= 4 || Self::b8_enabled())
10128        {
10129            // b16 tier (2026-07-11, spec K>7): Q4_0/Q6_K have base+_rp b16 kernels; Q8_0's
10130            // b16 exists only as the split-plane _rp twin, so it joins iff the q8rp mirror
10131            // is present (rp4) — the mirror pick below then routes to the _rp family.
10132            // QT_F8_E4M3 joins unconditionally (lane/rp-on-st): its b16 IS the base kernel,
10133            // because the native e4m3 row layout is already aligned and needs no mirror.
10134            // NVFP4/Q4_K/Q8_0 all join unconditionally now (lane/rp-on-st): each has base + _rp
10135            // b16 twins, so either residency layout has its aligned form at this width. Q8_0's
10136            // old `rp4.is_some()` precondition is GONE — the mirror is a bandwidth lever, not the
10137            // exact tier's admission ticket (it was refusing FP8-ST over 23.9 MiB of ssm_beta).
10138            let m_ok = m <= 8
10139                || matches!(w, GpuTensor::Quant { qtype, .. }
10140                if *qtype == QT_Q4_0 || *qtype == QT_Q6_K || *qtype == QT_F8_E4M3
10141                    || *qtype == QT_NVFP4 || *qtype == QT_Q4_K || *qtype == QT_Q5_K || *qtype == QT_Q8_0);
10142            if m_ok {
10143                if let GpuTensor::Quant {
10144                    bytes,
10145                    qtype,
10146                    row_bytes,
10147                    rp,
10148                    rp4,
10149                    ..
10150                } = w
10151                {
10152                    if self.batched_supports(*qtype) && self.mmvq_supports(*qtype) {
10153                        let (bytes, rp) = match rp4 {
10154                            Some(m4) => (m4, true),
10155                            None => (bytes, *rp),
10156                        };
10157                        let mcols = Self::batched_mcols(m);
10158                        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10159                        let mut y = self.qmatvec_mmvq_batched(
10160                            bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, mcols, 1.0, rp,
10161                        )?;
10162                        if let GpuTensor::Quant { scale, .. } = w {
10163                            if *scale != 1.0 {
10164                                self.scale_inplace(&mut y, *scale, m * out_f)?;
10165                            }
10166                        }
10167                        return Ok(y);
10168                    }
10169                }
10170            }
10171        }
10172        // F8-E4M3 (MEMRA_ST_E4M3) catch-all for the m<16 band the arms above didn't take (m=9..15,
10173        // the K=8 verify tier; or m=2..8 under MEMRA_NO_BATCHED/MEMRA_B8=0): grid.y=m e4m3 mmvq —
10174        // the SAME per-(token,row) program as the m=1 decode launch (bit-identical by construction),
10175        // weight re-read m times (rare tier; exactness over bandwidth here). There is no _dp4a twin
10176        // for this dtype, so the generic match below must never see it under `fast`.
10177        if fast {
10178            if let GpuTensor::Quant {
10179                bytes,
10180                qtype,
10181                row_bytes,
10182                scale,
10183                ..
10184            } = w
10185            {
10186                if *qtype == QT_F8_E4M3 {
10187                    let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10188                    return self.qmatvec_mmvq(
10189                        bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, *scale, false,
10190                    );
10191                }
10192            }
10193        }
10194        let mut y = match w {
10195            GpuTensor::Quant {
10196                bytes,
10197                qtype,
10198                row_bytes,
10199                ..
10200            } if fast && *qtype == QT_Q8_0 => {
10201                self.qmatvec_q8_0_fast(bytes, x, m, in_f, out_f, *row_bytes)?
10202            }
10203            GpuTensor::Quant {
10204                bytes,
10205                qtype,
10206                row_bytes,
10207                ..
10208            } if fast && *qtype == QT_Q4_K => {
10209                self.qmatvec_q4_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
10210            }
10211            GpuTensor::Quant {
10212                bytes,
10213                qtype,
10214                row_bytes,
10215                ..
10216            } if fast && *qtype == QT_Q6_K => {
10217                self.qmatvec_q6_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
10218            }
10219            GpuTensor::Quant {
10220                bytes,
10221                qtype,
10222                row_bytes,
10223                ..
10224            } if fast && *qtype == QT_Q5_K => {
10225                self.qmatvec_q5_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
10226            }
10227            GpuTensor::Quant {
10228                bytes,
10229                qtype,
10230                row_bytes,
10231                ..
10232            } if fast && *qtype == QT_Q3_K => {
10233                self.qmatvec_q3_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
10234            }
10235            GpuTensor::Quant {
10236                bytes,
10237                qtype,
10238                row_bytes,
10239                rp,
10240                ..
10241            } if fast && *qtype == QT_NVFP4 => self.qmatvec_dp4a_named(
10242                if *rp {
10243                    "qmatvec_nvfp4_dp4a_rp"
10244                } else {
10245                    "qmatvec_nvfp4_dp4a"
10246                },
10247                bytes,
10248                x,
10249                m,
10250                in_f,
10251                out_f,
10252                *row_bytes,
10253            )?,
10254            // IQ4_XS trunk fast path — DEFAULT ON since 2026-08-02 (MEMRA_IQ_FAST=0 reverts to
10255            // Stage-A; see iq_fast_enabled). The old opt-in default was the KAT-Coder decode
10256            // anomaly (research/kat-anomaly-20260802/).
10257            GpuTensor::Quant {
10258                bytes,
10259                qtype,
10260                row_bytes,
10261                ..
10262            } if fast && *qtype == QT_IQ4_XS && Self::iq_fast_enabled() => {
10263                self.qmatvec_iq4_XS_fast(bytes, x, m, in_f, out_f, *row_bytes)?
10264            }
10265            // B3: IQ3_S uses the Stage-A f32 dequant-in-kernel path. There is NO
10266            // qmatvec_iq3_s_dp4a kernel — do NOT add a `*qtype == QT_IQ3_S` fast guard here
10267            // without first writing the matching kernel, or func() will panic
10268            // "kernel ... not in any fatbin".
10269            GpuTensor::Quant {
10270                bytes,
10271                qtype,
10272                row_bytes,
10273                rp,
10274                ..
10275            } =>
10276            // Stage-A generic: repacked NVFP4 uses the device-side split-plane tag (the
10277            // deq(row,j) form cannot address the planes; same value/product order).
10278            {
10279                self.qmatvec(
10280                    bytes,
10281                    x,
10282                    m,
10283                    in_f,
10284                    out_f,
10285                    if *rp && *qtype == QT_NVFP4 {
10286                        QT_NVFP4_RP
10287                    } else {
10288                        *qtype
10289                    },
10290                    *row_bytes,
10291                )?
10292            }
10293            GpuTensor::Float { data, .. } => self.linear(x, data, m, in_f, out_f)?,
10294            // MEMRA_FULL_PREC bf16-resident weight: dequant-on-use to f32 scratch, then the same
10295            // cuBLASLt f32 GEMV as the Float arm.
10296            GpuTensor::FloatBf16 { data, .. } => {
10297                self.linear_bf16_chunked(x, data, m, in_f, out_f, false)?
10298            }
10299        };
10300        // NVFP4 per-tensor macro-scale (post-matmul). scale==1.0 for all other quants/float -> no-op.
10301        if let GpuTensor::Quant { scale, .. } = w {
10302            if *scale != 1.0 {
10303                self.scale_inplace(&mut y, *scale, m * out_f)?;
10304            }
10305        }
10306        Ok(y)
10307    }
10308
10309    /// True when `MEMRA_FAST=0`, i.e. the Stage-A f32 oracle is the requested arithmetic and every
10310    /// `matmul_pre` call will take the raw-f32 escape rather than the q8_1 pair.
10311    ///
10312    /// WHY THIS EXISTS AS ITS OWN PREDICATE: `uses_q8_1_fast` needs a weight, but the callers that
10313    /// have to DECIDE WHETHER TO MATERIALIZE an f32 activation sit one level above any weight — the
10314    /// gemma-4 decode trunk emits a q8_1 pair per layer and hands it down, so the f32 has to be
10315    /// produced (or not) by the loop that owns the residual. Cached in a OnceLock like every other
10316    /// `*_on()` flag here: this is read once per layer per token on the decode path, and the daily
10317    /// path must not pay an env lookup for a flag that is off.
10318    pub fn stage_a_raw_needed() -> bool {
10319        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
10320        *ON.get_or_init(|| std::env::var("MEMRA_FAST").as_deref() == Ok("0"))
10321    }
10322
10323    /// True if `w` would take the int8-dp4a fast path under MEMRA_FAST (so its activation can be
10324    /// pre-quantized once and shared across sibling matmuls via `matmul_pre`).
10325    pub fn uses_q8_1_fast(&self, w: &crate::model::GpuTensor) -> bool {
10326        use crate::model::GpuTensor;
10327        if std::env::var("MEMRA_FAST").as_deref() == Ok("0") {
10328            return false;
10329        }
10330        match w {
10331            // QT_F8_E4M3_BLK is admitted for the same reason QT_F8_E4M3 is: its ONLY kernel class
10332            // takes the shared q8_1 activation, so callers may pre-quantize once and share it
10333            // across siblings. It is NOT admitted to any of the fused/dual epilogue doors those
10334            // siblings can then open (`q8_fused_params`, `e4m3_fused_params` and
10335            // `matmul_pre_dual_noscale` all match on their own qtype and refuse this one) — the
10336            // block class has no fused twin yet, so each of its projections takes its own launch.
10337            GpuTensor::Quant { qtype, .. } => {
10338                matches!(
10339                    *qtype,
10340                    QT_Q8_0
10341                        | QT_Q4_K
10342                        | QT_Q6_K
10343                        | QT_Q5_K
10344                        | QT_Q3_K
10345                        | QT_NVFP4
10346                        | QT_F8_E4M3
10347                        | QT_F8_E4M3_BLK
10348                        | QT_Q4_0
10349                ) || (*qtype == QT_IQ4_XS && Self::iq_fast_enabled())
10350            }
10351            GpuTensor::Float { .. } | GpuTensor::FloatBf16 { .. } => false,
10352        }
10353    }
10354
10355    /// matmul with a PRE-QUANTIZED q8_1 activation (aq,ad from `quantize_q8_1`). Skips the
10356    /// per-matmul re-quantize so sibling matmuls that share an input (gate+up share `z`;
10357    /// q/k/v + wqkv/gate/beta/alpha share `h`) quantize ONCE. Caller MUST have checked
10358    /// `uses_q8_1_fast(w)`; falls back to plain `matmul` otherwise (Stage-A / Float / non-fast).
10359    pub fn matmul_pre(
10360        &self,
10361        w: &crate::model::GpuTensor,
10362        aq: &CudaSlice<i8>,
10363        ad: &CudaSlice<f32>,
10364        x_fallback: &CudaSlice<f32>,
10365        m: usize,
10366    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10367        use crate::model::GpuTensor;
10368        // Every raw-f32 arm below (fp8/f16/MMQ/fp4) reads m*in_f from x_fallback. Callers that
10369        // pre-quantized and dropped the f32 input pass an EMPTY x_fallback (E4B's fusion port:
10370        // h = zeros(0)) — the length guard keeps those on the aq/ad GEMM instead of feeding a
10371        // 0-byte buffer to a convert kernel (illegal address -> cublasLt status 13; the E4B
10372        // rc=30013 dig, 2026-07-31).
10373        let x_raw_ok = x_fallback.len() >= m * w.in_features();
10374        // FP8-ACT PREFILL (MEMRA_PP_FP8=1): same arm as `matmul` — the fp8 operand needs the RAW
10375        // f32 activation (per-batch e4m3 quant differs from q8_1), so x_fallback not aq/ad.
10376        if m >= 16 && x_raw_ok && !self.verify_exact_on() {
10377            if let Some(y) = self.try_fp8_gemm(w, x_fallback, m)? {
10378                return Ok(y);
10379            }
10380            // PER-BLOCK FP8 MMQ — same arm as `matmul` (stash opt-in, native-resident default ON);
10381            // its own quantizer wants the RAW f32 activation, so x_fallback not aq/ad.
10382            if let Some(y) = self.try_fp8_blk_mmq(w, x_fallback, m)? {
10383                return Ok(y);
10384            }
10385            // FP16-mirror prefill (same arm as `matmul` — fp16 wants the RAW f32 activation).
10386            if let Some(y) = self.try_f16_gemm(w, x_fallback, m)? {
10387                return Ok(y);
10388            }
10389        }
10390        // BLOCK-128 e4m3 (QT_F8_E4M3_BLK) — the same two arms as `matmul`, split at the same m, and
10391        // placed at the same point in the order (after the prefill GEMM hooks, before every arm
10392        // that refuses this qtype). The prefill arm needs the RAW f32 activation for the Q8_0
10393        // dispatch it recurses into, so it takes x_fallback and is skipped when that is empty
10394        // (a pre-quantized caller that dropped its f32 input never runs at prefill m anyway).
10395        if m >= 16 && x_raw_ok && !self.verify_exact_on() {
10396            if let Some(y) = self.try_e4m3_blk_prefill(w, x_fallback, m)? {
10397                return Ok(y);
10398            }
10399        }
10400        if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
10401            return Ok(y);
10402        }
10403        // VENDORED llama MMQ prefill GEMMs (NVFP4 W4A8 default-on; W4A4/k-quant behind MEMRA_MMQ=1
10404        // — policy in mmq_supports) — use the RAW f32 activation (their own internal quant:
10405        // q8_1 D4 for NVFP4 W4A8, FP8/UE4M3 for W4A4, q8_1 DS4 for Q4_K/Q5_K), so x_fallback not
10406        // aq/ad.
10407        if m >= 16
10408            && w.out_features() >= 128
10409            && self.mmq_supports(w)
10410            && !self.verify_exact_on()
10411            && x_raw_ok
10412        {
10413            return self.qmatvec_mmq(w, x_fallback, m);
10414        }
10415        // Stage-C FP4 prefill (MEMRA_FP4): native mxf4 GEMM needs the f32 activation (FP4-quant differs
10416        // from q8_1), so re-quantize from x_fallback rather than reuse aq/ad. NVFP4 only, m>=16.
10417        if m >= 16 && x_raw_ok && !self.verify_exact_on() {
10418            if let Some(y) =
10419                self.try_fp4_gemm(w, x_fallback, m, w.in_features(), w.out_features())?
10420            {
10421                return Ok(y);
10422            }
10423        }
10424        // Prefill GEMM root fix: if T>1 and the dtype has a GEMM kernel, batch via tensor cores
10425        // (reuses the already-quantized aq/ad — no extra quantize). m=1 falls through to dp4a.
10426        if m >= 16 && self.gemm_supports(w) && !self.verify_exact_on() {
10427            return self.qmatvec_gemm(w, aq, ad, m);
10428        }
10429        // THE SEVENTH ARM. Every raw-f32 arm above is guarded by `x_raw_ok`; this one — the
10430        // Stage-A / Float escape — was not, and it is the ONLY one `MEMRA_FAST=0` opens. So the
10431        // 2026-07-31 E4B rc=30013 fix (add the length guard, keep empty-fallback callers off the
10432        // raw-f32 arms) protected the six arms the FAST path can reach and missed the one the
10433        // ORACLE path reaches. Consequence before this guard: the gemma-4 decode arms, which pass
10434        // `e.zeros(0)` because their trunk emits only a q8_1 pair and never materializes an f32
10435        // attn-normed activation, fell through to Stage-A `qmatvec_f32` (cu/qmatvec.cu:5442),
10436        // which reads `m * in_f` floats out of a 0-byte allocation ->
10437        // CUDA_ERROR_ILLEGAL_ADDRESS at layer 0 of the first decode token. The fault is STICKY:
10438        // it poisons the context, so every LATER request in that process fails with an unrelated
10439        // message ("cache alloc failed: ...") and the true cause appears exactly once, in the
10440        // first failure. That is what made `MEMRA_FAST=0` — the reference named in
10441        // `tools/argmax-margin-gate.sh`'s own header — return an opaque HTTP 500 on the gemma
10442        // dense artifact and left the arm with no working truth instrument.
10443        //
10444        // Refuse loudly instead of reading out of bounds. A named error at the true call site is
10445        // strictly better than an illegal address surfacing later at an unrelated sync point, and
10446        // an oracle that cannot run must say so rather than corrupt the context it runs in.
10447        if !self.uses_q8_1_fast(w) {
10448            if !x_raw_ok {
10449                return Err(format!(
10450                    "matmul_pre: q8_1-fast is off for this weight but x_fallback holds {} f32 \
10451                     (need m*in_f = {}*{} = {}). This call site pre-quantized its activation and \
10452                     dropped the f32, so there is nothing to fall back to — pass the real f32 \
10453                     activation (see Engine::rms_norm_decode, which is bit-identical to \
10454                     rms_norm_q8_1's reduction) or keep the weight on the q8_1 path.",
10455                    x_fallback.len(),
10456                    m,
10457                    w.in_features(),
10458                    m * w.in_features()
10459                )
10460                .into());
10461            }
10462            return self.matmul(w, x_fallback, m);
10463        }
10464        let in_f = w.in_features();
10465        let out_f = w.out_features();
10466        let (bytes, qtype, row_bytes, scale, rp) = match w {
10467            GpuTensor::Quant {
10468                bytes,
10469                qtype,
10470                row_bytes,
10471                scale,
10472                rp,
10473                ..
10474            } => (bytes, *qtype, *row_bytes, *scale, *rp),
10475            _ => unreachable!("uses_q8_1_fast guaranteed Quant"),
10476        };
10477        // Q4_0 split-plane mirror: only the mmvq/batched decode arms read it (the _rp twins);
10478        // the dp4a/oracle tails below keep the raw GGUF bytes.
10479        let (mbytes, mrp) = match w {
10480            GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
10481            _ => (bytes, rp),
10482        };
10483        // PERF-3 decode-GEMV: warp-per-row MMVQ for the m=1 decode arm, gated behind MEMRA_MMVQ.
10484        // Only the 4 daily-hot dtypes have an _mmvq kernel (Q8_0/Q4_K/Q6_K/NVFP4); Q5_K/Q3_K/IQ4_XS
10485        // keep _dp4a (the oracle/fallback). Bit-equivalent to _dp4a up to f32 reduction order.
10486        if m == 1 && self.mmvq_supports(qtype) {
10487            return self.qmatvec_mmvq(mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, mrp);
10488        }
10489        // BATCHED weight-resident matvec for the m=2-4 band (the MTP/verify forward: full_attn_verify
10490        // and decode_step_t run their projections at m=T=k=2..4). The plain _dp4a path below launches
10491        // grid.y=m INDEPENDENT blocks per output row -> the weight row is re-read m times from HBM/L2.
10492        // The _b2/_b4 kernels walk the weight ONCE and dp4a vs all m activation columns, so m tokens
10493        // cost ~1 weight read instead of m (decode is weight-BW-bound). BIT-IDENTICAL per (token,row)
10494        // to the _mmvq path (32-thread warp reduce — NOT the dp4a 128-thread reduce below).
10495        // m=2 -> mcols=2; m∈{3,4} -> mcols=4; m∈{5..8} -> mcols=8 (kernel guards c>=m).
10496        // MEMRA_NO_BATCHED forces the per-m grid.y=m path (the A/B reference); MEMRA_B8=0 keeps
10497        // m=5..8 on the old per-m path (b8-tier-only seam).
10498        // DECODE-PARITY GATE (2026-07-07): batched iff mmvq_supports — see matmul's parity note.
10499        // Without MEMRA_MMVQ, m=1 decode rides dp4a (the arm below at m=1); the verify must ride
10500        // the SAME class per column (grid.y=m dp4a = the exact m=1 dp4a program per column).
10501        if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
10502            && std::env::var("MEMRA_NO_BATCHED").is_err()
10503            && (m <= 4 || Self::b8_enabled())
10504            // b16 tier: every class routed here now has base + _rp b16 kernels (Q4_0/Q6_K
10505            // pre-existing; NVFP4/Q4_K/Q8_0-base/F8_E4M3 added lane/rp-on-st 2026-08-06), so
10506            // there is no mirror precondition left — `mrp` still selects the LAYOUT below.
10507            && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || qtype == QT_NVFP4
10508                || qtype == QT_Q4_K || qtype == QT_Q5_K || qtype == QT_F8_E4M3 || qtype == QT_Q8_0)
10509        {
10510            let mcols = Self::batched_mcols(m);
10511            return self.qmatvec_mmvq_batched(
10512                mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, mrp,
10513            );
10514        }
10515        // F8-E4M3 catch-all (m=9..15 / batched-disabled seams): grid.y=m e4m3 mmvq — this dtype
10516        // has NO _dp4a twin, and per (token,row) the mmvq body is the exact m=1 decode program.
10517        // Q4_0 joins the catch-all (2026-07-11): adaptive-K cap 8 makes verify t=9 reachable
10518        // for the first time (past the b8 tier) and Q4_0 has no dp4a twin either. The mirror
10519        // (mbytes/mrp) keeps the rp layout consistent with the m=1 decode program.
10520        if qtype == QT_F8_E4M3 || qtype == QT_Q4_0 {
10521            let (b2, r2) = if qtype == QT_Q4_0 {
10522                (mbytes, mrp)
10523            } else {
10524                (bytes, rp)
10525            };
10526            return self.qmatvec_mmvq(b2, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, r2);
10527        }
10528        let name = match qtype {
10529            QT_Q8_0 => "qmatvec_q8_0_dp4a",
10530            QT_Q4_K => "qmatvec_q4_K_dp4a",
10531            QT_Q6_K => "qmatvec_q6_K_dp4a",
10532            QT_Q5_K => "qmatvec_q5_K_dp4a",
10533            QT_Q3_K => "qmatvec_q3_K_dp4a",
10534            QT_NVFP4 => {
10535                if rp {
10536                    "qmatvec_nvfp4_dp4a_rp"
10537                } else {
10538                    "qmatvec_nvfp4_dp4a"
10539                }
10540            }
10541            QT_IQ4_XS => "qmatvec_iq4_XS_dp4a",
10542            _ => unreachable!(),
10543        };
10544        let f = self.func(name);
10545        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output: skip memset
10546        let cfg = LaunchConfig {
10547            grid_dim: (out_f as u32, m as u32, 1),
10548            block_dim: (128, 1, 1),
10549            shared_mem_bytes: 0,
10550        };
10551        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
10552        let __s_b = self.gpu.stream();
10553        let mut b = __s_b.launch_builder(&f);
10554        b.arg(bytes)
10555            .arg(aq)
10556            .arg(ad)
10557            .arg(&mut y)
10558            .arg(&inf)
10559            .arg(&outf)
10560            .arg(&mi)
10561            .arg(&rb);
10562        unsafe {
10563            b.launch(cfg)?;
10564        }
10565        if scale != 1.0 {
10566            self.scale_inplace(&mut y, scale, m * out_f)?;
10567        }
10568        Ok(y)
10569    }
10570
10571    /// DECODE-EXACT matmul at any m: guarantees the SAME warp-per-row (MMVQ, 32-thread) FP
10572    /// accumulation order as the T=1 decode path for EVERY token row. The spec-decode verify MUST
10573    /// use this for linear-attn projections to be bit-identical to greedy decode. The dp4a kernel
10574    /// (128 threads, two-level reduction) used by `matmul`/`matmul_pre` at m>=5 has a different
10575    /// shfl-tree shape that produces ULP differences propagating through gdn_scan into argmax flips.
10576    /// The MMVQ kernel with grid.y=m already processes each row independently (same 32-thread warp
10577    /// reduce as m=1); this method just forces that path unconditionally.
10578    pub fn matmul_decode_exact(
10579        &self,
10580        w: &crate::model::GpuTensor,
10581        x: &CudaSlice<f32>,
10582        m: usize,
10583    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10584        use crate::model::GpuTensor;
10585        // FLOAT tensors (35B ssm_beta/ssm_alpha on every linear layer, F32 ne=[2048,32]): the
10586        // generic path is cuBLASLt, whose reduction splits are n-DEPENDENT — m=1 vs m=2 col-0
10587        // outputs differ in every bit (probe 2026-07-06: 32/32 bit-diff, maxdiff 3.5e-3), which
10588        // shifted 35B verify logits 0.26-0.56 vs eager and flipped greedy at tight margins (the
10589        // p3 spec FAIL). Decode-exact contract: per-COLUMN m=1 cuBLASLt calls — each column's
10590        // reduction is the exact kernel the T=1 decode path runs, so verify==decode bit-for-bit.
10591        // m<=10 here (K+2 verify tier), so the extra launches are a handful of 4us gemvs.
10592        if let GpuTensor::Float { data, .. } = w {
10593            return self.linear_decode_exact(x, data, m, w.in_features(), w.out_features());
10594        }
10595        // MEMRA_FULL_PREC bf16-resident weight: dequant-on-use, then the per-column decode-exact
10596        // float linear (same n-independent reduction contract as the Float arm above).
10597        if let GpuTensor::FloatBf16 { data, .. } = w {
10598            let (in_f, out_f) = (w.in_features(), w.out_features());
10599            return self.linear_bf16_chunked(x, data, m, in_f, out_f, true);
10600        }
10601        if !self.uses_q8_1_fast(w) {
10602            return self.matmul(w, x, m);
10603        }
10604        let in_f = w.in_features();
10605        let out_f = w.out_features();
10606        let (bytes, qtype, row_bytes, scale, rp) = match w {
10607            GpuTensor::Quant {
10608                bytes,
10609                qtype,
10610                row_bytes,
10611                scale,
10612                rp,
10613                ..
10614            } => (bytes, *qtype, *row_bytes, *scale, *rp),
10615            _ => return self.matmul(w, x, m),
10616        };
10617        // Q4_0 split-plane mirror for the mmvq/batched arms below (dp4a tail = matmul_pre,
10618        // which does its own mirror pick).
10619        let (bytes, rp) = match w {
10620            GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
10621            _ => (bytes, rp),
10622        };
10623        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10624        // BLOCK-128 e4m3 (QT_F8_E4M3_BLK): the same single kernel every other entry dispatches, so
10625        // the decode-exact contract needs nothing special — grid.y=m runs the m=1 program per
10626        // (token,row) by construction, which is exactly what this method exists to guarantee.
10627        if let Some(y) = self.try_e4m3_blk_pre(w, &aq, &ad, m)? {
10628            return Ok(y);
10629        }
10630        // Batched weight-resident matvec for m=2-8: BIT-IDENTICAL per (token,row) to MMVQ (exact
10631        // integer dp4a, same warp reduce — kernel-check gate rel=0.00e0), one weight read for m
10632        // tokens. The dispatch the divergence fix must avoid is dp4a's 128-thread two-level
10633        // reduce, NOT this. m=5..8 is the K=4..7 spec-verify tier (b8): pre-b8 T=5 fell to the
10634        // grid.y=m per-row MMVQ below = 5 full weight reads/launch — the measured 27B K=4 cliff.
10635        // DECODE-PARITY GATE (2026-07-07): batched (MMVQ-class order) only when the m=1 decode
10636        // chain rides MMVQ too — without MEMRA_MMVQ decode is dp4a, so the exact-contract here
10637        // must be per-column dp4a (matmul_pre fallthrough), not the MMVQ order.
10638        if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
10639            && std::env::var("MEMRA_NO_BATCHED").is_err()
10640            && (m <= 4 || Self::b8_enabled())
10641            // Every b16 class has base + _rp twins after lane/rp-on-st (see matmul_pre's note):
10642            // no mirror precondition, `rp` selects the layout only.
10643            && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || qtype == QT_F8_E4M3
10644                || qtype == QT_NVFP4 || qtype == QT_Q4_K || qtype == QT_Q5_K || qtype == QT_Q8_0)
10645        {
10646            let mcols = Self::batched_mcols(m);
10647            return self.qmatvec_mmvq_batched(
10648                bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, rp,
10649            );
10650        }
10651        if self.mmvq_supports(qtype) {
10652            // MMVQ at grid.y=m: each row is processed by its own warp independently — same 32-thread
10653            // accumulation + warp_reduce_sum as m=1 decode. Bit-identical per row.
10654            return self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, scale, rp);
10655        }
10656        // Fallback for non-MMVQ quant types (Q5_K, Q3_K): use dp4a (the only available kernel).
10657        // These types are not used in the 27B's linear-attn NVFP4+Q4_K layers.
10658        self.matmul_pre(w, &aq, &ad, x, m)
10659    }
10660
10661    /// DECODE-EXACT matmul from a PRE-QUANTIZED q8_1 activation (batched-verify epilogue
10662    /// re-fuse, lane/vt-fixes fix 2, 2026-08-03): the EXACT `matmul_decode_exact` dispatch for
10663    /// q8_1-fast Quant tensors, with the caller's (aq, ad) replacing the internal
10664    /// `quantize_q8_1`. quantize_q8_1 is deterministic (same input bytes -> same q8 bytes), so
10665    /// sharing one quantize across sibling matmuls of the same activation — or consuming the
10666    /// q8 emitted by a fused epilogue (rms_norm_q8_1 / add_rms_norm_q8_1 /
10667    /// silu_mul_scaled_q8_1 / gated_rmsnorm_q8_1, all kernel-check-pinned bit-identical to
10668    /// their unfused chains) — cannot change any dispatched kernel's input bytes.
10669    /// Caller MUST guarantee `uses_q8_1_fast(w)` (the fused epilogues only exist on that path).
10670    pub fn matmul_decode_exact_pre(
10671        &self,
10672        w: &crate::model::GpuTensor,
10673        aq: &CudaSlice<i8>,
10674        ad: &CudaSlice<f32>,
10675        m: usize,
10676    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10677        use crate::model::GpuTensor;
10678        debug_assert!(
10679            self.uses_q8_1_fast(w),
10680            "matmul_decode_exact_pre: caller must guarantee q8_1-fast"
10681        );
10682        // BLOCK-128 e4m3: same single kernel, all m — see matmul_decode_exact's note.
10683        if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
10684            return Ok(y);
10685        }
10686        let in_f = w.in_features();
10687        let out_f = w.out_features();
10688        let (bytes, qtype, row_bytes, scale, rp) = match w {
10689            GpuTensor::Quant {
10690                bytes,
10691                qtype,
10692                row_bytes,
10693                scale,
10694                rp,
10695                ..
10696            } => (bytes, *qtype, *row_bytes, *scale, *rp),
10697            _ => {
10698                return Err(
10699                    "matmul_decode_exact_pre: Quant tensor required (q8_1-fast contract)".into(),
10700                );
10701            }
10702        };
10703        // Q4_0 split-plane mirror — same pick as matmul_decode_exact.
10704        let (bytes, rp) = match w {
10705            GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
10706            _ => (bytes, rp),
10707        };
10708        // Dispatch mirror of matmul_decode_exact's q8_1-fast tail, condition for condition.
10709        if (2..=16).contains(&m)
10710            && self.batched_supports(qtype)
10711            && self.mmvq_supports(qtype)
10712            && std::env::var("MEMRA_NO_BATCHED").is_err()
10713            && (m <= 4 || Self::b8_enabled())
10714            && (m <= 8
10715                || qtype == QT_Q4_0
10716                || qtype == QT_Q6_K
10717                || qtype == QT_F8_E4M3
10718                || qtype == QT_NVFP4
10719                || qtype == QT_Q4_K
10720                || qtype == QT_Q5_K
10721                || qtype == QT_Q8_0)
10722        {
10723            let mcols = Self::batched_mcols(m);
10724            return self.qmatvec_mmvq_batched(
10725                bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, rp,
10726            );
10727        }
10728        if self.mmvq_supports(qtype) {
10729            return self.qmatvec_mmvq(bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, rp);
10730        }
10731        // Non-MMVQ quant types (Q5_K/Q3_K under MEMRA_MMVQ=0): dp4a via matmul_pre — the same
10732        // fallback matmul_decode_exact takes. m <= 16 on the verify tier never reads x_fallback.
10733        let x0 = self.zeros(0)?;
10734        self.matmul_pre(w, aq, ad, &x0, m)
10735    }
10736
10737    /// DUAL gate+up batched matvec from a PRE-QUANTIZED activation, macro-scales DEFERRED
10738    /// (lane/vt-fixes fix 2): same eligibility as `matmul_decode_exact_dual`, but the caller's
10739    /// (aq, ad) replaces the internal quantize and the NVFP4 per-tensor scales are RETURNED
10740    /// instead of applied via two `scale_inplace` launches — the fused SwiGLU epilogue
10741    /// (`silu_mul_scaled_q8_1`) folds them, exactly like the m=1 decode chain does. Deferring
10742    /// is value-exact: `y[i]*s` inline in the epilogue is the same IEEE multiply scale_inplace
10743    /// would store (f32 store/load round-trips are exact). None -> caller falls back to the
10744    /// per-tensor path.
10745    pub fn matmul_decode_exact_dual_pre(
10746        &self,
10747        w0: &crate::model::GpuTensor,
10748        w1: &crate::model::GpuTensor,
10749        aq: &CudaSlice<i8>,
10750        ad: &CudaSlice<f32>,
10751        m: usize,
10752    ) -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn std::error::Error>>
10753    {
10754        use crate::model::GpuTensor;
10755        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
10756        let on = *ON.get_or_init(|| {
10757            std::env::var("MEMRA_SPEC_DUAL_T")
10758                .map(|v| v != "0")
10759                .unwrap_or(true)
10760        });
10761        if !on
10762            || !(2..=7).contains(&m)
10763            || std::env::var("MEMRA_NO_BATCHED").is_ok()
10764            || !self.uses_q8_1_fast(w0)
10765            || !self.uses_q8_1_fast(w1)
10766        {
10767            return Ok(None);
10768        }
10769        // DECODE-PARITY GATE (lane/nvfp4-strict, 2026-08-05): batched iff MMVQ — the dual
10770        // kernels are the MMVQ warp-reduce family, and without MEMRA_MMVQ the m=1 decode
10771        // chain this verify must match bit-for-bit rides dp4a (see matmul_decode_exact's
10772        // note). The singles enforce this via `mmvq_supports`; the dual door skipped it.
10773        if !self.mmvq_supports(QT_NVFP4) {
10774            return Ok(None);
10775        }
10776        let (in_f, out_f) = (w0.in_features(), w0.out_features());
10777        if w1.in_features() != in_f || w1.out_features() != out_f {
10778            return Ok(None);
10779        }
10780        let (b0, b1, row_bytes, s0, s1, rp) = match (w0, w1) {
10781            (
10782                GpuTensor::Quant {
10783                    bytes: b0,
10784                    qtype: q0,
10785                    row_bytes: rb0,
10786                    scale: s0,
10787                    rp: rp0,
10788                    rp4: None,
10789                    ..
10790                },
10791                GpuTensor::Quant {
10792                    bytes: b1,
10793                    qtype: q1,
10794                    row_bytes: rb1,
10795                    scale: s1,
10796                    rp: rp1,
10797                    rp4: None,
10798                    ..
10799                },
10800            ) if *q0 == QT_NVFP4 && *q1 == QT_NVFP4 && rb0 == rb1 && rp0 == rp1 => {
10801                (b0, b1, *rb0, *s0, *s1, *rp0)
10802            }
10803            _ => return Ok(None),
10804        };
10805        // m=5..7: only the exact-width rp duals exist (vt-fixes fix 1b); GGUF layout keeps
10806        // the singles. The b8 dual (MCOLS=8 at m=5..8) measured FLAT and stays dead.
10807        if m > 4 && !(rp && Self::b8_enabled() && std::env::var("MEMRA_B567").as_deref() != Ok("0"))
10808        {
10809            return Ok(None);
10810        }
10811        let (y0, y1) =
10812            self.qmatvec_batched_dual_raw(b0, b1, aq, ad, m, in_f, out_f, row_bytes, rp)?;
10813        Ok(Some(((y0, s0), (y1, s1))))
10814    }
10815
10816    /// DUAL gate+up BATCHED matvec at verify t=2..8 (lane/verify-economics, 2026-08-02): ONE
10817    /// launch computes both FFN projections of a verify batch — same activation, same shape,
10818    /// blockIdx.y selects the tensor. Per (tensor, token, row) the kernel body is the single
10819    /// batched program on the SAME layout (split-plane rp: b2 rp / b4 rpr2 / b8 rpr2; GGUF:
10820    /// b2 base / b4 r2 / b8 r2) -> BIT-IDENTICAL to the two single `matmul_decode_exact`
10821    /// launches (kernel-check gates bitwise on both layouts; run-spec K=1..8 arbitrates e2e).
10822    /// The one activation quantize replaces two IDENTICAL quantizes of the same `x` (same
10823    /// kernel, same input -> same q8_1 bytes), and the two independent weight streams in one
10824    /// grid restore the memory-level parallelism the two-launch form loses to tail drain +
10825    /// launch gap (m=1 dual_mr2 precedent: DRAM 40% -> 47-50% on the 27B pair).
10826    /// `Some((y0, y1))` only when both tensors are NVFP4, the SAME layout (both rp or both
10827    /// GGUF, no rp4 mirror), identical (in_f, out_f, row_bytes), q8_1-fast, and m in 2..=4
10828    /// (the b2/b4 tiers = verify T for K=1..3, the profitable-K window — the b8 dual measured
10829    /// FLAT vs the rpsc singles x3 interleaved, research/verify-economics-20260802, and was
10830    /// killed per doctrine). None -> caller runs the two singles. MEMRA_SPEC_DUAL_T=0 rollback.
10831    pub fn matmul_decode_exact_dual(
10832        &self,
10833        w0: &crate::model::GpuTensor,
10834        w1: &crate::model::GpuTensor,
10835        x: &CudaSlice<f32>,
10836        m: usize,
10837    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
10838        use crate::model::GpuTensor;
10839        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
10840        let on = *ON.get_or_init(|| {
10841            std::env::var("MEMRA_SPEC_DUAL_T")
10842                .map(|v| v != "0")
10843                .unwrap_or(true)
10844        });
10845        if !on
10846            || !(2..=4).contains(&m)
10847            || std::env::var("MEMRA_NO_BATCHED").is_ok()
10848            || !self.uses_q8_1_fast(w0)
10849            || !self.uses_q8_1_fast(w1)
10850        {
10851            return Ok(None);
10852        }
10853        // DECODE-PARITY GATE (lane/nvfp4-strict, 2026-08-05): batched iff MMVQ — same law as
10854        // the singles' `batched_supports && mmvq_supports` check in matmul_decode_exact,
10855        // which this dual door bypassed. Without MEMRA_MMVQ the m=1 decode is dp4a; the
10856        // verify must ride the per-column dp4a class, not the MMVQ-family dual.
10857        if !self.mmvq_supports(QT_NVFP4) {
10858            return Ok(None);
10859        }
10860        let (in_f, out_f) = (w0.in_features(), w0.out_features());
10861        if w1.in_features() != in_f || w1.out_features() != out_f {
10862            return Ok(None);
10863        }
10864        let (b0, b1, row_bytes, s0, s1, rp) = match (w0, w1) {
10865            (
10866                GpuTensor::Quant {
10867                    bytes: b0,
10868                    qtype: q0,
10869                    row_bytes: rb0,
10870                    scale: s0,
10871                    rp: rp0,
10872                    rp4: None,
10873                    ..
10874                },
10875                GpuTensor::Quant {
10876                    bytes: b1,
10877                    qtype: q1,
10878                    row_bytes: rb1,
10879                    scale: s1,
10880                    rp: rp1,
10881                    rp4: None,
10882                    ..
10883                },
10884            ) if *q0 == QT_NVFP4 && *q1 == QT_NVFP4 && rb0 == rb1 && rp0 == rp1 => {
10885                (b0, b1, *rb0, *s0, *s1, *rp0)
10886            }
10887            _ => return Ok(None),
10888        };
10889        // Engagement receipt (MEMRA_DEBUG=1): the first dead-arm A/B lesson — a `rp: false`
10890        // gate silently no-op'd the whole experiment; prove the arm is live in the log.
10891        if std::env::var("MEMRA_DEBUG").is_ok() {
10892            static ONCE: std::sync::Once = std::sync::Once::new();
10893            ONCE.call_once(|| eprintln!("[memra] dual gate+up batched ENGAGED (m={m} rp={rp})"));
10894        }
10895        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10896        let (y0, y1) =
10897            self.qmatvec_batched_dual_raw(b0, b1, &aq, &ad, m, in_f, out_f, row_bytes, rp)?;
10898        let mut y0 = y0;
10899        let mut y1 = y1;
10900        if s0 != 1.0 {
10901            self.scale_inplace(&mut y0, s0, m * out_f)?;
10902        }
10903        if s1 != 1.0 {
10904            self.scale_inplace(&mut y1, s1, m * out_f)?;
10905        }
10906        Ok(Some((y0, y1)))
10907    }
10908
10909    /// Launch body of the dual batched twins from raw NVFP4 weight bytes + a pre-quantized q8_1
10910    /// activation (kernel-check's bit-equivalence entry; matmul_decode_exact_dual's core).
10911    /// mcols tier = batched_mcols(m); macro-scale NOT applied. `rp` selects the split-plane
10912    /// twins (both buffers must be the repacked layout).
10913    #[allow(clippy::too_many_arguments)]
10914    pub fn qmatvec_batched_dual_raw(
10915        &self,
10916        b0: &CudaSlice<u8>,
10917        b1: &CudaSlice<u8>,
10918        aq: &CudaSlice<i8>,
10919        ad: &CudaSlice<f32>,
10920        m: usize,
10921        in_f: usize,
10922        out_f: usize,
10923        row_bytes: usize,
10924        rp: bool,
10925    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10926        const ROWS_PER_BLOCK: u32 = 4;
10927        let mcols = Self::batched_mcols(m);
10928        // EXACT-WIDTH duals at m=5..7 (vt-fixes fix 1b): rp-only; bit-identical to the two
10929        // b5/b6/b7 singles (blockIdx.y selects the tensor, same template body).
10930        let tiny_rp1 = rp
10931            && mcols == 4
10932            && out_f <= 128
10933            && std::env::var("MEMRA_NVFP4_AUX_DUAL").as_deref() != Ok("0");
10934        let (name, rows_per_block) = if tiny_rp1 {
10935            ("qmatvec_nvfp4_mmvq_dual_b4_rp", ROWS_PER_BLOCK)
10936        } else {
10937            match (mcols, rp, m) {
10938                (2, false, _) => ("qmatvec_nvfp4_mmvq_dual_b2", ROWS_PER_BLOCK),
10939                (4, false, _) => ("qmatvec_nvfp4_mmvq_dual_b4_r2", ROWS_PER_BLOCK * 2),
10940                (2, true, _) => ("qmatvec_nvfp4_mmvq_dual_b2_rp", ROWS_PER_BLOCK),
10941                (4, true, _) => ("qmatvec_nvfp4_mmvq_dual_b4_rpr2", ROWS_PER_BLOCK * 2),
10942                (8, true, 5) => ("qmatvec_nvfp4_mmvq_dual_b5_rpr2", ROWS_PER_BLOCK * 2),
10943                (8, true, 6) => ("qmatvec_nvfp4_mmvq_dual_b6_rpr2", ROWS_PER_BLOCK * 2),
10944                (8, true, 7) => ("qmatvec_nvfp4_mmvq_dual_b7_rpr2", ROWS_PER_BLOCK * 2),
10945                _ => {
10946                    return Err(
10947                        format!("qmatvec_batched_dual_raw: no dual kernel for m {m}").into(),
10948                    );
10949                }
10950            }
10951        };
10952        let f = self.func(name);
10953        let mut y0 = self.alloc_uninit::<f32>(m * out_f)?;
10954        let mut y1 = self.alloc_uninit::<f32>(m * out_f)?;
10955        let cfg = LaunchConfig {
10956            grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 2, 1),
10957            block_dim: (32, ROWS_PER_BLOCK, 1),
10958            shared_mem_bytes: 0,
10959        };
10960        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
10961        let __s_b = self.gpu.stream();
10962        let mut b = __s_b.launch_builder(&f);
10963        b.arg(b0)
10964            .arg(b1)
10965            .arg(aq)
10966            .arg(ad)
10967            .arg(&mut y0)
10968            .arg(&mut y1)
10969            .arg(&inf)
10970            .arg(&outf)
10971            .arg(&mi)
10972            .arg(&rb);
10973        unsafe {
10974            b.launch(cfg)?;
10975        }
10976        Ok((y0, y1))
10977    }
10978
10979    /// Like `matmul_pre` but RETURNS THE RAW (un-macro-scaled) matmul output together with the
10980    /// per-tensor NVFP4 scale, instead of applying `scale_inplace` internally. Used by the fused
10981    /// SwiGLU epilogue (RANK3 LEVER 2) so the gate/up scales fold into one `silu_mul_scaled` launch.
10982    /// `Some((y_raw, scale))` only on the m==1 decode fast path (mmvq / dp4a) where the scale is a
10983    /// separate post-launch op we can defer; returns `None` for every other path (prefill GEMM, FP4
10984    /// GEMM, Stage-A, Float) so the caller falls back to the scaled `matmul_pre` + `silu_mul`.
10985    /// DUAL gate+up NVFP4 matvec (mm-fusion): ONE launch computes both projections (same
10986    /// activation, same shape) — grid.y selects the tensor. Bit-identical per element to two
10987    /// mr2 launches at m=1. Returns (gate_raw, up_raw) un-scaled (caller folds the two macro
10988    /// scales into the SwiGLU epilogue, same as the matmul_pre_noscale contract). None unless
10989    /// both tensors are NVFP4 q8_1-fast with identical (in_f, out_f, row_bytes) and m==1.
10990    pub fn matmul_pre_dual_noscale(
10991        &self,
10992        w0: &crate::model::GpuTensor,
10993        w1: &crate::model::GpuTensor,
10994        aq: &CudaSlice<i8>,
10995        ad: &CudaSlice<f32>,
10996        m: usize,
10997    ) -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn std::error::Error>>
10998    {
10999        use crate::model::GpuTensor;
11000        if m != 1 || !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) {
11001            return Ok(None);
11002        }
11003        // FP-ORDER LAW (lane/nvfp4-strict, 2026-08-05): every kernel this door can dispatch
11004        // (q8_0 fused2, nvfp4 dual_mr2) is the MMVQ family — 32-thread warp reduce. Without
11005        // MEMRA_MMVQ the m=1 singles ride dp4a (128-thread two-level reduce), so fusing here
11006        // would mix dispatch families across the pair — the exact class `q8_fused_params`
11007        // already refuses for Q8_0. The NVFP4 arm lacked this check, which is why
11008        // decode-batch-gate `--mode strict`'s equalizing env (MEMRA_MMVQ=0) never pinned
11009        // NVFP4 models: decode_step_h kept riding dual_mr2 while the batched body fell to
11010        // dp4a (gate1 maxdiff 1.639e-1 / gate2 step-8 divergence at the 2026-08-05 train
11011        // HEAD, research/nvfp4-strict-20260805/). Default env (MMVQ on) is dispatch-unchanged.
11012        if !self.mmvq_supports(QT_NVFP4) {
11013            return Ok(None);
11014        }
11015        let (in_f, out_f) = (w0.in_features(), w0.out_features());
11016        if w1.in_features() != in_f || w1.out_features() != out_f {
11017            return Ok(None);
11018        }
11019        // Q8_0 ARM (lane/q27-deepdive, 2026-08-05): the dense-FFN gate+up pair on a Q8_0 trunk fell
11020        // through this NVFP4-only gate to two `matmul_pre_noscale` launches — measured 128 of the
11021        // 1015 launches/token on q27-Q8_0 decode, the single largest un-fused class in the tick
11022        // (nsys `research/q27-deepdive-20260805/nsys/`). `q8_fused2_core` already serves the same
11023        // pair shape for the shared-expert gate/up, and its kernel body is `qmatvec_q8_0_mmvq`
11024        // VERBATIM per (tensor,row) -> BIT-IDENTICAL to the two separate launches. Q8_0 carries no
11025        // macro-scale (q8_fused_params requires scale==1.0), so the noscale contract is satisfied
11026        // by returning 1.0 for both: the SwiGLU epilogue's fold becomes the identity it already is
11027        // on this dtype today. Seam: MEMRA_Q8_FFN_FUSE2=0 rolls back to the two-launch pair.
11028        // rp4 guard: with MEMRA_Q8RP the singles route to the `_rp` split-plane twin over the
11029        // mirror buffer; the fused2 kernel has no `_rp` form, so fusing there would swap
11030        // dispatch families mid-model. Bail and let the two singles run (mirror lane unchanged).
11031        let no_mirror =
11032            |w: &crate::model::GpuTensor| !matches!(w, GpuTensor::Quant { rp4: Some(_), .. });
11033        if self.q8_ffn_fuse2_on()
11034            && no_mirror(w0)
11035            && no_mirror(w1)
11036            && let Some([p0, p1]) = self.q8_fused_params(&[w0, w1])
11037        {
11038            let (y0, y1) = self.q8_fused2_core(p0.0, p1.0, aq, ad, in_f, p0.1, p1.1, p0.2)?;
11039            return Ok(Some(((y0, 1.0), (y1, 1.0))));
11040        }
11041        // F8-E4M3 ARM (lane/fp8-decode-v1, 2026-08-05): with native e4m3 residency the FFN gate+up
11042        // pair (and the ssm beta+alpha dual, which routes through this same entry) fell through
11043        // both the NVFP4 gate below and the Q8_0 arm above to two `matmul_pre_noscale` launches —
11044        // native residency was UN-FUSING the trunk relative to the Q8_0 slab it replaces. The
11045        // fused2 kernel body is `qmatvec_e4m3_mmvq` VERBATIM per (tensor,row). Contract match:
11046        // `matmul_pre_noscale` on e4m3 launches with scale 1.0 and RETURNS the per-tensor
11047        // weight_scale for the caller to fold, so we pass ws=1.0 here and return (s0,s1) — same
11048        // bits, and the two macro-scale multiplies still fold into the SwiGLU epilogue.
11049        // MEMRA_E4M3_DUAL=0 rolls back to the two-launch pair.
11050        if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
11051            let (y0, y1) =
11052                self.e4m3_fused2_core(p0.0, p1.0, aq, ad, in_f, p0.1, p1.1, p0.2, 1.0, 1.0)?;
11053            return Ok(Some(((y0, p0.3), (y1, p1.3))));
11054        }
11055        let (b0, q0, rb0, s0, rp0) = match w0 {
11056            GpuTensor::Quant {
11057                bytes,
11058                qtype,
11059                row_bytes,
11060                scale,
11061                rp,
11062                ..
11063            } => (bytes, *qtype, *row_bytes, *scale, *rp),
11064            _ => return Ok(None),
11065        };
11066        let (b1, q1, rb1, s1, rp1) = match w1 {
11067            GpuTensor::Quant {
11068                bytes,
11069                qtype,
11070                row_bytes,
11071                scale,
11072                rp,
11073                ..
11074            } => (bytes, *qtype, *row_bytes, *scale, *rp),
11075            _ => return Ok(None),
11076        };
11077        if q0 != QT_NVFP4 || q1 != QT_NVFP4 || rb0 != rb1 || rp0 != rp1 {
11078            return Ok(None);
11079        }
11080        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
11081        const RPW: u32 = 2;
11082        let rows_per_block = ROWS_PER_BLOCK * RPW;
11083        let f = self.func(if rp0 {
11084            "qmatvec_nvfp4_mmvq_dual_mr2_rp"
11085        } else {
11086            "qmatvec_nvfp4_mmvq_dual_mr2"
11087        });
11088        let mut y0 = self.alloc_uninit::<f32>(out_f)?;
11089        let mut y1 = self.alloc_uninit::<f32>(out_f)?;
11090        let cfg = LaunchConfig {
11091            grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 2, 1),
11092            block_dim: (32, ROWS_PER_BLOCK, 1),
11093            shared_mem_bytes: 0,
11094        };
11095        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, 1i32, rb0 as i64);
11096        // noscale contract: the caller folds s0/s1 into the SwiGLU epilogue — the kernel's fused
11097        // yscale args stay 1.0 here (they exist for the single-tensor callers).
11098        let one = 1.0f32;
11099        let __s_b = self.gpu.stream();
11100        let mut b = __s_b.launch_builder(&f);
11101        b.arg(b0)
11102            .arg(b1)
11103            .arg(aq)
11104            .arg(ad)
11105            .arg(&mut y0)
11106            .arg(&mut y1)
11107            .arg(&inf)
11108            .arg(&outf)
11109            .arg(&mi)
11110            .arg(&rb)
11111            .arg(&one)
11112            .arg(&one);
11113        unsafe {
11114            b.launch(cfg)?;
11115        }
11116        Ok(Some(((y0, s0), (y1, s1))))
11117    }
11118
11119    /// FUSED NVFP4 matvec TRIPLE with unequal out_f (rig-native decode increment 1,
11120    /// lane/rig-native-nvfp4): wq+wk+wv in ONE launch via the q8_0 fused2 block-offset
11121    /// recipe. Per (tensor,row,t) the kernel body is nvfp4_mmvq_multirow_rp VERBATIM ->
11122    /// bit-identical to three separate `matmul_pre` launches; yscales fold in-kernel exactly
11123    /// as the singles do. grid.y = m (the t-parallel verify rows ride the same launch).
11124    /// None when ineligible (not all rp NVFP4 / in_f mismatch / mmvq off) — callers fall
11125    /// back to the three singles.
11126    #[allow(clippy::too_many_arguments)]
11127    pub fn matmul_nvfp4_fused3(
11128        &self,
11129        w0: &crate::model::GpuTensor,
11130        w1: &crate::model::GpuTensor,
11131        w2: &crate::model::GpuTensor,
11132        aq: &CudaSlice<i8>,
11133        ad: &CudaSlice<f32>,
11134        m: usize,
11135    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
11136    {
11137        use crate::model::GpuTensor;
11138        // m==1 ONLY: at m>1 the singles ride the _b16 weight-once column program (one weight
11139        // read serves all m rows); the fused segments would re-read the weight per row. The
11140        // fusion win is the B=1 decode tick.
11141        if m != 1
11142            || !self.mmvq_supports(QT_NVFP4)
11143            || !self.uses_q8_1_fast(w0)
11144            || !self.uses_q8_1_fast(w1)
11145            || !self.uses_q8_1_fast(w2)
11146        {
11147            return Ok(None);
11148        }
11149        let unpack = |w: &crate::model::GpuTensor| match w {
11150            GpuTensor::Quant {
11151                bytes,
11152                qtype,
11153                scale,
11154                rp,
11155                ..
11156            } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
11157            _ => None,
11158        };
11159        let (Some(p0), Some(p1), Some(p2)) = (unpack(w0), unpack(w1), unpack(w2)) else {
11160            return Ok(None);
11161        };
11162        let in_f = w0.in_features();
11163        if w1.in_features() != in_f || w2.in_features() != in_f {
11164            return Ok(None);
11165        }
11166        let (o0, o1, o2) = (w0.out_features(), w1.out_features(), w2.out_features());
11167        const ROWS_PER_BLOCK: u32 = 4; // MEMRA_MMVQ_ROWS
11168        const RPW: u32 = 2;
11169        let rows_pb = ROWS_PER_BLOCK * RPW;
11170        let nb = |o: usize| (o as u32).div_ceil(rows_pb);
11171        let f = self.func("qmatvec_nvfp4_mmvq_fused3_rp");
11172        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
11173        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
11174        let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
11175        let cfg = LaunchConfig {
11176            grid_dim: (nb(o0) + nb(o1) + nb(o2), m as u32, 1),
11177            block_dim: (32, ROWS_PER_BLOCK, 1),
11178            shared_mem_bytes: 0,
11179        };
11180        let (inf, oi0, oi1, oi2, mi) = (in_f as i32, o0 as i32, o1 as i32, o2 as i32, m as i32);
11181        // SAFETY: the raw pointers come straight from the &GpuTensor borrows above and are
11182        // only dereferenced for the launch-arg build inside this call.
11183        let (b0, b1, b2) = unsafe { (&*p0.0, &*p1.0, &*p2.0) };
11184        let __s_b = self.gpu.stream();
11185        let mut b = __s_b.launch_builder(&f);
11186        b.arg(b0)
11187            .arg(b1)
11188            .arg(b2)
11189            .arg(aq)
11190            .arg(ad)
11191            .arg(&mut y0)
11192            .arg(&mut y1)
11193            .arg(&mut y2)
11194            .arg(&inf)
11195            .arg(&oi0)
11196            .arg(&oi1)
11197            .arg(&oi2)
11198            .arg(&mi)
11199            .arg(&p0.1)
11200            .arg(&p1.1)
11201            .arg(&p2.1);
11202        unsafe {
11203            b.launch(cfg)?;
11204        }
11205        Ok(Some((y0, y1, y2)))
11206    }
11207
11208    /// fused2 twin of `matmul_nvfp4_fused3`, for MIXED-type weight groups: the gemma4
11209    /// dense NVFP4mix recipe keeps `attn_v` and `ffn_down` at Q8_0 (full-NVFP4 was
11210    /// measured garbage on this dense class), so its q/k/v trio and gate/up/down never
11211    /// satisfy an all-NVFP4 fused3 — the pairs that ARE uniformly NVFP4 (q,k and
11212    /// gate,up) fuse here instead. m==1 only, same law as fused3/fused4: per
11213    /// (tensor,row) the kernel seg body is VERBATIM, so the fusion is bit-identical to
11214    /// two separate launches. `MEMRA_NVFP4_FUSED2=0` is the rollback seam and the
11215    /// same-binary interleaved A/B arm.
11216    pub fn matmul_nvfp4_fused2(
11217        &self,
11218        w0: &crate::model::GpuTensor,
11219        w1: &crate::model::GpuTensor,
11220        aq: &CudaSlice<i8>,
11221        ad: &CudaSlice<f32>,
11222        m: usize,
11223    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
11224        use crate::model::GpuTensor;
11225        static FUSED2_OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11226        let off =
11227            *FUSED2_OFF.get_or_init(|| std::env::var("MEMRA_NVFP4_FUSED2").as_deref() == Ok("0"));
11228        // m==1 ONLY: at m>1 the singles ride the _b16 weight-once column program (one weight
11229        // read serves all m rows); the fused segments would re-read the weight per row.
11230        if off
11231            || m != 1
11232            || !self.mmvq_supports(QT_NVFP4)
11233            || !self.uses_q8_1_fast(w0)
11234            || !self.uses_q8_1_fast(w1)
11235        {
11236            return Ok(None);
11237        }
11238        let unpack = |w: &crate::model::GpuTensor| match w {
11239            GpuTensor::Quant {
11240                bytes,
11241                qtype,
11242                scale,
11243                rp,
11244                ..
11245            } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
11246            _ => None,
11247        };
11248        let (Some(p0), Some(p1)) = (unpack(w0), unpack(w1)) else {
11249            return Ok(None);
11250        };
11251        let in_f = w0.in_features();
11252        if w1.in_features() != in_f {
11253            return Ok(None);
11254        }
11255        let (o0, o1) = (w0.out_features(), w1.out_features());
11256        const ROWS_PER_BLOCK: u32 = 4; // MEMRA_MMVQ_ROWS
11257        const RPW: u32 = 2;
11258        let rows_pb = ROWS_PER_BLOCK * RPW;
11259        let nb = |o: usize| (o as u32).div_ceil(rows_pb);
11260        let f = self.func("qmatvec_nvfp4_mmvq_fused2_rp");
11261        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
11262        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
11263        let cfg = LaunchConfig {
11264            grid_dim: (nb(o0) + nb(o1), m as u32, 1),
11265            block_dim: (32, ROWS_PER_BLOCK, 1),
11266            shared_mem_bytes: 0,
11267        };
11268        let (inf, oi0, oi1, mi) = (in_f as i32, o0 as i32, o1 as i32, m as i32);
11269        // SAFETY: the raw pointers come straight from the &GpuTensor borrows above and are
11270        // only dereferenced for the launch-arg build inside this call.
11271        let (b0, b1) = unsafe { (&*p0.0, &*p1.0) };
11272        // PDL wave-B: the kernel carries MEMRA_PDL_ENTRY — launch overlaps the producer's
11273        // drain (quantize/norm emitting aq/ad). Same math, same order, bit-identical.
11274        if Self::pdl_on() && Self::pdl_mmvq_on() && Self::pdl_nvfp4q8_on() {
11275            {
11276                use cudarc::driver::{DevicePtr, DevicePtrMut};
11277                let s = &self.gpu.stream();
11278                let (pw0, _g0) = b0.device_ptr(s);
11279                let (pw1, _g1) = b1.device_ptr(s);
11280                let (paq, _g2) = aq.device_ptr(s);
11281                let (pad, _g3) = ad.device_ptr(s);
11282                let (py0, _g4) = y0.device_ptr_mut(s);
11283                let (py1, _g5) = y1.device_ptr_mut(s);
11284                let (s0, s1) = (p0.1, p1.1);
11285                let mut ps = [
11286                    &pw0 as *const _ as *mut std::ffi::c_void,
11287                    &pw1 as *const _ as *mut _,
11288                    &paq as *const _ as *mut _,
11289                    &pad as *const _ as *mut _,
11290                    &py0 as *const _ as *mut _,
11291                    &py1 as *const _ as *mut _,
11292                    &inf as *const _ as *mut _,
11293                    &oi0 as *const _ as *mut _,
11294                    &oi1 as *const _ as *mut _,
11295                    &mi as *const _ as *mut _,
11296                    &s0 as *const _ as *mut _,
11297                    &s1 as *const _ as *mut _,
11298                ];
11299                unsafe {
11300                    self.launch_pdl(
11301                        "qmatvec_nvfp4_mmvq_fused2_rp",
11302                        cfg.grid_dim,
11303                        cfg.block_dim,
11304                        &mut ps,
11305                    )?;
11306                }
11307            }
11308            return Ok(Some((y0, y1)));
11309        }
11310        let __s_b = self.gpu.stream();
11311        let mut b = __s_b.launch_builder(&f);
11312        b.arg(b0)
11313            .arg(b1)
11314            .arg(aq)
11315            .arg(ad)
11316            .arg(&mut y0)
11317            .arg(&mut y1)
11318            .arg(&inf)
11319            .arg(&oi0)
11320            .arg(&oi1)
11321            .arg(&mi)
11322            .arg(&p0.1)
11323            .arg(&p1.1);
11324        unsafe {
11325            b.launch(cfg)?;
11326        }
11327        Ok(Some((y0, y1)))
11328    }
11329
11330    /// Slot-fed fused2 twin (alloc-free capture lane): identical launch to
11331    /// `matmul_nvfp4_fused2`, caller-owned outputs — the gemma4 dc_slotted graph body
11332    /// needs zero mem nodes, so the allocating wrapper can't serve it. Returns false
11333    /// (decline) on any non-NVFP4/rp pair; the caller chains or refuses.
11334    pub fn matmul_nvfp4_fused2_into(
11335        &self,
11336        w0: &crate::model::GpuTensor,
11337        w1: &crate::model::GpuTensor,
11338        aq: &CudaSlice<i8>,
11339        ad: &CudaSlice<f32>,
11340        y0: &mut CudaSlice<f32>,
11341        y1: &mut CudaSlice<f32>,
11342    ) -> Result<bool, Box<dyn std::error::Error>> {
11343        use crate::model::GpuTensor;
11344        static FUSED2_OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11345        let off =
11346            *FUSED2_OFF.get_or_init(|| std::env::var("MEMRA_NVFP4_FUSED2").as_deref() == Ok("0"));
11347        if off
11348            || !self.mmvq_supports(QT_NVFP4)
11349            || !self.uses_q8_1_fast(w0)
11350            || !self.uses_q8_1_fast(w1)
11351        {
11352            return Ok(false);
11353        }
11354        let unpack = |w: &crate::model::GpuTensor| match w {
11355            GpuTensor::Quant {
11356                bytes,
11357                qtype,
11358                scale,
11359                rp,
11360                ..
11361            } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
11362            _ => None,
11363        };
11364        let (Some(p0), Some(p1)) = (unpack(w0), unpack(w1)) else {
11365            return Ok(false);
11366        };
11367        let in_f = w0.in_features();
11368        if w1.in_features() != in_f {
11369            return Ok(false);
11370        }
11371        let (o0, o1) = (w0.out_features(), w1.out_features());
11372        if y0.len() < o0 || y1.len() < o1 {
11373            return Ok(false);
11374        }
11375        const ROWS_PER_BLOCK: u32 = 4; // MEMRA_MMVQ_ROWS
11376        const RPW: u32 = 2;
11377        let rows_pb = ROWS_PER_BLOCK * RPW;
11378        let nb = |o: usize| (o as u32).div_ceil(rows_pb);
11379        let f = self.func("qmatvec_nvfp4_mmvq_fused2_rp");
11380        let cfg = LaunchConfig {
11381            grid_dim: (nb(o0) + nb(o1), 1, 1),
11382            block_dim: (32, ROWS_PER_BLOCK, 1),
11383            shared_mem_bytes: 0,
11384        };
11385        let (inf, oi0, oi1, mi) = (in_f as i32, o0 as i32, o1 as i32, 1i32);
11386        // SAFETY: the raw pointers come straight from the &GpuTensor borrows above and are
11387        // only dereferenced for the launch-arg build inside this call.
11388        let (b0, b1) = unsafe { (&*p0.0, &*p1.0) };
11389        let __s_b = self.gpu.stream();
11390        let mut b = __s_b.launch_builder(&f);
11391        b.arg(b0)
11392            .arg(b1)
11393            .arg(aq)
11394            .arg(ad)
11395            .arg(&mut *y0)
11396            .arg(&mut *y1)
11397            .arg(&inf)
11398            .arg(&oi0)
11399            .arg(&oi1)
11400            .arg(&mi)
11401            .arg(&p0.1)
11402            .arg(&p1.1);
11403        unsafe {
11404            b.launch(cfg)?;
11405        }
11406        Ok(true)
11407    }
11408
11409    /// fused4 twin of `matmul_nvfp4_fused3`: the Linear-mixer projection quartet
11410    /// (wqkv + wqkv_gate + ssm_beta + ssm_alpha) in one launch, m==1 only. Per
11411    /// (tensor,row) the kernel body is nvfp4_mmvq_multirow_rp VERBATIM — bit-identical
11412    /// to four separate launches (rig-native decode increment 2, RIG-NATIVE-DECODE.md).
11413    #[allow(clippy::type_complexity)]
11414    pub fn matmul_nvfp4_fused4(
11415        &self,
11416        w0: &crate::model::GpuTensor,
11417        w1: &crate::model::GpuTensor,
11418        w2: &crate::model::GpuTensor,
11419        w3: &crate::model::GpuTensor,
11420        aq: &CudaSlice<i8>,
11421        ad: &CudaSlice<f32>,
11422        m: usize,
11423    ) -> Result<
11424        Option<(
11425            CudaSlice<f32>,
11426            CudaSlice<f32>,
11427            CudaSlice<f32>,
11428            CudaSlice<f32>,
11429        )>,
11430        Box<dyn std::error::Error>,
11431    > {
11432        use crate::model::GpuTensor;
11433        // MEMRA_NVFP4_FUSED4=0: rollback seam + the same-binary interleaved A/B arm.
11434        if m != 1
11435            || std::env::var("MEMRA_NVFP4_FUSED4").as_deref() == Ok("0")
11436            || !self.mmvq_supports(QT_NVFP4)
11437            || !self.uses_q8_1_fast(w0)
11438            || !self.uses_q8_1_fast(w1)
11439            || !self.uses_q8_1_fast(w2)
11440            || !self.uses_q8_1_fast(w3)
11441        {
11442            return Ok(None);
11443        }
11444        let unpack = |w: &crate::model::GpuTensor| match w {
11445            GpuTensor::Quant {
11446                bytes,
11447                qtype,
11448                scale,
11449                rp,
11450                ..
11451            } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
11452            _ => None,
11453        };
11454        let (Some(p0), Some(p1), Some(p2), Some(p3)) =
11455            (unpack(w0), unpack(w1), unpack(w2), unpack(w3))
11456        else {
11457            return Ok(None);
11458        };
11459        let in_f = w0.in_features();
11460        if w1.in_features() != in_f || w2.in_features() != in_f || w3.in_features() != in_f {
11461            return Ok(None);
11462        }
11463        let (o0, o1, o2, o3) = (
11464            w0.out_features(),
11465            w1.out_features(),
11466            w2.out_features(),
11467            w3.out_features(),
11468        );
11469        const ROWS_PER_BLOCK: u32 = 4; // MEMRA_MMVQ_ROWS
11470        const RPW: u32 = 2;
11471        let rows_pb = ROWS_PER_BLOCK * RPW;
11472        let nb = |o: usize| (o as u32).div_ceil(rows_pb);
11473        let f = self.func("qmatvec_nvfp4_mmvq_fused4_rp");
11474        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
11475        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
11476        let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
11477        let mut y3 = self.alloc_uninit::<f32>(m * o3)?;
11478        let cfg = LaunchConfig {
11479            grid_dim: (nb(o0) + nb(o1) + nb(o2) + nb(o3), m as u32, 1),
11480            block_dim: (32, ROWS_PER_BLOCK, 1),
11481            shared_mem_bytes: 0,
11482        };
11483        let (inf, oi0, oi1, oi2, oi3, mi) = (
11484            in_f as i32,
11485            o0 as i32,
11486            o1 as i32,
11487            o2 as i32,
11488            o3 as i32,
11489            m as i32,
11490        );
11491        // SAFETY: the raw pointers come straight from the &GpuTensor borrows above and are
11492        // only dereferenced for the launch-arg build inside this call.
11493        let (b0, b1, b2, b3) = unsafe { (&*p0.0, &*p1.0, &*p2.0, &*p3.0) };
11494        let __s_b = self.gpu.stream();
11495        let mut b = __s_b.launch_builder(&f);
11496        b.arg(b0)
11497            .arg(b1)
11498            .arg(b2)
11499            .arg(b3)
11500            .arg(aq)
11501            .arg(ad)
11502            .arg(&mut y0)
11503            .arg(&mut y1)
11504            .arg(&mut y2)
11505            .arg(&mut y3)
11506            .arg(&inf)
11507            .arg(&oi0)
11508            .arg(&oi1)
11509            .arg(&oi2)
11510            .arg(&oi3)
11511            .arg(&mi)
11512            .arg(&p0.1)
11513            .arg(&p1.1)
11514            .arg(&p2.1)
11515            .arg(&p3.1);
11516        unsafe {
11517            b.launch(cfg)?;
11518        }
11519        Ok(Some((y0, y1, y2, y3)))
11520    }
11521
11522    /// FUSED Q8_0 m=1 matvec PAIR with UNEQUAL out_f (trunk launch-fusion, 2026-07-05). Folds two
11523    /// same-input q8_0 projections (35B trunk: wqkv+wqkv_gate 8192/4096, gate_shexp+up_shexp
11524    /// 512/512) into ONE launch via a block-offset split (blocks [0,nb0) -> w0, rest -> w1) — the
11525    /// dual-mr2 recipe with the same-out_f restriction lifted. Per (tensor,row) the kernel body is
11526    /// qmatvec_q8_0_mmvq VERBATIM -> BIT-IDENTICAL to two separate m=1 launches. Returns None when
11527    /// ineligible (not both Q8_0 / in_f mismatch / MEMRA_MMVQ off / MEMRA_Q8_DUAL=0) — caller falls
11528    /// back to the per-tensor path.
11529    pub fn matmul_q8_fused2(
11530        &self,
11531        w0: &crate::model::GpuTensor,
11532        w1: &crate::model::GpuTensor,
11533        aq: &CudaSlice<i8>,
11534        ad: &CudaSlice<f32>,
11535    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
11536        // e4m3 twin (lane/fp8-decode-v1): this entry is the trunk's generic m=1 pair door
11537        // (wqkv+wqkv_gate, ssm_beta+alpha, gate_shexp+up_shexp), so admitting QT_F8_E4M3 here
11538        // fuses the NATIVE-RESIDENCY FP8 trunk at every existing call site with no call-site
11539        // change. Scale is folded in-kernel per range -> the returned buffers are already scaled,
11540        // exactly like the per-tensor `matmul_pre` e4m3 dispatch this replaces.
11541        if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
11542            return Ok(Some(self.e4m3_fused2_core(
11543                p0.0,
11544                p1.0,
11545                aq,
11546                ad,
11547                w0.in_features(),
11548                p0.1,
11549                p1.1,
11550                p0.2,
11551                p0.3,
11552                p1.3,
11553            )?));
11554        }
11555        let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
11556            return Ok(None);
11557        };
11558        Ok(Some(self.q8_fused2_core(
11559            p0.0,
11560            p1.0,
11561            aq,
11562            ad,
11563            w0.in_features(),
11564            p0.1,
11565            p1.1,
11566            p0.2,
11567        )?))
11568    }
11569
11570    #[allow(clippy::too_many_arguments)]
11571    fn q8_fused2_core(
11572        &self,
11573        b0: &CudaSlice<u8>,
11574        b1: &CudaSlice<u8>,
11575        aq: &CudaSlice<i8>,
11576        ad: &CudaSlice<f32>,
11577        in_f: usize,
11578        out0: usize,
11579        out1: usize,
11580        row_bytes: usize,
11581    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11582        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
11583        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
11584        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
11585        let f = self.func("qmatvec_q8_0_mmvq_fused2");
11586        let mut y0 = self.alloc_uninit::<f32>(out0)?;
11587        let mut y1 = self.alloc_uninit::<f32>(out1)?;
11588        let cfg = LaunchConfig {
11589            grid_dim: (nb0 + nb1, 1, 1),
11590            block_dim: (32, ROWS_PER_BLOCK, 1),
11591            shared_mem_bytes: 0,
11592        };
11593        let (inf, o0, o1, rbl) = (in_f as i32, out0 as i32, out1 as i32, row_bytes as i64);
11594        let __s_b = self.gpu.stream();
11595        let mut b = __s_b.launch_builder(&f);
11596        b.arg(b0)
11597            .arg(b1)
11598            .arg(aq)
11599            .arg(ad)
11600            .arg(&mut y0)
11601            .arg(&mut y1)
11602            .arg(&inf)
11603            .arg(&o0)
11604            .arg(&o1)
11605            .arg(&rbl);
11606        unsafe {
11607            b.launch(cfg)?;
11608        }
11609        Ok((y0, y1))
11610    }
11611
11612    /// f32-activation entry for the fused2 pair: quantizes x to q8_1 ONCE then runs the fused
11613    /// launch — replaces two `matmul(w, x, 1)` calls that would each re-quantize the same x
11614    /// (35B shared-expert gate+up per MoE layer per token). Same bits: quantize_q8_1 is
11615    /// deterministic, the fused body is the MMVQ kernel verbatim. None when ineligible (the
11616    /// callers' m==1-under-MEMRA_FAST dispatch would take MMVQ; anything else falls back).
11617    pub fn matmul_q8_fused2_x(
11618        &self,
11619        w0: &crate::model::GpuTensor,
11620        w1: &crate::model::GpuTensor,
11621        x: &CudaSlice<f32>,
11622    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
11623        if !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) {
11624            return Ok(None);
11625        }
11626        if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
11627            let (aq, ad) = self.quantize_q8_1(x, 1, w0.in_features())?;
11628            return Ok(Some(self.e4m3_fused2_core(
11629                p0.0,
11630                p1.0,
11631                &aq,
11632                &ad,
11633                w0.in_features(),
11634                p0.1,
11635                p1.1,
11636                p0.2,
11637                p0.3,
11638                p1.3,
11639            )?));
11640        }
11641        let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
11642            return Ok(None);
11643        };
11644        let (aq, ad) = self.quantize_q8_1(x, 1, w0.in_features())?;
11645        Ok(Some(self.q8_fused2_core(
11646            p0.0,
11647            p1.0,
11648            &aq,
11649            &ad,
11650            w0.in_features(),
11651            p0.1,
11652            p1.1,
11653            p0.2,
11654        )?))
11655    }
11656
11657    /// Test entry for the kernel_check gate: launch the fused2 kernel from raw weight bytes,
11658    /// quantizing the f32 activation internally (mirrors qmatvec_mmvq_raw; no env gating).
11659    #[allow(clippy::too_many_arguments)]
11660    pub fn qmatvec_q8_fused2_raw(
11661        &self,
11662        b0: &CudaSlice<u8>,
11663        b1: &CudaSlice<u8>,
11664        x: &CudaSlice<f32>,
11665        in_f: usize,
11666        out0: usize,
11667        out1: usize,
11668        row_bytes: usize,
11669    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11670        let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
11671        self.q8_fused2_core(b0, b1, &aq, &ad, in_f, out0, out1, row_bytes)
11672    }
11673
11674    /// FUSED Q8_0 m=1 matvec TRIPLE (wq+wk+wv on the 35B full-attn layers: out_f 8192/512/512).
11675    /// Same block-offset recipe as `matmul_q8_fused2` with three ranges. BIT-IDENTICAL per
11676    /// (tensor,row) to three separate m=1 MMVQ launches.
11677    /// FUSED Q4_0 m=1 TRIPLE (gemma q/k/v — same quantized input; per (tensor,row) chain
11678    /// identical to the mr2 kernel). Returns None unless all three are Q4_0 with equal in_f.
11679    pub fn matmul_q4_fused3(
11680        &self,
11681        w0: &crate::model::GpuTensor,
11682        w1: &crate::model::GpuTensor,
11683        w2: &crate::model::GpuTensor,
11684        aq: &CudaSlice<i8>,
11685        ad: &CudaSlice<f32>,
11686    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
11687    {
11688        use crate::model::GpuTensor;
11689        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
11690            match w {
11691                GpuTensor::Quant {
11692                    qtype, row_bytes, ..
11693                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
11694                _ => None,
11695            }
11696        };
11697        let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (q4(w0), q4(w1), q4(w2)) else {
11698            return Ok(None);
11699        };
11700        if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
11701            return Ok(None);
11702        }
11703        // Effective (bytes, rp) per tensor: mirror (rp4) OR the in-place swap (rp flag,
11704        // bytes already split). Mixed layouts cannot share one fused launch -> fall back to
11705        // the separate matvecs (each routes its own rp).
11706        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
11707            match w {
11708                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
11709                    Some(m) => (m, true),
11710                    None => (bytes, *rp),
11711                },
11712                _ => unreachable!(),
11713            }
11714        }
11715        let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
11716        if rp0 != rp1 || rp1 != rp2 {
11717            return Ok(None);
11718        }
11719        let rp = rp0;
11720        let rpb: u32 = 4;
11721        // mr1 (one row/warp, 2026-07-14): follows the singles' MEMRA_Q40_MR default — the
11722        // fused t=1 kernels were left on mr2 when the singles flipped (DRAM-duty map:
11723        // fused3 57% / fused2 86%; small qkv segments starve under mr2's half grid).
11724        let mr1 = rp && Self::q40_mr1_on();
11725        let nb = |o: usize| {
11726            if mr1 {
11727                (o as u32).div_ceil(rpb)
11728            } else {
11729                (o as u32).div_ceil(2).div_ceil(rpb)
11730            }
11731        };
11732        let grid = nb(o0) + nb(o1) + nb(o2);
11733        let mut y0 = self.alloc_uninit::<f32>(o0)?;
11734        let mut y1 = self.alloc_uninit::<f32>(o1)?;
11735        let mut y2 = self.alloc_uninit::<f32>(o2)?;
11736        let f = self.func(if mr1 {
11737            "qmatvec_q4_0_mmvq_fused3_mr1_rp"
11738        } else if rp {
11739            "qmatvec_q4_0_mmvq_fused3_rp"
11740        } else {
11741            "qmatvec_q4_0_mmvq_fused3"
11742        });
11743        let cfg = LaunchConfig {
11744            grid_dim: (grid, 1, 1),
11745            block_dim: (32, rpb, 1),
11746            shared_mem_bytes: 0,
11747        };
11748        let inf = w0.in_features() as i32;
11749        let (oo0, oo1, oo2) = (o0 as i32, o1 as i32, o2 as i32);
11750        let (r0, r1, r2) = (rb0 as i64, rb1 as i64, rb2 as i64);
11751        // PDL wave-A (2026-07-23): the mr1 kernel carries MEMRA_PDL_ENTRY; only that
11752        // variant may take the programmatic-serialization launch.
11753        if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
11754            {
11755                use cudarc::driver::{DevicePtr, DevicePtrMut};
11756                let s = &self.gpu.stream();
11757                let (p0, _g0) = b0.device_ptr(s);
11758                let (p1, _g1) = b1.device_ptr(s);
11759                let (p2, _g2) = b2.device_ptr(s);
11760                let (paq, _g3) = aq.device_ptr(s);
11761                let (pad, _g4) = ad.device_ptr(s);
11762                let (py0, _g5) = y0.device_ptr_mut(s);
11763                let (py1, _g6) = y1.device_ptr_mut(s);
11764                let (py2, _g7) = y2.device_ptr_mut(s);
11765                let mut ps = [
11766                    &p0 as *const _ as *mut std::ffi::c_void,
11767                    &p1 as *const _ as *mut _,
11768                    &p2 as *const _ as *mut _,
11769                    &paq as *const _ as *mut _,
11770                    &pad as *const _ as *mut _,
11771                    &py0 as *const _ as *mut _,
11772                    &py1 as *const _ as *mut _,
11773                    &py2 as *const _ as *mut _,
11774                    &inf as *const _ as *mut _,
11775                    &oo0 as *const _ as *mut _,
11776                    &oo1 as *const _ as *mut _,
11777                    &oo2 as *const _ as *mut _,
11778                    &r0 as *const _ as *mut _,
11779                    &r1 as *const _ as *mut _,
11780                    &r2 as *const _ as *mut _,
11781                ];
11782                unsafe {
11783                    self.launch_pdl(
11784                        "qmatvec_q4_0_mmvq_fused3_mr1_rp",
11785                        (grid, 1, 1),
11786                        (32, rpb, 1),
11787                        &mut ps,
11788                    )?;
11789                }
11790            }
11791            return Ok(Some((y0, y1, y2)));
11792        }
11793        let __s_b = self.gpu.stream();
11794        let mut b = __s_b.launch_builder(&f);
11795        b.arg(b0)
11796            .arg(b1)
11797            .arg(b2)
11798            .arg(aq)
11799            .arg(ad)
11800            .arg(&mut y0)
11801            .arg(&mut y1)
11802            .arg(&mut y2)
11803            .arg(&inf)
11804            .arg(&oo0)
11805            .arg(&oo1)
11806            .arg(&oo2)
11807            .arg(&r0)
11808            .arg(&r1)
11809            .arg(&r2);
11810        unsafe {
11811            b.launch(cfg)?;
11812        }
11813        Ok(Some((y0, y1, y2)))
11814    }
11815
11816    /// Slot-fed fused3 twin (alloc-free capture lane): identical launch, caller-owned outputs.
11817    /// Returns Ok(false) when the fused path is unavailable (caller falls back).
11818    #[allow(clippy::too_many_arguments)]
11819    pub fn matmul_q4_fused3_into(
11820        &self,
11821        w0: &crate::model::GpuTensor,
11822        w1: &crate::model::GpuTensor,
11823        w2: &crate::model::GpuTensor,
11824        aq: &CudaSlice<i8>,
11825        ad: &CudaSlice<f32>,
11826        y0: &mut CudaSlice<f32>,
11827        y1: &mut CudaSlice<f32>,
11828        y2: &mut CudaSlice<f32>,
11829    ) -> Result<bool, Box<dyn std::error::Error>> {
11830        use crate::model::GpuTensor;
11831        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
11832            match w {
11833                GpuTensor::Quant {
11834                    qtype, row_bytes, ..
11835                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
11836                _ => None,
11837            }
11838        };
11839        let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (q4(w0), q4(w1), q4(w2)) else {
11840            return Ok(false);
11841        };
11842        if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
11843            return Ok(false);
11844        }
11845        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
11846            match w {
11847                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
11848                    Some(m) => (m, true),
11849                    None => (bytes, *rp),
11850                },
11851                _ => unreachable!(),
11852            }
11853        }
11854        let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
11855        if rp0 != rp1 || rp1 != rp2 {
11856            return Ok(false);
11857        }
11858        let rp = rp0;
11859        let rpb: u32 = 4;
11860        let mr1 = rp && Self::q40_mr1_on();
11861        let nb = |o: usize| {
11862            if mr1 {
11863                (o as u32).div_ceil(rpb)
11864            } else {
11865                (o as u32).div_ceil(2).div_ceil(rpb)
11866            }
11867        };
11868        let grid = nb(o0) + nb(o1) + nb(o2);
11869        debug_assert!(y0.len() >= o0 && y1.len() >= o1 && y2.len() >= o2);
11870        let f = self.func(if mr1 {
11871            "qmatvec_q4_0_mmvq_fused3_mr1_rp"
11872        } else if rp {
11873            "qmatvec_q4_0_mmvq_fused3_rp"
11874        } else {
11875            "qmatvec_q4_0_mmvq_fused3"
11876        });
11877        let cfg = LaunchConfig {
11878            grid_dim: (grid, 1, 1),
11879            block_dim: (32, rpb, 1),
11880            shared_mem_bytes: 0,
11881        };
11882        let inf = w0.in_features() as i32;
11883        let (oo0, oo1, oo2) = (o0 as i32, o1 as i32, o2 as i32);
11884        let (r0, r1, r2) = (rb0 as i64, rb1 as i64, rb2 as i64);
11885        // PDL wave-A: identical to the owned twin (capture-lane parity).
11886        if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
11887            use cudarc::driver::{DevicePtr, DevicePtrMut};
11888            let s = &self.gpu.stream();
11889            let (p0, _g0) = b0.device_ptr(s);
11890            let (p1, _g1) = b1.device_ptr(s);
11891            let (p2, _g2) = b2.device_ptr(s);
11892            let (paq, _g3) = aq.device_ptr(s);
11893            let (pad, _g4) = ad.device_ptr(s);
11894            let (py0, _g5) = y0.device_ptr_mut(s);
11895            let (py1, _g6) = y1.device_ptr_mut(s);
11896            let (py2, _g7) = y2.device_ptr_mut(s);
11897            let mut ps = [
11898                &p0 as *const _ as *mut std::ffi::c_void,
11899                &p1 as *const _ as *mut _,
11900                &p2 as *const _ as *mut _,
11901                &paq as *const _ as *mut _,
11902                &pad as *const _ as *mut _,
11903                &py0 as *const _ as *mut _,
11904                &py1 as *const _ as *mut _,
11905                &py2 as *const _ as *mut _,
11906                &inf as *const _ as *mut _,
11907                &oo0 as *const _ as *mut _,
11908                &oo1 as *const _ as *mut _,
11909                &oo2 as *const _ as *mut _,
11910                &r0 as *const _ as *mut _,
11911                &r1 as *const _ as *mut _,
11912                &r2 as *const _ as *mut _,
11913            ];
11914            unsafe {
11915                self.launch_pdl(
11916                    "qmatvec_q4_0_mmvq_fused3_mr1_rp",
11917                    (grid, 1, 1),
11918                    (32, rpb, 1),
11919                    &mut ps,
11920                )?;
11921            }
11922            return Ok(true);
11923        }
11924        let __s_b = self.gpu.stream();
11925        let mut b = __s_b.launch_builder(&f);
11926        b.arg(b0)
11927            .arg(b1)
11928            .arg(b2)
11929            .arg(aq)
11930            .arg(ad)
11931            .arg(&mut *y0)
11932            .arg(&mut *y1)
11933            .arg(&mut *y2)
11934            .arg(&inf)
11935            .arg(&oo0)
11936            .arg(&oo1)
11937            .arg(&oo2)
11938            .arg(&r0)
11939            .arg(&r1)
11940            .arg(&r2);
11941        unsafe {
11942            b.launch(cfg)?;
11943        }
11944        Ok(true)
11945    }
11946
11947    /// FUSED Q4_0 m=1 PAIR (gemma shared gate+up).
11948    pub fn matmul_q4_fused2(
11949        &self,
11950        w0: &crate::model::GpuTensor,
11951        w1: &crate::model::GpuTensor,
11952        aq: &CudaSlice<i8>,
11953        ad: &CudaSlice<f32>,
11954    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
11955        use crate::model::GpuTensor;
11956        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
11957            match w {
11958                GpuTensor::Quant {
11959                    qtype, row_bytes, ..
11960                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
11961                _ => None,
11962            }
11963        };
11964        let (Some((rb0, o0)), Some((rb1, o1))) = (q4(w0), q4(w1)) else {
11965            return Ok(None);
11966        };
11967        if w0.in_features() != w1.in_features() {
11968            return Ok(None);
11969        }
11970        // Effective (bytes, rp) per tensor (mirror or in-place swap); mixed -> separate matvecs.
11971        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
11972            match w {
11973                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
11974                    Some(m) => (m, true),
11975                    None => (bytes, *rp),
11976                },
11977                _ => unreachable!(),
11978            }
11979        }
11980        let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
11981        if rp0 != rp1 {
11982            return Ok(None);
11983        }
11984        let rp = rp0;
11985        let rpb: u32 = 4;
11986        // mr1 twin — see matmul_q4_fused3.
11987        let mr1 = rp && Self::q40_mr1_on();
11988        let nb = |o: usize| {
11989            if mr1 {
11990                (o as u32).div_ceil(rpb)
11991            } else {
11992                (o as u32).div_ceil(2).div_ceil(rpb)
11993            }
11994        };
11995        let grid = nb(o0) + nb(o1);
11996        let mut y0 = self.alloc_uninit::<f32>(o0)?;
11997        let mut y1 = self.alloc_uninit::<f32>(o1)?;
11998        let f = self.func(if mr1 {
11999            "qmatvec_q4_0_mmvq_fused2_mr1_rp"
12000        } else if rp {
12001            "qmatvec_q4_0_mmvq_fused2_rp"
12002        } else {
12003            "qmatvec_q4_0_mmvq_fused2"
12004        });
12005        let cfg = LaunchConfig {
12006            grid_dim: (grid, 1, 1),
12007            block_dim: (32, rpb, 1),
12008            shared_mem_bytes: 0,
12009        };
12010        let inf = w0.in_features() as i32;
12011        let (oo0, oo1) = (o0 as i32, o1 as i32);
12012        let (r0, r1) = (rb0 as i64, rb1 as i64);
12013        // PDL wave-A: mr1 kernel carries MEMRA_PDL_ENTRY.
12014        if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
12015            {
12016                use cudarc::driver::{DevicePtr, DevicePtrMut};
12017                let s = &self.gpu.stream();
12018                let (p0, _g0) = b0.device_ptr(s);
12019                let (p1, _g1) = b1.device_ptr(s);
12020                let (paq, _g2) = aq.device_ptr(s);
12021                let (pad, _g3) = ad.device_ptr(s);
12022                let (py0, _g4) = y0.device_ptr_mut(s);
12023                let (py1, _g5) = y1.device_ptr_mut(s);
12024                let mut ps = [
12025                    &p0 as *const _ as *mut std::ffi::c_void,
12026                    &p1 as *const _ as *mut _,
12027                    &paq as *const _ as *mut _,
12028                    &pad as *const _ as *mut _,
12029                    &py0 as *const _ as *mut _,
12030                    &py1 as *const _ as *mut _,
12031                    &inf as *const _ as *mut _,
12032                    &oo0 as *const _ as *mut _,
12033                    &oo1 as *const _ as *mut _,
12034                    &r0 as *const _ as *mut _,
12035                    &r1 as *const _ as *mut _,
12036                ];
12037                unsafe {
12038                    self.launch_pdl(
12039                        "qmatvec_q4_0_mmvq_fused2_mr1_rp",
12040                        (grid, 1, 1),
12041                        (32, rpb, 1),
12042                        &mut ps,
12043                    )?;
12044                }
12045            }
12046            return Ok(Some((y0, y1)));
12047        }
12048        let __s_b = self.gpu.stream();
12049        let mut b = __s_b.launch_builder(&f);
12050        b.arg(b0)
12051            .arg(b1)
12052            .arg(aq)
12053            .arg(ad)
12054            .arg(&mut y0)
12055            .arg(&mut y1)
12056            .arg(&inf)
12057            .arg(&oo0)
12058            .arg(&oo1)
12059            .arg(&r0)
12060            .arg(&r1);
12061        unsafe {
12062            b.launch(cfg)?;
12063        }
12064        Ok(Some((y0, y1)))
12065    }
12066
12067    /// Slot-fed fused2 twin (alloc-free capture lane): identical launch, caller-owned outputs.
12068    pub fn matmul_q4_fused2_into(
12069        &self,
12070        w0: &crate::model::GpuTensor,
12071        w1: &crate::model::GpuTensor,
12072        aq: &CudaSlice<i8>,
12073        ad: &CudaSlice<f32>,
12074        y0: &mut CudaSlice<f32>,
12075        y1: &mut CudaSlice<f32>,
12076    ) -> Result<bool, Box<dyn std::error::Error>> {
12077        use crate::model::GpuTensor;
12078        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
12079            match w {
12080                GpuTensor::Quant {
12081                    qtype, row_bytes, ..
12082                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
12083                _ => None,
12084            }
12085        };
12086        let (Some((rb0, o0)), Some((rb1, o1))) = (q4(w0), q4(w1)) else {
12087            return Ok(false);
12088        };
12089        if w0.in_features() != w1.in_features() {
12090            return Ok(false);
12091        }
12092        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
12093            match w {
12094                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
12095                    Some(m) => (m, true),
12096                    None => (bytes, *rp),
12097                },
12098                _ => unreachable!(),
12099            }
12100        }
12101        let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
12102        if rp0 != rp1 {
12103            return Ok(false);
12104        }
12105        let rp = rp0;
12106        let rpb: u32 = 4;
12107        let mr1 = rp && Self::q40_mr1_on();
12108        let nb = |o: usize| {
12109            if mr1 {
12110                (o as u32).div_ceil(rpb)
12111            } else {
12112                (o as u32).div_ceil(2).div_ceil(rpb)
12113            }
12114        };
12115        let grid = nb(o0) + nb(o1);
12116        debug_assert!(y0.len() >= o0 && y1.len() >= o1);
12117        let f = self.func(if mr1 {
12118            "qmatvec_q4_0_mmvq_fused2_mr1_rp"
12119        } else if rp {
12120            "qmatvec_q4_0_mmvq_fused2_rp"
12121        } else {
12122            "qmatvec_q4_0_mmvq_fused2"
12123        });
12124        let cfg = LaunchConfig {
12125            grid_dim: (grid, 1, 1),
12126            block_dim: (32, rpb, 1),
12127            shared_mem_bytes: 0,
12128        };
12129        let inf = w0.in_features() as i32;
12130        let (oo0, oo1) = (o0 as i32, o1 as i32);
12131        let (r0, r1) = (rb0 as i64, rb1 as i64);
12132        // PDL wave-A: identical to the owned twin (capture-lane parity).
12133        if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
12134            use cudarc::driver::{DevicePtr, DevicePtrMut};
12135            let s = &self.gpu.stream();
12136            let (p0, _g0) = b0.device_ptr(s);
12137            let (p1, _g1) = b1.device_ptr(s);
12138            let (paq, _g2) = aq.device_ptr(s);
12139            let (pad, _g3) = ad.device_ptr(s);
12140            let (py0, _g4) = y0.device_ptr_mut(s);
12141            let (py1, _g5) = y1.device_ptr_mut(s);
12142            let mut ps = [
12143                &p0 as *const _ as *mut std::ffi::c_void,
12144                &p1 as *const _ as *mut _,
12145                &paq as *const _ as *mut _,
12146                &pad as *const _ as *mut _,
12147                &py0 as *const _ as *mut _,
12148                &py1 as *const _ as *mut _,
12149                &inf as *const _ as *mut _,
12150                &oo0 as *const _ as *mut _,
12151                &oo1 as *const _ as *mut _,
12152                &r0 as *const _ as *mut _,
12153                &r1 as *const _ as *mut _,
12154            ];
12155            unsafe {
12156                self.launch_pdl(
12157                    "qmatvec_q4_0_mmvq_fused2_mr1_rp",
12158                    (grid, 1, 1),
12159                    (32, rpb, 1),
12160                    &mut ps,
12161                )?;
12162            }
12163            return Ok(true);
12164        }
12165        let __s_b = self.gpu.stream();
12166        let mut b = __s_b.launch_builder(&f);
12167        b.arg(b0)
12168            .arg(b1)
12169            .arg(aq)
12170            .arg(ad)
12171            .arg(&mut *y0)
12172            .arg(&mut *y1)
12173            .arg(&inf)
12174            .arg(&oo0)
12175            .arg(&oo1)
12176            .arg(&r0)
12177            .arg(&r1);
12178        unsafe {
12179            b.launch(cfg)?;
12180        }
12181        Ok(true)
12182    }
12183
12184    /// BATCHED fused2 (2026-07-13, megakernel-microcosm probe): gate+up b-tier matvecs in
12185    /// ONE segmented-grid launch — the up segment fills SMs as the gate segment drains
12186    /// (the per-launch tail waves behind the 6x-falsified b-tier plateau). Bit-identical
12187    /// per row to two mr2_rp launches. rp layout required; m in 2..=8 (b16 has no twin).
12188    pub fn matmul_q4_fused2_batched(
12189        &self,
12190        w0: &crate::model::GpuTensor,
12191        w1: &crate::model::GpuTensor,
12192        aq: &CudaSlice<i8>,
12193        ad: &CudaSlice<f32>,
12194        m: usize,
12195    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
12196        use crate::model::GpuTensor;
12197        if m < 2 || m > 8 {
12198            return Ok(None);
12199        }
12200        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
12201            match w {
12202                GpuTensor::Quant {
12203                    qtype, row_bytes, ..
12204                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
12205                _ => None,
12206            }
12207        };
12208        let (Some((rb0, o0)), Some((_rb1, o1))) = (q4(w0), q4(w1)) else {
12209            return Ok(None);
12210        };
12211        if w0.in_features() != w1.in_features() {
12212            return Ok(None);
12213        }
12214        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
12215            match w {
12216                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
12217                    Some(mr) => (mr, true),
12218                    None => (bytes, *rp),
12219                },
12220                _ => unreachable!(),
12221            }
12222        }
12223        let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
12224        if !rp0 || !rp1 {
12225            return Ok(None);
12226        }
12227        let mcols = Self::batched_mcols(m);
12228        let rpb: u32 = 4;
12229        let nb = |o: usize| (o as u32).div_ceil(2 * rpb);
12230        let grid = nb(o0) + nb(o1);
12231        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
12232        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
12233        let f = self.func(match mcols {
12234            2 => "qmatvec_q4_0_mmvq_b2_f2_rp",
12235            4 => "qmatvec_q4_0_mmvq_b4_f2_rp",
12236            _ => "qmatvec_q4_0_mmvq_b8_f2_rp",
12237        });
12238        let cfg = LaunchConfig {
12239            grid_dim: (grid, 1, 1),
12240            block_dim: (32, rpb, 1),
12241            shared_mem_bytes: 0,
12242        };
12243        let inf = w0.in_features() as i32;
12244        let (oo0, oo1, mi) = (o0 as i32, o1 as i32, m as i32);
12245        let rb = rb0 as i64;
12246        let __s_b = self.gpu.stream();
12247        let mut b = __s_b.launch_builder(&f);
12248        b.arg(b0)
12249            .arg(b1)
12250            .arg(aq)
12251            .arg(ad)
12252            .arg(&mut y0)
12253            .arg(&mut y1)
12254            .arg(&inf)
12255            .arg(&oo0)
12256            .arg(&oo1)
12257            .arg(&mi)
12258            .arg(&rb);
12259        unsafe {
12260            b.launch(cfg)?;
12261        }
12262        Ok(Some((y0, y1)))
12263    }
12264
12265    /// BATCHED fused3 (see matmul_q4_fused2_batched): three-segment single launch for the
12266    /// verify qkv triple. Same-in_f q4_0 rp tensors, m in 2..=8. Bit-identical per row.
12267    #[allow(clippy::too_many_arguments)]
12268    pub fn matmul_q4_fused3_batched(
12269        &self,
12270        w0: &crate::model::GpuTensor,
12271        w1: &crate::model::GpuTensor,
12272        w2: &crate::model::GpuTensor,
12273        aq: &CudaSlice<i8>,
12274        ad: &CudaSlice<f32>,
12275        m: usize,
12276    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
12277    {
12278        use crate::model::GpuTensor;
12279        if m < 2 || m > 8 {
12280            return Ok(None);
12281        }
12282        let q4 = |w: &GpuTensor| -> Option<usize> {
12283            match w {
12284                GpuTensor::Quant { qtype, .. } if *qtype == QT_Q4_0 => Some(w.out_features()),
12285                _ => None,
12286            }
12287        };
12288        let (Some(o0), Some(o1), Some(o2)) = (q4(w0), q4(w1), q4(w2)) else {
12289            return Ok(None);
12290        };
12291        if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
12292            return Ok(None);
12293        }
12294        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
12295            match w {
12296                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
12297                    Some(mr) => (mr, true),
12298                    None => (bytes, *rp),
12299                },
12300                _ => unreachable!(),
12301            }
12302        }
12303        let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
12304        if !rp0 || !rp1 || !rp2 {
12305            return Ok(None);
12306        }
12307        let mcols = Self::batched_mcols(m);
12308        let rpb: u32 = 4;
12309        let nb = |o: usize| (o as u32).div_ceil(2 * rpb);
12310        let grid = nb(o0) + nb(o1) + nb(o2);
12311        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
12312        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
12313        let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
12314        let f = self.func(match mcols {
12315            2 => "qmatvec_q4_0_mmvq_b2_f3_rp",
12316            4 => "qmatvec_q4_0_mmvq_b4_f3_rp",
12317            _ => "qmatvec_q4_0_mmvq_b8_f3_rp",
12318        });
12319        let cfg = LaunchConfig {
12320            grid_dim: (grid, 1, 1),
12321            block_dim: (32, rpb, 1),
12322            shared_mem_bytes: 0,
12323        };
12324        let inf = w0.in_features() as i32;
12325        let (oo0, oo1, oo2, mi) = (o0 as i32, o1 as i32, o2 as i32, m as i32);
12326        let rb = 0i64;
12327        let __s_b = self.gpu.stream();
12328        let mut b = __s_b.launch_builder(&f);
12329        b.arg(b0)
12330            .arg(b1)
12331            .arg(b2)
12332            .arg(aq)
12333            .arg(ad)
12334            .arg(&mut y0)
12335            .arg(&mut y1)
12336            .arg(&mut y2)
12337            .arg(&inf)
12338            .arg(&oo0)
12339            .arg(&oo1)
12340            .arg(&oo2)
12341            .arg(&mi)
12342            .arg(&rb);
12343        unsafe {
12344            b.launch(cfg)?;
12345        }
12346        Ok(Some((y0, y1, y2)))
12347    }
12348
12349    pub fn matmul_q8_fused3(
12350        &self,
12351        w0: &crate::model::GpuTensor,
12352        w1: &crate::model::GpuTensor,
12353        w2: &crate::model::GpuTensor,
12354        aq: &CudaSlice<i8>,
12355        ad: &CudaSlice<f32>,
12356    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
12357    {
12358        // e4m3 twin (lane/fp8-decode-v1): the full-attn wq/wk/wv triple — on the NV-27B those three
12359        // are per-tensor FP8, so native residency without this arm meant three separate launches.
12360        if let Some([p0, p1, p2]) = self.e4m3_fused_params(&[w0, w1, w2]) {
12361            return Ok(Some(self.e4m3_fused3_core(
12362                p0.0,
12363                p1.0,
12364                p2.0,
12365                aq,
12366                ad,
12367                w0.in_features(),
12368                p0.1,
12369                p1.1,
12370                p2.1,
12371                p0.2,
12372                p0.3,
12373                p1.3,
12374                p2.3,
12375            )?));
12376        }
12377        let Some([p0, p1, p2]) = self.q8_fused_params(&[w0, w1, w2]) else {
12378            return Ok(None);
12379        };
12380        Ok(Some(self.q8_fused3_core(
12381            p0.0,
12382            p1.0,
12383            p2.0,
12384            aq,
12385            ad,
12386            w0.in_features(),
12387            p0.1,
12388            p1.1,
12389            p2.1,
12390            p0.2,
12391        )?))
12392    }
12393
12394    #[allow(clippy::too_many_arguments)]
12395    fn q8_fused3_core(
12396        &self,
12397        b0: &CudaSlice<u8>,
12398        b1: &CudaSlice<u8>,
12399        b2: &CudaSlice<u8>,
12400        aq: &CudaSlice<i8>,
12401        ad: &CudaSlice<f32>,
12402        in_f: usize,
12403        out0: usize,
12404        out1: usize,
12405        out2: usize,
12406        row_bytes: usize,
12407    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12408        const ROWS_PER_BLOCK: u32 = 4;
12409        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
12410        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
12411        let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
12412        let f = self.func("qmatvec_q8_0_mmvq_fused3");
12413        let mut y0 = self.alloc_uninit::<f32>(out0)?;
12414        let mut y1 = self.alloc_uninit::<f32>(out1)?;
12415        let mut y2 = self.alloc_uninit::<f32>(out2)?;
12416        let cfg = LaunchConfig {
12417            grid_dim: (nb0 + nb1 + nb2, 1, 1),
12418            block_dim: (32, ROWS_PER_BLOCK, 1),
12419            shared_mem_bytes: 0,
12420        };
12421        let (inf, o0, o1, o2, rbl) = (
12422            in_f as i32,
12423            out0 as i32,
12424            out1 as i32,
12425            out2 as i32,
12426            row_bytes as i64,
12427        );
12428        let __s_b = self.gpu.stream();
12429        let mut b = __s_b.launch_builder(&f);
12430        b.arg(b0)
12431            .arg(b1)
12432            .arg(b2)
12433            .arg(aq)
12434            .arg(ad)
12435            .arg(&mut y0)
12436            .arg(&mut y1)
12437            .arg(&mut y2)
12438            .arg(&inf)
12439            .arg(&o0)
12440            .arg(&o1)
12441            .arg(&o2)
12442            .arg(&rbl);
12443        unsafe {
12444            b.launch(cfg)?;
12445        }
12446        Ok((y0, y1, y2))
12447    }
12448
12449    /// Test entry for the kernel_check gate: fused3 from raw weight bytes (internal q8_1 quant).
12450    #[allow(clippy::too_many_arguments)]
12451    pub fn qmatvec_q8_fused3_raw(
12452        &self,
12453        b0: &CudaSlice<u8>,
12454        b1: &CudaSlice<u8>,
12455        b2: &CudaSlice<u8>,
12456        x: &CudaSlice<f32>,
12457        in_f: usize,
12458        out0: usize,
12459        out1: usize,
12460        out2: usize,
12461        row_bytes: usize,
12462    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12463        let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
12464        self.q8_fused3_core(b0, b1, b2, &aq, &ad, in_f, out0, out1, out2, row_bytes)
12465    }
12466
12467    /// BATCHED twin of `matmul_q8_fused2` for the verify t=2-4 tier (MEMRA_SPEC_FUSED_T call
12468    /// sites, lane/close35b): ONE launch computes both same-input Q8_0 projections for m tokens.
12469    /// Per (tensor,token,row) the kernel body is q8_0_mmvq_batched VERBATIM with the identical
12470    /// row mapping (Q8_0's batched_variant is always "base") -> BIT-IDENTICAL to the two
12471    /// per-tensor _b2/_b4 launches `matmul_decode_exact` dispatches at m=2-4, with the caller's
12472    /// single shared q8_1 activation replacing two per-call re-quantizes (quantize_q8_1 is
12473    /// deterministic -> same bytes). None when ineligible (m outside 2..=4 / not both Q8_0 /
12474    /// in_f mismatch / MEMRA_MMVQ=0 / MEMRA_Q8_DUAL=0 / MEMRA_NO_BATCHED set — the last keeps
12475    /// dispatch parity: without batched kernels decode-exact runs grid.y=m MMVQ, and the fused
12476    /// twin must not introduce a batched program the reference path would not run).
12477    pub fn matmul_q8_fused2_t(
12478        &self,
12479        w0: &crate::model::GpuTensor,
12480        w1: &crate::model::GpuTensor,
12481        aq: &CudaSlice<i8>,
12482        ad: &CudaSlice<f32>,
12483        m: usize,
12484    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
12485        // m<=8 (lane/q27-deepdive, 2026-08-05): was 2..=4 (the verify tier's mcols 2/4). The
12486        // serving tick's mcols-8 tier now has its fused2_b8 wrapper, so c=5..8 batched decode
12487        // fuses too — same template body, still bit-identical to the two _b8 launches.
12488        if !(2..=8).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok() {
12489            return Ok(None);
12490        }
12491        // e4m3 twin: MEMRA_B8 parity — without it m=5..8 e4m3 decode runs the per-m grid.y=m path,
12492        // so the fused b8 launch would introduce a batched program the reference path would not run.
12493        if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
12494            if m > 4 && !Self::b8_enabled() {
12495                return Ok(None);
12496            }
12497            return Ok(Some(self.e4m3_fused2_t_core(
12498                p0.0,
12499                p1.0,
12500                aq,
12501                ad,
12502                m,
12503                w0.in_features(),
12504                p0.1,
12505                p1.1,
12506                p0.2,
12507                p0.3,
12508                p1.3,
12509            )?));
12510        }
12511        let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
12512            return Ok(None);
12513        };
12514        Ok(Some(self.q8_fused2_t_core(
12515            p0.0,
12516            p1.0,
12517            aq,
12518            ad,
12519            m,
12520            w0.in_features(),
12521            p0.1,
12522            p1.1,
12523            p0.2,
12524        )?))
12525    }
12526
12527    #[allow(clippy::too_many_arguments)]
12528    fn q8_fused2_t_core(
12529        &self,
12530        b0: &CudaSlice<u8>,
12531        b1: &CudaSlice<u8>,
12532        aq: &CudaSlice<i8>,
12533        ad: &CudaSlice<f32>,
12534        m: usize,
12535        in_f: usize,
12536        out0: usize,
12537        out1: usize,
12538        row_bytes: usize,
12539    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12540        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
12541        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
12542        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
12543        let f = self.func(match Self::batched_mcols(m) {
12544            2 => "qmatvec_q8_0_mmvq_fused2_b2",
12545            4 => "qmatvec_q8_0_mmvq_fused2_b4",
12546            // b8 = the SERVING tier (lane/q27-deepdive): c=5..8 batched decode.
12547            _ => "qmatvec_q8_0_mmvq_fused2_b8",
12548        });
12549        let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
12550        let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
12551        let cfg = LaunchConfig {
12552            grid_dim: (nb0 + nb1, 1, 1),
12553            block_dim: (32, ROWS_PER_BLOCK, 1),
12554            shared_mem_bytes: 0,
12555        };
12556        let (inf, o0, o1, mi, rbl) = (
12557            in_f as i32,
12558            out0 as i32,
12559            out1 as i32,
12560            m as i32,
12561            row_bytes as i64,
12562        );
12563        let __s_b = self.gpu.stream();
12564        let mut b = __s_b.launch_builder(&f);
12565        b.arg(b0)
12566            .arg(b1)
12567            .arg(aq)
12568            .arg(ad)
12569            .arg(&mut y0)
12570            .arg(&mut y1)
12571            .arg(&inf)
12572            .arg(&o0)
12573            .arg(&o1)
12574            .arg(&mi)
12575            .arg(&rbl);
12576        unsafe {
12577            b.launch(cfg)?;
12578        }
12579        Ok((y0, y1))
12580    }
12581
12582    /// Test entry for the kernel_check gate: fused2 batched from raw weight bytes (internal
12583    /// q8_1 quant of the [m, in_f] activation), no env gating.
12584    #[allow(clippy::too_many_arguments)]
12585    pub fn qmatvec_q8_fused2_t_raw(
12586        &self,
12587        b0: &CudaSlice<u8>,
12588        b1: &CudaSlice<u8>,
12589        x: &CudaSlice<f32>,
12590        m: usize,
12591        in_f: usize,
12592        out0: usize,
12593        out1: usize,
12594        row_bytes: usize,
12595    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12596        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
12597        self.q8_fused2_t_core(b0, b1, &aq, &ad, m, in_f, out0, out1, row_bytes)
12598    }
12599
12600    /// BATCHED twin of `matmul_q8_fused3` (wq+wk+wv at verify t=2-4). Same contract as
12601    /// `matmul_q8_fused2_t` with three ranges.
12602    #[allow(clippy::too_many_arguments)]
12603    pub fn matmul_q8_fused3_t(
12604        &self,
12605        w0: &crate::model::GpuTensor,
12606        w1: &crate::model::GpuTensor,
12607        w2: &crate::model::GpuTensor,
12608        aq: &CudaSlice<i8>,
12609        ad: &CudaSlice<f32>,
12610        m: usize,
12611    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
12612    {
12613        if !(2..=4).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok() {
12614            return Ok(None);
12615        }
12616        if let Some([p0, p1, p2]) = self.e4m3_fused_params(&[w0, w1, w2]) {
12617            return Ok(Some(self.e4m3_fused3_t_core(
12618                p0.0,
12619                p1.0,
12620                p2.0,
12621                aq,
12622                ad,
12623                m,
12624                w0.in_features(),
12625                p0.1,
12626                p1.1,
12627                p2.1,
12628                p0.2,
12629                p0.3,
12630                p1.3,
12631                p2.3,
12632            )?));
12633        }
12634        let Some([p0, p1, p2]) = self.q8_fused_params(&[w0, w1, w2]) else {
12635            return Ok(None);
12636        };
12637        Ok(Some(self.q8_fused3_t_core(
12638            p0.0,
12639            p1.0,
12640            p2.0,
12641            aq,
12642            ad,
12643            m,
12644            w0.in_features(),
12645            p0.1,
12646            p1.1,
12647            p2.1,
12648            p0.2,
12649        )?))
12650    }
12651
12652    #[allow(clippy::too_many_arguments)]
12653    fn q8_fused3_t_core(
12654        &self,
12655        b0: &CudaSlice<u8>,
12656        b1: &CudaSlice<u8>,
12657        b2: &CudaSlice<u8>,
12658        aq: &CudaSlice<i8>,
12659        ad: &CudaSlice<f32>,
12660        m: usize,
12661        in_f: usize,
12662        out0: usize,
12663        out1: usize,
12664        out2: usize,
12665        row_bytes: usize,
12666    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12667        const ROWS_PER_BLOCK: u32 = 4;
12668        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
12669        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
12670        let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
12671        let f = self.func(if Self::batched_mcols(m) == 2 {
12672            "qmatvec_q8_0_mmvq_fused3_b2"
12673        } else {
12674            "qmatvec_q8_0_mmvq_fused3_b4"
12675        });
12676        let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
12677        let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
12678        let mut y2 = self.alloc_uninit::<f32>(m * out2)?;
12679        let cfg = LaunchConfig {
12680            grid_dim: (nb0 + nb1 + nb2, 1, 1),
12681            block_dim: (32, ROWS_PER_BLOCK, 1),
12682            shared_mem_bytes: 0,
12683        };
12684        let (inf, o0, o1, o2, mi, rbl) = (
12685            in_f as i32,
12686            out0 as i32,
12687            out1 as i32,
12688            out2 as i32,
12689            m as i32,
12690            row_bytes as i64,
12691        );
12692        let __s_b = self.gpu.stream();
12693        let mut b = __s_b.launch_builder(&f);
12694        b.arg(b0)
12695            .arg(b1)
12696            .arg(b2)
12697            .arg(aq)
12698            .arg(ad)
12699            .arg(&mut y0)
12700            .arg(&mut y1)
12701            .arg(&mut y2)
12702            .arg(&inf)
12703            .arg(&o0)
12704            .arg(&o1)
12705            .arg(&o2)
12706            .arg(&mi)
12707            .arg(&rbl);
12708        unsafe {
12709            b.launch(cfg)?;
12710        }
12711        Ok((y0, y1, y2))
12712    }
12713
12714    /// Test entry for the kernel_check gate: fused3 batched from raw weight bytes.
12715    #[allow(clippy::too_many_arguments)]
12716    pub fn qmatvec_q8_fused3_t_raw(
12717        &self,
12718        b0: &CudaSlice<u8>,
12719        b1: &CudaSlice<u8>,
12720        b2: &CudaSlice<u8>,
12721        x: &CudaSlice<f32>,
12722        m: usize,
12723        in_f: usize,
12724        out0: usize,
12725        out1: usize,
12726        out2: usize,
12727        row_bytes: usize,
12728    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12729        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
12730        self.q8_fused3_t_core(b0, b1, b2, &aq, &ad, m, in_f, out0, out1, out2, row_bytes)
12731    }
12732
12733    /// Rollback seam for the Q8_0 dense-FFN gate+up fusion arm in `matmul_pre_dual_noscale`
12734    /// (lane/q27-deepdive, 2026-08-05). Default ON; `MEMRA_Q8_FFN_FUSE2=0` restores the
12735    /// two-`matmul_pre_noscale` pair. Read once — the dispatch must not vary within a run.
12736    pub fn q8_ffn_fuse2_on(&self) -> bool {
12737        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12738        *ON.get_or_init(|| std::env::var("MEMRA_Q8_FFN_FUSE2").as_deref() != Ok("0"))
12739    }
12740
12741    /// Eligibility + param extraction for the fused q8_0 launches: every tensor must be Quant Q8_0
12742    /// with macro-scale 1.0 (always true for GGUF q8_0; only NVFP4 carries scale) and share w[0]'s
12743    /// in_f (q8_0 row_bytes is a pure function of in_f, so equal in_f => equal row_bytes). MEMRA_MMVQ
12744    /// must be on: the fused body is the MMVQ kernel; without it decode m=1 runs dp4a and fusing
12745    /// would mix dispatch families (FP-order law). MEMRA_Q8_DUAL=0 = rollback seam.
12746    #[allow(clippy::type_complexity)]
12747    fn q8_fused_params<'w, const N: usize>(
12748        &self,
12749        ws: &[&'w crate::model::GpuTensor; N],
12750    ) -> Option<[(&'w CudaSlice<u8>, usize, usize); N]> {
12751        use crate::model::GpuTensor;
12752        if std::env::var("MEMRA_MMVQ").as_deref() == Ok("0") {
12753            return None;
12754        }
12755        if std::env::var("MEMRA_Q8_DUAL").is_ok_and(|v| v == "0") {
12756            return None;
12757        }
12758        let in_f = ws[0].in_features();
12759        let mut out: [Option<(&CudaSlice<u8>, usize, usize)>; N] = [None; N];
12760        for (i, w) in ws.iter().enumerate() {
12761            match w {
12762                GpuTensor::Quant {
12763                    bytes,
12764                    qtype,
12765                    row_bytes,
12766                    scale,
12767                    ..
12768                } if *qtype == QT_Q8_0 && *scale == 1.0 && w.in_features() == in_f => {
12769                    out[i] = Some((bytes, w.out_features(), *row_bytes))
12770                }
12771                _ => return None,
12772            }
12773        }
12774        Some(out.map(|o| o.unwrap()))
12775    }
12776
12777    /// Rollback seam for the F8-E4M3 launch-fusion arm (lane/fp8-decode-v1, 2026-08-05).
12778    /// Default ON; `MEMRA_E4M3_DUAL=0` restores the per-tensor m=1/batched launches.
12779    pub fn e4m3_dual_on(&self) -> bool {
12780        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12781        *ON.get_or_init(|| std::env::var("MEMRA_E4M3_DUAL").as_deref() != Ok("0"))
12782    }
12783
12784    /// Eligibility + param extraction for the FUSED e4m3 launches — the QT_F8_E4M3 twin of
12785    /// `q8_fused_params`. Differences that are inherent to the dtype, not policy:
12786    ///   * each tensor carries its OWN per-tensor `weight_scale` (returned as the 4th field);
12787    ///     Q8_0 hard-requires scale==1.0 because it has no macro-scale at all.
12788    ///   * no MEMRA_MMVQ gate: `mmvq_supports` exempts QT_F8_E4M3 (the e4m3 mmvq family is that
12789    ///     dtype's ONLY int8-act kernel class), so the per-tensor fallback these fused kernels
12790    ///     replace is ALWAYS the same mmvq body under every env — the FP-order law holds.
12791    ///   * `row_bytes == in_f` is asserted rather than derived: the native-residency load arm keeps
12792    ///     the checkpoint's raw [out_f, in_f] rows, and a re-encoded slab must never reach here.
12793    /// Rejects any split-plane mirror (`rp`/`rp4`): there is no `_rp` e4m3 fused form, so fusing
12794    /// there would swap dispatch families mid-model. MEMRA_E4M3_DUAL=0 = rollback seam.
12795    #[allow(clippy::type_complexity)]
12796    fn e4m3_fused_params<'w, const N: usize>(
12797        &self,
12798        ws: &[&'w crate::model::GpuTensor; N],
12799    ) -> Option<[(&'w CudaSlice<u8>, usize, usize, f32); N]> {
12800        use crate::model::GpuTensor;
12801        if !self.e4m3_dual_on() {
12802            return None;
12803        }
12804        let in_f = ws[0].in_features();
12805        let mut out: [Option<(&CudaSlice<u8>, usize, usize, f32)>; N] = [None; N];
12806        for (i, w) in ws.iter().enumerate() {
12807            match w {
12808                GpuTensor::Quant {
12809                    bytes,
12810                    qtype,
12811                    row_bytes,
12812                    scale,
12813                    rp,
12814                    rp4,
12815                    ..
12816                } if *qtype == QT_F8_E4M3
12817                    && w.in_features() == in_f
12818                    && *row_bytes == in_f
12819                    && !*rp
12820                    && rp4.is_none() =>
12821                {
12822                    out[i] = Some((bytes, w.out_features(), *row_bytes, *scale))
12823                }
12824                _ => return None,
12825            }
12826        }
12827        Some(out.map(|o| o.unwrap()))
12828    }
12829
12830    /// FUSED e4m3 m=1 PAIR. Block-offset split (`qmatvec_e4m3_mmvq_fused2`), per-tensor
12831    /// weight_scale folded at the write like the single-tensor `qmatvec_e4m3_mmvq` — so per
12832    /// (tensor,row) this is BIT-IDENTICAL to two separate m=1 launches, scale included.
12833    #[allow(clippy::too_many_arguments)]
12834    fn e4m3_fused2_core(
12835        &self,
12836        b0: &CudaSlice<u8>,
12837        b1: &CudaSlice<u8>,
12838        aq: &CudaSlice<i8>,
12839        ad: &CudaSlice<f32>,
12840        in_f: usize,
12841        out0: usize,
12842        out1: usize,
12843        row_bytes: usize,
12844        ws0: f32,
12845        ws1: f32,
12846    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12847        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
12848        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
12849        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
12850        let f = self.func("qmatvec_e4m3_mmvq_fused2");
12851        let mut y0 = self.alloc_uninit::<f32>(out0)?;
12852        let mut y1 = self.alloc_uninit::<f32>(out1)?;
12853        let cfg = LaunchConfig {
12854            grid_dim: (nb0 + nb1, 1, 1),
12855            block_dim: (32, ROWS_PER_BLOCK, 1),
12856            shared_mem_bytes: 0,
12857        };
12858        let (inf, o0, o1, rbl) = (in_f as i32, out0 as i32, out1 as i32, row_bytes as i64);
12859        let __s_b = self.gpu.stream();
12860        let mut b = __s_b.launch_builder(&f);
12861        b.arg(b0)
12862            .arg(b1)
12863            .arg(aq)
12864            .arg(ad)
12865            .arg(&mut y0)
12866            .arg(&mut y1)
12867            .arg(&inf)
12868            .arg(&o0)
12869            .arg(&o1)
12870            .arg(&rbl)
12871            .arg(&ws0)
12872            .arg(&ws1);
12873        unsafe {
12874            b.launch(cfg)?;
12875        }
12876        Ok((y0, y1))
12877    }
12878
12879    /// FUSED e4m3 m=1 TRIPLE (`qmatvec_e4m3_mmvq_fused3`). Same contract as the pair.
12880    #[allow(clippy::too_many_arguments)]
12881    fn e4m3_fused3_core(
12882        &self,
12883        b0: &CudaSlice<u8>,
12884        b1: &CudaSlice<u8>,
12885        b2: &CudaSlice<u8>,
12886        aq: &CudaSlice<i8>,
12887        ad: &CudaSlice<f32>,
12888        in_f: usize,
12889        out0: usize,
12890        out1: usize,
12891        out2: usize,
12892        row_bytes: usize,
12893        ws0: f32,
12894        ws1: f32,
12895        ws2: f32,
12896    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12897        const ROWS_PER_BLOCK: u32 = 4;
12898        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
12899        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
12900        let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
12901        let f = self.func("qmatvec_e4m3_mmvq_fused3");
12902        let mut y0 = self.alloc_uninit::<f32>(out0)?;
12903        let mut y1 = self.alloc_uninit::<f32>(out1)?;
12904        let mut y2 = self.alloc_uninit::<f32>(out2)?;
12905        let cfg = LaunchConfig {
12906            grid_dim: (nb0 + nb1 + nb2, 1, 1),
12907            block_dim: (32, ROWS_PER_BLOCK, 1),
12908            shared_mem_bytes: 0,
12909        };
12910        let (inf, o0, o1, o2, rbl) = (
12911            in_f as i32,
12912            out0 as i32,
12913            out1 as i32,
12914            out2 as i32,
12915            row_bytes as i64,
12916        );
12917        let __s_b = self.gpu.stream();
12918        let mut b = __s_b.launch_builder(&f);
12919        b.arg(b0)
12920            .arg(b1)
12921            .arg(b2)
12922            .arg(aq)
12923            .arg(ad)
12924            .arg(&mut y0)
12925            .arg(&mut y1)
12926            .arg(&mut y2)
12927            .arg(&inf)
12928            .arg(&o0)
12929            .arg(&o1)
12930            .arg(&o2)
12931            .arg(&rbl)
12932            .arg(&ws0)
12933            .arg(&ws1)
12934            .arg(&ws2);
12935        unsafe {
12936            b.launch(cfg)?;
12937        }
12938        Ok((y0, y1, y2))
12939    }
12940
12941    /// BATCHED FUSED e4m3 pair (m=2..8). The batched kernels carry no `ws` arg (every batched
12942    /// kernel in the tree is scale-free), so each output takes its own `scale_inplace` — the
12943    /// SAME post-op the per-tensor batched dispatch applies, hence still bit-identical.
12944    #[allow(clippy::too_many_arguments)]
12945    fn e4m3_fused2_t_core(
12946        &self,
12947        b0: &CudaSlice<u8>,
12948        b1: &CudaSlice<u8>,
12949        aq: &CudaSlice<i8>,
12950        ad: &CudaSlice<f32>,
12951        m: usize,
12952        in_f: usize,
12953        out0: usize,
12954        out1: usize,
12955        row_bytes: usize,
12956        ws0: f32,
12957        ws1: f32,
12958    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12959        const ROWS_PER_BLOCK: u32 = 4;
12960        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
12961        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
12962        let f = self.func(match Self::batched_mcols(m) {
12963            2 => "qmatvec_e4m3_mmvq_fused2_b2",
12964            4 => "qmatvec_e4m3_mmvq_fused2_b4",
12965            _ => "qmatvec_e4m3_mmvq_fused2_b8",
12966        });
12967        let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
12968        let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
12969        let cfg = LaunchConfig {
12970            grid_dim: (nb0 + nb1, 1, 1),
12971            block_dim: (32, ROWS_PER_BLOCK, 1),
12972            shared_mem_bytes: 0,
12973        };
12974        let (inf, o0, o1, mi, rbl) = (
12975            in_f as i32,
12976            out0 as i32,
12977            out1 as i32,
12978            m as i32,
12979            row_bytes as i64,
12980        );
12981        let __s_b = self.gpu.stream();
12982        let mut b = __s_b.launch_builder(&f);
12983        b.arg(b0)
12984            .arg(b1)
12985            .arg(aq)
12986            .arg(ad)
12987            .arg(&mut y0)
12988            .arg(&mut y1)
12989            .arg(&inf)
12990            .arg(&o0)
12991            .arg(&o1)
12992            .arg(&mi)
12993            .arg(&rbl);
12994        unsafe {
12995            b.launch(cfg)?;
12996        }
12997        if ws0 != 1.0 {
12998            self.scale_inplace(&mut y0, ws0, m * out0)?;
12999        }
13000        if ws1 != 1.0 {
13001            self.scale_inplace(&mut y1, ws1, m * out1)?;
13002        }
13003        Ok((y0, y1))
13004    }
13005
13006    /// BATCHED FUSED e4m3 triple (m=2..4). Same contract as the batched pair.
13007    #[allow(clippy::too_many_arguments)]
13008    fn e4m3_fused3_t_core(
13009        &self,
13010        b0: &CudaSlice<u8>,
13011        b1: &CudaSlice<u8>,
13012        b2: &CudaSlice<u8>,
13013        aq: &CudaSlice<i8>,
13014        ad: &CudaSlice<f32>,
13015        m: usize,
13016        in_f: usize,
13017        out0: usize,
13018        out1: usize,
13019        out2: usize,
13020        row_bytes: usize,
13021        ws0: f32,
13022        ws1: f32,
13023        ws2: f32,
13024    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13025        const ROWS_PER_BLOCK: u32 = 4;
13026        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
13027        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
13028        let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
13029        let f = self.func(if Self::batched_mcols(m) == 2 {
13030            "qmatvec_e4m3_mmvq_fused3_b2"
13031        } else {
13032            "qmatvec_e4m3_mmvq_fused3_b4"
13033        });
13034        let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
13035        let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
13036        let mut y2 = self.alloc_uninit::<f32>(m * out2)?;
13037        let cfg = LaunchConfig {
13038            grid_dim: (nb0 + nb1 + nb2, 1, 1),
13039            block_dim: (32, ROWS_PER_BLOCK, 1),
13040            shared_mem_bytes: 0,
13041        };
13042        let (inf, o0, o1, o2, mi, rbl) = (
13043            in_f as i32,
13044            out0 as i32,
13045            out1 as i32,
13046            out2 as i32,
13047            m as i32,
13048            row_bytes as i64,
13049        );
13050        let __s_b = self.gpu.stream();
13051        let mut b = __s_b.launch_builder(&f);
13052        b.arg(b0)
13053            .arg(b1)
13054            .arg(b2)
13055            .arg(aq)
13056            .arg(ad)
13057            .arg(&mut y0)
13058            .arg(&mut y1)
13059            .arg(&mut y2)
13060            .arg(&inf)
13061            .arg(&o0)
13062            .arg(&o1)
13063            .arg(&o2)
13064            .arg(&mi)
13065            .arg(&rbl);
13066        unsafe {
13067            b.launch(cfg)?;
13068        }
13069        if ws0 != 1.0 {
13070            self.scale_inplace(&mut y0, ws0, m * out0)?;
13071        }
13072        if ws1 != 1.0 {
13073            self.scale_inplace(&mut y1, ws1, m * out1)?;
13074        }
13075        if ws2 != 1.0 {
13076            self.scale_inplace(&mut y2, ws2, m * out2)?;
13077        }
13078        Ok((y0, y1, y2))
13079    }
13080
13081    /// BLOCK-128 e4m3 MMVQ launcher (`qmatvec_e4m3_blk_mmvq`, lane/fp8-blk128-decode 2026-08-05).
13082    /// The per-block-dequant twin of `qmatvec_mmvq`'s QT_F8_E4M3 arm: same grid/block decomposition
13083    /// (warp per output row, ROWS_PER_BLOCK warps per block, grid.y = m), same q8_1 activation, but
13084    /// the weight scale is a resident [rows, cols] f32 grid read per k128 block inside the kernel
13085    /// instead of one scalar folded at the write. It cannot share `qmatvec_mmvq`'s body because
13086    /// that launcher's arg list is fixed at (bytes, aq, ad, y, in_f, out_f, m, row_bytes [, scale]).
13087    ///
13088    /// `mr` and `rp` have no analogue here (no split-plane e4m3 layout exists), so there is exactly
13089    /// one kernel and no name table — a shape this cannot serve must be refused at LOAD, not here.
13090    pub fn qmatvec_e4m3_blk_mmvq(
13091        &self,
13092        bytes: &CudaSlice<u8>,
13093        aq: &CudaSlice<i8>,
13094        ad: &CudaSlice<f32>,
13095        scales: &CudaSlice<f32>,
13096        m: usize,
13097        in_f: usize,
13098        out_f: usize,
13099        row_bytes: usize,
13100        scale_cols: usize,
13101    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13102        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
13103        self.qmatvec_e4m3_blk_mmvq_into(
13104            bytes, aq, ad, scales, m, in_f, out_f, row_bytes, scale_cols, &mut y,
13105        )?;
13106        Ok(y)
13107    }
13108
13109    /// Slot-fed twin of `qmatvec_e4m3_blk_mmvq` (caller-owned output; the alloc-free capture lane).
13110    #[allow(clippy::too_many_arguments)]
13111    pub fn qmatvec_e4m3_blk_mmvq_into(
13112        &self,
13113        bytes: &CudaSlice<u8>,
13114        aq: &CudaSlice<i8>,
13115        ad: &CudaSlice<f32>,
13116        scales: &CudaSlice<f32>,
13117        m: usize,
13118        in_f: usize,
13119        out_f: usize,
13120        row_bytes: usize,
13121        scale_cols: usize,
13122        y: &mut CudaSlice<f32>,
13123    ) -> Result<(), Box<dyn std::error::Error>> {
13124        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
13125        let f = self.func("qmatvec_e4m3_blk_mmvq");
13126        let cfg = LaunchConfig {
13127            grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), m as u32, 1),
13128            block_dim: (32, ROWS_PER_BLOCK, 1), // warp-per-row
13129            shared_mem_bytes: 0,                // warp-only reduce
13130        };
13131        let (inf, outf, mi, rb, sc) = (
13132            in_f as i32,
13133            out_f as i32,
13134            m as i32,
13135            row_bytes as i64,
13136            scale_cols as i32,
13137        );
13138        let __s_b = self.gpu.stream();
13139        let mut b = __s_b.launch_builder(&f);
13140        b.arg(bytes)
13141            .arg(aq)
13142            .arg(ad)
13143            .arg(scales)
13144            .arg(&mut *y)
13145            .arg(&inf)
13146            .arg(&outf)
13147            .arg(&mi)
13148            .arg(&rb)
13149            .arg(&sc);
13150        unsafe {
13151            b.launch(cfg)?;
13152        }
13153        Ok(())
13154    }
13155
13156    /// BLOCK-128 e4m3 BATCHED matvec (lane/rp-on-st, 2026-08-06): the weight-read-once twin of
13157    /// `qmatvec_e4m3_blk_mmvq` for m=2..16. Per (token,row) BIT-IDENTICAL to the grid.y=m launch
13158    /// (same fmaf chain, same per-k32 `s * ad` fold, same warp reduce), so it inherits the
13159    /// decode-exactness contract while reading the weight ONCE for up to `mcols` columns instead
13160    /// of `m` times. `mcols` must be one of {2,4,8,16} and satisfy `mcols >= m`.
13161    #[allow(clippy::too_many_arguments)]
13162    pub fn qmatvec_e4m3_blk_mmvq_batched(
13163        &self,
13164        bytes: &CudaSlice<u8>,
13165        aq: &CudaSlice<i8>,
13166        ad: &CudaSlice<f32>,
13167        scales: &CudaSlice<f32>,
13168        m: usize,
13169        in_f: usize,
13170        out_f: usize,
13171        row_bytes: usize,
13172        scale_cols: usize,
13173        mcols: usize,
13174    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13175        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
13176        debug_assert!(mcols >= m, "blk batched: mcols {mcols} < m {m}");
13177        let name = match mcols {
13178            2 => "qmatvec_e4m3_blk_mmvq_b2",
13179            4 => "qmatvec_e4m3_blk_mmvq_b4",
13180            8 => "qmatvec_e4m3_blk_mmvq_b8",
13181            16 => "qmatvec_e4m3_blk_mmvq_b16",
13182            _ => {
13183                return Err(
13184                    format!("qmatvec_e4m3_blk_mmvq_batched: no kernel for mcols {mcols}").into(),
13185                );
13186            }
13187        };
13188        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
13189        let f = self.func(name);
13190        let cfg = LaunchConfig {
13191            grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
13192            block_dim: (32, ROWS_PER_BLOCK, 1),
13193            shared_mem_bytes: 0,
13194        };
13195        let (inf, outf, mi, rb, sc) = (
13196            in_f as i32,
13197            out_f as i32,
13198            m as i32,
13199            row_bytes as i64,
13200            scale_cols as i32,
13201        );
13202        let __s_b = self.gpu.stream();
13203        let mut b = __s_b.launch_builder(&f);
13204        b.arg(bytes)
13205            .arg(aq)
13206            .arg(ad)
13207            .arg(scales)
13208            .arg(&mut y)
13209            .arg(&inf)
13210            .arg(&outf)
13211            .arg(&mi)
13212            .arg(&rb)
13213            .arg(&sc);
13214        unsafe {
13215            b.launch(cfg)?;
13216        }
13217        Ok(y)
13218    }
13219
13220    /// Test entry for the kernel_check exactness gate: the block-128 e4m3 batched MMVQ from raw
13221    /// bytes with an internal q8_1 quantize (mirrors `qmatvec_batched_raw`).
13222    #[allow(clippy::too_many_arguments)]
13223    pub fn qmatvec_e4m3_blk_batched_raw(
13224        &self,
13225        bytes: &CudaSlice<u8>,
13226        x: &CudaSlice<f32>,
13227        scales: &CudaSlice<f32>,
13228        m: usize,
13229        in_f: usize,
13230        out_f: usize,
13231        row_bytes: usize,
13232        scale_cols: usize,
13233        mcols: usize,
13234    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13235        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
13236        self.qmatvec_e4m3_blk_mmvq_batched(
13237            bytes, &aq, &ad, scales, m, in_f, out_f, row_bytes, scale_cols, mcols,
13238        )
13239    }
13240
13241    /// Test entry for the kernel_check exactness gate: the block-128 e4m3 MMVQ from raw bytes with
13242    /// an internal q8_1 quantize (mirrors `qmatvec_mmvq_raw`).
13243    #[allow(clippy::too_many_arguments)]
13244    pub fn qmatvec_e4m3_blk_mmvq_raw(
13245        &self,
13246        bytes: &CudaSlice<u8>,
13247        x: &CudaSlice<f32>,
13248        scales: &CudaSlice<f32>,
13249        m: usize,
13250        in_f: usize,
13251        out_f: usize,
13252        row_bytes: usize,
13253        scale_cols: usize,
13254    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13255        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
13256        self.qmatvec_e4m3_blk_mmvq(
13257            bytes, &aq, &ad, scales, m, in_f, out_f, row_bytes, scale_cols,
13258        )
13259    }
13260
13261    /// Test entries for the kernel_check bit-parity gate: fused e4m3 launches from raw weight
13262    /// bytes with internal q8_1 quantize, no env gating (mirrors `qmatvec_q8_fused*_raw`).
13263    #[allow(clippy::too_many_arguments)]
13264    pub fn qmatvec_e4m3_fused2_raw(
13265        &self,
13266        b0: &CudaSlice<u8>,
13267        b1: &CudaSlice<u8>,
13268        x: &CudaSlice<f32>,
13269        in_f: usize,
13270        out0: usize,
13271        out1: usize,
13272        row_bytes: usize,
13273        ws0: f32,
13274        ws1: f32,
13275    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13276        let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
13277        self.e4m3_fused2_core(b0, b1, &aq, &ad, in_f, out0, out1, row_bytes, ws0, ws1)
13278    }
13279
13280    #[allow(clippy::too_many_arguments)]
13281    pub fn qmatvec_e4m3_fused3_raw(
13282        &self,
13283        b0: &CudaSlice<u8>,
13284        b1: &CudaSlice<u8>,
13285        b2: &CudaSlice<u8>,
13286        x: &CudaSlice<f32>,
13287        in_f: usize,
13288        out0: usize,
13289        out1: usize,
13290        out2: usize,
13291        row_bytes: usize,
13292        ws0: f32,
13293        ws1: f32,
13294        ws2: f32,
13295    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13296        let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
13297        self.e4m3_fused3_core(
13298            b0, b1, b2, &aq, &ad, in_f, out0, out1, out2, row_bytes, ws0, ws1, ws2,
13299        )
13300    }
13301
13302    #[allow(clippy::too_many_arguments)]
13303    pub fn qmatvec_e4m3_fused2_t_raw(
13304        &self,
13305        b0: &CudaSlice<u8>,
13306        b1: &CudaSlice<u8>,
13307        x: &CudaSlice<f32>,
13308        m: usize,
13309        in_f: usize,
13310        out0: usize,
13311        out1: usize,
13312        row_bytes: usize,
13313        ws0: f32,
13314        ws1: f32,
13315    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13316        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
13317        self.e4m3_fused2_t_core(b0, b1, &aq, &ad, m, in_f, out0, out1, row_bytes, ws0, ws1)
13318    }
13319
13320    #[allow(clippy::too_many_arguments)]
13321    pub fn qmatvec_e4m3_fused3_t_raw(
13322        &self,
13323        b0: &CudaSlice<u8>,
13324        b1: &CudaSlice<u8>,
13325        b2: &CudaSlice<u8>,
13326        x: &CudaSlice<f32>,
13327        m: usize,
13328        in_f: usize,
13329        out0: usize,
13330        out1: usize,
13331        out2: usize,
13332        row_bytes: usize,
13333        ws0: f32,
13334        ws1: f32,
13335        ws2: f32,
13336    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13337        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
13338        self.e4m3_fused3_t_core(
13339            b0, b1, b2, &aq, &ad, m, in_f, out0, out1, out2, row_bytes, ws0, ws1, ws2,
13340        )
13341    }
13342
13343    /// THE single dispatch point for `QT_F8_E4M3_BLK` from a PRE-QUANTIZED q8_1 activation
13344    /// (lane/fp8-blk128-decode). Every `matmul_pre`-family entry calls this first, so the block-128
13345    /// class has exactly ONE code path across `matmul`, `matmul_pre`, `matmul_pre_noscale`,
13346    /// `matmul_decode_exact` and `matmul_decode_exact_pre` — the same kernel at the same grid for
13347    /// every m, which is what makes verify == decode bit-for-bit at every tier for free.
13348    ///
13349    /// Returns None for any other qtype (the caller continues its normal dispatch). The `blk: Some`
13350    /// pattern is part of the match, not an unwrap: qtype and grid presence are set together in the
13351    /// one residency arm that builds this tensor, and a qtype-without-grid would be a construction
13352    /// bug — better to fall through and hit a loud refusal than to unwrap a None here.
13353    fn try_e4m3_blk_pre(
13354        &self,
13355        w: &crate::model::GpuTensor,
13356        aq: &CudaSlice<i8>,
13357        ad: &CudaSlice<f32>,
13358        m: usize,
13359    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
13360        use crate::model::GpuTensor;
13361        if let GpuTensor::Quant {
13362            bytes,
13363            qtype,
13364            row_bytes,
13365            blk: Some(g),
13366            ..
13367        } = w
13368        {
13369            if *qtype == QT_F8_E4M3_BLK {
13370                // BATCHED tier m=2..16 (lane/rp-on-st): weight read ONCE for up to mcols columns
13371                // instead of m grid.y re-reads. Bit-identical per (token,row) to the grid.y=m form
13372                // below, so the decode-exactness contract is preserved at every width. Gated by
13373                // the same seams the other batched families honor (MEMRA_NO_BATCHED, MEMRA_B8) so
13374                // one rollback door covers every dtype's batched tier.
13375                if (2..=16).contains(&m)
13376                    && std::env::var("MEMRA_NO_BATCHED").is_err()
13377                    && (m <= 4 || Self::b8_enabled())
13378                {
13379                    let mcols = Self::batched_mcols(m);
13380                    return Ok(Some(self.qmatvec_e4m3_blk_mmvq_batched(
13381                        bytes,
13382                        aq,
13383                        ad,
13384                        &g.scales,
13385                        m,
13386                        w.in_features(),
13387                        w.out_features(),
13388                        *row_bytes,
13389                        g.cols,
13390                        mcols,
13391                    )?));
13392                }
13393                return Ok(Some(self.qmatvec_e4m3_blk_mmvq(
13394                    bytes,
13395                    aq,
13396                    ad,
13397                    &g.scales,
13398                    m,
13399                    w.in_features(),
13400                    w.out_features(),
13401                    *row_bytes,
13402                    g.cols,
13403                )?));
13404            }
13405        }
13406        Ok(None)
13407    }
13408
13409    /// PREFILL (m >= GEMM_M_THRESHOLD) for `QT_F8_E4M3_BLK` — DEQUANT-PER-CALL to the Q8_0 slab
13410    /// this class's residency replaced, then the ordinary Q8_0 prefill dispatch on the transient.
13411    ///
13412    /// WHY THIS EXISTS AT ALL, i.e. the regression it prevents: the decode kernel is a warp-per-row
13413    /// GEMV. At grid.y=m it re-reads the whole weight once PER TOKEN, so letting a 512-token prefill
13414    /// chunk reach it would be a ~500x weight-traffic blowup on the single most bandwidth-bound part
13415    /// of the forward. Native residency is a DECODE win and must not be paid for in prefill, so
13416    /// prefill keeps the floor's arithmetic and the floor's kernels.
13417    ///
13418    /// WHY DEQUANT-PER-CALL rather than a second resident slab: a resident slab is dual residency —
13419    /// it gives back the entire 1.0-vs-1.0625 B/weight win this lane exists to capture (and then
13420    /// some, since the e4m3 copy stays too). The transient costs one linear device pass per
13421    /// (projection, prefill call) and frees immediately.
13422    ///
13423    /// NUMERICALLY IT IS THE FLOOR, EXACTLY: `fp8_blk_dequant_q8_0` is the merged ARM B' kernel,
13424    /// gate-proven BYTE-IDENTICAL to the host dequant+re-encode (kernel-check `fp8-blk-gpu`). So the
13425    /// slab these bytes form is bit-for-bit the slab the `MEMRA_ST_E4M3_BLK=0` arm makes resident,
13426    /// and every prefill kernel downstream sees identical input — prefill logits under this lane are
13427    /// bit-identical to prefill logits under the floor, which is what makes the decode A/B a clean
13428    /// single-variable comparison instead of a two-variable one.
13429    ///
13430    /// WHAT IT COSTS, MEASURED, AND WHY THAT COST IS MOSTLY STRUCTURAL (27B block-128 ckpt, pp512,
13431    /// this rig = RTX 5090 Laptop, ~896 GB/s GDDR7). This arm makes prefill move the weight THREE
13432    /// times instead of once: read 6.88 GB of e4m3, write 7.31 GB of Q8_0, then the MMQ reads that
13433    /// 7.31 GB back. The two extra passes are 14.19 GB = 15.8 ms at this card's roofline against a
13434    /// ~332 ms pp512, i.e. **~-4.5% pp is a floor no kernel tuning can remove** — only deleting the
13435    /// dequant can. Measured: the dequant kernel costs 27.9 ms/pass (nsys, 208 projections) after
13436    /// the 2026-08-05 vector rewrite (was 66.5 ms at one byte per thread), and e2e pp512 is
13437    /// 1451.4 vs the slab arm's 1541.6 tok/s = -5.8% (N=3 interleaved pairs). So ~1.3pp of the
13438    /// -5.8% is residual kernel inefficiency and ~4.5pp is the extra traffic itself.
13439    ///
13440    /// SO THE DEQUANT IS NO LONGER THE DEFAULT ROUTE — it is the FALLBACK. The per-block FP8 MMQ
13441    /// tile (`try_fp8_blk_mmq`) consumes the resident e4m3 bytes + grid DIRECTLY, deleting both extra
13442    /// passes, and since 2026-08-05 it runs FIRST and by default for the native-resident source
13443    /// (`fp8_blk_mmq_native_enabled`; `MEMRA_FP8_MMQ=0` is the seam back to this dequant). On paper
13444    /// the trade was unassumable — lane/fp8-mmq-v2 measured that tile at 0.85-1.09x the Q8_0 MMQ
13445    /// floor GEMM-only, so it swapped a -4.5% traffic cost for a 0-to-15% GEMM cost of unknown sign.
13446    /// Measured on the 27B (3 arms interleaved, N=3, research/fp8blk-20260805/VERDICT.md): slab
13447    /// 1540.5 / this dequant 1449.1 / the tile 1553.3 tok/s, min(tile) > max(slab). The tile wins
13448    /// because v2's denominator had its slab already resident while this class's floor must build it
13449    /// every call; same tile, opposite sign, because the question changed.
13450    ///
13451    /// THIS ARM STILL RUNS, and is not dead code: every `try_fp8_blk_mmq` precondition (in_f % 16,
13452    /// grid dims vs shape, per-tensor scale == 1.0, the e4m3-NaN scan) refuses by falling through to
13453    /// here, so a checkpoint the tile cannot take keeps exact prefill on the floor's own bits rather
13454    /// than losing the class. It is also what `MEMRA_FP8_MMQ=0` reverts to.
13455    fn try_e4m3_blk_prefill(
13456        &self,
13457        w: &crate::model::GpuTensor,
13458        x: &CudaSlice<f32>,
13459        m: usize,
13460    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
13461        use crate::model::GpuTensor;
13462        let GpuTensor::Quant {
13463            bytes,
13464            qtype,
13465            blk: Some(g),
13466            ..
13467        } = w
13468        else {
13469            return Ok(None);
13470        };
13471        if *qtype != QT_F8_E4M3_BLK {
13472            return Ok(None);
13473        }
13474        // NO-DEQUANT ROUTE, THE DEFAULT (MEMRA_FP8_MMQ=0 reverts): the per-block MMQ tile eats the
13475        // resident e4m3 bytes and grid as-is, so neither extra weight pass happens. Its own
13476        // preconditions (in_f % 16, grid dims, scale == 1.0, no e4m3 NaN code) can refuse — fall
13477        // through to the dequant below when they do, never silently produce nothing.
13478        if let Some(y) = self.try_fp8_blk_mmq(w, x, m)? {
13479            return Ok(Some(y));
13480        }
13481        let (in_f, out_f) = (w.in_features(), w.out_features());
13482        let slab = self.fp8_blk_dequant_q8_0_dev(bytes, &g.scales, out_f, in_f)?;
13483        let tmp = GpuTensor::Quant {
13484            bytes: slab,
13485            qtype: QT_Q8_0,
13486            row_bytes: in_f / 32 * 34,
13487            ne: vec![in_f as u64, out_f as u64],
13488            scale: 1.0,
13489            rp: false,
13490            #[cfg(memra_cutlass)]
13491            cutlass: None,
13492            fp8: None,
13493            blk: None,
13494            f16: None,
13495            rp4: None,
13496        };
13497        // Recursion terminates: `tmp` is QT_Q8_0 with `blk: None`, so it cannot re-enter this arm.
13498        Ok(Some(self.matmul(&tmp, x, m)?))
13499    }
13500
13501    pub fn matmul_pre_noscale(
13502        &self,
13503        w: &crate::model::GpuTensor,
13504        aq: &CudaSlice<i8>,
13505        ad: &CudaSlice<f32>,
13506        m: usize,
13507    ) -> Result<Option<(CudaSlice<f32>, f32)>, Box<dyn std::error::Error>> {
13508        use crate::model::GpuTensor;
13509        // BLOCK-128 e4m3: every scale factor is folded inside the kernel per k128, so the
13510        // "separable post-op scale" this entry exists to defer is 1.0 — return it explicitly
13511        // rather than let the tail below refuse and cost the caller a re-dispatch.
13512        if m == 1 {
13513            if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
13514                return Ok(Some((y, 1.0)));
13515            }
13516        }
13517        // Only the m==1 fast path applies the scale as a separable post-op; bail everywhere else.
13518        if m != 1 || !self.uses_q8_1_fast(w) {
13519            return Ok(None);
13520        }
13521        let in_f = w.in_features();
13522        let out_f = w.out_features();
13523        let (bytes, qtype, row_bytes, scale, rp) = match w {
13524            GpuTensor::Quant {
13525                bytes,
13526                qtype,
13527                row_bytes,
13528                scale,
13529                rp,
13530                ..
13531            } => (bytes, *qtype, *row_bytes, *scale, *rp),
13532            _ => return Ok(None),
13533        };
13534        // MMVQ warp-per-row (scale==1.0 passed -> kernel skips its internal scale; we return scale).
13535        if self.mmvq_supports(qtype) {
13536            // Q4_0 split-plane mirror (dp4a fallback below keeps the raw GGUF bytes).
13537            let (mbytes, mrp) = match w {
13538                GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
13539                _ => (bytes, rp),
13540            };
13541            let y = self.qmatvec_mmvq(
13542                mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, /*scale*/ 1.0, mrp,
13543            )?;
13544            return Ok(Some((y, scale)));
13545        }
13546        // dp4a fallback: same launch as matmul_pre but WITHOUT the post scale_inplace.
13547        let name = match qtype {
13548            QT_Q8_0 => "qmatvec_q8_0_dp4a",
13549            QT_Q4_K => "qmatvec_q4_K_dp4a",
13550            QT_Q6_K => "qmatvec_q6_K_dp4a",
13551            QT_Q5_K => "qmatvec_q5_K_dp4a",
13552            QT_Q3_K => "qmatvec_q3_K_dp4a",
13553            QT_NVFP4 => {
13554                if rp {
13555                    "qmatvec_nvfp4_dp4a_rp"
13556                } else {
13557                    "qmatvec_nvfp4_dp4a"
13558                }
13559            }
13560            QT_IQ4_XS => "qmatvec_iq4_XS_dp4a",
13561            _ => return Ok(None),
13562        };
13563        let f = self.func(name);
13564        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
13565        let cfg = LaunchConfig {
13566            grid_dim: (out_f as u32, m as u32, 1),
13567            block_dim: (128, 1, 1),
13568            shared_mem_bytes: 0,
13569        };
13570        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
13571        let __s_b = self.gpu.stream();
13572        let mut b = __s_b.launch_builder(&f);
13573        b.arg(bytes)
13574            .arg(aq)
13575            .arg(ad)
13576            .arg(&mut y)
13577            .arg(&inf)
13578            .arg(&outf)
13579            .arg(&mi)
13580            .arg(&rb);
13581        unsafe {
13582            b.launch(cfg)?;
13583        }
13584        Ok(Some((y, scale)))
13585    }
13586
13587    /// True if `qtype` has a warp-per-row MMVQ decode kernel AND MEMRA_MMVQ is set. Only the 4
13588    /// daily-hot dtypes (Q8_0, Q4_K, Q6_K, NVFP4) — others keep the _dp4a matvec (oracle/fallback).
13589    pub fn mmvq_supports(&self, qtype: i32) -> bool {
13590        // DEFAULT ON since 2026-07-08 (MEMRA_MMVQ=0 reverts to the _dp4a matvec class).
13591        // QT_F8_E4M3 is exempt from the MEMRA_MMVQ=0 escape: the e4m3 mmvq family is that dtype's
13592        // ONLY int8-act kernel class (there is no _dp4a twin), so its m=1/verify/batched dispatch
13593        // is a pure function of the dtype — the decode-parity law holds under every env.
13594        if qtype == QT_F8_E4M3 {
13595            return true;
13596        }
13597        if std::env::var("MEMRA_MMVQ").as_deref() == Ok("0") {
13598            return false;
13599        }
13600        matches!(
13601            qtype,
13602            QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q6_K | QT_NVFP4 | QT_Q4_0
13603        )
13604    }
13605
13606    /// PERF-3 warp-per-row MMVQ launcher (decode m=1 hot path). block=(32,ROWS_PER_BLOCK,1):
13607    /// one warp owns one output row, warp-only __shfl reduction (no smem barrier). Bit-equivalent
13608    /// to qmatvec_*_dp4a up to f32 reduction order. Pre-quantized q8_1 activation (aq,ad). NVFP4
13609    /// per-tensor macro-scale applied post (scale==1.0 for other dtypes -> no-op).
13610    pub fn qmatvec_mmvq(
13611        &self,
13612        bytes: &CudaSlice<u8>,
13613        aq: &CudaSlice<i8>,
13614        ad: &CudaSlice<f32>,
13615        m: usize,
13616        in_f: usize,
13617        out_f: usize,
13618        qtype: i32,
13619        row_bytes: usize,
13620        scale: f32,
13621        rp: bool,
13622    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13623        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output: skip memset
13624        self.qmatvec_mmvq_into(
13625            bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, rp, &mut y,
13626        )?;
13627        Ok(y)
13628    }
13629
13630    /// Slot-fed MMVQ twin (alloc-free capture lane): full policy body, caller-owned output.
13631    #[allow(clippy::too_many_arguments)]
13632    pub fn qmatvec_mmvq_into(
13633        &self,
13634        bytes: &CudaSlice<u8>,
13635        aq: &CudaSlice<i8>,
13636        ad: &CudaSlice<f32>,
13637        m: usize,
13638        in_f: usize,
13639        out_f: usize,
13640        qtype: i32,
13641        row_bytes: usize,
13642        scale: f32,
13643        rp: bool,
13644        y: &mut CudaSlice<f32>,
13645    ) -> Result<(), Box<dyn std::error::Error>> {
13646        debug_assert!(y.len() >= m * out_f);
13647        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
13648        // SMALL-SHAPE GRID FILL (H100 lane, 2026-07-26 microbench: attn qkv out_f=2048 =
13649        // 0.97 waves at the 4-warp block -> 66% of peak). The g2 twin (2 warps/block)
13650        // doubles the grid when the 4-warp launch would be sub-wave; per-row program
13651        // identical -> bit-identical. MEMRA_Q80_G2=0 reverts.
13652        if qtype == QT_Q8_0
13653            && rp
13654            && m == 1
13655            && out_f >= 64
13656            && (out_f as u32).div_ceil(ROWS_PER_BLOCK) < 4 * self.sm_count() as u32
13657            && {
13658                static G2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13659                *G2.get_or_init(|| std::env::var("MEMRA_Q80_G2").as_deref() != Ok("0"))
13660            }
13661        {
13662            let f = self.func("qmatvec_q8_0_mmvq_rp_g2");
13663            let cfg = LaunchConfig {
13664                grid_dim: ((out_f as u32).div_ceil(2), 1, 1),
13665                block_dim: (32, 2, 1),
13666                shared_mem_bytes: 0,
13667            };
13668            let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, 1i32, row_bytes as i64);
13669            let __s_b = self.gpu.stream();
13670            let mut b = __s_b.launch_builder(&f);
13671            b.arg(bytes)
13672                .arg(aq)
13673                .arg(ad)
13674                .arg(&mut *y)
13675                .arg(&inf)
13676                .arg(&outf)
13677                .arg(&mi)
13678                .arg(&rb);
13679            unsafe {
13680                b.launch(cfg)?;
13681            }
13682            if scale != 1.0 {
13683                self.scale_inplace(y, scale, out_f)?;
13684            }
13685            return Ok(());
13686        }
13687        // Multi-row-per-warp (mr2) policy, fixed since the 2026-07 sweeps (the MEMRA_MMVQ_MR
13688        // override + mr4 kernel were retired 2026-07-08 — mr4 regressed on register pressure and
13689        // crashed under rp; q4_K/q6_K mr2 measured flat, "no gain = no change"):
13690        //   NVFP4 m=1 -> mr2 (clean +1-2% on 9B: RPW acc chains hide the weight-load latency
13691        //     that pins the single-row kernel at 30-46% DRAM). Bit-identical per row.
13692        //   Q5_K m=1 -> mr2 (2026-07-05: the FR-Spec trimmed draft head is Q5_K 32768 rows = 8%
13693        //     of the 27B p3 spec wall; latency-bound like the other k-quants pre-fix).
13694        //   Q4_K/Q6_K m=1 -> single-row (mr2 measured +0.7% / flat — weight-bandwidth-bound).
13695        let mut mr: u32 = if m == 1 && (qtype == QT_NVFP4 || qtype == QT_Q5_K) {
13696            2
13697        } else {
13698            1
13699        };
13700        // Q4_0 mr (gemma trunk): DEFAULT 1 since 2026-07-13 (MEMRA_Q40_MR=2 reverts) — the
13701        // mr1 rp twin doubles the block count and wins the tail-quantization/latency battle
13702        // on every gemma model (E4B +3.75%: 198.9 vs 191.7; 26B +0.7%; 31B +0.9%; N=2-3
13703        // valid-window interleaved, bit-identical per row — same dot program).
13704        if m == 1 && qtype == QT_Q4_0 {
13705            static Q40MR: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
13706            // shape policy PROBED NEGATIVE (2026-07-13): tall-only mr1 197.2 vs
13707            // mr1-everywhere 198.7 — mr1 wins wide-output shapes too; arm removed.
13708            mr = *Q40MR.get_or_init(|| {
13709                std::env::var("MEMRA_Q40_MR")
13710                    .ok()
13711                    .and_then(|v| v.parse().ok())
13712                    .unwrap_or(1)
13713            });
13714        }
13715        // q5issue lane (2026-07-08): MEMRA_Q5K_ISSUE swaps the q5_K m=1 mmvq kernels for the
13716        // issue-reduced `_il` bodies (uint4 header/qh/qs loads + branchless scale decode —
13717        // cuts ~34 LDG.U16 + ~5 LDG.U8 + a warp-divergent scale branch per 32-elem group-row
13718        // to 5 LDG.128). Bit-identical per (token,row) to the reference kernels.
13719        // `1` = shape-aware policy (N=3 clock-locked micro-bench, mem P0, synthetic real shapes):
13720        //   out_f <= 65536 (trunk/frspec regime): il at the default mr — mr2_il -9.5%/-10.5%
13721        //     on 4096x4096/4096x8192, -3.1% on the 32768 frspec head vs the mr2-ref default;
13722        //   out_f > 65536 (the 248320-row 27B lm_head, already ~97% of the mem wall): mr2_il
13723        //     REGRESSES +22% there but mr1_il wins -2.1% vs the mr2-ref default -> force mr=1.
13724        // `2` = force il at the current mr for EVERY shape (A/B probe seam). Default OFF.
13725        let q5_mode = std::env::var("MEMRA_Q5K_ISSUE").ok();
13726        let q5_force = q5_mode.as_deref() == Some("2");
13727        // DEFAULT ON since 2026-07-08 (MEMRA_Q5K_ISSUE=0 reverts): +1.8% 9B plain e2e N=3
13728        // (128.2 -> 130.4), 27B flat (its big head is already at the mem wall), all gates green.
13729        let q5_il = qtype == QT_Q5_K
13730            && m == 1
13731            && (q5_force || q5_mode.as_deref().map(|v| v != "0").unwrap_or(true));
13732        if q5_il && !q5_force && out_f > 65536 {
13733            mr = 1;
13734        }
13735        // Q4_0 split-plane rp: mr2 default; MEMRA_Q40_MR=1 reaches the mr1 rp twin
13736        // (2026-07-13 — the tall-input/short-output tail-quantization probe).
13737        if qtype == QT_Q4_0 && rp && mr != 1 {
13738            mr = 2;
13739        }
13740        // Q8_0 rp (H100 lane): mr1 default — the q4_0 mr2 recipe MEASURED NEGATIVE on H100
13741        // (2026-07-26 N=3: mr1 186.2 vs mr2 171.5 tok/s; halving the grid on 132 SMs costs
13742        // more than 2-row ILP buys). mr2 kernel stays behind MEMRA_Q80_MR=2 for the corpus.
13743        if qtype == QT_Q8_0 && rp {
13744            static Q80MR: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
13745            mr = *Q80MR.get_or_init(|| {
13746                std::env::var("MEMRA_Q80_MR")
13747                    .ok()
13748                    .and_then(|v| v.parse().ok())
13749                    .unwrap_or(1)
13750            });
13751        }
13752        let name = match (qtype, mr, rp) {
13753            (QT_NVFP4, 2, false) => "qmatvec_nvfp4_mmvq_mr2",
13754            (QT_NVFP4, 2, true) => "qmatvec_nvfp4_mmvq_mr2_rp",
13755            (QT_NVFP4, _, true) => "qmatvec_nvfp4_mmvq_rp",
13756            (QT_Q4_0, 1, true) => "qmatvec_q4_0_mmvq_rp",
13757            (QT_Q4_0, _, true) => "qmatvec_q4_0_mmvq_mr2_rp",
13758            (QT_Q5_K, 2, _) => {
13759                if q5_il {
13760                    "qmatvec_q5_K_mmvq_mr2_il"
13761                } else {
13762                    "qmatvec_q5_K_mmvq_mr2"
13763                }
13764            }
13765            (QT_Q8_0, 2, true) => "qmatvec_q8_0_mmvq_mr2_rp",
13766            // rpca (cp.async-staged weight ring): MEASURED NEGATIVE on H100 for Q8_0
13767            // (2026-07-26 N=3: 181.8 vs plain rp 185.5 — the smem round-trip exceeds the
13768            // latency it hides for 8-bit direct-dp4a; the NVFP4 win case overlaps table
13769            // decode with half the bytes). OPT-IN via MEMRA_Q80_CA=1 for the corpus.
13770            (QT_Q8_0, _, true)
13771                if in_f % 1024 == 0 && {
13772                    static CA: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13773                    *CA.get_or_init(|| std::env::var("MEMRA_Q80_CA").as_deref() == Ok("1"))
13774                } =>
13775            {
13776                "qmatvec_q8_0_mmvq_rpca"
13777            }
13778            (QT_Q8_0, _, true) => "qmatvec_q8_0_mmvq_rp",
13779            (QT_Q8_0, _, _) => "qmatvec_q8_0_mmvq",
13780            // K-quant split-plane twins (H100 K-quant coalescing fix, 2026-08-01): the rp4
13781            // mirror routes here; GGUF layout keeps the plain kernels. rp bytes MUST never
13782            // reach a GGUF-layout kernel or vice versa.
13783            (QT_Q4_K, _, true) => "qmatvec_q4_K_mmvq_rp",
13784            (QT_Q6_K, _, true) => "qmatvec_q6_K_mmvq_rp",
13785            (QT_Q4_K, _, _) => "qmatvec_q4_K_mmvq",
13786            (QT_Q4_0, 2, false) => "qmatvec_q4_0_mmvq_mr2",
13787            (QT_Q4_0, _, false) => "qmatvec_q4_0_mmvq",
13788            (QT_Q5_K, _, _) => {
13789                if q5_il {
13790                    "qmatvec_q5_K_mmvq_il"
13791                } else {
13792                    "qmatvec_q5_K_mmvq"
13793                }
13794            }
13795            (QT_Q6_K, _, _) => "qmatvec_q6_K_mmvq",
13796            (QT_NVFP4, _, false) => "qmatvec_nvfp4_mmvq",
13797            (QT_F8_E4M3, _, _) => "qmatvec_e4m3_mmvq",
13798            _ => panic!("qmatvec_mmvq: qtype {qtype} has no MMVQ kernel"),
13799        };
13800        let f = self.func(name);
13801        // each block still has ROWS_PER_BLOCK warps; with mr rows/warp it covers ROWS_PER_BLOCK*mr rows.
13802        let rows_per_block = ROWS_PER_BLOCK * mr;
13803        let cfg = LaunchConfig {
13804            grid_dim: (
13805                (out_f as u32 + rows_per_block - 1) / rows_per_block,
13806                m as u32,
13807                1,
13808            ),
13809            block_dim: (32, ROWS_PER_BLOCK, 1), // warp-per-row (x mr rows each)
13810            shared_mem_bytes: 0,                // warp-only reduce at m=1
13811        };
13812        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
13813        let __s_b = self.gpu.stream();
13814        let mut b = __s_b.launch_builder(&f);
13815        // NVFP4 + e4m3 mmvq kernels take the macro-scale as a fused epilogue arg (applied at the
13816        // write — bit-identical to the old separate scale_inplace pass, minus one launch per matvec:
13817        // 53 scale launches/token on the 9B; for e4m3 the scale is the checkpoint's per-tensor f32
13818        // weight_scale). Other mmvq kernels keep the 8-arg signature.
13819        if qtype == QT_NVFP4 || qtype == QT_F8_E4M3 {
13820            // PDL wave-B: the nvfp4 mr2_rp single (gemma wo / generic rp singles) joins
13821            // the wave-A launch class — 9-arg flavor (fused macro-scale epilogue).
13822            if Self::pdl_on()
13823                && Self::pdl_mmvq_on()
13824                && Self::pdl_nvfp4q8_on()
13825                && name == "qmatvec_nvfp4_mmvq_mr2_rp"
13826            {
13827                use cudarc::driver::{DevicePtr, DevicePtrMut};
13828                let s = &self.gpu.stream();
13829                let (pw, _g0) = bytes.device_ptr(s);
13830                let (paq, _g1) = aq.device_ptr(s);
13831                let (pad, _g2) = ad.device_ptr(s);
13832                let (py, _g3) = y.device_ptr_mut(s);
13833                let mut ps = [
13834                    &pw as *const _ as *mut std::ffi::c_void,
13835                    &paq as *const _ as *mut _,
13836                    &pad as *const _ as *mut _,
13837                    &py as *const _ as *mut _,
13838                    &inf as *const _ as *mut _,
13839                    &outf as *const _ as *mut _,
13840                    &mi as *const _ as *mut _,
13841                    &rb as *const _ as *mut _,
13842                    &scale as *const _ as *mut _,
13843                ];
13844                unsafe {
13845                    self.launch_pdl(name, cfg.grid_dim, cfg.block_dim, &mut ps)?;
13846                }
13847                return Ok(());
13848            }
13849            b.arg(bytes)
13850                .arg(aq)
13851                .arg(ad)
13852                .arg(&mut *y)
13853                .arg(&inf)
13854                .arg(&outf)
13855                .arg(&mi)
13856                .arg(&rb)
13857                .arg(&scale);
13858            unsafe {
13859                b.launch(cfg)?;
13860            }
13861        } else if Self::pdl_on()
13862            && Self::pdl_mmvq_on()
13863            && (matches!(
13864                name,
13865                "qmatvec_q4_0_mmvq_rp" | "qmatvec_q6_K_mmvq" | "qmatvec_q6_K_mmvq_rp"
13866            ) || (Self::pdl_nvfp4q8_on()
13867                && matches!(name, "qmatvec_q8_0_mmvq_rp" | "qmatvec_q8_0_mmvq_mr2_rp")))
13868        {
13869            // PDL wave-A (2026-07-23): the two decode-hot single-matvec kernels carry
13870            // MEMRA_PDL_ENTRY — grid launches while the producer drains. ONLY the marked
13871            // names may take this launch (unmarked kernels would read unordered).
13872            {
13873                use cudarc::driver::{DevicePtr, DevicePtrMut};
13874                let s = &self.gpu.stream();
13875                let (pw, _g0) = bytes.device_ptr(s);
13876                let (paq, _g1) = aq.device_ptr(s);
13877                let (pad, _g2) = ad.device_ptr(s);
13878                let (py, _g3) = y.device_ptr_mut(s);
13879                let mut ps = [
13880                    &pw as *const _ as *mut std::ffi::c_void,
13881                    &paq as *const _ as *mut _,
13882                    &pad as *const _ as *mut _,
13883                    &py as *const _ as *mut _,
13884                    &inf as *const _ as *mut _,
13885                    &outf as *const _ as *mut _,
13886                    &mi as *const _ as *mut _,
13887                    &rb as *const _ as *mut _,
13888                ];
13889                unsafe {
13890                    self.launch_pdl(name, cfg.grid_dim, cfg.block_dim, &mut ps)?;
13891                }
13892            }
13893            if scale != 1.0 {
13894                self.scale_inplace(y, scale, m * out_f)?;
13895            }
13896        } else {
13897            b.arg(bytes)
13898                .arg(aq)
13899                .arg(ad)
13900                .arg(&mut *y)
13901                .arg(&inf)
13902                .arg(&outf)
13903                .arg(&mi)
13904                .arg(&rb);
13905            unsafe {
13906                b.launch(cfg)?;
13907            }
13908            if scale != 1.0 {
13909                self.scale_inplace(y, scale, m * out_f)?;
13910            }
13911        }
13912        Ok(())
13913    }
13914
13915    /// Test entry for the kernel_check bit-equivalence gate: run the warp-per-row MMVQ directly
13916    /// from raw weight bytes (quantize the f32 activation `x` to q8_1 internally). NVFP4 per-tensor
13917    /// macro-scale is NOT applied (caller compares bare, like qmatvec_*_fast). Mirrors qmatvec_gemm_raw.
13918    pub fn qmatvec_mmvq_raw(
13919        &self,
13920        bytes: &CudaSlice<u8>,
13921        x: &CudaSlice<f32>,
13922        m: usize,
13923        in_f: usize,
13924        out_f: usize,
13925        qtype: i32,
13926        row_bytes: usize,
13927        rp: bool,
13928    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13929        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
13930        self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, 1.0, rp)
13931    }
13932
13933    /// True if `qtype` has a batched weight-resident (`_b2`/`_b4`) matvec kernel. These mirror the
13934    /// `_mmvq` kernels but iterate the m token columns INSIDE one warp/row, so the weight bytes leave
13935    /// HBM/L2 once for m tokens (vs grid.y=m re-reading m times). The 5 daily-hot dtypes have them.
13936    pub fn batched_supports(&self, qtype: i32) -> bool {
13937        matches!(
13938            qtype,
13939            QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q6_K | QT_NVFP4 | QT_F8_E4M3 | QT_Q4_0
13940        )
13941    }
13942
13943    /// IQ4_XS trunk fast seam: MEMRA_IQ_FAST=0 reverts non-expert IQ4_XS matmuls to the Stage-A
13944    /// f32 oracle path. Default ON since 2026-08-02 (research/kat-anomaly-20260802/): the old
13945    /// opt-in default left every IQ4_XS-trunk artifact (KAT-Coder IQ4_XS: attn_qkv/attn_gate/
13946    /// ssm_out/shexp, ~0.52GB re-read per decode tick) on the oracle kernel — decode 106.7 ->
13947    /// 193.4 tok/s (x5 interleaved), pp512 228 -> 697, same bytes, via qmatvec_iq4_XS_dp4a. The
13948    /// supported artifacts carry IQ4_XS only in EXPERT banks (their own dispatch, not this seam),
13949    /// so this admission is dispatch-unchanged for every non-IQ4_XS-trunk model.
13950    pub fn iq_fast_enabled() -> bool {
13951        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13952        *ON.get_or_init(|| {
13953            std::env::var("MEMRA_IQ_FAST")
13954                .map(|v| v != "0")
13955                .unwrap_or(true)
13956        })
13957    }
13958
13959    /// b8 tier seam: MEMRA_B8=0 keeps m=5..8 on the per-m grid.y=m path (m=2..4 batched dispatch
13960    /// unaffected). Default ON — the K=4..7 spec-verify weight-read-once fix.
13961    pub fn b8_enabled() -> bool {
13962        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13963        *ON.get_or_init(|| std::env::var("MEMRA_B8").map(|v| v != "0").unwrap_or(true))
13964    }
13965
13966    /// Compile-time column batch for a runtime m: 2 -> b2, 3..4 -> b4, 5..8 -> b8.
13967    pub fn batched_mcols(m: usize) -> usize {
13968        if m == 2 {
13969            2
13970        } else if m <= 4 {
13971            4
13972        } else if m <= 8 {
13973            8
13974        } else {
13975            16
13976        }
13977    }
13978
13979    /// Kernel name for the batched matvec of `(qtype, mcols)`. mcols ∈ {2,4,8}. The b8 tier is the
13980    /// K=4..7 spec-verify fix (T=5..8): pre-b8 those T fell to grid.y=m per-row MMVQ = m full
13981    /// weight reads/launch — the measured 27B K=4 cliff (101 -> 73 tok/s at p3 despite acceptance
13982    /// holding 54%). One b8 launch reads the weight ONCE for up to 8 columns (c >= m masked).
13983    fn batched_kernel_name(qtype: i32, mcols: usize) -> Option<&'static str> {
13984        Some(match (qtype, mcols) {
13985            (QT_Q8_0, 2) => "qmatvec_q8_0_mmvq_b2",
13986            (QT_Q8_0, 4) => "qmatvec_q8_0_mmvq_b4",
13987            (QT_Q8_0, 8) => "qmatvec_q8_0_mmvq_b8",
13988            // b16 now has BOTH forms (lane/rp-on-st, 2026-08-06). It used to be rp-ONLY, which
13989            // made the q8rp mirror the exact-16 tier's admission ticket for any model carrying a
13990            // single Q8_0 matmul — measured as the FP8-ST refusal (`L0.ssm_beta qtype=0
13991            // rp4=false`, 96 t / 23.9 MiB = 0.143% of resident weight). The mirror stays a
13992            // BANDWIDTH lever on Q8_0-dominant GGUFs; it is no longer a correctness prerequisite.
13993            (QT_Q8_0, 16) => "qmatvec_q8_0_mmvq_b16",
13994            (QT_Q4_K, 2) => "qmatvec_q4_K_mmvq_b2",
13995            (QT_Q4_K, 4) => "qmatvec_q4_K_mmvq_b4",
13996            (QT_Q4_K, 8) => "qmatvec_q4_K_mmvq_b8",
13997            // b16 base + _rp (lane/rp-on-st): the 9B NVFP4 GGUF's blocker — real NVFP4 GGUFs keep
13998            // Q4_K attention next to NVFP4 MLP, and the tier's predicate is an ALL.
13999            (QT_Q4_K, 16) => "qmatvec_q4_K_mmvq_b16",
14000            (QT_Q5_K, 2) => "qmatvec_q5_K_mmvq_b2",
14001            (QT_Q5_K, 4) => "qmatvec_q5_K_mmvq_b4",
14002            (QT_Q5_K, 8) => "qmatvec_q5_K_mmvq_b8",
14003            // b16 base only (lane/rp-on-st): Q5_K has no rp twins at any width, so there is
14004            // nothing to mirror. Named by the diagnostic as `L0.wqkv_gate qtype=3` on the 9B.
14005            (QT_Q5_K, 16) => "qmatvec_q5_K_mmvq_b16",
14006            (QT_Q6_K, 2) => "qmatvec_q6_K_mmvq_b2",
14007            (QT_Q6_K, 4) => "qmatvec_q6_K_mmvq_b4",
14008            (QT_Q6_K, 8) => "qmatvec_q6_K_mmvq_b8",
14009            (QT_Q6_K, 16) => "qmatvec_q6_K_mmvq_b16",
14010            (QT_NVFP4, 2) => "qmatvec_nvfp4_mmvq_b2",
14011            (QT_NVFP4, 4) => "qmatvec_nvfp4_mmvq_b4",
14012            (QT_NVFP4, 8) => "qmatvec_nvfp4_mmvq_b8",
14013            // b16 (lane/rp-on-st): no mirror needed — NVFP4's 36 B/k32 block is already the
14014            // aligned form its own kernel walks. Unlocks the exact-16 tier for every NVFP4 model
14015            // AND for the mixed FP8-ST artifact, whose 193 NVFP4 tensors were refusing it.
14016            (QT_NVFP4, 16) => "qmatvec_nvfp4_mmvq_b16",
14017            (QT_F8_E4M3, 2) => "qmatvec_e4m3_mmvq_b2",
14018            (QT_F8_E4M3, 4) => "qmatvec_e4m3_mmvq_b4",
14019            (QT_F8_E4M3, 8) => "qmatvec_e4m3_mmvq_b8",
14020            // b16 tier (lane/rp-on-st): e4m3 needs NO split-plane mirror to reach it — its native
14021            // row-major layout is already 32B-aligned per k32 block, so the base kernel IS the
14022            // aligned form. Contrast Q8_0, whose b16 exists only as the `_rp` twin (hence q8rp).
14023            (QT_F8_E4M3, 16) => "qmatvec_e4m3_mmvq_b16",
14024            (QT_Q4_0, 2) => "qmatvec_q4_0_mmvq_b2",
14025            (QT_Q4_0, 4) => "qmatvec_q4_0_mmvq_b4",
14026            (QT_Q4_0, 8) => "qmatvec_q4_0_mmvq_b8",
14027            (QT_Q4_0, 16) => "qmatvec_q4_0_mmvq_b16",
14028            _ => return None,
14029        })
14030    }
14031
14032    /// BATCHED weight-tile-resident matvec from a PRE-QUANTIZED q8_1 activation (the m=2-8 verify/MTP
14033    /// win). One warp walks the weight row ONCE, dp4a vs all m activation columns -> weight HBM/L2
14034    /// traffic 1x for m tokens (vs grid.y=m re-reading it m times). `mcols` ∈ {2,4,8} is the
14035    /// compile-time batch; m must be <= mcols (the c >= m columns are masked in-kernel). y is
14036    /// [m, out_f] token-major. NVFP4 per-tensor macro-scale applied post
14037    /// (scale==1.0 for other dtypes -> no-op). BIT-IDENTICAL per (token,row) to qmatvec_*_mmvq.
14038    ///
14039    /// NVFP4 VARIANT DISPATCH: the batched NVFP4 kernel measured memory-LATENCY bound on the real
14040    /// 27B verify (ncu --set full, 12 steady launches: long_scoreboard 18-30 stalls/issue vs <=1.7
14041    /// for every other reason, DRAM only 41-51% active, lg_throttle 0.7, L1 hit 94% — ONE 6-LDG
14042    /// weight wavefront in flight per warp is the binding constraint, NOT bandwidth and NOT the
14043    /// column-unroll break). Two exactness-free fixes, chosen PER SHAPE from the DRAM-cold 8-copy
14044    /// msweep on all six 27B shapes (2026-07-03):
14045    ///   `pf` = next-g weight-prefetch double-buffer (48 regs, occupancy intact) — wins everywhere
14046    ///          it applies for b4 (-3..-14%), never loses;
14047    ///   `r2` = two rows/warp (67 regs -> 7 resident blocks/SM) — the bigger win (-8.5..-30%) but
14048    ///          wave-quantization-sensitive: with the grid halved to ceil(out_f/8) blocks, a
14049    ///          fractional straggler wave (waves in ~1.05-1.5) costs a full extra latency round on
14050    ///          a latency-bound kernel (27B ffn_down 640 blocks / 574 resident = 1.11 waves: +17%),
14051    ///          while <=1 wave (9B ffn_down 0.89: -30%) or >=2 waves (tail amortized; qkv 2.2:
14052    ///          -8.5%, ffn_gate 3.8: -12.5%) win. For b2, r2 wins on DEEP k-loops (in_f>=6144:
14053    ///          -8..-19%) where the 2-col body starves weight MLP hardest; pf measured negative.
14054    /// b4: r2 when waves(out_f) <= 1 (and grid fills >=half the SMs) or >= 2, else pf.
14055    /// b2: in_f>=6144 -> r2, else base.
14056    /// MEMRA_MMVQ_BV=base|pf|r2|pfr2 forces one variant everywhere (A/B + rollback seam).
14057    /// All variants BIT-IDENTICAL per (token,row): same dp4a order, scales, adg factor, reduce —
14058    /// only load issue time and the row->warp mapping change (kernel-check gates all of them).
14059    /// `rp` = the weight buffer is the A6 SPLIT-PLANE repacked layout (NVFP4 only): the same
14060    /// wave-aware auto rule applies, mapped onto the `_rp` twins (rp/rpr2/rpr2w8 mirror
14061    /// pf/r2/r2w8 — regs 44/67/64 land in the same residency classes).
14062    /// The variant the batched dispatch will pick for this (shape, m, mcols, layout) — exposed so
14063    /// gates can distinguish bit-identical variants (bit-bad==0 required) from the k-split family
14064    /// (deterministic but k-reduce-order-shifted: rel<1e-3 + run-to-run bit-identity required).
14065    /// Device SM count (cached) — grid-fill policy input.
14066    pub fn sm_count(&self) -> i32 {
14067        static SMS: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
14068        *SMS.get_or_init(|| {
14069            use cudarc::driver::sys::CUdevice_attribute_enum as A;
14070            self.gpu
14071                .ctx
14072                .attribute(A::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)
14073                .unwrap_or(82)
14074        })
14075    }
14076
14077    pub fn batched_variant(
14078        &self,
14079        _m: usize,
14080        in_f: usize,
14081        out_f: usize,
14082        qtype: i32,
14083        row_bytes: usize,
14084        mcols: usize,
14085        rp: bool,
14086    ) -> &'static str {
14087        // Q8_0 never joined the auto variant machinery (on sm_120 its only batched shapes
14088        // were tiny aux tensors). On Q8_0-trunk models the layout is the whole game: the
14089        // split-plane mirror (rp) routes to the _rp twins (H100 coalescing fix, 2026-07-26);
14090        // GGUF layout stays "base". rp bytes MUST never reach the base kernel or vice versa.
14091        if qtype == QT_Q8_0 {
14092            return if rp { "rp" } else { "base" };
14093        }
14094        static BV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
14095        let bv = *BV.get_or_init(|| match std::env::var("MEMRA_MMVQ_BV").as_deref() {
14096            Ok("base") => "base",
14097            Ok("pf") => "pf",
14098            Ok("r2") => "r2",
14099            Ok("r2w8") => "r2w8",
14100            Ok("pfr2") => "pfr2",
14101            Ok("ca") => "ca",
14102            Ok("car2") => "car2",
14103            // rp* = SPLIT-PLANE REPACKED layout kernels (A6 prototype): W must already be the
14104            // repacked buffer (msweep MSWEEP_RP harness) — never valid on GGUF-layout weights.
14105            Ok("rp") => "rp",
14106            Ok("rpr2") => "rpr2",
14107            Ok("rpr2w8") => "rpr2w8",
14108            // rpca* = cp.async software-pipelined split-plane (2026-07-05): hides the _rp
14109            // long_scoreboard load stall. rp-layout only; b4/b2 (no b8 twin).
14110            Ok("rpca") => "rpca",
14111            Ok("rpcar2") => "rpcar2",
14112            // 2026-07-06 m-small latency arc: rpsc = rpr2 + per-warp smem scale prestage (kills
14113            // the scale-plane global dependency, zero reg growth); rpms/rpmsc = m-split x2
14114            // across warp pairs (2x blocks of rpr2, column halves per warp, BIT-identical to
14115            // _rp); rpks/rpksc = k-split x2 (fastest microbench cells but k-reduce-order-shifted:
14116            // run-spec self-consistency FAILED on the 27B daily driver — verify logits must be
14117            // bit-identical to the decode path — measurement corpus ONLY, never auto).
14118            Ok("rpsc") => "rpsc",
14119            Ok("rpms") => "rpms",
14120            Ok("rpmsc") => "rpmsc",
14121            Ok("rpks") => "rpks",
14122            Ok("rpksc") => "rpksc",
14123            _ => "auto",
14124        });
14125        // cp.async ring variants need 16B-aligned rows (in_f%256==0 -> (in_f/64)*36 % 16 == 0)
14126        // and whole 32-group warp iterations (nsb%32==0 <=> in_f%1024==0). All 27B/9B trunk
14127        // shapes qualify; anything else falls back to the register variants.
14128        let ca_ok = qtype == QT_NVFP4 && (row_bytes % 16 == 0) && (in_f % 1024 == 0);
14129        // rpsc: smem scale plane fits (nsb64 <= 272) + int4-aligned staging (nsb64 % 4 == 0).
14130        // rpks/rpksc: half-plane staging alignment needs nsb64 % 8 == 0 (in_f % 512 == 0).
14131        // MEMRA_KS=0 removes the 2026-07-06 rpsc/rpks/rpksc entries from AUTO (rollback seam;
14132        // forced MEMRA_MMVQ_BV values still work).
14133        static KS_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
14134        let ks_on = *KS_ON.get_or_init(|| std::env::var("MEMRA_KS").as_deref() != Ok("0"));
14135        let sc_ok = ks_on && qtype == QT_NVFP4 && (in_f % 256 == 0) && (in_f / 64 <= 272);
14136        let ks_ok = ks_on && qtype == QT_NVFP4 && (in_f % 512 == 0) && (in_f / 64 <= 272);
14137        static SMS: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
14138        let sms = *SMS.get_or_init(|| {
14139            use cudarc::driver::sys::CUdevice_attribute_enum as A;
14140            self.gpu
14141                .ctx
14142                .attribute(A::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)
14143                .unwrap_or(82)
14144        });
14145        // k-quant r2 port (2026-07-04): q4_K/q5_K/q6_K have _r2/_r2w8 twins. ncu on the DRAM-cold
14146        // 9B msweep showed q4_K/q5_K b4 memory-latency bound like NVFP4 pre-fix (long_scoreboard
14147        // 19.6/16.4 per issue, DRAM 47.7/38.2%, L2 weight hit ~13%); q6_K lm_head is the exception
14148        // at DRAM 90-91% = wall-bound (yet r2 still wins -8%: deeper MLP raises achieved DRAM).
14149        // No _pf port (a k-quant group stages 10+ words vs NVFP4's 5 — register cost outweighs;
14150        // r2 covers the same MLP) and no rp (GGUF layout only). Q8_0 stays base: its only real
14151        // batched shapes are the tiny out_f=32 ssm_alpha/beta (8-block grids never fill one SM).
14152        // AUTO RULE = the measured winners table (differs from NVFP4's!):
14153        //   r2w8 NEVER in auto — the reg squeeze (72 -> 64 regs = stack spill) loses to unbounded
14154        //     r2 on every measured k-quant cell, incl. the wave-crossing lm_heads (q6_K 1316 vs
14155        //     r2 1258us) — kernels kept behind the force seam for the corpus;
14156        //   q4_K: r2 whenever the halved grid fills the SMs (blocks >= 4*SMs), INCLUDING the
14157        //     1.05-2.0 straggler window where NVFP4's r2 lost (qkv 1.78 waves: r2 -15% here; the
14158        //     k-quant base kernel leaves more latency on the table than a straggler wave costs);
14159        //   q5_K/q6_K: r2 only at waves >= 2 (the 248320-row lm_heads, 48+ waves: q6_K -8%, q5_K
14160        //     -2%); mid shapes measured base-or-flat (q5_K qkv 49.1 base vs 49.7 r2, attn_gate
14161        //     flat, attn_k base) — the 5/6-bit two-stream unpack makes r2's staging pricier.
14162        //   b2 same table with 8-row blocks: q4_K r2 when filled (-3..-22% all measured shapes),
14163        //     q5_K/q6_K r2 at waves >= 2 (27B lm_head -2.9%; 9B q6_K flat, harmless).
14164        let kq_r2 = matches!(qtype, QT_Q4_K | QT_Q5_K | QT_Q6_K);
14165        // MEMRA_KQ_BV=base|r2|r2w8 forces the k-quant variant WITHOUT touching the NVFP4 dispatch
14166        // (MEMRA_MMVQ_BV is global — an interleaved k-quant-only e2e A/B needs this narrower seam).
14167        static KQBV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
14168        let kq_bv = *KQBV.get_or_init(|| match std::env::var("MEMRA_KQ_BV").as_deref() {
14169            Ok("base") => "base",
14170            Ok("r2") => "r2",
14171            Ok("r2w8") => "r2w8",
14172            _ => "auto",
14173        });
14174        let variant: &'static str = if qtype == QT_Q4_0 {
14175            // Q4_0 r2 (gemma verify trunk, 2026-07-10): shared activation loads + the
14176            // row-independent ones-sum computed once per (col,group) for 2 rows. Same
14177            // fill rule as q4_K: r2 when the halved grid still fills the SMs.
14178            static Q40BV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
14179            let q40 = *Q40BV.get_or_init(|| match std::env::var("MEMRA_Q40_BV").as_deref() {
14180                // ms/sm/la = force-only measurement seams (ALL FLAT/NEGATIVE 2026-07-13,
14181                // never auto): m-split flat (nvcc keeps 72 regs); smem-slab −11% (staging
14182                // + syncs cost more than the stalls, bank-pad made no difference);
14183                // register load-ahead flat (nvcc already reorders). The b-tier limiter
14184                // is still unidentified — see the jsonl row.
14185                Ok("base") => "base",
14186                Ok("r2") => "r2",
14187                Ok("ms") => "ms",
14188                Ok("sm") => "sm",
14189                Ok("la") => "la",
14190                _ => "auto",
14191            });
14192            let v = if q40 != "auto" {
14193                q40
14194            } else if (out_f as u32).div_ceil(8) >= 4 * sms as u32 {
14195                "r2"
14196            } else {
14197                "base"
14198            };
14199            // split-plane mirror twins (2026-07-10): same fill rule, _rp names.
14200            // (m-split r2 pair twin PROBED FLAT 2026-07-13 — nvcc kept 72 regs either way
14201            // and the limiter is the per-column activation load chain (long_scoreboard
14202            // 42.5%), not occupancy; arm killed per doctrine, jsonl row is the record.)
14203            if rp {
14204                match v {
14205                    "ms" => "r2ms_rp",
14206                    "sm" => "r2sm_rp",
14207                    "la" => "r2la_rp",
14208                    "r2" => "r2_rp",
14209                    _ => "rp",
14210                }
14211            } else if matches!(v, "ms" | "sm" | "la") {
14212                "r2"
14213            } else {
14214                v
14215            }
14216        } else if qtype != QT_NVFP4 && !kq_r2 {
14217            "base"
14218        } else if kq_r2 && rp {
14219            // K-quant split-plane mirror (2026-08-01): only the plain _rp batched twins are
14220            // compiled for q4_K/q6_K — rp is a LAYOUT, it must survive every heuristic
14221            // (split-plane bytes through a GGUF-layout kernel = NaN). q5_K never mirrors.
14222            "rp"
14223        } else if kq_r2 {
14224            // k-quant r2w8 only exists at b4 (b2_r2 already 8-resident; b8 has no w8 twin) ->
14225            // mcols != 4 forced r2w8 falls to unbounded r2.
14226            if kq_bv != "auto" {
14227                if kq_bv == "r2w8" && mcols != 4 {
14228                    "r2"
14229                } else {
14230                    kq_bv
14231                }
14232            } else if bv != "auto" {
14233                match bv {
14234                    "r2" | "pfr2" | "rpr2" | "car2" => "r2",
14235                    "r2w8" | "rpr2w8" => {
14236                        if mcols != 4 {
14237                            "r2"
14238                        } else {
14239                            "r2w8"
14240                        }
14241                    }
14242                    _ => "base", // base/pf/ca/rp forced -> base (no such k-quant kernels)
14243                }
14244            } else {
14245                let blocks = (out_f + 7) / 8;
14246                let waves = blocks as f64 / (7 * sms as usize) as f64;
14247                let filled = blocks >= 4 * sms as usize;
14248                let use_r2 = if qtype == QT_Q4_K {
14249                    filled
14250                } else {
14251                    waves >= 2.0
14252                };
14253                if use_r2 { "r2" } else { "base" }
14254            }
14255        } else if bv != "auto" {
14256            // r2w8 only exists for b4/b8 (the b2_r2 kernel is already 8-blocks-resident at 60 regs).
14257            // ca/car2 need the alignment gate AND have no b8 twins; pfr2 has no b8 twin either —
14258            // unsupported (shape, mcols) combos fall back to pf/r2.
14259            // On rp buffers, forced legacy names map to their rp twins (layout law).
14260            let v = if bv == "r2w8" && mcols == 2 {
14261                "r2"
14262            } else if bv == "ca" && (!ca_ok || mcols == 8) {
14263                "pf"
14264            } else if bv == "car2" && (!ca_ok || mcols == 8) {
14265                "r2"
14266            } else if bv == "pfr2" && mcols == 8 {
14267                "r2"
14268            } else if (bv == "rpr2w8" || bv == "rpr2") && mcols == 2 {
14269                "rpr2"
14270            }
14271            // rpca* has no b8 twin (falls to rpr2w8/rpr2); needs the ca alignment gate.
14272            else if (bv == "rpca" || bv == "rpcar2") && (!ca_ok || mcols == 8) {
14273                if mcols == 8 { "rpr2w8" } else { "rpr2" }
14274            } else if bv == "rpcar2" && mcols == 2 {
14275                "rpca"
14276            }
14277            // rpsc/rpmsc/rpks* gate on smem-fit + alignment; fall to rpr2 outside it
14278            // (rpms has no smem and no alignment need — always valid on rp buffers).
14279            else if (bv == "rpsc" || bv == "rpmsc") && !sc_ok {
14280                "rpr2"
14281            } else if (bv == "rpks" || bv == "rpksc") && !ks_ok {
14282                "rpr2"
14283            } else {
14284                bv
14285            };
14286            if rp {
14287                match v {
14288                    "base" | "pf" | "ca" | "rp" => "rp",
14289                    "r2" | "pfr2" | "car2" | "rpr2" => "rpr2",
14290                    "r2w8" | "rpr2w8" => {
14291                        if mcols == 2 {
14292                            "rpr2"
14293                        } else {
14294                            "rpr2w8"
14295                        }
14296                    }
14297                    other => other, // rpca/rpcar2/rpsc/rpks/rpksc pass through (already rp-layout)
14298                }
14299            } else {
14300                v
14301            }
14302        } else if mcols == 8 {
14303            // b8 AUTO (2026-07-06 m-small latency arc, g7e DRAM-cold rp msweep m=5/6/8 all five
14304            // 27B shapes): rpsc — the rpr2w8 schedule with the warp's scale rows prestaged to
14305            // smem, leaving ONE global dependency (the quant stream) in the k-loop at zero reg
14306            // growth. BIT-identical to rpr2w8 and wins or ties EVERY b8 cell: ffn_gate m5
14307            // 50.7->46.9 m8 64.1->57.1 (-11%), qkv m8 34.6->33.0, ssm_out m8 29.7->28.8,
14308            // attn_gate m8 26.9->26.1, ffn_down m5 58.2->56.9. The faster split-grid twins are
14309            // OUT: rpksc (k-split, ffn_down m5 -21%) broke run-spec self-consistency (k-reduce
14310            // order shifts verify argmax at tie margins — verify must stay bit-identical to the
14311            // m=1 decode chain); rpmsc (m-split, bit-identical) measured NEGATIVE everywhere
14312            // (twin warp's duplicated weight stream: ffn_down m5 85.7 vs 56.9).
14313            if rp {
14314                if sc_ok { "rpsc" } else { "rpr2w8" }
14315            } else {
14316                "r2w8"
14317            }
14318        } else if mcols >= 4 {
14319            // r2 runs 7 resident blocks/SM (67 regs); its __launch_bounds__(128,8) twin `r2w8`
14320            // (64 regs) runs 8. grid = ceil(out_f/8) for both. rp twins land in the same
14321            // residency classes (rp 44 regs ~ pf-class occupancy, rpr2 67, rpr2w8 64).
14322            let blocks = (out_f + 7) / 8;
14323            let r7 = 7 * sms as usize;
14324            let r8 = 8 * sms as usize;
14325            let waves = blocks as f64 / r7 as f64;
14326            let filled = blocks >= 4 * sms as usize;
14327            // 2026-07-06 m-small latency arc: b4 keeps the wave rule (rpms/rpmsc measured
14328            // flat-to-negative at m=3/4 on every shape — the m-split twin duplicates the weight
14329            // stream; rpsc b4 also negative on r2-class picks, ffn_down m4 51.1 vs 46.5).
14330            if filled && blocks.div_ceil(r8) < blocks.div_ceil(r7) {
14331                // the extra residency drops the INTEGER wave count -> the straggler wave a
14332                // latency-bound kernel pays in full disappears (ffn_down 1.11 -> 0.98 waves:
14333                // 112.5 -> 81.6us, beats pf 90.1; qkv 2.23 -> 1.95: 58.1 -> 51.1).
14334                if rp { "rpr2w8" } else { "r2w8" }
14335            } else if waves >= 2.0 || (waves <= 1.0 && filled) {
14336                // tail amortized (>=2 waves) or single wave: unbounded r2 (no reg-squeeze tax —
14337                // gate/up 81.1 vs 83.9 bounded, attn_q 61.0 vs 63.4).
14338                if rp { "rpr2" } else { "r2" }
14339            } else {
14340                // fractional straggler-wave window with no crossing, or grid too small to fill
14341                // the SMs (tiny out_f<=1024 shapes want max row-parallelism): prefetch variant
14342                // (rp = the r1 split-plane twin — measured the attn_gate winner, 35.4 vs pf 36.4).
14343                if rp { "rp" } else { "pf" }
14344            }
14345        } else if in_f >= 6144 {
14346            // b2 deep-k (2026-07-06): every new twin measured flat-to-negative here (rpms 44.1
14347            // vs rpr2 40.8 ffn_down; rpsc 43.6; the winning rpks is banned on k-order) — rpr2
14348            // stays.
14349            if rp { "rpr2" } else { "r2" }
14350        } else if rp {
14351            // b2 shallow-k: qkv (out_f=10240, 0.97 waves at 7-resident) is the one measured cell
14352            // where the r2-schedule scale-prestage twin beats the r1 rp pick (24.7 vs 28.9us
14353            // -15%); the wider (ffn_gate 1.65 waves) and smaller (attn_gate 0.58) shapes LOSE
14354            // (41.8 vs 38.2 / 16.6 vs 14.6) — gate on the single-wave window.
14355            let waves = ((out_f + 7) / 8) as f64 / (7 * sms as usize) as f64;
14356            if sc_ok && waves >= 0.9 && waves <= 1.1 {
14357                "rpsc"
14358            } else {
14359                "rp"
14360            }
14361        } else {
14362            "base"
14363        };
14364        variant
14365    }
14366
14367    pub fn qmatvec_mmvq_batched(
14368        &self,
14369        bytes: &CudaSlice<u8>,
14370        aq: &CudaSlice<i8>,
14371        ad: &CudaSlice<f32>,
14372        m: usize,
14373        in_f: usize,
14374        out_f: usize,
14375        qtype: i32,
14376        row_bytes: usize,
14377        mcols: usize,
14378        scale: f32,
14379        rp: bool,
14380    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14381        const ROWS_PER_BLOCK: u32 = 4;
14382        // TUNE SEAM (H100 lane): MEMRA_BVAR forces the batched-variant pick for the whole
14383        // process — the auto heuristics were tuned on sm_120 (82 SMs / 858 GB/s) and the
14384        // sm_90a re-tune sweeps this seam empirically. Layout variants stay safe: an rp
14385        // weight keeps its rp-layout kernel family regardless of the override.
14386        let forced: Option<&'static str> = {
14387            static V: std::sync::OnceLock<Option<String>> = std::sync::OnceLock::new();
14388            V.get_or_init(|| std::env::var("MEMRA_BVAR").ok())
14389                .as_deref()
14390                .map(|s| Box::leak(s.to_string().into_boxed_str()) as &'static str)
14391        };
14392        let variant = match forced {
14393            Some(v) if !rp || v.contains("rp") => v,
14394            _ => self.batched_variant(m, in_f, out_f, qtype, row_bytes, mcols, rp),
14395        };
14396        let base_name = Self::batched_kernel_name(qtype, mcols).ok_or_else(|| {
14397            format!("qmatvec_mmvq_batched: no kernel for qtype {qtype} mcols {mcols}")
14398        })?;
14399        // b16 tier (t=9..16 verify): only base/_rp b16 kernels are compiled — the b2..b8
14400        // per-shape perf variants (r2/pf/...) do not apply at this width. rp is a LAYOUT,
14401        // not a perf variant: it must survive (base kernel on split-plane bytes = NaN).
14402        let variant = if mcols == 16 {
14403            if rp { "rp" } else { "base" }
14404        } else {
14405            variant
14406        };
14407        // EXACT-WIDTH b5/b6/b7 twins (lane/vt-fixes fix 1, 2026-08-03): the b8 kernels
14408        // allocate acc[WROWS][8] at ANY m, so T=5..7 verify paid the full 8-wide register
14409        // tax — the measured T=4->5 cliff. The same template at MCOLS=m runs the identical
14410        // per-(token,row) chain (columns c >= m never execute in either form) ->
14411        // BIT-IDENTICAL to the b8 launch. NVFP4 split-plane only (the sm_120 default trunk);
14412        // covers both b8 auto schedules (rpsc, rpr2w8). MEMRA_B567=0 rollback.
14413        static B567: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
14414        let b567 = *B567.get_or_init(|| std::env::var("MEMRA_B567").as_deref() != Ok("0"));
14415        if b567
14416            && qtype == QT_NVFP4
14417            && rp
14418            && mcols == 8
14419            && (5..=7).contains(&m)
14420            && matches!(variant, "rpsc" | "rpr2w8")
14421        {
14422            let f = self.func(&format!("qmatvec_nvfp4_mmvq_b{m}_{variant}"));
14423            let rows_per_block = ROWS_PER_BLOCK * 2; // r2-class schedules: 2 rows/warp
14424            let mut y = self.alloc_uninit::<f32>(m * out_f)?;
14425            let cfg = LaunchConfig {
14426                grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
14427                block_dim: (32, ROWS_PER_BLOCK, 1),
14428                shared_mem_bytes: 0,
14429            };
14430            let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
14431            let __s_b = self.gpu.stream();
14432            let mut b = __s_b.launch_builder(&f);
14433            b.arg(bytes)
14434                .arg(aq)
14435                .arg(ad)
14436                .arg(&mut y)
14437                .arg(&inf)
14438                .arg(&outf)
14439                .arg(&mi)
14440                .arg(&rb);
14441            unsafe {
14442                b.launch(cfg)?;
14443            }
14444            if scale != 1.0 {
14445                self.scale_inplace(&mut y, scale, m * out_f)?;
14446            }
14447            return Ok(y);
14448        }
14449        let (name, rows_per_block): (std::borrow::Cow<'static, str>, u32) = match variant {
14450            "base" => (base_name.into(), ROWS_PER_BLOCK),
14451            "pf" => (format!("{base_name}_pf").into(), ROWS_PER_BLOCK),
14452            "ca" => (format!("{base_name}_ca").into(), ROWS_PER_BLOCK),
14453            "rp" => (format!("{base_name}_rp").into(), ROWS_PER_BLOCK),
14454            "rpca" => (format!("{base_name}_rpca").into(), ROWS_PER_BLOCK), // 1 row/warp cp.async
14455            // split families: 2 warp-pairs x 2 rows = 4 rows/block (the k-range or column set
14456            // splits across the pair's two warps; grid.x doubles vs rpr2 at the same regs).
14457            "rpks" => (format!("{base_name}_rpks").into(), ROWS_PER_BLOCK),
14458            "rpksc" => (format!("{base_name}_rpksc").into(), ROWS_PER_BLOCK),
14459            "rpms" => (format!("{base_name}_rpms").into(), ROWS_PER_BLOCK),
14460            "rpmsc" => (format!("{base_name}_rpmsc").into(), ROWS_PER_BLOCK),
14461            "r2ms_rp" => (format!("{base_name}_r2ms_rp").into(), ROWS_PER_BLOCK),
14462            "r2sm_rp" => (format!("{base_name}_r2sm_rp").into(), ROWS_PER_BLOCK * 2),
14463            "r2la_rp" => (format!("{base_name}_r2la_rp").into(), ROWS_PER_BLOCK * 2),
14464            v => (format!("{base_name}_{v}").into(), ROWS_PER_BLOCK * 2), // r2-class: 2 rows/warp
14465        };
14466        debug_assert!(
14467            !rp || name.contains("_rp"),
14468            "rp weight dispatched to a GGUF-layout kernel"
14469        );
14470        let f = self.func(&name);
14471        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
14472        // r2sm_rp: [MCOLS][32 blk][8 int] activation slab + [MCOLS][32] f32 scales.
14473        let smem = if name.contains("_r2sm_rp") {
14474            (mcols * 32 * 9 * 4 + mcols * 32 * 4) as u32
14475        } else {
14476            0
14477        };
14478        let cfg = LaunchConfig {
14479            grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
14480            block_dim: (32, ROWS_PER_BLOCK, 1),
14481            shared_mem_bytes: smem,
14482        };
14483        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
14484        let __s_b = self.gpu.stream();
14485        let mut b = __s_b.launch_builder(&f);
14486        b.arg(bytes)
14487            .arg(aq)
14488            .arg(ad)
14489            .arg(&mut y)
14490            .arg(&inf)
14491            .arg(&outf)
14492            .arg(&mi)
14493            .arg(&rb);
14494        unsafe {
14495            b.launch(cfg)?;
14496        }
14497        if scale != 1.0 {
14498            self.scale_inplace(&mut y, scale, m * out_f)?;
14499        }
14500        Ok(y)
14501    }
14502
14503    /// BATCHED weight-tile-resident matvec from raw weight bytes (quantizes the f32 activation `x` to
14504    /// q8_1 internally; macro-scale NOT applied — caller compares bare, like qmatvec_*_fast). For the
14505    /// kernel_check bit-equivalence gate. `mcols` ∈ {2,4,8}. Works for Q8_0/Q4_K/Q5_K/Q6_K/NVFP4.
14506    pub fn qmatvec_batched_raw(
14507        &self,
14508        bytes: &CudaSlice<u8>,
14509        x: &CudaSlice<f32>,
14510        m: usize,
14511        in_f: usize,
14512        out_f: usize,
14513        qtype: i32,
14514        row_bytes: usize,
14515        mcols: usize,
14516        rp: bool,
14517    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14518        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
14519        self.qmatvec_mmvq_batched(
14520            bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, mcols, 1.0, rp,
14521        )
14522    }
14523
14524    /// Back-compat NVFP4-only batched raw launcher (used by older gates). Delegates to the generic one.
14525    pub fn qmatvec_nvfp4_batched_raw(
14526        &self,
14527        bytes: &CudaSlice<u8>,
14528        x: &CudaSlice<f32>,
14529        m: usize,
14530        in_f: usize,
14531        out_f: usize,
14532        row_bytes: usize,
14533        mcols: usize,
14534        rp: bool,
14535    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14536        self.qmatvec_batched_raw(bytes, x, m, in_f, out_f, QT_NVFP4, row_bytes, mcols, rp)
14537    }
14538
14539    /// Stage-C FP4 gate (MEMRA_FP4): if `w` is an NVFP4 weight with in_f%64==0, run the native mxf4
14540    /// block-scale GEMM and apply the per-tensor macro-scale, returning Some(y). Else None (caller
14541    /// falls through to the int8 GEMM / dp4a). Strict opt-in over the proven int8 path; m>=16 only.
14542    fn try_fp4_gemm(
14543        &self,
14544        w: &crate::model::GpuTensor,
14545        x: &CudaSlice<f32>,
14546        m: usize,
14547        in_f: usize,
14548        out_f: usize,
14549    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
14550        use crate::model::GpuTensor;
14551        if cfg!(memra_portable_cuda) {
14552            return Ok(None);
14553        }
14554        if std::env::var("MEMRA_FP4").is_err() {
14555            return Ok(None);
14556        }
14557        // CUTLASS prefill branch (m>=128 + MEMRA_FP4_CUTLASS + a repacked CutlassWeight present): route
14558        // to the CUTLASS sm120 NVFP4 GEMM, folding the per-tensor macro-scale into the epilogue alpha
14559        // (1/scale) — no post-matmul scale_inplace. Decode (m<128) and the m∈[16,128) middle band keep
14560        // the hand-roll below: CUTLASS's 128-row M-tile wastes work under 128.
14561        // The hand-roll applies the per-tensor macro-scale as a POST-matmul MULTIPLY (scale_inplace(y,
14562        // scale)); CUTLASS's epilogue does D = alpha * (A@B^T), so alpha == scale reproduces it exactly
14563        // (NOT 1/scale — the plan sketch had this inverted; the kernel_check arm gates it). scale==1.0
14564        // for the common no-macro-scale case.
14565        #[cfg(memra_cutlass)]
14566        if m >= 128 && std::env::var("MEMRA_FP4_CUTLASS").is_ok() {
14567            if let GpuTensor::Quant {
14568                bytes,
14569                qtype,
14570                scale,
14571                row_bytes,
14572                cutlass,
14573                ..
14574            } = w
14575            {
14576                if *qtype == QT_NVFP4 && in_f % 64 == 0 {
14577                    if let Some(cw) = cutlass {
14578                        // Resident fast path: load-time-repacked B + swizzled SFB (no per-call repack).
14579                        let y = self.cutlass_fp4_gemm(
14580                            &cw.b_packed,
14581                            &cw.sfb_swizzled,
14582                            x,
14583                            *scale,
14584                            m,
14585                            out_f,
14586                            in_f,
14587                        )?;
14588                        return Ok(Some(y));
14589                    } else if std::env::var("MEMRA_FP4_CUTLASS_OTF").is_ok() {
14590                        // On-the-fly repack (MEMRA_FP4_CUTLASS_OTF): de-interleave + swizzle the B operand
14591                        // from raw bytes per prefill call. No resident doubling of the NVFP4 weight VRAM
14592                        // (the load-time repack ~doubles it) — needed for models that don't fit the
14593                        // resident path (e.g. the 27B on 24GB). Slower (per-call repack) but argmax-exact.
14594                        let (b_packed, sfb_sw) =
14595                            self.build_cutlass_weight(bytes, out_f, in_f, *row_bytes)?;
14596                        let y =
14597                            self.cutlass_fp4_gemm(&b_packed, &sfb_sw, x, *scale, m, out_f, in_f)?;
14598                        return Ok(Some(y));
14599                    }
14600                }
14601            }
14602        }
14603        if let GpuTensor::Quant {
14604            bytes,
14605            qtype,
14606            row_bytes,
14607            scale,
14608            rp,
14609            ..
14610        } = w
14611        {
14612            // A6: the hand-rolled W4A4 mxf4 GEMM reads 36B GGUF blocks — no rp port (MEMRA_FP4 is
14613            // an opt-in accuracy tradeoff); repacked tensors fall through to the int8 GEMM.
14614            if *qtype == QT_NVFP4 && in_f % 64 == 0 && !*rp {
14615                let y =
14616                    self.qmatvec_gemm_nvfp4_fp4(bytes, x, m, in_f, out_f, *row_bytes, *scale)?;
14617                return Ok(Some(y));
14618            }
14619        }
14620        Ok(None)
14621    }
14622
14623    /// rms_norm + fused fp16 twin (task #14): f32 output verbatim `rms_norm` + the fp16
14624    /// copy the f16-mirror GEMM group would otherwise produce with a standalone convert
14625    /// launch. BIT-IDENTICAL end-to-end (same reduction, same __float2half values).
14626    pub fn rms_norm_f16out(
14627        &self,
14628        x: &CudaSlice<f32>,
14629        w: &CudaSlice<f32>,
14630        dst: &mut CudaSlice<f32>,
14631        dst16: &mut CudaSlice<u8>,
14632        ncols: usize,
14633        nrows: usize,
14634        eps: f32,
14635    ) -> Result<(), Box<dyn std::error::Error>> {
14636        let f = self.func("rms_norm_f16out_f32");
14637        let cfg = LaunchConfig {
14638            grid_dim: (nrows as u32, 1, 1),
14639            block_dim: (rms_block(), 1, 1),
14640            shared_mem_bytes: 0,
14641        };
14642        let (nc, e) = (ncols as i32, eps);
14643        let __s_b = self.gpu.stream();
14644        let mut b = __s_b.launch_builder(&f);
14645        b.arg(x).arg(w).arg(dst).arg(dst16).arg(&nc).arg(&e);
14646        unsafe {
14647            b.launch(cfg)?;
14648        }
14649        Ok(())
14650    }
14651
14652    /// add+norm(+f16out) fusion for the prefill trunk (round 28; add_rms_norm precedent —
14653    /// bit-identical to add_f32 -> rms_norm_f16out). block_dim matches rms_norm_f16out's.
14654    #[allow(clippy::too_many_arguments)]
14655    pub fn add_rms_norm_f16out(
14656        &self,
14657        a: &CudaSlice<f32>,
14658        b: &CudaSlice<f32>,
14659        w: &CudaSlice<f32>,
14660        res: &mut CudaSlice<f32>,
14661        dst: &mut CudaSlice<f32>,
14662        dst16: &mut CudaSlice<u8>,
14663        ncols: usize,
14664        nrows: usize,
14665        eps: f32,
14666    ) -> Result<(), Box<dyn std::error::Error>> {
14667        let f = self.func("add_rms_norm_f16out_f32");
14668        let cfg = LaunchConfig {
14669            grid_dim: (nrows as u32, 1, 1),
14670            block_dim: (rms_block(), 1, 1),
14671            shared_mem_bytes: 0,
14672        };
14673        let (nc, e) = (ncols as i32, eps);
14674        let __s_lb = self.gpu.stream();
14675        let mut lb = __s_lb.launch_builder(&f);
14676        lb.arg(a)
14677            .arg(b)
14678            .arg(w)
14679            .arg(res)
14680            .arg(dst)
14681            .arg(dst16)
14682            .arg(&nc)
14683            .arg(&e);
14684        unsafe {
14685            lb.launch(cfg)?;
14686        }
14687        Ok(())
14688    }
14689
14690    /// matmul_group with a PRE-EMITTED fp16 activation (task #14: the producer norm fused
14691    /// the convert). Mirror-less members fall back to `matmul` on the f32 activation.
14692    pub fn matmul_group_xh(
14693        &self,
14694        ws: &[&crate::model::GpuTensor],
14695        x: &CudaSlice<f32>,
14696        xh: &CudaSlice<u8>,
14697        m: usize,
14698    ) -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
14699        let mut out = Vec::with_capacity(ws.len());
14700        let in_f = ws[0].in_features();
14701        for w in ws {
14702            if w.in_features() == in_f && m >= 16 && !self.verify_exact_on() {
14703                if let Some(y) = self.try_f16_gemm_pre(w, xh, m)? {
14704                    out.push(y);
14705                    continue;
14706                }
14707            }
14708            out.push(self.matmul(w, x, m)?);
14709        }
14710        Ok(out)
14711    }
14712
14713    /// task #14 pad-proofing: zero beta/g_log at rows >= len_d[0] (pads become identity
14714    /// GDN steps). Layouts [T, H].
14715    pub fn gdn_pad_mask(
14716        &self,
14717        beta: &mut CudaSlice<f32>,
14718        g_log: &mut CudaSlice<f32>,
14719        len_d: &CudaSlice<i32>,
14720        h: usize,
14721        t: usize,
14722    ) -> Result<(), Box<dyn std::error::Error>> {
14723        let f = self.func("gdn_pad_mask_f32");
14724        let cfg = LaunchConfig::for_num_elems((t * h) as u32);
14725        let (hi, ti) = (h as i32, t as i32);
14726        let __s_b = self.gpu.stream();
14727        let mut b = __s_b.launch_builder(&f);
14728        b.arg(beta).arg(g_log).arg(len_d).arg(&hi).arg(&ti);
14729        unsafe {
14730            b.launch(cfg)?;
14731        }
14732        Ok(())
14733    }
14734
14735    /// task #14 pad-proofing: dst[ncols] = src row (len_d[0]-1) — device-indexed last-row
14736    /// gather for the padded prime graph's h_seed/hlast.
14737    pub fn row_gather_dev(
14738        &self,
14739        src: &CudaSlice<f32>,
14740        dst: &mut CudaSlice<f32>,
14741        len_d: &CudaSlice<i32>,
14742        ncols: usize,
14743    ) -> Result<(), Box<dyn std::error::Error>> {
14744        let f = self.func("row_gather_dev_f32");
14745        let cfg = LaunchConfig::for_num_elems(ncols as u32);
14746        let nc = ncols as i32;
14747        let __s_b = self.gpu.stream();
14748        let mut b = __s_b.launch_builder(&f);
14749        b.arg(src).arg(dst).arg(len_d).arg(&nc);
14750        unsafe {
14751            b.launch(cfg)?;
14752        }
14753        Ok(())
14754    }
14755
14756    /// Grouped matmul: several weights consuming ONE activation (hybrid layers: the GDN
14757    /// 4-tuple wqkv/gate/beta/alpha, attention q/k/v, ffn gate/up). Semantics identical to
14758    /// calling `matmul` per weight; the f16-mirror arm converts the activation ONCE for the
14759    /// whole group instead of once per GEMM (the standalone converts were ~250 launches/prime
14760    /// of small-kernel gap fuel — nsys 2026-07-26). Any member without a mirror (or with a
14761    /// different in_f) falls back to its own `matmul` — behavior unchanged.
14762    pub fn matmul_group(
14763        &self,
14764        ws: &[&crate::model::GpuTensor],
14765        x: &CudaSlice<f32>,
14766        m: usize,
14767    ) -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
14768        use crate::model::GpuTensor;
14769        let mut out = Vec::with_capacity(ws.len());
14770        let any_mirror = ws
14771            .iter()
14772            .any(|w| matches!(w, GpuTensor::Quant { f16: Some(_), .. }));
14773        if m >= 16 && any_mirror && !self.verify_exact_on() {
14774            let in_f = ws[0].in_features();
14775            let xh = self.f16_act(x, m * in_f, in_f)?;
14776            for w in ws {
14777                if w.in_features() == in_f {
14778                    if let Some(y) = self.try_f16_gemm_pre(w, &xh, m)? {
14779                        out.push(y);
14780                        continue;
14781                    }
14782                }
14783                out.push(self.matmul(w, x, m)?);
14784            }
14785            return Ok(out);
14786        }
14787        for w in ws {
14788            out.push(self.matmul(w, x, m)?);
14789        }
14790        Ok(out)
14791    }
14792
14793    /// Cross-request grouped matmul (task #13): run ONE projection group over the
14794    /// CONCATENATION of several sequences' activations (m = sum of per-seq rows — the
14795    /// GEMM-batch win vLLM gets from continuous batching), then split each output back
14796    /// into per-seq buffers. Zero view plumbing: gather/scatter are stream-ordered D2D
14797    /// copies (~us at prime sizes). NUMERIC CONFIG NOTE: a GEMM at m=sum tiles K
14798    /// differently than per-seq GEMMs — argmax-gated like every prefill GEMM change.
14799    pub fn matmul_group_multi(
14800        &self,
14801        ws: &[&crate::model::GpuTensor],
14802        xs: &[&CudaSlice<f32>],
14803        ms: &[usize],
14804    ) -> Result<Vec<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
14805        assert_eq!(xs.len(), ms.len());
14806        let in_f = ws[0].in_features();
14807        let total: usize = ms.iter().sum();
14808        let mut xcat = self.uninit(total * in_f)?;
14809        let mut off = 0usize;
14810        for (x, &m) in xs.iter().zip(ms) {
14811            self.copy_into(&mut xcat, off * in_f, x, m * in_f)?;
14812            off += m;
14813        }
14814        let ys = self.matmul_group(ws, &xcat, total)?;
14815        let mut out: Vec<Vec<CudaSlice<f32>>> = (0..xs.len()).map(|_| Vec::new()).collect();
14816        for (w, y) in ws.iter().zip(ys) {
14817            let out_f = w.out_features();
14818            let mut off = 0usize;
14819            for (s, &m) in ms.iter().enumerate() {
14820                let mut ys_s = self.uninit(m * out_f)?;
14821                let src = y.slice(off * out_f..(off + m) * out_f);
14822                self.gpu.stream().memcpy_dtod(&src, &mut ys_s)?;
14823                out[s].push(ys_s);
14824                off += m;
14825            }
14826        }
14827        Ok(out)
14828    }
14829
14830    /// True if `w`'s qtype has a batched tensor-core GEMM kernel (the prefill T>1 root fix).
14831    /// Only the 4 daily-hot dtypes: Q8_0, Q4_K, Q6_K, NVFP4. NVFP4 needs in_f % 64 == 0.
14832    /// DEFAULT-ON (2026-06-28): measured pp512 9B-NVFP4 = 1413 tok/s WITH this GEMM vs 298 with the
14833    /// dp4a fallback (4.7x) AND MORE accurate (prefill logit maxdiff 0.159 vs dp4a 0.55, both argmax
14834    /// MATCH). The int8 tensor-core GEMM is unconditional (its historical MEMRA_GEMM opt-in gate
14835    /// shipped with Phase 0 — mma + smem swizzle + cp.async — and was removed). Prefill-only
14836    /// (m>=GEMM_M_THRESHOLD); m=1 decode keeps dp4a/MMVQ (this returns true but matmul only calls it
14837    /// at m>=threshold). Portable CUDA targets always use the correctness fallback; on sm_120a,
14838    /// MEMRA_NO_GEMM forces that same dp4a fallback (the bit-reference).
14839    pub fn gemm_supports(&self, w: &crate::model::GpuTensor) -> bool {
14840        use crate::model::GpuTensor;
14841        if !legacy_quant_gemm_allowed(
14842            cfg!(memra_portable_cuda),
14843            cfg!(memra_hopper_mma),
14844            std::env::var_os("MEMRA_NO_GEMM").is_some(),
14845        ) {
14846            return false;
14847        }
14848        match w {
14849            GpuTensor::Quant { qtype, .. } => {
14850                matches!(*qtype, QT_Q8_0 | QT_Q4_K | QT_Q6_K | QT_Q5_K | QT_Q4_0)
14851                    || (*qtype == QT_NVFP4 && w.in_features() % 64 == 0)
14852            }
14853            GpuTensor::Float { .. } | GpuTensor::FloatBf16 { .. } => false,
14854        }
14855    }
14856
14857    /// Batched tensor-core int8 GEMM with a PRE-QUANTIZED q8_1 activation (aq,ad). The prefill
14858    /// (T>1) root fix: decode each weight 32-block to int8 in shared memory ONCE per (row-tile,
14859    /// K-step) and reuse it across all BN tokens via mma.sync.m16n8k32.s8 — amortizing the weight
14860    /// read/decode N-fold (vs the dp4a matvec's per-token re-read). s32 accumulate is exact vs
14861    /// dp4a; only the final f32 block-scale rounding differs. Caller MUST have checked
14862    /// `gemm_supports(w)`. y[m,out] token-major. NVFP4 per-tensor macro-scale applied post.
14863    pub fn qmatvec_gemm(
14864        &self,
14865        w: &crate::model::GpuTensor,
14866        aq: &CudaSlice<i8>,
14867        ad: &CudaSlice<f32>,
14868        m: usize,
14869    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14870        use crate::model::GpuTensor;
14871        let in_f = w.in_features();
14872        let out_f = w.out_features();
14873        let (bytes, qtype, row_bytes, scale, rp) = match w {
14874            GpuTensor::Quant {
14875                bytes,
14876                qtype,
14877                row_bytes,
14878                scale,
14879                rp,
14880                ..
14881            } => (bytes, *qtype, *row_bytes, *scale, *rp),
14882            _ => unreachable!("gemm_supports guaranteed Quant"),
14883        };
14884        // wgmma arm (sm_90a, task 8): the m64n64k32 warpgroup kernel reads the rp4 split-plane
14885        // mirror AS-IS (qplane rows = its A operand, the half dplane its scales) and the same
14886        // (aq, ad) activation planes. Same numeric class as the mma kernel below (exact s32 per
14887        // 32-block, one f32 scale fold per block, ascending K) — argmax/tolerance gated like
14888        // every prefill GEMM, not bit-gated. MEMRA_WGMMA=0 restores the portable kernel.
14889        if cfg!(memra_hopper_mma) && qtype == QT_Q8_0 && out_f % 64 == 0 && wgmma_gemm_enabled() {
14890            if let GpuTensor::Quant { rp4: Some(m4), .. } = w {
14891                let mut y = self.qmatvec_gemm_q8_0_wgmma_raw(m4, aq, ad, m, in_f, out_f)?;
14892                if scale != 1.0 {
14893                    self.scale_inplace(&mut y, scale, m * out_f)?;
14894                }
14895                return Ok(y);
14896            }
14897        }
14898        let name = match qtype {
14899            QT_Q8_0 => "qmatvec_gemm_q8_0",
14900            QT_Q4_K => "qmatvec_gemm_q4_K",
14901            QT_Q4_0 => {
14902                if rp {
14903                    "qmatvec_gemm_q4_0_rp"
14904                } else {
14905                    "qmatvec_gemm_q4_0"
14906                }
14907            }
14908            QT_Q5_K => "qmatvec_gemm_q5_K",
14909            QT_Q6_K => "qmatvec_gemm_q6_K",
14910            QT_NVFP4 => {
14911                if rp {
14912                    "qmatvec_gemm_nvfp4_rp"
14913                } else {
14914                    "qmatvec_gemm_nvfp4"
14915                }
14916            }
14917            _ => unreachable!(),
14918        };
14919        let f = self.func(name);
14920        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output: skip memset
14921        // CTA tile MUST match the .cu per-kernel tile. MMQ-PORT: kernel1 (Q8_0/Q4_K/Q5_K) runs llama's
14922        // 128x128 SQUARE tile (K1_BM=128 x K1_BN=128, 8 warps); kernel2 (Q6_K/NVFP4) keeps 64x256, 4 warps
14923        // (the macro BM/BN in the .cu). Grid dims are selected by qtype so each launches its own tile.
14924        let is_k1 = matches!(qtype, QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q4_0);
14925        // TUNE SEAM: MEMRA_GEMM_K1_LAUNCH overrides kernel1's launch tile to match a -D-swept fatbin.
14926        let k1_tile = if is_k1 {
14927            k1_launch_override().unwrap_or((128, 128, 8))
14928        } else {
14929            (128, 128, 8)
14930        };
14931        let (bm, bn): (u32, u32) = if is_k1 {
14932            (k1_tile.0, k1_tile.1)
14933        } else {
14934            (64, 256)
14935        };
14936        let warps: u32 = if is_k1 {
14937            k1_tile.2
14938        } else {
14939            match qtype {
14940                QT_NVFP4 => 8,
14941                _ => 4,
14942            }
14943        };
14944        let cfg = LaunchConfig {
14945            grid_dim: ((out_f as u32 + bm - 1) / bm, (m as u32 + bn - 1) / bn, 1),
14946            block_dim: (32, warps, 1),
14947            shared_mem_bytes: 0,
14948        };
14949        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
14950        let __s_b = self.gpu.stream();
14951        let mut b = __s_b.launch_builder(&f);
14952        b.arg(bytes)
14953            .arg(aq)
14954            .arg(ad)
14955            .arg(&mut y)
14956            .arg(&inf)
14957            .arg(&outf)
14958            .arg(&mi)
14959            .arg(&rb);
14960        unsafe {
14961            b.launch(cfg)?;
14962        }
14963        if scale != 1.0 {
14964            self.scale_inplace(&mut y, scale, m * out_f)?;
14965        }
14966        Ok(y)
14967    }
14968
14969    /// Test entry: run the GEMM directly from raw weight bytes + qtype (no GpuTensor). Quantizes
14970    /// the f32 activation `x` to q8_1 internally then launches the tensor-core GEMM. NVFP4 per-tensor
14971    /// macro-scale is NOT applied here (caller passes it separately, like the dp4a path). Used by
14972    /// kernel_check for the bit-equivalence gate vs qmatvec_*_dp4a.
14973    pub fn qmatvec_gemm_raw(
14974        &self,
14975        bytes: &CudaSlice<u8>,
14976        x: &CudaSlice<f32>,
14977        m: usize,
14978        in_f: usize,
14979        out_f: usize,
14980        qtype: i32,
14981        row_bytes: usize,
14982    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14983        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
14984        let name = match qtype {
14985            QT_Q8_0 => "qmatvec_gemm_q8_0",
14986            QT_Q4_K => "qmatvec_gemm_q4_K",
14987            QT_Q4_0 => "qmatvec_gemm_q4_0",
14988            QT_Q5_K => "qmatvec_gemm_q5_K",
14989            QT_Q6_K => "qmatvec_gemm_q6_K",
14990            QT_NVFP4 => "qmatvec_gemm_nvfp4",
14991            QT_NVFP4_RP => "qmatvec_gemm_nvfp4_rp",
14992            _ => panic!("qmatvec_gemm_raw: qtype {qtype} has no GEMM kernel"),
14993        };
14994        let f = self.func(name);
14995        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output: skip memset
14996        // MMQ-PORT: kernel1 (Q8_0/Q4_K/Q5_K) = llama 128x128 tile, 8 warps; kernel2 (Q6_K/NVFP4) = 64x256,
14997        // 4/8 warps. Grid tile per qtype (must match the .cu K1_BM/K1_BN vs BM/BN). KEEP IN SYNC w/ qmatvec_gemm.
14998        let is_k1 = matches!(qtype, QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q4_0);
14999        // TUNE SEAM: MEMRA_GEMM_K1_LAUNCH overrides kernel1's launch tile to match a -D-swept fatbin.
15000        let k1_tile = if is_k1 {
15001            k1_launch_override().unwrap_or((128, 128, 8))
15002        } else {
15003            (128, 128, 8)
15004        };
15005        let (bm, bn): (u32, u32) = if is_k1 {
15006            (k1_tile.0, k1_tile.1)
15007        } else {
15008            (64, 256)
15009        };
15010        let warps: u32 = if is_k1 {
15011            k1_tile.2
15012        } else {
15013            match qtype {
15014                QT_NVFP4 | QT_NVFP4_RP => 8,
15015                _ => 4,
15016            }
15017        };
15018        let cfg = LaunchConfig {
15019            grid_dim: ((out_f as u32 + bm - 1) / bm, (m as u32 + bn - 1) / bn, 1),
15020            block_dim: (32, warps, 1),
15021            shared_mem_bytes: 0,
15022        };
15023        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
15024        let __s_b = self.gpu.stream();
15025        let mut b = __s_b.launch_builder(&f);
15026        b.arg(bytes)
15027            .arg(&aq)
15028            .arg(&ad)
15029            .arg(&mut y)
15030            .arg(&inf)
15031            .arg(&outf)
15032            .arg(&mi)
15033            .arg(&rb);
15034        unsafe {
15035            b.launch(cfg)?;
15036        }
15037        Ok(y)
15038    }
15039
15040    /// H100 warpgroup GEMM raw entry (task 8): launch `qmatvec_gemm_q8_0_wgmma` on an rp4
15041    /// split-plane mirror + pre-quantized (aq, ad) activation planes. One warpgroup (128 thr)
15042    /// owns a 64x64 C tile; grid (out_f/64, ceil(m/64)). out_f % 64 == 0 REQUIRED (row loads
15043    /// and dplane scale reads are unguarded); the token edge is guarded in-kernel.
15044    /// Standalone harness verdict (tools/bench_q8_gemm_wgmma.cu, 4096x4096x512): rel 1.6e-05
15045    /// vs CPU ref, 179us vs the portable mma kernel's 688us (3.84x, unpipelined).
15046    pub fn qmatvec_gemm_q8_0_wgmma_raw(
15047        &self,
15048        rp4: &CudaSlice<u8>,
15049        aq: &CudaSlice<i8>,
15050        ad: &CudaSlice<f32>,
15051        m: usize,
15052        in_f: usize,
15053        out_f: usize,
15054    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15055        assert!(
15056            out_f % 64 == 0 && in_f % 32 == 0,
15057            "wgmma GEMM needs out_f%64==0, in_f%32==0"
15058        );
15059        let f = self.func("qmatvec_gemm_q8_0_wgmma");
15060        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output
15061        let cfg = LaunchConfig {
15062            grid_dim: ((out_f / 64) as u32, (m as u32).div_ceil(64), 1),
15063            block_dim: (128, 1, 1),
15064            shared_mem_bytes: 0,
15065        };
15066        let (inf, outf, mi) = (in_f as i32, out_f as i32, m as i32);
15067        let __s_b = self.gpu.stream();
15068        let mut b = __s_b.launch_builder(&f);
15069        b.arg(rp4)
15070            .arg(aq)
15071            .arg(ad)
15072            .arg(&mut y)
15073            .arg(&inf)
15074            .arg(&outf)
15075            .arg(&mi);
15076        unsafe {
15077            b.launch(cfg)?;
15078        }
15079        Ok(y)
15080    }
15081
15082    /// y[i] *= s. NVFP4 per-tensor macro-scale broadcast over the whole output.
15083    pub fn scale_inplace(
15084        &self,
15085        y: &mut CudaSlice<f32>,
15086        s: f32,
15087        n: usize,
15088    ) -> Result<(), Box<dyn std::error::Error>> {
15089        let f = self.func("scale_f32");
15090        let cfg = LaunchConfig::for_num_elems(n as u32);
15091        let (sf, ni) = (s, n as i32);
15092        let __s_b = self.gpu.stream();
15093        let mut b = __s_b.launch_builder(&f);
15094        b.arg(y).arg(&sf).arg(&ni);
15095        unsafe {
15096            b.launch(cfg)?;
15097        }
15098        Ok(())
15099    }
15100
15101    /// MEMRA_FULL_PREC dequant-on-use: expand a bf16-resident weight (`GpuTensor::FloatBf16`, raw
15102    /// bf16 bytes) to a transient f32 scratch of `n` elements, which then feeds the existing f32
15103    /// cuBLASLt GEMV. The scratch is freed when the caller drops it, so peak VRAM = resident bf16
15104    /// weights + ONE (largest) weight's f32 expansion + activations. SLOW IS FINE (research mode).
15105    pub fn bf16_to_f32(
15106        &self,
15107        data: &cudarc::driver::CudaView<'_, u8>,
15108        n: usize,
15109    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15110        let mut out = self.alloc_uninit::<f32>(n)?;
15111        let f = self.func("bf16_to_f32");
15112        let cfg = LaunchConfig::for_num_elems(n as u32);
15113        let ni = n as i32;
15114        let __s_b = self.gpu.stream();
15115        let mut b = __s_b.launch_builder(&f);
15116        b.arg(data).arg(&mut out).arg(&ni);
15117        unsafe {
15118            b.launch(cfg)?;
15119        }
15120        Ok(out)
15121    }
15122
15123    /// Chunked bf16 linear (MEMRA_FULL_PREC): y[m,out] = x @ W_bf16^T with the f32 dequant scratch
15124    /// bounded to CHUNK_ROWS rows (256MB at in_f=4096) instead of the whole weight — the 4GB
15125    /// lm_head expansion OOM'd the 24GB budget. Row-chunking partitions OUTPUT rows; each row's
15126    /// dot is computed by the identical kernel on identical bytes, so per-(token,row) results are
15127    /// bit-identical to the unchunked form. `exact` selects linear_decode_exact (per-column m=1
15128    /// calls, the spec-verify contract) vs plain linear.
15129    fn linear_bf16_chunked(
15130        &self,
15131        x: &CudaSlice<f32>,
15132        data: &CudaSlice<u8>,
15133        m: usize,
15134        in_f: usize,
15135        out_f: usize,
15136        exact: bool,
15137    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15138        const CHUNK_BYTES: usize = 256 << 20;
15139        let chunk_rows = (CHUNK_BYTES / (in_f * 4)).max(1).min(out_f);
15140        if chunk_rows >= out_f {
15141            let wf32 = self.bf16_to_f32(&data.slice(0..in_f * out_f * 2), in_f * out_f)?;
15142            return if exact {
15143                self.linear_decode_exact(x, &wf32, m, in_f, out_f)
15144            } else {
15145                self.linear(x, &wf32, m, in_f, out_f)
15146            };
15147        }
15148        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
15149        let mut r0 = 0usize;
15150        while r0 < out_f {
15151            let rows = chunk_rows.min(out_f - r0);
15152            let wslice = data.slice(r0 * in_f * 2..(r0 + rows) * in_f * 2);
15153            let wf32 = self.bf16_to_f32(&wslice, in_f * rows)?;
15154            let yc = if exact {
15155                self.linear_decode_exact(x, &wf32, m, in_f, rows)?
15156            } else {
15157                self.linear(x, &wf32, m, in_f, rows)?
15158            };
15159            // scatter [m, rows] into y[m, out_f] at column offset r0 (m is tiny in decode/verify)
15160            for mi in 0..m {
15161                let src = yc.slice(mi * rows..(mi + 1) * rows);
15162                let mut dst = y.slice_mut(mi * out_f + r0..mi * out_f + r0 + rows);
15163                self.gpu.stream().memcpy_dtod(&src, &mut dst)?;
15164            }
15165            r0 += rows;
15166        }
15167        Ok(y)
15168    }
15169
15170    /// On-device linear: y[m,out] = x[m,in] @ W[out,in]^T, weights row-major [out,in] (ggml).
15171    /// cuBLASLt col-major mapping (see memra_runtime::Gpu::linear_f32 for the derivation).
15172    /// DECODE-EXACT float linear: per-column m=1 cuBLASLt calls. cuBLASLt's reduction split is
15173    /// n-dependent (lt_ndep probe: m=1 vs m=2 col0 differs every bit), so spec-verify batches
15174    /// must not batch float matmuls the T=1 decode chain runs at m=1. Used by the small-t MoE
15175    /// router/shexp sites and matmul_decode_exact's Float arm.
15176    pub fn linear_decode_exact(
15177        &self,
15178        x: &CudaSlice<f32>,
15179        w: &CudaSlice<f32>,
15180        m_tokens: usize,
15181        in_f: usize,
15182        out_f: usize,
15183    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15184        if m_tokens == 1 {
15185            return self.linear(x, w, 1, in_f, out_f);
15186        }
15187        let xv = self.view(x, m_tokens * in_f);
15188        let mut y = self.alloc_uninit::<f32>(m_tokens * out_f)?;
15189        for t in 0..m_tokens {
15190            let row = xv.slice(t * in_f..(t + 1) * in_f);
15191            let mut xr = self.alloc_uninit::<f32>(in_f)?;
15192            self.copy_view_into(&mut xr, 0, &row, in_f)?;
15193            let yr = self.linear(&xr, w, 1, in_f, out_f)?;
15194            self.copy_into(&mut y, t * out_f, &yr, out_f)?;
15195        }
15196        Ok(y)
15197    }
15198
15199    pub fn linear(
15200        &self,
15201        x: &CudaSlice<f32>,
15202        w: &CudaSlice<f32>,
15203        m_tokens: usize,
15204        in_f: usize,
15205        out_f: usize,
15206    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15207        use cudarc::cublaslt::{Matmul, MatmulConfig};
15208        let mut c = self.alloc_uninit::<f32>(m_tokens * out_f)?; // cuBLASLt beta=0: C fully written
15209        let cfg = MatmulConfig {
15210            transa: true,
15211            transb: false,
15212            transc: false,
15213            m: out_f as u64,
15214            n: m_tokens as u64,
15215            k: in_f as u64,
15216            alpha: 1.0,
15217            lda: in_f as i64,
15218            ldb: in_f as i64,
15219            beta: 0.0,
15220            ldc: out_f as i64,
15221            stride_a: None,
15222            stride_b: None,
15223            stride_c: None,
15224            stride_bias: None,
15225            batch_size: None,
15226        };
15227        unsafe {
15228            self.gpu.blas.matmul(cfg, w, x, &mut c, None, None)?;
15229        }
15230        Ok(c)
15231    }
15232
15233    /// Naive SDPA. Q:[head_dim,n_head,T], K/V:[head_dim,n_head_kv,T_kv] -> O:[head_dim,n_head,T].
15234    pub fn sdpa_naive(
15235        &self,
15236        q: &CudaSlice<f32>,
15237        k: &CudaSlice<f32>,
15238        v: &CudaSlice<f32>,
15239        o: &mut CudaSlice<f32>,
15240        head_dim: usize,
15241        n_head: usize,
15242        n_head_kv: usize,
15243        t: usize,
15244        t_kv: usize,
15245        scale: f32,
15246        causal: bool,
15247    ) -> Result<(), Box<dyn std::error::Error>> {
15248        let f = self.func("sdpa_naive_f32");
15249        let cfg = LaunchConfig {
15250            grid_dim: (n_head as u32, t as u32, 1),
15251            block_dim: (128, 1, 1),
15252            shared_mem_bytes: (t_kv * 4) as u32,
15253        };
15254        let (hd, nh, nhkv, ti, tkvi, cz) = (
15255            head_dim as i32,
15256            n_head as i32,
15257            n_head_kv as i32,
15258            t as i32,
15259            t_kv as i32,
15260            causal as i32,
15261        );
15262        let __s_b = self.gpu.stream();
15263        let mut b = __s_b.launch_builder(&f);
15264        b.arg(q)
15265            .arg(k)
15266            .arg(v)
15267            .arg(o)
15268            .arg(&hd)
15269            .arg(&nh)
15270            .arg(&nhkv)
15271            .arg(&ti)
15272            .arg(&tkvi)
15273            .arg(&scale)
15274            .arg(&cz);
15275        unsafe {
15276            b.launch(cfg)?;
15277        }
15278        Ok(())
15279    }
15280
15281    /// Island twin (lane/gemma-vision): causal + sliding-window attention with
15282    /// bidirectional image islands. `span_id` labels each absolute kv position
15283    /// (-1 text, >=0 island id); same-island keys are visible unconditionally,
15284    /// reproducing the reference's non-causal image batch. window 0 = no window.
15285    #[allow(clippy::too_many_arguments)]
15286    pub fn sdpa_naive_island(
15287        &self,
15288        q: &CudaSlice<f32>,
15289        k: &CudaSlice<f32>,
15290        v: &CudaSlice<f32>,
15291        o: &mut CudaSlice<f32>,
15292        span_id: &CudaSlice<i32>,
15293        head_dim: usize,
15294        n_head: usize,
15295        n_head_kv: usize,
15296        t: usize,
15297        t_kv: usize,
15298        scale: f32,
15299        window: usize,
15300    ) -> Result<(), Box<dyn std::error::Error>> {
15301        let f = self.func("sdpa_naive_island_f32");
15302        let cfg = LaunchConfig {
15303            grid_dim: (n_head as u32, t as u32, 1),
15304            block_dim: (128, 1, 1),
15305            shared_mem_bytes: (t_kv * 4) as u32,
15306        };
15307        let (hd, nh, nhkv, ti, tkvi, wi) = (
15308            head_dim as i32,
15309            n_head as i32,
15310            n_head_kv as i32,
15311            t as i32,
15312            t_kv as i32,
15313            window as i32,
15314        );
15315        let __s_b = self.gpu.stream();
15316        let mut b = __s_b.launch_builder(&f);
15317        b.arg(q)
15318            .arg(k)
15319            .arg(v)
15320            .arg(o)
15321            .arg(span_id)
15322            .arg(&hd)
15323            .arg(&nh)
15324            .arg(&nhkv)
15325            .arg(&ti)
15326            .arg(&tkvi)
15327            .arg(&scale)
15328            .arg(&wi);
15329        unsafe {
15330            b.launch(cfg)?;
15331        }
15332        Ok(())
15333    }
15334
15335    /// Windowed sdpa_naive twin (gemma4 R6): masks keys older than q_pos-(window-1).
15336    #[allow(clippy::too_many_arguments)]
15337    pub fn sdpa_naive_w(
15338        &self,
15339        q: &CudaSlice<f32>,
15340        k: &CudaSlice<f32>,
15341        v: &CudaSlice<f32>,
15342        o: &mut CudaSlice<f32>,
15343        head_dim: usize,
15344        n_head: usize,
15345        n_head_kv: usize,
15346        t: usize,
15347        t_kv: usize,
15348        scale: f32,
15349        causal: bool,
15350        window: usize,
15351    ) -> Result<(), Box<dyn std::error::Error>> {
15352        let f = self.func("sdpa_naive_w_f32");
15353        let cfg = LaunchConfig {
15354            grid_dim: (n_head as u32, t as u32, 1),
15355            block_dim: (128, 1, 1),
15356            shared_mem_bytes: (t_kv * 4) as u32,
15357        };
15358        let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
15359            head_dim as i32,
15360            n_head as i32,
15361            n_head_kv as i32,
15362            t as i32,
15363            t_kv as i32,
15364            causal as i32,
15365            window as i32,
15366        );
15367        let __s_b = self.gpu.stream();
15368        let mut b = __s_b.launch_builder(&f);
15369        b.arg(q)
15370            .arg(k)
15371            .arg(v)
15372            .arg(o)
15373            .arg(&hd)
15374            .arg(&nh)
15375            .arg(&nhkv)
15376            .arg(&ti)
15377            .arg(&tkvi)
15378            .arg(&scale)
15379            .arg(&cz)
15380            .arg(&wi);
15381        unsafe {
15382            b.launch(cfg)?;
15383        }
15384        Ok(())
15385    }
15386
15387    /// SDPA where K/V are CudaViews into a resident KV cache (decode hot path, no host round-trip).
15388    pub fn sdpa_naive_view(
15389        &self,
15390        q: &CudaSlice<f32>,
15391        k: &cudarc::driver::CudaView<f32>,
15392        v: &cudarc::driver::CudaView<f32>,
15393        o: &mut CudaSlice<f32>,
15394        head_dim: usize,
15395        n_head: usize,
15396        n_head_kv: usize,
15397        t: usize,
15398        t_kv: usize,
15399        scale: f32,
15400        causal: bool,
15401    ) -> Result<(), Box<dyn std::error::Error>> {
15402        let f = self.func("sdpa_naive_f32");
15403        let cfg = LaunchConfig {
15404            grid_dim: (n_head as u32, t as u32, 1),
15405            block_dim: (128, 1, 1),
15406            shared_mem_bytes: (t_kv * 4) as u32,
15407        };
15408        let (hd, nh, nhkv, ti, tkvi, cz) = (
15409            head_dim as i32,
15410            n_head as i32,
15411            n_head_kv as i32,
15412            t as i32,
15413            t_kv as i32,
15414            causal as i32,
15415        );
15416        let __s_b = self.gpu.stream();
15417        let mut b = __s_b.launch_builder(&f);
15418        b.arg(q)
15419            .arg(k)
15420            .arg(v)
15421            .arg(o)
15422            .arg(&hd)
15423            .arg(&nh)
15424            .arg(&nhkv)
15425            .arg(&ti)
15426            .arg(&tkvi)
15427            .arg(&scale)
15428            .arg(&cz);
15429        unsafe {
15430            b.launch(cfg)?;
15431        }
15432        Ok(())
15433    }
15434
15435    /// Correctness fallback for quantized resident K/V views. Dequantizes K and V once into f32
15436    /// workspaces, then calls `sdpa_naive`. This is an explicit API: the optimized prefill view
15437    /// dispatch remains unchanged, so callers can use it as a reference or compatibility path.
15438    /// Dequant a quantized KV view into caller-owned f32 buffers (one grid-stride launch).
15439    /// `g` picks the kf8vf8-module stamp for e4m3 caches (same flag contract as fa_decode/
15440    /// fa_prefill_view). Used by the E4B shared-KV prefill arms (2026-07-31) to feed the
15441    /// f32 fa_prefill_w / fa_prefill_hd512 twins from the target layer's quantized rows.
15442    #[allow(clippy::too_many_arguments)]
15443    pub fn fa_dequant_kv_view_f32(
15444        &self,
15445        k: &cudarc::driver::CudaView<u8>,
15446        v: &cudarc::driver::CudaView<u8>,
15447        kf: &mut CudaSlice<f32>,
15448        vf: &mut CudaSlice<f32>,
15449        kv_dim_k: usize,
15450        kv_dim_v: usize,
15451        t_kv: usize,
15452        k_tok_bytes: usize,
15453        v_tok_bytes: usize,
15454        g: bool,
15455    ) -> Result<(), Box<dyn std::error::Error>> {
15456        let f = if g {
15457            self.func_g("fa_dequant_kv_ws_f32")
15458        } else {
15459            self.func("fa_dequant_kv_ws_f32")
15460        };
15461        let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
15462        let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
15463        let cfg = LaunchConfig {
15464            grid_dim: (nblk.max(1), 1, 1),
15465            block_dim: (256, 1, 1),
15466            shared_mem_bytes: 0,
15467        };
15468        let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
15469        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
15470        let __s_b = self.gpu.stream();
15471        let mut b = __s_b.launch_builder(&f);
15472        b.arg(k)
15473            .arg(v)
15474            .arg(&mut *kf)
15475            .arg(&mut *vf)
15476            .arg(&kdk)
15477            .arg(&kdv)
15478            .arg(&tkvi)
15479            .arg(&ktb)
15480            .arg(&vtb);
15481        unsafe {
15482            b.launch(cfg)?;
15483        }
15484        Ok(())
15485    }
15486
15487    #[allow(clippy::too_many_arguments)]
15488    pub fn sdpa_naive_quantized_view(
15489        &self,
15490        q: &CudaSlice<f32>,
15491        k: &cudarc::driver::CudaView<u8>,
15492        v: &cudarc::driver::CudaView<u8>,
15493        o: &mut CudaSlice<f32>,
15494        head_dim: usize,
15495        n_head: usize,
15496        n_head_kv: usize,
15497        t: usize,
15498        t_kv: usize,
15499        scale: f32,
15500        causal: bool,
15501        k_tok_bytes: usize,
15502        v_tok_bytes: usize,
15503    ) -> Result<(), Box<dyn std::error::Error>> {
15504        let kv_dim = n_head_kv * head_dim;
15505        let mut kf = self.uninit(t_kv * kv_dim)?;
15506        let mut vf = self.uninit(t_kv * kv_dim)?;
15507        let f = self.func("fa_dequant_kv_ws_f32");
15508        let total = (2 * t_kv * kv_dim) as u64;
15509        let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
15510        let cfg = LaunchConfig {
15511            grid_dim: (nblk.max(1), 1, 1),
15512            block_dim: (256, 1, 1),
15513            shared_mem_bytes: 0,
15514        };
15515        let (kv_dim_i, t_kv_i) = (kv_dim as i32, t_kv as i32);
15516        let (k_tok_bytes_i, v_tok_bytes_i) = (k_tok_bytes as i64, v_tok_bytes as i64);
15517        let __s_b = self.gpu.stream();
15518        let mut b = __s_b.launch_builder(&f);
15519        b.arg(k)
15520            .arg(v)
15521            .arg(&mut kf)
15522            .arg(&mut vf)
15523            .arg(&kv_dim_i)
15524            .arg(&kv_dim_i)
15525            .arg(&t_kv_i)
15526            .arg(&k_tok_bytes_i)
15527            .arg(&v_tok_bytes_i);
15528        unsafe { b.launch(cfg)? };
15529        self.sdpa_naive(
15530            q, &kf, &vf, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
15531        )
15532    }
15533
15534    /// WINDOWED twin of `sdpa_naive_quantized_view` (step35 SWA prefill): dequant the KV byte
15535    /// view into f32 workspaces with the SAME `fa_dequant_kv_ws_f32` launch, then run
15536    /// `sdpa_naive_w` instead of `sdpa_naive`. `window == 0` is the unwindowed form (the kernel
15537    /// treats a non-positive window as "no window mask"), so this is a strict superset of the
15538    /// unwindowed function above and produces bit-identical output at window == 0.
15539    ///
15540    /// Why this exists: EVERY windowed FlashAttention stamp in flash_attn.cu is head_dim-256
15541    /// only (`fa_prefill_w_f32` == `fa_prefill_f32_body<256>`, and the quantized-view windowed
15542    /// twins likewise), while step35 is head_dim 128. Its SWA layers therefore have no windowed
15543    /// FA path and take this f32 floor in v0 — same cache bytes, same numeric class as the
15544    /// unwindowed quantized-view fallback, so the chunk-invariance contract holds on both.
15545    #[allow(clippy::too_many_arguments)]
15546    pub fn sdpa_naive_w_quantized_view(
15547        &self,
15548        q: &CudaSlice<f32>,
15549        k: &cudarc::driver::CudaView<u8>,
15550        v: &cudarc::driver::CudaView<u8>,
15551        o: &mut CudaSlice<f32>,
15552        head_dim: usize,
15553        n_head: usize,
15554        n_head_kv: usize,
15555        t: usize,
15556        t_kv: usize,
15557        scale: f32,
15558        causal: bool,
15559        window: usize,
15560        k_tok_bytes: usize,
15561        v_tok_bytes: usize,
15562    ) -> Result<(), Box<dyn std::error::Error>> {
15563        let kv_dim = n_head_kv * head_dim;
15564        let mut kf = self.uninit(t_kv * kv_dim)?;
15565        let mut vf = self.uninit(t_kv * kv_dim)?;
15566        let f = self.func("fa_dequant_kv_ws_f32");
15567        let total = (2 * t_kv * kv_dim) as u64;
15568        let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
15569        let cfg = LaunchConfig {
15570            grid_dim: (nblk.max(1), 1, 1),
15571            block_dim: (256, 1, 1),
15572            shared_mem_bytes: 0,
15573        };
15574        let (kv_dim_i, t_kv_i) = (kv_dim as i32, t_kv as i32);
15575        let (k_tok_bytes_i, v_tok_bytes_i) = (k_tok_bytes as i64, v_tok_bytes as i64);
15576        let __s_b = self.gpu.stream();
15577        let mut b = __s_b.launch_builder(&f);
15578        b.arg(k)
15579            .arg(v)
15580            .arg(&mut kf)
15581            .arg(&mut vf)
15582            .arg(&kv_dim_i)
15583            .arg(&kv_dim_i)
15584            .arg(&t_kv_i)
15585            .arg(&k_tok_bytes_i)
15586            .arg(&v_tok_bytes_i);
15587        unsafe { b.launch(cfg)? };
15588        self.sdpa_naive_w(
15589            q, &kf, &vf, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal, window,
15590        )
15591    }
15592
15593    /// Hand-written FlashAttention prefill (sm_120, FA-2 online softmax on validated mma.sync,
15594    /// head_dim 256 or 128 (template-stamped twins), GQA, causal). Replaces sdpa_naive for T>1.
15595    /// Q/K/V/O [head_dim, n_head(_kv), T].
15596    pub fn fa_prefill(
15597        &self,
15598        q: &CudaSlice<f32>,
15599        k: &CudaSlice<f32>,
15600        v: &CudaSlice<f32>,
15601        o: &mut CudaSlice<f32>,
15602        head_dim: usize,
15603        n_head: usize,
15604        n_head_kv: usize,
15605        t: usize,
15606        t_kv: usize,
15607        scale: f32,
15608        causal: bool,
15609    ) -> Result<(), Box<dyn std::error::Error>> {
15610        if portable_mma_gated() {
15611            return self.sdpa_naive(
15612                q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
15613            );
15614        }
15615        // FA3 v10 arm (task #20, OPT-IN MEMRA_FA3=1 — harness-proven 883us vs the shipped
15616        // kernel's 993us at T=2048): TMA-swizzled wgmma FA, fresh causal hd256 only.
15617        // NEW NUMERIC CONFIG (GDN-mma precedent): online softmax / bf16-P class — the
15618        // run-gen argmax + greedy-stream batteries arbitrate; not bit-paired.
15619        // PROMOTED default-ON hopper (2026-07-27): 3-seed 2048-prime -> 128-decode
15620        // streams MATCH vs mma, full battery green, lane interleaved 5/5 (+2.4%).
15621        // MEMRA_FA3=0 reverts; kernel-check pins the mma config regardless.
15622        let fa3_on = head_dim == 256
15623            && causal
15624            && t == t_kv
15625            && match std::env::var("MEMRA_FA3").as_deref() {
15626                Ok("0") => false,
15627                Ok("1") => true,
15628                _ => cfg!(memra_hopper_mma),
15629            };
15630        if fa3_on {
15631            let n = t * n_head * head_dim;
15632            let nkv = t * n_head_kv * head_dim;
15633            let mut q16 = self.alloc_u8_uninit(n * 2)?;
15634            let mut k16 = self.alloc_u8_uninit(nkv * 2)?;
15635            let mut v16 = self.alloc_u8_uninit(nkv * 2)?;
15636            self.f32_to_bf16_into(q, &mut q16, n)?;
15637            self.f32_to_bf16_into(k, &mut k16, nkv)?;
15638            self.f32_to_bf16_into(v, &mut v16, nkv)?;
15639            let rc = {
15640                use cudarc::driver::{DevicePtr, DevicePtrMut};
15641                let stream = self.gpu.stream();
15642                let (qp, _g1) = q16.device_ptr(&stream);
15643                let (kp, _g2) = k16.device_ptr(&stream);
15644                let (vp, _g3) = v16.device_ptr(&stream);
15645                let (op, _g4) = o.device_ptr_mut(&stream);
15646                unsafe {
15647                    memra_fa3_prefill(
15648                        qp as *const core::ffi::c_void,
15649                        kp as *const core::ffi::c_void,
15650                        vp as *const core::ffi::c_void,
15651                        op as *mut f32,
15652                        t as i32,
15653                        n_head as i32,
15654                        n_head_kv as i32,
15655                        head_dim as i32,
15656                        scale,
15657                        stream.cu_stream() as *mut core::ffi::c_void,
15658                    )
15659                }
15660            };
15661            if rc != 0 {
15662                return Err(format!("memra_fa3_prefill rc={rc}").into());
15663            }
15664            return Ok(());
15665        }
15666        // FLOOR PORT (P2+P0a+P0b+P1): 4 warps/CTA, BLOCK_Q=64 query rows, BK=32 KV tile,
15667        // Q-in-reg + register-O, grid.y=n_head_kv (4 Q-heads share staged K/V).
15668        // P1 plain arm (MEMRA_FA_P1=1 opt-in until the qwen battery): the engine-study body
15669        // (FA2 schedule + boundary split + swizzle) on the non-windowed lane. bf16 pre-convert.
15670        static FA_P1: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15671        let fa_p1 = *FA_P1.get_or_init(|| std::env::var("MEMRA_FA_P1").as_deref() == Ok("1"));
15672        if fa_p1 && head_dim == 256 && !std::env::var("MEMRA_FA_FLOOR").is_ok() {
15673            const BLOCK_Q: usize = 64;
15674            const BKX: usize = 32;
15675            let f = self.func("fa_prefill_bf16_p1");
15676            let shmem = (2 * (2 * BKX * head_dim + BLOCK_Q * BKX)
15677                + 4 * (BLOCK_Q * BKX + 2 * BLOCK_Q)) as u32;
15678            use cudarc::driver::sys::CUfunction_attribute_enum as A;
15679            f.set_attribute(
15680                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
15681                shmem as i32,
15682            )?;
15683            let cfg = LaunchConfig {
15684                grid_dim: (
15685                    (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
15686                    n_head as u32,
15687                    1,
15688                ),
15689                block_dim: (32, 4, 1),
15690                shared_mem_bytes: shmem,
15691            };
15692            let (hd, nh, nhkv, ti, tkvi, cz) = (
15693                head_dim as i32,
15694                n_head as i32,
15695                n_head_kv as i32,
15696                t as i32,
15697                t_kv as i32,
15698                causal as i32,
15699            );
15700            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
15701            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
15702            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
15703            let __s_b = self.gpu.stream();
15704            let mut b = __s_b.launch_builder(&f);
15705            b.arg(&qb)
15706                .arg(&kb)
15707                .arg(&vb)
15708                .arg(o)
15709                .arg(&hd)
15710                .arg(&nh)
15711                .arg(&nhkv)
15712                .arg(&ti)
15713                .arg(&tkvi)
15714                .arg(&scale)
15715                .arg(&cz);
15716            unsafe {
15717                b.launch(cfg)?;
15718            }
15719            return Ok(());
15720        }
15721        // Edge 5a (DEFAULT): fa_prefill_f32_pp — register-resident softmax (no sSw smem
15722        // round-trip), the FA3 softmax-GEMM overlap variant. ncu (pp512): short_scoreboard
15723        // 4.32->3.47, wait 1.99->1.45, per-call ~577us->~440us (1.31x) at flat 12.1% warps /
15724        // 255 regs / 2 CTAs (occupancy preserved). Bit-safe: 9B+27B argmax MATCH, rel 2.55e-3
15725        // vs floor 3.03e-3. MEMRA_FA_FLOOR reverts to the serialized-softmax floor kernel.
15726        const BK: usize = 32;
15727        // W2 lane (MEMRA_FA_PP_W2=1, ncu 2026-07-26): 2-warp/32-row CTA tile doubles grid.x —
15728        // bit-identical per-row math, pure coverage trade for the 6.25%-occupancy starvation.
15729        let w2 = std::env::var("MEMRA_FA_PP_W2").as_deref() == Ok("1");
15730        let (block_q, warps, w2_sfx): (usize, u32, &str) =
15731            if w2 { (32, 2, "_w2") } else { (64, 4, "") };
15732        // hd128 twins (2026-07-07): the prefill kernels are template-stamped at 256 (original
15733        // names, dispatch unchanged) and 128 (`_hd128`, the MiniMax-M3 class). Callers gate
15734        // other head_dims to sdpa_naive before reaching here.
15735        let hd_sfx = fa_hd_suffix(head_dim)?;
15736        let floor = std::env::var("MEMRA_FA_FLOOR").is_ok();
15737        // BF16-KV staging lane (2026-07-26, default ON): the kernel converts K/V to bf16
15738        // during staging anyway — pre-converting to bf16 mirrors is BIT-IDENTICAL (same
15739        // __float2bfloat16 values into the same mma) and turns the 67%-of-stalls scalar
15740        // staging into int4 vector copies. MEMRA_FA_BF16KV=0 reverts.
15741        let bf16kv = !floor && !w2 && std::env::var("MEMRA_FA_BF16KV").as_deref() != Ok("0");
15742        let (kb16, vb16) = if bf16kv {
15743            let n = t_kv * n_head_kv * head_dim;
15744            let mut kb = self.alloc_u8_uninit(n * 2)?;
15745            let mut vb = self.alloc_u8_uninit(n * 2)?;
15746            let fcv = self.func("f32_to_bf16_bulk");
15747            let ni = n as i64;
15748            let cfgc = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
15749            let __s_b = self.gpu.stream();
15750            let mut b = __s_b.launch_builder(&fcv);
15751            b.arg(k).arg(&mut kb).arg(&ni);
15752            unsafe {
15753                b.launch(cfgc)?;
15754            }
15755            let __s_b = self.gpu.stream();
15756            let mut b = __s_b.launch_builder(&fcv);
15757            b.arg(v).arg(&mut vb).arg(&ni);
15758            unsafe {
15759                b.launch(cfgc)?;
15760            }
15761            (Some(kb), Some(vb))
15762        } else {
15763            (None, None)
15764        };
15765        let f = self.func(&if bf16kv {
15766            format!("fa_prefill_bf16kv_pp{hd_sfx}")
15767        } else {
15768            format!(
15769                "fa_prefill_f32{}{}{hd_sfx}",
15770                if floor { "" } else { "_pp" },
15771                if floor { "" } else { w2_sfx }
15772            )
15773        });
15774        // persistent smem: bf16*(KV_STAGES*(sK + sV) + sP) + f32*(sS + sM + sL);
15775        // the bf16kv ring doubles the K/V stages (KV_STAGES=2).
15776        let kv_stages = if bf16kv { 2 } else { 1 };
15777        let shmem = (2 * (kv_stages * 2 * BK * head_dim + block_q * BK)
15778            + 4 * (block_q * BK + 2 * block_q)) as u32;
15779        use cudarc::driver::sys::CUfunction_attribute_enum as A;
15780        f.set_attribute(
15781            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
15782            shmem as i32,
15783        )?;
15784        let cfg = LaunchConfig {
15785            grid_dim: (
15786                (t as u32 + block_q as u32 - 1) / block_q as u32,
15787                n_head as u32,
15788                1,
15789            ),
15790            block_dim: (32, warps, 1),
15791            shared_mem_bytes: shmem,
15792        };
15793        let (hd, nh, nhkv, ti, tkvi, cz) = (
15794            head_dim as i32,
15795            n_head as i32,
15796            n_head_kv as i32,
15797            t as i32,
15798            t_kv as i32,
15799            causal as i32,
15800        );
15801        let __s_b = self.gpu.stream();
15802        let mut b = __s_b.launch_builder(&f);
15803        b.arg(q);
15804        match (&kb16, &vb16) {
15805            (Some(kb), Some(vb)) => {
15806                b.arg(kb).arg(vb);
15807            }
15808            _ => {
15809                b.arg(k).arg(v);
15810            }
15811        }
15812        b.arg(o)
15813            .arg(&hd)
15814            .arg(&nh)
15815            .arg(&nhkv)
15816            .arg(&ti)
15817            .arg(&tkvi)
15818            .arg(&scale)
15819            .arg(&cz);
15820        unsafe {
15821            b.launch(cfg)?;
15822        }
15823        Ok(())
15824    }
15825
15826    /// Windowed FA prefill (gemma4 SWA layers past the sliding window, hd256): fa_prefill's
15827    /// exact dispatch (pp default, MEMRA_FA_FLOOR seam) with the sliding-window mask + tile
15828    /// skip in-kernel. Replaces the O(T*T_kv) scalar sdpa_naive_w on the prime path.
15829    #[allow(clippy::too_many_arguments)]
15830    pub fn fa_prefill_w(
15831        &self,
15832        q: &CudaSlice<f32>,
15833        k: &CudaSlice<f32>,
15834        v: &CudaSlice<f32>,
15835        o: &mut CudaSlice<f32>,
15836        head_dim: usize,
15837        n_head: usize,
15838        n_head_kv: usize,
15839        t: usize,
15840        t_kv: usize,
15841        scale: f32,
15842        causal: bool,
15843        window: usize,
15844    ) -> Result<(), Box<dyn std::error::Error>> {
15845        // sm_90a rides the mma twins (portable_mma_gated, 2026-07-31 — the raw
15846        // portable_cuda gate was stale-conservative on Hopper; fa_prefill already flipped).
15847        if portable_mma_gated() {
15848            return self.sdpa_naive_w(
15849                q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal, window,
15850            );
15851        }
15852        // Default: bf16-prestaged twin (same treatment as hd512 — Q/K/V pre-converted once,
15853        // int4 stage copies; bit-identical, kernel_check-gated). MEMRA_FAW_STAGE=f32 reverts;
15854        // MEMRA_FA_FLOOR keeps the f32 floor stamp untouched.
15855        static FAW_F32: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15856        let faw_f32 =
15857            *FAW_F32.get_or_init(|| std::env::var("MEMRA_FAW_STAGE").as_deref() == Ok("f32"));
15858        let floor = std::env::var("MEMRA_FA_FLOOR").is_ok();
15859        self.fa_prefill_w_arm(
15860            q,
15861            k,
15862            v,
15863            o,
15864            head_dim,
15865            n_head,
15866            n_head_kv,
15867            t,
15868            t_kv,
15869            scale,
15870            causal,
15871            window,
15872            floor || faw_f32,
15873            floor,
15874        )
15875    }
15876
15877    /// Windowed FA prefill with PRE-CONVERTED bf16 operands (producer-emitted; 31B glue lane).
15878    /// Launches the P1 stamp directly — callers guarantee qb/kb/vb hold the exact bf16 of q/k/v.
15879    #[allow(clippy::too_many_arguments)]
15880    pub fn fa_prefill_w_pre(
15881        &self,
15882        qb: &CudaSlice<u8>,
15883        kb: &CudaSlice<u8>,
15884        vb: &CudaSlice<u8>,
15885        o: &mut CudaSlice<f32>,
15886        head_dim: usize,
15887        n_head: usize,
15888        n_head_kv: usize,
15889        t: usize,
15890        t_kv: usize,
15891        scale: f32,
15892        causal: bool,
15893        window: usize,
15894        v_f16: bool,
15895    ) -> Result<(), Box<dyn std::error::Error>> {
15896        const BLOCK_Q: usize = 64;
15897        const BK: usize = 32;
15898        debug_assert_eq!(head_dim, 256);
15899        let hp = fa_f16pv_on() && faw_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
15900        debug_assert!(!v_f16 || hp, "f16 V emitted but the SWA hp arm is off");
15901        if hp {
15902            const BLOCK_QH: usize = 32;
15903            // V bytes must be f16 for the h2 stamp; producer normally emits f16 (v_f16),
15904            // else re-encode through the pooled scratch (stream-ordered reuse).
15905            let mut vguard = self.fa_vf16_scratch.lock().unwrap();
15906            let vh: &CudaSlice<u8> = if v_f16 {
15907                vb
15908            } else {
15909                let n = t_kv * n_head_kv * head_dim;
15910                if vguard.as_ref().map(|b| b.len() < n * 2).unwrap_or(true) {
15911                    *vguard = Some(self.alloc_uninit::<u8>(n * 2)?);
15912                }
15913                self.bf16_to_f16_into(vb, n, vguard.as_mut().unwrap())?;
15914                vguard.as_ref().unwrap()
15915            };
15916            let f = self.func("fa_prefill_w_bf16_p1h2");
15917            let shmem = (2 * (2 * BK * head_dim + 2 * BLOCK_QH * BK) + 4 * (2 * BLOCK_QH)) as u32;
15918            use cudarc::driver::sys::CUfunction_attribute_enum as A;
15919            f.set_attribute(
15920                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
15921                shmem as i32,
15922            )?;
15923            let cfg = LaunchConfig {
15924                grid_dim: ((t as u32).div_ceil(BLOCK_QH as u32), (n_head / 2) as u32, 1),
15925                block_dim: (32, 4, 1),
15926                shared_mem_bytes: shmem,
15927            };
15928            let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
15929                head_dim as i32,
15930                n_head as i32,
15931                n_head_kv as i32,
15932                t as i32,
15933                t_kv as i32,
15934                causal as i32,
15935                window as i32,
15936            );
15937            let __s_b = self.gpu.stream();
15938            let mut b = __s_b.launch_builder(&f);
15939            b.arg(qb)
15940                .arg(kb)
15941                .arg(vh)
15942                .arg(o)
15943                .arg(&hd)
15944                .arg(&nh)
15945                .arg(&nhkv)
15946                .arg(&ti)
15947                .arg(&tkvi)
15948                .arg(&scale)
15949                .arg(&cz)
15950                .arg(&wi);
15951            unsafe {
15952                b.launch(cfg)?;
15953            }
15954            return Ok(());
15955        }
15956        let f = self.func("fa_prefill_w_bf16_p1");
15957        let shmem =
15958            (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
15959        use cudarc::driver::sys::CUfunction_attribute_enum as A;
15960        f.set_attribute(
15961            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
15962            shmem as i32,
15963        )?;
15964        let cfg = LaunchConfig {
15965            grid_dim: (
15966                (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
15967                n_head as u32,
15968                1,
15969            ),
15970            block_dim: (32, 4, 1),
15971            shared_mem_bytes: shmem,
15972        };
15973        let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
15974            head_dim as i32,
15975            n_head as i32,
15976            n_head_kv as i32,
15977            t as i32,
15978            t_kv as i32,
15979            causal as i32,
15980            window as i32,
15981        );
15982        let __s_b = self.gpu.stream();
15983        let mut b = __s_b.launch_builder(&f);
15984        b.arg(qb)
15985            .arg(kb)
15986            .arg(vb)
15987            .arg(o)
15988            .arg(&hd)
15989            .arg(&nh)
15990            .arg(&nhkv)
15991            .arg(&ti)
15992            .arg(&tkvi)
15993            .arg(&scale)
15994            .arg(&cz)
15995            .arg(&wi);
15996        unsafe {
15997            b.launch(cfg)?;
15998        }
15999        Ok(())
16000    }
16001
16002    /// Windowed FA prefill with the stage arm FORCED — the kernel_check bit-identity entry.
16003    #[allow(clippy::too_many_arguments)]
16004    pub fn fa_prefill_w_arm(
16005        &self,
16006        q: &CudaSlice<f32>,
16007        k: &CudaSlice<f32>,
16008        v: &CudaSlice<f32>,
16009        o: &mut CudaSlice<f32>,
16010        head_dim: usize,
16011        n_head: usize,
16012        n_head_kv: usize,
16013        t: usize,
16014        t_kv: usize,
16015        scale: f32,
16016        causal: bool,
16017        window: usize,
16018        f32_stage: bool,
16019        floor: bool,
16020    ) -> Result<(), Box<dyn std::error::Error>> {
16021        const BLOCK_Q: usize = 64;
16022        const BK: usize = 32;
16023        debug_assert_eq!(head_dim, 256, "fa_prefill_w is stamped hd256 only");
16024        // P1 (2026-07-22 engine study): per-head Br=64 stamp with the FA2 schedule (V-copy
16025        // over GEMM0, next-K over softmax+GEMM1) + boundary/interior mask split. FP order
16026        // preserved -> bit-identical (gated). MEMRA_FAW_P1=0 reverts to the g4/o2 arms.
16027        static P1_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16028        let p1 = !floor
16029            && !f32_stage
16030            && *P1_ON.get_or_init(|| {
16031                std::env::var("MEMRA_FAW_P1")
16032                    .map(|v| v != "0")
16033                    .unwrap_or(true)
16034            });
16035        let hp =
16036            p1 && fa_f16pv_on() && faw_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
16037        if hp {
16038            const BLOCK_QH: usize = 32;
16039            let f = self.func("fa_prefill_w_bf16_p1h2");
16040            let shmem = (2 * (2 * BK * head_dim + 2 * BLOCK_QH * BK) + 4 * (2 * BLOCK_QH)) as u32;
16041            use cudarc::driver::sys::CUfunction_attribute_enum as A;
16042            f.set_attribute(
16043                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16044                shmem as i32,
16045            )?;
16046            let cfg = LaunchConfig {
16047                grid_dim: ((t as u32).div_ceil(BLOCK_QH as u32), (n_head / 2) as u32, 1),
16048                block_dim: (32, 4, 1),
16049                shared_mem_bytes: shmem,
16050            };
16051            let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
16052                head_dim as i32,
16053                n_head as i32,
16054                n_head_kv as i32,
16055                t as i32,
16056                t_kv as i32,
16057                causal as i32,
16058                window as i32,
16059            );
16060            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
16061            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
16062            let vh = self.f32_to_f16(v, t_kv * n_head_kv * head_dim)?;
16063            let __s_b = self.gpu.stream();
16064            let mut b = __s_b.launch_builder(&f);
16065            b.arg(&qb)
16066                .arg(&kb)
16067                .arg(&vh)
16068                .arg(o)
16069                .arg(&hd)
16070                .arg(&nh)
16071                .arg(&nhkv)
16072                .arg(&ti)
16073                .arg(&tkvi)
16074                .arg(&scale)
16075                .arg(&cz)
16076                .arg(&wi);
16077            unsafe {
16078                b.launch(cfg)?;
16079            }
16080            return Ok(());
16081        }
16082        if p1 {
16083            let f = self.func("fa_prefill_w_bf16_p1");
16084            let shmem =
16085                (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
16086            use cudarc::driver::sys::CUfunction_attribute_enum as A;
16087            f.set_attribute(
16088                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16089                shmem as i32,
16090            )?;
16091            let cfg = LaunchConfig {
16092                grid_dim: (
16093                    (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
16094                    n_head as u32,
16095                    1,
16096                ),
16097                block_dim: (32, 4, 1),
16098                shared_mem_bytes: shmem,
16099            };
16100            let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
16101                head_dim as i32,
16102                n_head as i32,
16103                n_head_kv as i32,
16104                t as i32,
16105                t_kv as i32,
16106                causal as i32,
16107                window as i32,
16108            );
16109            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
16110            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
16111            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
16112            let __s_b = self.gpu.stream();
16113            let mut b = __s_b.launch_builder(&f);
16114            b.arg(&qb)
16115                .arg(&kb)
16116                .arg(&vb)
16117                .arg(o)
16118                .arg(&hd)
16119                .arg(&nh)
16120                .arg(&nhkv)
16121                .arg(&ti)
16122                .arg(&tkvi)
16123                .arg(&scale)
16124                .arg(&cz)
16125                .arg(&wi);
16126            unsafe {
16127                b.launch(cfg)?;
16128            }
16129            return Ok(());
16130        }
16131        // MQA head-grouping (MEMRA_FAW_G4=0 reverts): 4 heads/CTA share the staged K/V —
16132        // per-(head,row) FP chain identical to the per-head stamp -> bit-identical (gated).
16133        static G4_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16134        let g4 = !floor
16135            && !f32_stage
16136            && n_head_kv == 1
16137            && n_head % 4 == 0
16138            && *G4_ON.get_or_init(|| {
16139                std::env::var("MEMRA_FAW_G4")
16140                    .map(|v| v != "0")
16141                    .unwrap_or(true)
16142            });
16143        if g4 {
16144            const SP_M: usize = 16;
16145            // Occupancy-2 twin (MEMRA_FAW_O2=0 reverts): one shared K/V buffer inside the dead
16146            // Q-stage region -> ~36.5KB smem, 2 CTA/SM (the llama hd256 mechanism). Bit-identical.
16147            static O2_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16148            let o2 = *O2_ON.get_or_init(|| {
16149                std::env::var("MEMRA_FAW_O2")
16150                    .map(|v| v != "0")
16151                    .unwrap_or(true)
16152            });
16153            let f = self.func(if o2 {
16154                "fa_prefill_w_bf16_g4o2"
16155            } else {
16156                "fa_prefill_w_bf16_g4"
16157            });
16158            let shmem = if o2 {
16159                (2 * (4 * SP_M * head_dim + 4 * SP_M * BK) + 4 * (4 * SP_M)) as u32
16160            } else {
16161                (2 * (2 * BK * head_dim + 4 * SP_M * head_dim + 4 * SP_M * BK) + 4 * (4 * SP_M))
16162                    as u32
16163            };
16164            use cudarc::driver::sys::CUfunction_attribute_enum as A;
16165            f.set_attribute(
16166                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16167                shmem as i32,
16168            )?;
16169            let cfg = LaunchConfig {
16170                grid_dim: ((t as u32).div_ceil(SP_M as u32), (n_head / 4) as u32, 1),
16171                block_dim: (32, 4, 1),
16172                shared_mem_bytes: shmem,
16173            };
16174            let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
16175                head_dim as i32,
16176                n_head as i32,
16177                n_head_kv as i32,
16178                t as i32,
16179                t_kv as i32,
16180                causal as i32,
16181                window as i32,
16182            );
16183            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
16184            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
16185            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
16186            let __s_b = self.gpu.stream();
16187            let mut b = __s_b.launch_builder(&f);
16188            b.arg(&qb)
16189                .arg(&kb)
16190                .arg(&vb)
16191                .arg(o)
16192                .arg(&hd)
16193                .arg(&nh)
16194                .arg(&nhkv)
16195                .arg(&ti)
16196                .arg(&tkvi)
16197                .arg(&scale)
16198                .arg(&cz)
16199                .arg(&wi);
16200            unsafe {
16201                b.launch(cfg)?;
16202            }
16203            return Ok(());
16204        }
16205        let f = self.func(if floor {
16206            "fa_prefill_w_f32"
16207        } else if f32_stage {
16208            "fa_prefill_w_f32_pp"
16209        } else {
16210            "fa_prefill_w_bf16_pp"
16211        });
16212        let shmem =
16213            (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
16214        use cudarc::driver::sys::CUfunction_attribute_enum as A;
16215        f.set_attribute(
16216            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16217            shmem as i32,
16218        )?;
16219        let cfg = LaunchConfig {
16220            grid_dim: (
16221                (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
16222                n_head as u32,
16223                1,
16224            ),
16225            block_dim: (32, 4, 1),
16226            shared_mem_bytes: shmem,
16227        };
16228        let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
16229            head_dim as i32,
16230            n_head as i32,
16231            n_head_kv as i32,
16232            t as i32,
16233            t_kv as i32,
16234            causal as i32,
16235            window as i32,
16236        );
16237        if f32_stage {
16238            let __s_b = self.gpu.stream();
16239            let mut b = __s_b.launch_builder(&f);
16240            b.arg(q)
16241                .arg(k)
16242                .arg(v)
16243                .arg(o)
16244                .arg(&hd)
16245                .arg(&nh)
16246                .arg(&nhkv)
16247                .arg(&ti)
16248                .arg(&tkvi)
16249                .arg(&scale)
16250                .arg(&cz)
16251                .arg(&wi);
16252            unsafe {
16253                b.launch(cfg)?;
16254            }
16255        } else {
16256            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
16257            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
16258            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
16259            let __s_b = self.gpu.stream();
16260            let mut b = __s_b.launch_builder(&f);
16261            b.arg(&qb)
16262                .arg(&kb)
16263                .arg(&vb)
16264                .arg(o)
16265                .arg(&hd)
16266                .arg(&nh)
16267                .arg(&nhkv)
16268                .arg(&ti)
16269                .arg(&tkvi)
16270                .arg(&scale)
16271                .arg(&cz)
16272                .arg(&wi);
16273            unsafe {
16274                b.launch(cfg)?;
16275            }
16276        }
16277        Ok(())
16278    }
16279
16280    /// hd512 FA prefill (gemma4 GLOBAL layers): BLOCK_Q=32 x 2 warps, Q staged in smem,
16281    /// grid.z = 2 O-halves (each CTA computes the full 512-dim scores, accumulates half the
16282    /// V dims). Replaces the scalar sdpa_naive on the prime path's globals.
16283    #[allow(clippy::too_many_arguments)]
16284    pub fn fa_prefill_hd512(
16285        &self,
16286        q: &CudaSlice<f32>,
16287        k: &CudaSlice<f32>,
16288        v: &CudaSlice<f32>,
16289        o: &mut CudaSlice<f32>,
16290        head_dim: usize,
16291        n_head: usize,
16292        n_head_kv: usize,
16293        t: usize,
16294        t_kv: usize,
16295        scale: f32,
16296        causal: bool,
16297    ) -> Result<(), Box<dyn std::error::Error>> {
16298        // sm_90a rides the mma twins (portable_mma_gated, 2026-07-31 — same flip as _w).
16299        if portable_mma_gated() {
16300            return self.sdpa_naive(
16301                q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
16302            );
16303        }
16304        // Default: pre-convert Q/K/V to bf16 once and stage int4 (8 bf16/copy) — at 1 CTA/SM the
16305        // synchronous stage serializes with compute and MQA re-stages the same K/V per head CTA;
16306        // pre-converting halves staged bytes and cuts stage instructions 8x. BIT-IDENTICAL to the
16307        // f32-staged kernel (the converter applies the same __float2bfloat16 the stage applied;
16308        // kernel_check gates the identity). MEMRA_FA512_STAGE=f32 = rollback to the f32 kernel.
16309        static F32_STAGE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16310        let f32_stage =
16311            *F32_STAGE.get_or_init(|| std::env::var("MEMRA_FA512_STAGE").as_deref() == Ok("f32"));
16312        // Single-pass arm (MEMRA_FA512_SP=0 reverts to the z=2 bf16 kernel): GEMM0 split-K across
16313        // the 2 warps instead of recomputed per O-half CTA — the 2026-07-22 kernel-diff excess.
16314        // Own numeric config (partial-sum order) — battery-gated.
16315        static SP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16316        let sp = !f32_stage
16317            && *SP_ON.get_or_init(|| {
16318                std::env::var("MEMRA_FA512_SP")
16319                    .map(|v| v != "0")
16320                    .unwrap_or(true)
16321            });
16322        self.fa_prefill_hd512_arm(
16323            q,
16324            k,
16325            v,
16326            o,
16327            head_dim,
16328            n_head,
16329            n_head_kv,
16330            t,
16331            t_kv,
16332            scale,
16333            causal,
16334            f32_stage,
16335            sp,
16336            sp && fa_f16pv_on(),
16337        )
16338    }
16339
16340    /// hd512 single-pass FA with PRE-CONVERTED bf16 operands (producer-emitted).
16341    #[allow(clippy::too_many_arguments)]
16342    pub fn fa_prefill_hd512_pre(
16343        &self,
16344        qb: &CudaSlice<u8>,
16345        kb: &CudaSlice<u8>,
16346        vb: &CudaSlice<u8>,
16347        o: &mut CudaSlice<f32>,
16348        head_dim: usize,
16349        n_head: usize,
16350        n_head_kv: usize,
16351        t: usize,
16352        t_kv: usize,
16353        scale: f32,
16354        causal: bool,
16355        v_f16: bool,
16356    ) -> Result<(), Box<dyn std::error::Error>> {
16357        debug_assert_eq!(head_dim, 512);
16358        const SP_M: usize = 16;
16359        const BKS: usize = 32;
16360        // f16-P/V door (MEMRA_FA_F16PV=1): P and the P@V accumulation in f16 (llama's fa=1 VKQ
16361        // class); KQ/softmax/rescale-band/final-normalize stay f32. Own numeric config —
16362        // battery-gated. V bytes must be f16 for the sp16 kernel (stage/ldmatrix are typeless).
16363        let f16pv = fa_f16pv_on();
16364        let nw = if f16pv { fa512_wide_warps() } else { 2 };
16365        let hp = f16pv && fa512_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
16366        debug_assert!(!v_f16 || f16pv, "f16 V emitted without the door on");
16367        let mut vguard = self.fa_vf16_scratch.lock().unwrap();
16368        let vref: &CudaSlice<u8> = if f16pv && !v_f16 {
16369            // Fallback re-encode (producer emitted bf16); the emit lane normally hands f16.
16370            let n = t_kv * n_head_kv * head_dim;
16371            let need = n * 2;
16372            if vguard.as_ref().map(|b| b.len() < need).unwrap_or(true) {
16373                *vguard = Some(self.alloc_uninit::<u8>(need)?);
16374            }
16375            let dst = vguard.as_mut().unwrap();
16376            self.bf16_to_f16_into(vb, n, dst)?;
16377            vguard.as_ref().unwrap()
16378        } else {
16379            vb
16380        };
16381        let f = self.func(if hp {
16382            "fa_prefill_bf16_hd512_sp16h2"
16383        } else {
16384            match (f16pv, nw) {
16385                (true, 4) => "fa_prefill_bf16_hd512_sp16w4",
16386                (true, _) => "fa_prefill_bf16_hd512_sp16",
16387                _ => "fa_prefill_bf16_hd512_sp",
16388            }
16389        });
16390        let (nwarp, npart) = if hp {
16391            (4usize, 4usize)
16392        } else if nw > 2 {
16393            (nw, nw)
16394        } else {
16395            (2, 1)
16396        };
16397        // h2 drops sQ (Q register-resident) and doubles sP/sS/sL for the head pair.
16398        let shmem = if hp {
16399            (2 * (2 * BKS * head_dim + 2 * SP_M * BKS) + 4 * (2 * npart * SP_M * BKS + 2 * SP_M))
16400                as u32
16401        } else {
16402            (2 * (SP_M * head_dim + 2 * BKS * head_dim + SP_M * BKS)
16403                + 4 * (npart * SP_M * BKS + SP_M)) as u32
16404        };
16405        use cudarc::driver::sys::CUfunction_attribute_enum as A;
16406        f.set_attribute(
16407            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16408            shmem as i32,
16409        )?;
16410        let grid_y = if hp {
16411            (n_head / 2) as u32
16412        } else {
16413            n_head as u32
16414        };
16415        let cfg = LaunchConfig {
16416            grid_dim: ((t as u32).div_ceil(SP_M as u32), grid_y, 1),
16417            block_dim: (32, nwarp as u32, 1),
16418            shared_mem_bytes: shmem,
16419        };
16420        let (hd, nh, nhkv, ti, tkvi, cz) = (
16421            head_dim as i32,
16422            n_head as i32,
16423            n_head_kv as i32,
16424            t as i32,
16425            t_kv as i32,
16426            causal as i32,
16427        );
16428        let __s_b = self.gpu.stream();
16429        let mut b = __s_b.launch_builder(&f);
16430        b.arg(qb)
16431            .arg(kb)
16432            .arg(vref)
16433            .arg(o)
16434            .arg(&hd)
16435            .arg(&nh)
16436            .arg(&nhkv)
16437            .arg(&ti)
16438            .arg(&tkvi)
16439            .arg(&scale)
16440            .arg(&cz);
16441        unsafe {
16442            b.launch(cfg)?;
16443        }
16444        Ok(())
16445    }
16446
16447    /// hd512 FA prefill with the stage/sp arms FORCED — the kernel_check gate entry
16448    /// (`fa_prefill_hd512` picks the arms from MEMRA_FA512_STAGE / MEMRA_FA512_SP).
16449    #[allow(clippy::too_many_arguments)]
16450    pub fn fa_prefill_hd512_arm(
16451        &self,
16452        q: &CudaSlice<f32>,
16453        k: &CudaSlice<f32>,
16454        v: &CudaSlice<f32>,
16455        o: &mut CudaSlice<f32>,
16456        head_dim: usize,
16457        n_head: usize,
16458        n_head_kv: usize,
16459        t: usize,
16460        t_kv: usize,
16461        scale: f32,
16462        causal: bool,
16463        f32_stage: bool,
16464        sp: bool,
16465        f16pv: bool,
16466    ) -> Result<(), Box<dyn std::error::Error>> {
16467        debug_assert_eq!(head_dim, 512, "fa_prefill_hd512 is hd512 only");
16468        if sp && !f32_stage {
16469            // Single-pass: 16 q-rows/CTA, 2 warps, grid (ceil(T/16), n_head, 1).
16470            // smem: sQ[16][512] + sK[32][512] + sV[32][512] + sP[16][32] (bf16) + sS[16][32]+sL f32.
16471            // f16pv: sp16 kernel — f16 P + f16 P@V accum, V operand encoded f16.
16472            const SP_M: usize = 16;
16473            const BKS: usize = 32;
16474            let nw = if f16pv { fa512_wide_warps() } else { 2 };
16475            let hp = f16pv && fa512_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
16476            let f = self.func(if hp {
16477                "fa_prefill_bf16_hd512_sp16h2"
16478            } else {
16479                match (f16pv, nw) {
16480                    (true, 4) => "fa_prefill_bf16_hd512_sp16w4",
16481                    (true, _) => "fa_prefill_bf16_hd512_sp16",
16482                    _ => "fa_prefill_bf16_hd512_sp",
16483                }
16484            });
16485            let (nwarp, npart) = if hp {
16486                (4usize, 4usize)
16487            } else if nw > 2 {
16488                (nw, nw)
16489            } else {
16490                (2, 1)
16491            };
16492            let shmem = if hp {
16493                (2 * (2 * BKS * head_dim + 2 * SP_M * BKS)
16494                    + 4 * (2 * npart * SP_M * BKS + 2 * SP_M)) as u32
16495            } else {
16496                (2 * (SP_M * head_dim + 2 * BKS * head_dim + SP_M * BKS)
16497                    + 4 * (npart * SP_M * BKS + SP_M)) as u32
16498            };
16499            use cudarc::driver::sys::CUfunction_attribute_enum as A;
16500            f.set_attribute(
16501                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16502                shmem as i32,
16503            )?;
16504            let grid_y = if hp {
16505                (n_head / 2) as u32
16506            } else {
16507                n_head as u32
16508            };
16509            let cfg = LaunchConfig {
16510                grid_dim: ((t as u32).div_ceil(SP_M as u32), grid_y, 1),
16511                block_dim: (32, nwarp as u32, 1),
16512                shared_mem_bytes: shmem,
16513            };
16514            let (hd, nh, nhkv, ti, tkvi, cz) = (
16515                head_dim as i32,
16516                n_head as i32,
16517                n_head_kv as i32,
16518                t as i32,
16519                t_kv as i32,
16520                causal as i32,
16521            );
16522            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
16523            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
16524            let vb = if f16pv {
16525                self.f32_to_f16(v, t_kv * n_head_kv * head_dim)?
16526            } else {
16527                self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?
16528            };
16529            let __s_b = self.gpu.stream();
16530            let mut b = __s_b.launch_builder(&f);
16531            b.arg(&qb)
16532                .arg(&kb)
16533                .arg(&vb)
16534                .arg(o)
16535                .arg(&hd)
16536                .arg(&nh)
16537                .arg(&nhkv)
16538                .arg(&ti)
16539                .arg(&tkvi)
16540                .arg(&scale)
16541                .arg(&cz);
16542            unsafe {
16543                b.launch(cfg)?;
16544            }
16545            return Ok(());
16546        }
16547        const BLOCK_Q: usize = 32;
16548        const BK: usize = 32;
16549        const HALF: usize = 256;
16550        let f = self.func(if f32_stage {
16551            "fa_prefill_f32_hd512"
16552        } else {
16553            "fa_prefill_bf16_hd512"
16554        });
16555        // sQ[32][512] + sK[BK][512] + sV[BK][256] + sP[32][BK] (bf16) + sL[32] f32
16556        let shmem = (2 * (BLOCK_Q * head_dim + BK * head_dim + BK * HALF + BLOCK_Q * BK)
16557            + 4 * BLOCK_Q) as u32;
16558        use cudarc::driver::sys::CUfunction_attribute_enum as A;
16559        f.set_attribute(
16560            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16561            shmem as i32,
16562        )?;
16563        let cfg = LaunchConfig {
16564            grid_dim: (
16565                (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
16566                n_head as u32,
16567                2,
16568            ),
16569            block_dim: (32, 2, 1),
16570            shared_mem_bytes: shmem,
16571        };
16572        let (hd, nh, nhkv, ti, tkvi, cz) = (
16573            head_dim as i32,
16574            n_head as i32,
16575            n_head_kv as i32,
16576            t as i32,
16577            t_kv as i32,
16578            causal as i32,
16579        );
16580        if f32_stage {
16581            let __s_b = self.gpu.stream();
16582            let mut b = __s_b.launch_builder(&f);
16583            b.arg(q)
16584                .arg(k)
16585                .arg(v)
16586                .arg(o)
16587                .arg(&hd)
16588                .arg(&nh)
16589                .arg(&nhkv)
16590                .arg(&ti)
16591                .arg(&tkvi)
16592                .arg(&scale)
16593                .arg(&cz);
16594            unsafe {
16595                b.launch(cfg)?;
16596            }
16597        } else {
16598            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
16599            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
16600            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
16601            let __s_b = self.gpu.stream();
16602            let mut b = __s_b.launch_builder(&f);
16603            b.arg(&qb)
16604                .arg(&kb)
16605                .arg(&vb)
16606                .arg(o)
16607                .arg(&hd)
16608                .arg(&nh)
16609                .arg(&nhkv)
16610                .arg(&ti)
16611                .arg(&tkvi)
16612                .arg(&scale)
16613                .arg(&cz);
16614            unsafe {
16615                b.launch(cfg)?;
16616            }
16617        }
16618        Ok(())
16619    }
16620
16621    /// rope_neox2 with bf16 EMIT (31B glue lane): identical rope math/stores plus the post-rope
16622    /// values written as bf16 — the FA q/k operands come from this launch (bit-identical to the
16623    /// separate f32_to_bf16 the FA entries would run).
16624    #[allow(clippy::too_many_arguments)]
16625    pub fn rope_neox2_bf16e(
16626        &self,
16627        q: &mut CudaSlice<f32>,
16628        k: &mut CudaSlice<f32>,
16629        qb: &mut CudaSlice<u8>,
16630        kb: &mut CudaSlice<u8>,
16631        pos: &CudaSlice<i32>,
16632        head_dim: usize,
16633        n_dims: usize,
16634        nh_q: usize,
16635        nh_k: usize,
16636        n_tokens: usize,
16637        base: f32,
16638        freq_scale: f32,
16639        ff: Option<&CudaSlice<f32>>,
16640    ) -> Result<(), Box<dyn std::error::Error>> {
16641        let f = self.func("rope_neox2_bf16e_f32");
16642        let rows = ((nh_q + nh_k) * n_tokens) as u32;
16643        let cfg = LaunchConfig {
16644            grid_dim: (rows, 1, 1),
16645            block_dim: ((head_dim / 2) as u32, 1, 1),
16646            shared_mem_bytes: 0,
16647        };
16648        let theta_scale = base.powf(-2.0 / n_dims as f32);
16649        let (hd, nd, nhq, nhk, nt) = (
16650            head_dim as i32,
16651            n_dims as i32,
16652            nh_q as i32,
16653            nh_k as i32,
16654            n_tokens as i32,
16655        );
16656        let __s_b = self.gpu.stream();
16657        let mut b = __s_b.launch_builder(&f);
16658        match ff {
16659            Some(t) => {
16660                b.arg(&mut *q)
16661                    .arg(&mut *k)
16662                    .arg(&mut *qb)
16663                    .arg(&mut *kb)
16664                    .arg(pos)
16665                    .arg(&hd)
16666                    .arg(&nd)
16667                    .arg(&nhq)
16668                    .arg(&nhk)
16669                    .arg(&nt)
16670                    .arg(&theta_scale)
16671                    .arg(&freq_scale)
16672                    .arg(t);
16673                unsafe {
16674                    b.launch(cfg)?;
16675                }
16676            }
16677            None => {
16678                let null: u64 = 0;
16679                b.arg(&mut *q)
16680                    .arg(&mut *k)
16681                    .arg(&mut *qb)
16682                    .arg(&mut *kb)
16683                    .arg(pos)
16684                    .arg(&hd)
16685                    .arg(&nd)
16686                    .arg(&nhq)
16687                    .arg(&nhk)
16688                    .arg(&nt)
16689                    .arg(&theta_scale)
16690                    .arg(&freq_scale)
16691                    .arg(&null);
16692                unsafe {
16693                    b.launch(cfg)?;
16694                }
16695            }
16696        }
16697        Ok(())
16698    }
16699
16700    /// Flat f32 -> bf16 conversion into a fresh scratch buffer (2 bytes/elem). `n % 4 == 0`
16701    /// (float4 in, 4x bf16 out). Feeds the bf16-staged hd512 FA prefill.
16702    pub fn f32_to_bf16(
16703        &self,
16704        x: &CudaSlice<f32>,
16705        n: usize,
16706    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
16707        assert!(n % 4 == 0, "f32_to_bf16 requires n % 4 == 0, got {n}");
16708        let mut y = self.alloc_uninit::<u8>(n * 2)?;
16709        let f = self.func("f32_to_bf16_flat");
16710        let n_i = n as i64;
16711        let cfg = LaunchConfig {
16712            grid_dim: (((n / 4) as u32).div_ceil(256), 1, 1),
16713            block_dim: (256, 1, 1),
16714            shared_mem_bytes: 0,
16715        };
16716        let __s_b = self.gpu.stream();
16717        let mut b = __s_b.launch_builder(&f);
16718        b.arg(x).arg(&mut y).arg(&n_i);
16719        unsafe {
16720            b.launch(cfg)?;
16721        }
16722        Ok(y)
16723    }
16724
16725    pub fn f32_to_f16(
16726        &self,
16727        x: &CudaSlice<f32>,
16728        n: usize,
16729    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
16730        assert!(n % 4 == 0, "f32_to_f16 requires n % 4 == 0, got {n}");
16731        let mut y = self.alloc_uninit::<u8>(n * 2)?;
16732        let f = self.func("f32_to_f16_flat");
16733        let n_i = n as i64;
16734        let cfg = LaunchConfig {
16735            grid_dim: (((n / 4) as u32).div_ceil(256), 1, 1),
16736            block_dim: (256, 1, 1),
16737            shared_mem_bytes: 0,
16738        };
16739        let __s_b = self.gpu.stream();
16740        let mut b = __s_b.launch_builder(&f);
16741        b.arg(x).arg(&mut y).arg(&n_i);
16742        unsafe {
16743            b.launch(cfg)?;
16744        }
16745        Ok(y)
16746    }
16747
16748    /// bf16 bytes -> f16 bytes, n elements (the f16-P/V door's V re-encode on the emit lane).
16749    pub fn bf16_to_f16(
16750        &self,
16751        xb: &CudaSlice<u8>,
16752        n: usize,
16753    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
16754        let mut y = self.alloc_uninit::<u8>(n * 2)?;
16755        self.bf16_to_f16_into(xb, n, &mut y)?;
16756        Ok(y)
16757    }
16758
16759    /// Same conversion into a caller-owned (pooled) buffer; `y.len() >= n*2`.
16760    pub fn bf16_to_f16_into(
16761        &self,
16762        xb: &CudaSlice<u8>,
16763        n: usize,
16764        y: &mut CudaSlice<u8>,
16765    ) -> Result<(), Box<dyn std::error::Error>> {
16766        assert!(n % 2 == 0, "bf16_to_f16 requires n % 2 == 0, got {n}");
16767        assert!(y.len() >= n * 2);
16768        let f = self.func("bf16_to_f16_flat");
16769        let n2 = (n / 2) as i64;
16770        let cfg = LaunchConfig {
16771            grid_dim: (((n / 2) as u32).div_ceil(256), 1, 1),
16772            block_dim: (256, 1, 1),
16773            shared_mem_bytes: 0,
16774        };
16775        let __s_b = self.gpu.stream();
16776        let mut b = __s_b.launch_builder(&f);
16777        b.arg(xb).arg(y).arg(&n2);
16778        unsafe {
16779            b.launch(cfg)?;
16780        }
16781        Ok(())
16782    }
16783
16784    /// task #18 (attn side): varlen FA — bf16 K/V mirrors (2 launches) + ONE
16785    /// fa_prefill_bf16kv launch for every fresh sequence. Same per-block math as the
16786    /// per-seq path (bit-gateable). Caller guarantees: fresh causal (T_kv == T),
16787    /// head_dim in {256, 128}, bf16kv lane on.
16788    #[allow(clippy::too_many_arguments)]
16789    pub fn fa_prefill_vl8(
16790        &self,
16791        seqs: &[FaSeqVl],
16792        head_dim: usize,
16793        n_head: usize,
16794        n_head_kv: usize,
16795        scale: f32,
16796    ) -> Result<(), Box<dyn std::error::Error>> {
16797        const BK: usize = 32;
16798        let b = seqs.len();
16799        assert!(b >= 1 && b <= 8);
16800        let mut packed = [FaSeqVl::default(); 8];
16801        packed[..b].copy_from_slice(seqs);
16802        let v = FaVl8(packed);
16803        let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
16804        let ept = (n_head_kv * head_dim) as i32;
16805        {
16806            let f = self.func("fa_mirror_vl");
16807            let max_n = (max_t as i64) * ept as i64;
16808            let blocks = ((max_n as u32).div_ceil(4)).div_ceil(256);
16809            for which in 0..2i32 {
16810                let cfg = LaunchConfig {
16811                    grid_dim: (blocks, 1, b as u32),
16812                    block_dim: (256, 1, 1),
16813                    shared_mem_bytes: 0,
16814                };
16815                let __s_lb = self.gpu.stream();
16816                let mut lb = __s_lb.launch_builder(&f);
16817                lb.arg(&v).arg(&ept).arg(&which);
16818                unsafe {
16819                    lb.launch(cfg)?;
16820                }
16821            }
16822        }
16823        let hd_sfx = fa_hd_suffix(head_dim)?;
16824        let f = self.func(&format!("fa_prefill_bf16kv_vl{hd_sfx}"));
16825        let block_q = 64usize;
16826        let kv_stages = 2usize;
16827        let shmem = (2 * (kv_stages * 2 * BK * head_dim + block_q * BK)
16828            + 4 * (block_q * BK + 2 * block_q)) as u32;
16829        use cudarc::driver::sys::CUfunction_attribute_enum as A;
16830        f.set_attribute(
16831            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16832            shmem as i32,
16833        )?;
16834        let cfg = LaunchConfig {
16835            grid_dim: (max_t.div_ceil(block_q as u32), n_head as u32, b as u32),
16836            block_dim: (32, 4, 1),
16837            shared_mem_bytes: shmem,
16838        };
16839        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
16840        let __s_lb = self.gpu.stream();
16841        let mut lb = __s_lb.launch_builder(&f);
16842        lb.arg(&v).arg(&hd).arg(&nh).arg(&nhkv).arg(&scale);
16843        unsafe {
16844            lb.launch(cfg)?;
16845        }
16846        Ok(())
16847    }
16848
16849    /// task #18 (attn pre-FA): varlen split + QK-norm + RoPE + KV-append — FOUR launches
16850    /// for every fresh sequence (was 6 x B, plus the q/k/v split copies which the view
16851    /// inputs remove entirely). Fresh-only (append at t0=0, RoPE pos = token index).
16852    #[allow(clippy::too_many_arguments)]
16853    pub fn attn_pre_vl8(
16854        &self,
16855        seqs: &[AttnPreVl],
16856        wq: &CudaSlice<f32>,
16857        wk: &CudaSlice<f32>,
16858        head_dim: usize,
16859        rope_dims: usize,
16860        n_head: usize,
16861        n_head_kv: usize,
16862        eps: f32,
16863        freq_base: f32,
16864        freq_scale: f32,
16865        kv_dim_k: usize,
16866        kv_dim_v: usize,
16867        k_tok_bytes: usize,
16868        v_tok_bytes: usize,
16869    ) -> Result<(), Box<dyn std::error::Error>> {
16870        let b = seqs.len();
16871        assert!(b >= 1 && b <= 8);
16872        let mut packed = [AttnPreVl::default(); 8];
16873        packed[..b].copy_from_slice(seqs);
16874        let v = AttnPreVl8(packed);
16875        let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
16876        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
16877        {
16878            let f = self.func("q_gate_split_vl");
16879            let n = max_t * (n_head * head_dim) as u32;
16880            let cfg = LaunchConfig {
16881                grid_dim: (n.div_ceil(256), 1, b as u32),
16882                block_dim: (256, 1, 1),
16883                shared_mem_bytes: 0,
16884            };
16885            let __s_lb = self.gpu.stream();
16886            let mut lb = __s_lb.launch_builder(&f);
16887            lb.arg(&v).arg(&hd).arg(&nh);
16888            unsafe {
16889                lb.launch(cfg)?;
16890            }
16891        }
16892        {
16893            let f = self.func("attn_rms_vl");
16894            let cfg = LaunchConfig {
16895                grid_dim: (max_t * n_head as u32, 2, b as u32),
16896                block_dim: (rms_block(), 1, 1),
16897                shared_mem_bytes: 0,
16898            };
16899            let __s_lb = self.gpu.stream();
16900            let mut lb = __s_lb.launch_builder(&f);
16901            lb.arg(&v)
16902                .arg(wq)
16903                .arg(wk)
16904                .arg(&hd)
16905                .arg(&nh)
16906                .arg(&nhkv)
16907                .arg(&eps);
16908            unsafe {
16909                lb.launch(cfg)?;
16910            }
16911        }
16912        {
16913            let f = self.func("attn_rope_vl");
16914            let theta_scale = freq_base.powf(-2.0 / rope_dims as f32);
16915            let nd = rope_dims as i32;
16916            let cfg = LaunchConfig {
16917                grid_dim: (max_t * n_head as u32, 2, b as u32),
16918                block_dim: ((head_dim / 2) as u32, 1, 1),
16919                shared_mem_bytes: 0,
16920            };
16921            let __s_lb = self.gpu.stream();
16922            let mut lb = __s_lb.launch_builder(&f);
16923            lb.arg(&v)
16924                .arg(&hd)
16925                .arg(&nd)
16926                .arg(&nh)
16927                .arg(&nhkv)
16928                .arg(&theta_scale)
16929                .arg(&freq_scale);
16930            unsafe {
16931                lb.launch(cfg)?;
16932            }
16933        }
16934        {
16935            let f = self.func("append_kv_vl");
16936            let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
16937            let cfg = LaunchConfig {
16938                grid_dim: (nblk, max_t, b as u32),
16939                block_dim: (32, 1, 1),
16940                shared_mem_bytes: 0,
16941            };
16942            let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
16943            let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
16944            let __s_lb = self.gpu.stream();
16945            let mut lb = __s_lb.launch_builder(&f);
16946            lb.arg(&v).arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
16947            unsafe {
16948                lb.launch(cfg)?;
16949            }
16950        }
16951        Ok(())
16952    }
16953
16954    /// FA prefill where K/V are QUANTIZED CudaViews into the resident byte KV cache (the T=K verify
16955    /// path, MTP-PLAN §D.3). Uses `fa_prefill_q` (inline-dequant during stage-to-smem). The view's
16956    /// base+offset pointer is honored; the kernel reads [0..t_kv*tok_bytes). Q is the T fresh query
16957    /// rows; t = T, t_kv = cache len. k_tok_bytes/v_tok_bytes are the per-token byte strides.
16958    pub fn fa_prefill_view(
16959        &self,
16960        q: &CudaSlice<f32>,
16961        k: &cudarc::driver::CudaView<u8>,
16962        v: &cudarc::driver::CudaView<u8>,
16963        o: &mut CudaSlice<f32>,
16964        head_dim: usize,
16965        n_head: usize,
16966        n_head_kv: usize,
16967        t: usize,
16968        t_kv: usize,
16969        scale: f32,
16970        causal: bool,
16971        k_tok_bytes: usize,
16972        v_tok_bytes: usize,
16973        g: bool,
16974    ) -> Result<(), Box<dyn std::error::Error>> {
16975        if portable_mma_gated() {
16976            return self.sdpa_naive_quantized_view(
16977                q,
16978                k,
16979                v,
16980                o,
16981                head_dim,
16982                n_head,
16983                n_head_kv,
16984                t,
16985                t_kv,
16986                scale,
16987                causal,
16988                k_tok_bytes,
16989                v_tok_bytes,
16990            );
16991        }
16992        const BLOCK_Q: usize = 64;
16993        const BK: usize = 32;
16994        // g = e4m3 cache: the kernel parses via DQ_K_ELEM/DQ_V_ELEM (format macros) — the
16995        // kf8vf8-module stamp reads fp8 with the identical MMA/softmax/PV body.
16996        let name = format!("fa_prefill_q{}", fa_hd_suffix(head_dim)?);
16997        let f = if g {
16998            self.func_g(&name)
16999        } else {
17000            self.func(&name)
17001        };
17002        let shmem =
17003            (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
17004        use cudarc::driver::sys::CUfunction_attribute_enum as A;
17005        f.set_attribute(
17006            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
17007            shmem as i32,
17008        )?;
17009        let cfg = LaunchConfig {
17010            grid_dim: (
17011                (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
17012                n_head as u32,
17013                1,
17014            ),
17015            block_dim: (32, 4, 1),
17016            shared_mem_bytes: shmem,
17017        };
17018        let (hd, nh, nhkv, ti, tkvi, cz) = (
17019            head_dim as i32,
17020            n_head as i32,
17021            n_head_kv as i32,
17022            t as i32,
17023            t_kv as i32,
17024            causal as i32,
17025        );
17026        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
17027        let __s_b = self.gpu.stream();
17028        let mut b = __s_b.launch_builder(&f);
17029        b.arg(q)
17030            .arg(k)
17031            .arg(v)
17032            .arg(o)
17033            .arg(&hd)
17034            .arg(&nh)
17035            .arg(&nhkv)
17036            .arg(&ti)
17037            .arg(&tkvi)
17038            .arg(&scale)
17039            .arg(&cz)
17040            .arg(&ktb)
17041            .arg(&vtb);
17042        unsafe {
17043            b.launch(cfg)?;
17044        }
17045        Ok(())
17046    }
17047
17048    /// ARC B (2026-07-05): dequant-once chunk-prime FA. Same contract as `fa_prefill_view`, but
17049    /// instead of every (q-block, head) CTA re-dequanting the whole quantized KV stream inline
17050    /// (T/64 x n_head redundant at chunk prime — 30.5% of the 32k prime wall), dequant the full
17051    /// [t_kv, kv_dim] K and V ONCE into a resident bf16 workspace (fa_dequant_kv_ws_bf16), then
17052    /// run `fa_prefill_qw` (the bf16-workspace twin) over it. EXACT: the workspace holds the same
17053    /// __float2bfloat16(dq_*_elem(...)) values fa_prefill_q stages to smem, and the twin's MMA/
17054    /// softmax/PV code is byte-identical -> bit-identical O (kernel_check pins bitdiff=0).
17055    /// The workspace allocation is REUSED across layers/chunks (grown to the largest shape);
17056    /// contents are rewritten per call. MEMRA_PRIME_DEQW=0 falls back to fa_prefill_view (callers gate).
17057    #[allow(clippy::too_many_arguments)]
17058    pub fn fa_prefill_view_ws(
17059        &self,
17060        q: &CudaSlice<f32>,
17061        k: &cudarc::driver::CudaView<u8>,
17062        v: &cudarc::driver::CudaView<u8>,
17063        o: &mut CudaSlice<f32>,
17064        head_dim: usize,
17065        n_head: usize,
17066        n_head_kv: usize,
17067        t: usize,
17068        t_kv: usize,
17069        scale: f32,
17070        causal: bool,
17071        k_tok_bytes: usize,
17072        v_tok_bytes: usize,
17073        g: bool,
17074    ) -> Result<(), Box<dyn std::error::Error>> {
17075        if portable_mma_gated() {
17076            return self.sdpa_naive_quantized_view(
17077                q,
17078                k,
17079                v,
17080                o,
17081                head_dim,
17082                n_head,
17083                n_head_kv,
17084                t,
17085                t_kv,
17086                scale,
17087                causal,
17088                k_tok_bytes,
17089                v_tok_bytes,
17090            );
17091        }
17092        const BLOCK_Q: usize = 64;
17093        const BK: usize = 32;
17094        let kv_dim_k = n_head_kv * head_dim;
17095        let kv_dim_v = n_head_kv * head_dim;
17096        let k_ws_bytes = t_kv * kv_dim_k * 2; // bf16
17097        let v_ws_bytes = t_kv * kv_dim_v * 2;
17098        // Lock held across BOTH launches: enqueue-only (µs), all compute serializes on gpu.stream.
17099        let mut guard = self.prime_deqw_ws.lock().unwrap();
17100        let need_grow = match guard.as_ref() {
17101            Some((kw, vw)) => kw.len() < k_ws_bytes || vw.len() < v_ws_bytes,
17102            None => true,
17103        };
17104        if need_grow {
17105            let grow = |cur: usize, need: usize| if cur >= need { cur } else { need };
17106            let (ck, cv) = guard
17107                .as_ref()
17108                .map(|(a, b)| (a.len(), b.len()))
17109                .unwrap_or((0, 0));
17110            *guard = Some((
17111                self.alloc_u8(grow(ck, k_ws_bytes))?,
17112                self.alloc_u8(grow(cv, v_ws_bytes))?,
17113            ));
17114        }
17115        let (kw, vw) = guard.as_mut().unwrap();
17116        // pass 1: dequant K+V once into the bf16 workspace (grid-stride, 1 thread/elem)
17117        {
17118            // only THIS pass parses KV bytes — pass 2 reads the bf16 workspace (format-free).
17119            let f = if g {
17120                self.func_g("fa_dequant_kv_ws_bf16")
17121            } else {
17122                self.func("fa_dequant_kv_ws_bf16")
17123            };
17124            let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
17125            let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
17126            let cfg = LaunchConfig {
17127                grid_dim: (nblk.max(1), 1, 1),
17128                block_dim: (256, 1, 1),
17129                shared_mem_bytes: 0,
17130            };
17131            let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
17132            let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
17133            let __s_b = self.gpu.stream();
17134            let mut b = __s_b.launch_builder(&f);
17135            b.arg(k)
17136                .arg(v)
17137                .arg(&mut *kw)
17138                .arg(&mut *vw)
17139                .arg(&kdk)
17140                .arg(&kdv)
17141                .arg(&tkvi)
17142                .arg(&ktb)
17143                .arg(&vtb);
17144            unsafe {
17145                b.launch(cfg)?;
17146            }
17147        }
17148        // pass 2: the bf16-workspace prefill twin (same tile sizes/loop structure as fa_prefill_q).
17149        // DEFAULT: cp.async double-buffered staging twin (fa_prefill_qw_db, +32KB smem for the
17150        // second K/V tile pair, 1 CTA/SM): overlaps tile n+1's L2->smem copy with tile n's MMA.
17151        // Bit-identical output (staging is a pure byte copy; kernel_check pins bitdiff=0 under
17152        // both twins). A/B (27B g7e, N=3): 32k prime 17.10->16.51s, 16k 9.09->8.65s — the copy
17153        // latency hides behind the MMA pipe and beats the 2-CTA/SM occupancy of the sync twin.
17154        // MEMRA_PRIME_DEQW_DB=0 falls back to the single-buffer twin.
17155        let db = std::env::var("MEMRA_PRIME_DEQW_DB")
17156            .map(|v| v != "0")
17157            .unwrap_or(true);
17158        {
17159            let hd_sfx = fa_hd_suffix(head_dim)?;
17160            let f = self.func(&format!(
17161                "fa_prefill_qw{}{hd_sfx}",
17162                if db { "_db" } else { "" }
17163            ));
17164            let shmem = if db {
17165                // 4x KV tile buffers (bf16) + sP (bf16) + sL (f32)
17166                (2 * (4 * BK * head_dim + BLOCK_Q * BK) + 4 * BLOCK_Q) as u32
17167            } else {
17168                (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32
17169            };
17170            use cudarc::driver::sys::CUfunction_attribute_enum as A;
17171            f.set_attribute(
17172                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
17173                shmem as i32,
17174            )?;
17175            let cfg = LaunchConfig {
17176                grid_dim: (
17177                    (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
17178                    n_head as u32,
17179                    1,
17180                ),
17181                block_dim: (32, 4, 1),
17182                shared_mem_bytes: shmem,
17183            };
17184            let (hd, nh, nhkv, ti, tkvi, cz) = (
17185                head_dim as i32,
17186                n_head as i32,
17187                n_head_kv as i32,
17188                t as i32,
17189                t_kv as i32,
17190                causal as i32,
17191            );
17192            let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
17193            let __s_b = self.gpu.stream();
17194            let mut b = __s_b.launch_builder(&f);
17195            b.arg(q)
17196                .arg(&*kw)
17197                .arg(&*vw)
17198                .arg(o)
17199                .arg(&hd)
17200                .arg(&nh)
17201                .arg(&nhkv)
17202                .arg(&ti)
17203                .arg(&tkvi)
17204                .arg(&scale)
17205                .arg(&cz)
17206                .arg(&kdk)
17207                .arg(&kdv);
17208            unsafe {
17209                b.launch(cfg)?;
17210            }
17211        }
17212        Ok(())
17213    }
17214
17215    /// WINDOWED `fa_prefill_view_ws` twin at head_dim 128 (lane/pp-prefill 2026-08-07):
17216    /// step35's SWA prefill (win=512, 33 of 45 layers) previously had NO windowed FA prefill
17217    /// stamp — every windowed twin was hd256-only — and took `sdpa_naive_w_quantized_view`,
17218    /// the f32 floor, at 565 ms/layer on a pp4096 where the hd128 FA family does the harder
17219    /// causal-4096 in 3.3 ms (41% of the whole prime; research/pp-prefill-20260807 anatomy).
17220    /// Same two-pass shape as the unwindowed function: dequant K/V ONCE into the resident
17221    /// bf16 workspace, then the windowed qw kernel (`fa_prefill_qw_db_w_hd128`, cp.async
17222    /// double-buffered; MEMRA_PRIME_DEQW_DB=0 selects the single-buffer twin). The window
17223    /// mask is `fa_prefill_f32_body`'s exact predicate; `window == 0` is bit-identical to
17224    /// `fa_prefill_view_ws` by construction (default-arg body). NEW NUMERIC CLASS vs the
17225    /// f32 floor on SWA rows (bf16 MMA online-softmax vs f32 serial softmax) — adoption is
17226    /// gated by the full battery, and the class must change UNIFORMLY for a whole request
17227    /// (kernel selection keys on seq_end, never per chunk — the chunkfix law).
17228    /// hd128-only deliberately: the only windowed-prefill consumer at another head_dim is
17229    /// gemma4 (hd256), which already has `fa_prefill_w_f32`.
17230    #[allow(clippy::too_many_arguments)]
17231    pub fn fa_prefill_view_ws_w_hd128(
17232        &self,
17233        q: &CudaSlice<f32>,
17234        k: &cudarc::driver::CudaView<u8>,
17235        v: &cudarc::driver::CudaView<u8>,
17236        o: &mut CudaSlice<f32>,
17237        head_dim: usize,
17238        n_head: usize,
17239        n_head_kv: usize,
17240        t: usize,
17241        t_kv: usize,
17242        scale: f32,
17243        causal: bool,
17244        window: usize,
17245        k_tok_bytes: usize,
17246        v_tok_bytes: usize,
17247    ) -> Result<(), Box<dyn std::error::Error>> {
17248        assert_eq!(
17249            head_dim, 128,
17250            "fa_prefill_view_ws_w_hd128: only the hd128 twin is stamped"
17251        );
17252        if portable_mma_gated() {
17253            return self.sdpa_naive_w_quantized_view(
17254                q,
17255                k,
17256                v,
17257                o,
17258                head_dim,
17259                n_head,
17260                n_head_kv,
17261                t,
17262                t_kv,
17263                scale,
17264                causal,
17265                window,
17266                k_tok_bytes,
17267                v_tok_bytes,
17268            );
17269        }
17270        const BLOCK_Q: usize = 64;
17271        const BK: usize = 32;
17272        let kv_dim_k = n_head_kv * head_dim;
17273        let kv_dim_v = n_head_kv * head_dim;
17274        let k_ws_bytes = t_kv * kv_dim_k * 2; // bf16
17275        let v_ws_bytes = t_kv * kv_dim_v * 2;
17276        let mut guard = self.prime_deqw_ws.lock().unwrap();
17277        let need_grow = match guard.as_ref() {
17278            Some((kw, vw)) => kw.len() < k_ws_bytes || vw.len() < v_ws_bytes,
17279            None => true,
17280        };
17281        if need_grow {
17282            let grow = |cur: usize, need: usize| if cur >= need { cur } else { need };
17283            let (ck, cv) = guard
17284                .as_ref()
17285                .map(|(a, b)| (a.len(), b.len()))
17286                .unwrap_or((0, 0));
17287            *guard = Some((
17288                self.alloc_u8(grow(ck, k_ws_bytes))?,
17289                self.alloc_u8(grow(cv, v_ws_bytes))?,
17290            ));
17291        }
17292        let (kw, vw) = guard.as_mut().unwrap();
17293        // pass 1: dequant K+V once into the bf16 workspace (identical to fa_prefill_view_ws —
17294        // the workspace bytes are the SAME __float2bfloat16(dq(...)) values either way).
17295        {
17296            let f = self.func("fa_dequant_kv_ws_bf16");
17297            let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
17298            let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
17299            let cfg = LaunchConfig {
17300                grid_dim: (nblk.max(1), 1, 1),
17301                block_dim: (256, 1, 1),
17302                shared_mem_bytes: 0,
17303            };
17304            let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
17305            let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
17306            let __s_b = self.gpu.stream();
17307            let mut b = __s_b.launch_builder(&f);
17308            b.arg(k)
17309                .arg(v)
17310                .arg(&mut *kw)
17311                .arg(&mut *vw)
17312                .arg(&kdk)
17313                .arg(&kdv)
17314                .arg(&tkvi)
17315                .arg(&ktb)
17316                .arg(&vtb);
17317            unsafe {
17318                b.launch(cfg)?;
17319            }
17320        }
17321        // pass 2: the WINDOWED qw twin (db default, same as the unwindowed wrapper).
17322        let db = std::env::var("MEMRA_PRIME_DEQW_DB")
17323            .map(|v| v != "0")
17324            .unwrap_or(true);
17325        {
17326            let f = self.func(if db {
17327                "fa_prefill_qw_db_w_hd128"
17328            } else {
17329                "fa_prefill_qw_w_hd128"
17330            });
17331            let shmem = if db {
17332                (2 * (4 * BK * head_dim + BLOCK_Q * BK) + 4 * BLOCK_Q) as u32
17333            } else {
17334                (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32
17335            };
17336            use cudarc::driver::sys::CUfunction_attribute_enum as A;
17337            f.set_attribute(
17338                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
17339                shmem as i32,
17340            )?;
17341            let cfg = LaunchConfig {
17342                grid_dim: (
17343                    (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
17344                    n_head as u32,
17345                    1,
17346                ),
17347                block_dim: (32, 4, 1),
17348                shared_mem_bytes: shmem,
17349            };
17350            let (hd, nh, nhkv, ti, tkvi, cz) = (
17351                head_dim as i32,
17352                n_head as i32,
17353                n_head_kv as i32,
17354                t as i32,
17355                t_kv as i32,
17356                causal as i32,
17357            );
17358            let (kdk, kdv, wnd) = (kv_dim_k as i32, kv_dim_v as i32, window as i32);
17359            let __s_b = self.gpu.stream();
17360            let mut b = __s_b.launch_builder(&f);
17361            b.arg(q)
17362                .arg(&*kw)
17363                .arg(&*vw)
17364                .arg(o)
17365                .arg(&hd)
17366                .arg(&nh)
17367                .arg(&nhkv)
17368                .arg(&ti)
17369                .arg(&tkvi)
17370                .arg(&scale)
17371                .arg(&cz)
17372                .arg(&kdk)
17373                .arg(&kdv)
17374                .arg(&wnd);
17375            unsafe {
17376                b.launch(cfg)?;
17377            }
17378        }
17379        Ok(())
17380    }
17381
17382    /// FA decode (T=1 split-K) over the resident QUANTIZED KV cache (q8_0 K / q5_1 V) as u8 views.
17383    /// Replaces sdpa_naive_view for decode; inline-dequants per element. k_tok_bytes/v_tok_bytes are
17384    /// the per-token byte strides (differ: q8_0=34*nblk, q5_1=24*nblk per token).
17385    pub fn fa_decode(
17386        &self,
17387        q: &CudaSlice<f32>,
17388        k: &cudarc::driver::CudaView<u8>,
17389        v: &cudarc::driver::CudaView<u8>,
17390        o: &mut CudaSlice<f32>,
17391        head_dim: usize,
17392        n_head: usize,
17393        n_head_kv: usize,
17394        t_kv: usize,
17395        scale: f32,
17396        k_tok_bytes: usize,
17397        v_tok_bytes: usize,
17398    ) -> Result<(), Box<dyn std::error::Error>> {
17399        self.fa_decode_kvmod(
17400            q,
17401            k,
17402            v,
17403            o,
17404            head_dim,
17405            n_head,
17406            n_head_kv,
17407            t_kv,
17408            scale,
17409            k_tok_bytes,
17410            v_tok_bytes,
17411            false,
17412        )
17413    }
17414
17415    /// `fa_decode` with an explicit fp8-module flag (`g`): gemma windowed layers under
17416    /// MEMRA_GEMMA_WKV read an e4m3 cache — every kernel must come from the kf8vf8 module
17417    /// and the v4 lane (q8_0-hardcoded staging) is excluded.
17418    #[allow(clippy::too_many_arguments)]
17419    /// UNIFIED scalar decode launch (fa_decode_f32, nullable-ctr): ONE symbol for host-len
17420    /// (kvmod eager) and device-len (graph/stream) callers — the textually-identical f32_dc
17421    /// twin compiled apart and its ULP drift flipped 31B verify argmaxes (2026-07-12).
17422    #[allow(clippy::too_many_arguments)]
17423    #[allow(clippy::too_many_arguments)]
17424    fn fa_decode_scalar_unified(
17425        &self,
17426        q: &cudarc::driver::CudaView<f32>,
17427        k: &cudarc::driver::CudaView<u8>,
17428        v: &cudarc::driver::CudaView<u8>,
17429        o: &mut cudarc::driver::CudaViewMut<f32>,
17430        head_dim: usize,
17431        n_head: usize,
17432        n_head_kv: usize,
17433        t_kv_host: usize,
17434        t_kv_dev: Option<&CudaSlice<i32>>,
17435        scale: f32,
17436        n_splits: usize,
17437        split_keys: usize,
17438        k_tok_bytes: usize,
17439        v_tok_bytes: usize,
17440        g: bool,
17441        part_o: &mut CudaSlice<f32>,
17442        part_m: &mut CudaSlice<f32>,
17443        part_l: &mut CudaSlice<f32>,
17444        q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
17445    ) -> Result<(), Box<dyn std::error::Error>> {
17446        let f = if g {
17447            self.func_g("fa_decode_f32")
17448        } else {
17449            self.fa_func("fa_decode_f32", head_dim)
17450        };
17451        let cfg = LaunchConfig {
17452            grid_dim: (n_head as u32, n_splits as u32, 1),
17453            block_dim: (head_dim as u32, 1, 1),
17454            shared_mem_bytes: (4 * (head_dim + 32)) as u32,
17455        };
17456        let (hd, nh, nhkv, nsp) = (
17457            head_dim as i32,
17458            n_head as i32,
17459            n_head_kv as i32,
17460            n_splits as i32,
17461        );
17462        let (ktb, vtb, tkvi, ski) = (
17463            k_tok_bytes as i64,
17464            v_tok_bytes as i64,
17465            t_kv_host as i32,
17466            split_keys as i32,
17467        );
17468        let __s_b = self.gpu.stream();
17469        let mut b = __s_b.launch_builder(&f);
17470        match t_kv_dev {
17471            Some(d) => {
17472                b.arg(q)
17473                    .arg(k)
17474                    .arg(v)
17475                    .arg(&mut *part_o)
17476                    .arg(&mut *part_m)
17477                    .arg(&mut *part_l)
17478                    .arg(&hd)
17479                    .arg(&nh)
17480                    .arg(&nhkv)
17481                    .arg(&tkvi)
17482                    .arg(d)
17483                    .arg(&scale)
17484                    .arg(&nsp)
17485                    .arg(&ski)
17486                    .arg(&ktb)
17487                    .arg(&vtb);
17488                unsafe {
17489                    b.launch(cfg)?;
17490                }
17491            }
17492            None => {
17493                let null: u64 = 0;
17494                b.arg(q)
17495                    .arg(k)
17496                    .arg(v)
17497                    .arg(&mut *part_o)
17498                    .arg(&mut *part_m)
17499                    .arg(&mut *part_l)
17500                    .arg(&hd)
17501                    .arg(&nh)
17502                    .arg(&nhkv)
17503                    .arg(&tkvi)
17504                    .arg(&null)
17505                    .arg(&scale)
17506                    .arg(&nsp)
17507                    .arg(&ski)
17508                    .arg(&ktb)
17509                    .arg(&vtb);
17510                unsafe {
17511                    b.launch(cfg)?;
17512                }
17513            }
17514        }
17515        let cfg2 = LaunchConfig {
17516            grid_dim: (n_head as u32, 1, 1),
17517            block_dim: (head_dim as u32, 1, 1),
17518            shared_mem_bytes: 0,
17519        };
17520        if let Some((oq, od)) = q8_out {
17521            // wave-5b: q8-emitting combine — the wo matmul_pre consumes the pair directly.
17522            let fc = if g {
17523                self.func_g("fa_decode_combine_q8_1")
17524            } else {
17525                self.fa_func("fa_decode_combine_q8_1", head_dim)
17526            };
17527            let __s_b2 = self.gpu.stream();
17528            let mut b2 = __s_b2.launch_builder(&fc);
17529            b2.arg(&*part_o)
17530                .arg(&*part_m)
17531                .arg(&*part_l)
17532                .arg(oq)
17533                .arg(od)
17534                .arg(&hd)
17535                .arg(&nh)
17536                .arg(&nsp);
17537            unsafe {
17538                b2.launch(cfg2)?;
17539            }
17540            return Ok(());
17541        }
17542        let fc = if g {
17543            self.func_g("fa_decode_combine_f32")
17544        } else {
17545            self.fa_func("fa_decode_combine_f32", head_dim)
17546        };
17547        let __s_b2 = self.gpu.stream();
17548        let mut b2 = __s_b2.launch_builder(&fc);
17549        b2.arg(&*part_o)
17550            .arg(&*part_m)
17551            .arg(&*part_l)
17552            .arg(o)
17553            .arg(&hd)
17554            .arg(&nh)
17555            .arg(&nsp);
17556        unsafe {
17557            b2.launch(cfg2)?;
17558        }
17559        Ok(())
17560    }
17561
17562    pub fn fa_decode_kvmod(
17563        &self,
17564        q: &CudaSlice<f32>,
17565        k: &cudarc::driver::CudaView<u8>,
17566        v: &cudarc::driver::CudaView<u8>,
17567        o: &mut CudaSlice<f32>,
17568        head_dim: usize,
17569        n_head: usize,
17570        n_head_kv: usize,
17571        t_kv: usize,
17572        scale: f32,
17573        k_tok_bytes: usize,
17574        v_tok_bytes: usize,
17575        g: bool,
17576    ) -> Result<(), Box<dyn std::error::Error>> {
17577        let q_view = q.as_view();
17578        let mut o_view = o.as_view_mut();
17579        self.fa_decode_kvmod_view(
17580            &q_view,
17581            k,
17582            v,
17583            &mut o_view,
17584            head_dim,
17585            n_head,
17586            n_head_kv,
17587            t_kv,
17588            scale,
17589            k_tok_bytes,
17590            v_tok_bytes,
17591            g,
17592        )
17593    }
17594
17595    /// Row-view entry into `fa_decode_kvmod`. The kernel sees the selected Q/output rows as its
17596    /// base pointers, so the launch geometry and arithmetic are identical to the owned-slice entry.
17597    /// Batched fallback callers use this to avoid materializing rows around an otherwise unchanged
17598    /// per-session KV view and FA launch.
17599    #[allow(clippy::too_many_arguments)]
17600    pub fn fa_decode_kvmod_view(
17601        &self,
17602        q: &cudarc::driver::CudaView<f32>,
17603        k: &cudarc::driver::CudaView<u8>,
17604        v: &cudarc::driver::CudaView<u8>,
17605        o: &mut cudarc::driver::CudaViewMut<f32>,
17606        head_dim: usize,
17607        n_head: usize,
17608        n_head_kv: usize,
17609        t_kv: usize,
17610        scale: f32,
17611        k_tok_bytes: usize,
17612        v_tok_bytes: usize,
17613        g: bool,
17614    ) -> Result<(), Box<dyn std::error::Error>> {
17615        // PERF-4: the warp-per-token vec path replaces the scalar element-per-thread fa_decode_f32 —
17616        // warp-per-token fa_decode_vec_q (grid=(n_head_kv,n_splits), block=(32,gqa_ratio)).
17617        // The block dequants each KV tile ONCE into smem (bf16) and broadcasts to all gqa Q-head
17618        // warps -> each KV byte leaves HBM/L2 ~1x/group (vs 4x). ARGS identical; func/grid/block/
17619        // smem/n_splits differ. fa_decode_f32 stays the bit-reference fallback. Combine is shared.
17620        //
17621        // SPLIT-K: the scalar path has grid.x=n_head (32) blocks; the vec path only has
17622        // grid.x=n_head_kv (8). To avoid starving the GPU at mid ctx, the vec path splits MORE
17623        // aggressively (64 keys/split vs 256) so grid.y rises and 8*n_splits fills the SMs.
17624        // At VERY short ctx (t_kv<96) even 1 split can't fill the GPU from 8 KV heads, so the
17625        // broadcast can't beat the scalar path's 4x-more-blocks latency hiding — fall back to
17626        // scalar there (measured crossover: vec 0.68x at t_kv=64, 1.23x at t_kv=96, 2.2x at 256).
17627        // DEFAULT-ON (2026-06-28): clean clock-locked sweep proved vec beats scalar at every
17628        // t_kv>=96 and the gain WIDENS with ctx (graph decode: +9.5% @128, +11.6% @512, +11.8%
17629        // @2048) — the KV-byte-broadcast (4x fewer HBM reads/group) compounds as attention grows.
17630        // MEMRA_NO_FA_VEC forces the scalar bit-reference. Below FA_VEC_MIN_TKV the scalar path's
17631        // 4x-more-blocks (grid.x=n_head=32 vs n_head_kv=8) hides latency better, so keep scalar there.
17632        // g + no-v4: the g-module REGISTER twin mis-decodes the gemma windowed shape
17633        // (root-cause open, jsonl) — only reachable by forcing v4 off (MEMRA_FA_V4_MAX);
17634        // fall to the exact scalar there instead of the broken register arm.
17635        let mut fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && t_kv >= fa_vec_min_tkv();
17636        // hd256 under g: v4-or-scalar. hd128/other under g ride the REGISTER g-lane (the
17637        // dq_K_lane/dq_V_lane macros are format-aware; v3 is format-gated off, v2/smem
17638        // arms are excluded under g). Mirrored in kvmod / fa_decode_dc / fa_geom_eager.
17639        if g && head_dim == 256 && !fa_v4_at(t_kv) {
17640            fa_vec = false;
17641        }
17642        let sp = fa_split_keys(t_kv, n_head_kv);
17643        let n_splits = if fa_vec {
17644            ((t_kv + sp - 1) / sp).max(1)
17645        } else {
17646            ((t_kv + 255) / 256).max(1)
17647        };
17648        let o_len = n_head * n_splits * head_dim;
17649        let ml_len = n_head * n_splits;
17650        let mut part_guard = self.fa_part_pool.lock().unwrap();
17651        if part_guard
17652            .as_ref()
17653            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
17654            .unwrap_or(true)
17655        {
17656            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
17657            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
17658            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
17659            // later live allocations land at those addresses, and the next graph REPLAY writes
17660            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
17661            // output corruption began the burst after the trunk's t_kv growth first realloc'd
17662            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
17663            // the baked addresses alive (single-stream: eager writes the new buffers, replays
17664            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
17665            // (total retired < final size).
17666            let old = part_guard.take();
17667            let (co, cm) = old
17668                .as_ref()
17669                .map(|pp| (pp.0.len(), pp.1.len()))
17670                .unwrap_or((0, 0));
17671            if let Some(old) = old {
17672                self.fa_part_retired.lock().unwrap().push(old);
17673            }
17674            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
17675                eprintln!(
17676                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
17677                    co, o_len, cm, ml_len
17678                );
17679            }
17680            *part_guard = Some((
17681                self.alloc_uninit::<f32>(o_len.max(2 * co))?,
17682                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
17683                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
17684            ));
17685        }
17686        let pg = part_guard.as_mut().unwrap();
17687        self.gpu
17688            .stream()
17689            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
17690        self.gpu
17691            .stream()
17692            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
17693        self.gpu
17694            .stream()
17695            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
17696        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
17697        let (part_o, part_m, part_l) = (&mut *part_o, &mut *part_m, &mut *part_l);
17698        let (hd, nh, nhkv, tkvi, nsp) = (
17699            head_dim as i32,
17700            n_head as i32,
17701            n_head_kv as i32,
17702            t_kv as i32,
17703            n_splits as i32,
17704        );
17705        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
17706        // The vec kernel holds head_dim/32 register accumulators (FA_DEC_MAX_DPL=8 -> head_dim<=256).
17707        // All shipped models use head_dim=256; fall back to scalar for anything wider rather than
17708        // silently truncating the accumulator.
17709        let fa_vec = fa_vec && head_dim <= 512 && head_dim % 32 == 0;
17710        // hd-512 vec crossover (MEMRA_FA512_MIN, default 512): the DPL16 twin wins at depth
17711        // (82.5 -> vec at 1736) but the scalar's more-blocks latency hiding wins at tiny t_kv
17712        // (the same scalar-floor physics as hd256's old 96 floor; short-ctx plain regressed
17713        // 178.4 -> 173.7 when 512 rode vec unconditionally).
17714        let fa512_min = fa512_min_tkv();
17715        // FA-DEEP pick (bit-identical twins, see fa_deep_at): default module only — the
17716        // g-module keeps the v4 pick (its class is not the depth-decay class).
17717        let deep = fa_vec
17718            && head_dim == 256
17719            && fa_v4_at(t_kv)
17720            && !g
17721            && fa_deep_at(t_kv)
17722            && !matches!(fa_v4_mode(), "noB3" | "stage");
17723        let (f, cfg) = if fa_vec && head_dim == 512 && t_kv >= fa512_min {
17724            // gemma4 globals (hd 512): the DPL16 register twin (fa_decode_vec_q body with a
17725            // 16-slot accumulator ceiling). Scalar fallback measured 82.5us/layer at 1736 ctx.
17726            let gqa = (n_head / n_head_kv).max(1) as u32;
17727            let fv = self.fa_func("fa_decode_vec_q_dpl16", head_dim);
17728            (
17729                fv,
17730                LaunchConfig {
17731                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
17732                    block_dim: (32, gqa, 1),
17733                    shared_mem_bytes: 0,
17734                },
17735            )
17736        } else if fa_vec && head_dim <= 256 {
17737            let gqa = (n_head / n_head_kv).max(1) as u32;
17738            // DEEP-CTX smem twin (2026-07-05): the register-dequant path's GQA reuse rides L2,
17739            // which holds to ~8k ctx but dies at 40k (layer KV ~37MB) — the 4 GQA warps then
17740            // re-read every KV byte from DRAM (4x traffic). Above MEMRA_FA_SMEM_TKV (default
17741            // 1024 — the 2026-07-05 crossover re-sweep on real prompts: p3 spec 73.8->79.2 at
17742            // 2048, flat down to 512, p2 +5%, p1/9B unchanged; the ARC-A probe's synthetic
17743            // 2.1x smem-at-all-depths pointed here; 0=never) dispatch the smem-broadcast twin:
17744            // dequant each tile ONCE per block.
17745            // Bit-identical per (token,split): same bf16 round-trip, same accumulation order,
17746            // same partial layout -> same combine. Short/mid ctx keeps the register path (it won
17747            // there by 12x — latency, not bandwidth, rules small KV).
17748            static SMEM_TKV: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
17749            let smem_tkv = *SMEM_TKV.get_or_init(|| {
17750                std::env::var("MEMRA_FA_SMEM_TKV")
17751                    .ok()
17752                    .and_then(|v| v.parse().ok())
17753                    .unwrap_or_else(|| {
17754                        FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
17755                    })
17756            });
17757            if fa_v4_at(t_kv) && head_dim == 256 {
17758                // FA v4 lane (2026-07-10): key-per-lane score phase, zero shuffles per key.
17759                // NEW NUMERIC CONFIG (chunk-serial per-key dot) — battery-arbitrated.
17760                // g (fp8-windowed): the v4 staging is format-aware (2026-07-12) — kf8vf8 module.
17761                let v4name = match fa_v4_mode() {
17762                    "noB3" => "fa_decode_vec_q_v4_noB3", // phase probe (WRONG OUTPUT)
17763                    "stage" => "fa_decode_vec_q_v4_stage", // phase probe (WRONG OUTPUT)
17764                    _ if deep => "fa_decode_vec_q_v4_deep",
17765                    _ => "fa_decode_vec_q_v4",
17766                };
17767                let fv = if g {
17768                    self.func_g(v4name)
17769                } else {
17770                    self.func(v4name)
17771                };
17772                // fa_v4_smem (deep: fa_v4_deep_smem, +640B row pads) + sV (g: raw e4m3 sV
17773                // tile = 1B/elem — half the smem, 3->5 blocks/SM)
17774                let shmem = (if deep { 12160 } else { 11520 }
17775                    + 32 * head_dim * if g { 1 } else { 2 }) as u32;
17776                use cudarc::driver::sys::CUfunction_attribute_enum as A;
17777                fv.set_attribute(
17778                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
17779                    shmem as i32,
17780                )?;
17781                (
17782                    fv,
17783                    LaunchConfig {
17784                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
17785                        block_dim: (32, gqa, 1),
17786                        shared_mem_bytes: shmem,
17787                    },
17788                )
17789            } else if fa_v3_active(head_dim) {
17790                // FA v3 lane: dp4a-K hybrid (register-quantized Q, raw q8_0 K, staged-V kept).
17791                // smem = sV only (half of v2's).
17792                let fv = if g {
17793                    self.func_g("fa_decode_vec_q_v3")
17794                } else {
17795                    self.func("fa_decode_vec_q_v3")
17796                };
17797                let shmem = (32 * head_dim * 2) as u32; // sV bf16 [FA_DEC_TILE=32][hd]
17798                (
17799                    fv,
17800                    LaunchConfig {
17801                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
17802                        block_dim: (32, gqa, 1),
17803                        shared_mem_bytes: shmem,
17804                    },
17805                )
17806            } else if fa_v2_on() {
17807                // FAVENDOR lane: llama fattn-vec tile-batched softmax + wide-load staging on
17808                // OUR smem KV broadcast. Replaces BOTH per-key twins when on; same grid/block/
17809                // partials; same 32KB sK+sV tile as the smem twin.
17810                let fv = if g {
17811                    self.func_g("fa_decode_vec_q_v2")
17812                } else {
17813                    self.func("fa_decode_vec_q_v2")
17814                };
17815                let shmem = (2 * 32 * head_dim * 2) as u32; // sK+sV bf16 [FA_DEC_TILE=32][hd]
17816                (
17817                    fv,
17818                    LaunchConfig {
17819                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
17820                        block_dim: (32, gqa, 1),
17821                        shared_mem_bytes: shmem,
17822                    },
17823                )
17824            } else if smem_tkv > 0 && t_kv >= smem_tkv && !g && !(head_dim == 512 && Self::gkv_on())
17825            {
17826                // (fp8 exclusions: the smem twin's V-stage is q5_1-hardcoded — neither the wkv
17827                // windowed layers (g) nor the gkv globals (hd512) may be forced onto it via
17828                // MEMRA_FA_SMEM_TKV; they fall through to the format-clean register/scalar arms.)
17829                let fv = if g {
17830                    self.func_g("fa_decode_vec_q_smem")
17831                } else {
17832                    self.func("fa_decode_vec_q_smem")
17833                };
17834                let shmem = (2 * 32 * head_dim * 2) as u32; // sK+sV bf16 [FA_DEC_TILE=32][hd]
17835                use cudarc::driver::sys::CUfunction_attribute_enum as A;
17836                fv.set_attribute(
17837                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
17838                    shmem as i32,
17839                )?;
17840                (
17841                    fv,
17842                    LaunchConfig {
17843                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
17844                        block_dim: (32, gqa, 1),
17845                        shared_mem_bytes: shmem,
17846                    },
17847                )
17848            } else {
17849                // REGISTER-DEQUANT kernel (2026-07-03): per-warp direct q8_0/q5_1 register
17850                // dequant, zero dynamic shared memory.
17851                let fv = if g {
17852                    self.func_g("fa_decode_vec_q")
17853                } else {
17854                    self.func("fa_decode_vec_q")
17855                };
17856                (
17857                    fv,
17858                    LaunchConfig {
17859                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
17860                        block_dim: (32, gqa, 1),
17861                        shared_mem_bytes: 0,
17862                    },
17863                )
17864            }
17865        } else {
17866            // UNIFIED scalar (nullable-ctr symbol shared with graph/stream callers). The
17867            // split ladder value rides along so ns_eff reproduces THIS n_splits in-kernel.
17868            return self.fa_decode_scalar_unified(
17869                q,
17870                k,
17871                v,
17872                o,
17873                head_dim,
17874                n_head,
17875                n_head_kv,
17876                t_kv,
17877                None,
17878                scale,
17879                n_splits,
17880                if fa_vec { sp } else { 256 },
17881                k_tok_bytes,
17882                v_tok_bytes,
17883                g,
17884                part_o,
17885                part_m,
17886                part_l,
17887                None,
17888            );
17889        };
17890        let __s_b = self.gpu.stream();
17891        let mut b = __s_b.launch_builder(&f);
17892        b.arg(q)
17893            .arg(k)
17894            .arg(v)
17895            .arg(&mut *part_o)
17896            .arg(&mut *part_m)
17897            .arg(&mut *part_l)
17898            .arg(&hd)
17899            .arg(&nh)
17900            .arg(&nhkv)
17901            .arg(&tkvi)
17902            .arg(&scale)
17903            .arg(&nsp)
17904            .arg(&ktb)
17905            .arg(&vtb);
17906        unsafe {
17907            b.launch(cfg)?;
17908        }
17909        // (combine re-tile refuted in the fa-deep lane — flat/worse both shapes; the v4
17910        // combine stays for all arms. Receipts research/fa-decode-deep-20260802/.)
17911        let (fc, cfg2) = (
17912            if g {
17913                self.func_g("fa_decode_combine_f32")
17914            } else {
17915                self.fa_func("fa_decode_combine_f32", head_dim)
17916            },
17917            LaunchConfig {
17918                grid_dim: (n_head as u32, 1, 1),
17919                block_dim: (head_dim as u32, 1, 1),
17920                shared_mem_bytes: 0,
17921            },
17922        );
17923        let __s_b2 = self.gpu.stream();
17924        let mut b2 = __s_b2.launch_builder(&fc);
17925        b2.arg(&*part_o)
17926            .arg(&*part_m)
17927            .arg(&*part_l)
17928            .arg(o)
17929            .arg(&hd)
17930            .arg(&nh)
17931            .arg(&nsp);
17932        unsafe {
17933            b2.launch(cfg2)?;
17934        }
17935        Ok(())
17936    }
17937
17938    /// BATCHED-TICK increment 2: ONE fa_decode launch covering ALL B sequences of the
17939    /// batched decode step (blockIdx.z = sequence). Per-seq K/V cache bases ride a device
17940    /// pointer table (`kv_ptrs`, [2B] interleaved k0,v0,...); per-seq key bounds ride the
17941    /// tick's position table (`pos_seq`, T_kv = pos+1). v4-lane only: the CALLER
17942    /// (decode_batch) gates every row through `fa_seqs_eligible` AND one `fa_split_keys`
17943    /// rung (`split_keys`), so each sequence's split partition, key walk and combine order
17944    /// reproduce its per-seq eager v4 program exactly (kernel-check pins seqs-vs-loop bit
17945    /// identity; decode-batch-gate strict pins the whole tick vs decode_step_h).
17946    /// q is the stacked [B, n_head, head_dim] tick buffer read in place (no per-seq q
17947    /// copies); o is written [B, n_head, head_dim] in place (no per-seq a copies).
17948    #[allow(clippy::too_many_arguments)]
17949    pub fn fa_decode_batch_seqs_v4(
17950        &self,
17951        q: &CudaSlice<f32>,
17952        kv_ptrs: &cudarc::driver::CudaView<u64>,
17953        pos_seq: &CudaSlice<i32>,
17954        o: &mut CudaSlice<f32>,
17955        head_dim: usize,
17956        n_head: usize,
17957        n_head_kv: usize,
17958        b_n: usize,
17959        t_kv_max: usize,
17960        scale: f32,
17961        split_keys: usize,
17962        k_tok_bytes: usize,
17963        v_tok_bytes: usize,
17964    ) -> Result<(), Box<dyn std::error::Error>> {
17965        debug_assert!(head_dim == 256, "seqs twin is v4-stamped (hd256 only)");
17966        let n_splits_max = (t_kv_max + split_keys - 1) / split_keys;
17967        let o_len = b_n * n_head * n_splits_max * head_dim;
17968        let ml_len = b_n * n_head * n_splits_max;
17969        let mut part_guard = self.fa_part_pool.lock().unwrap();
17970        if part_guard
17971            .as_ref()
17972            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
17973            .unwrap_or(true)
17974        {
17975            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
17976            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
17977            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
17978            // later live allocations land at those addresses, and the next graph REPLAY writes
17979            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
17980            // output corruption began the burst after the trunk's t_kv growth first realloc'd
17981            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
17982            // the baked addresses alive (single-stream: eager writes the new buffers, replays
17983            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
17984            // (total retired < final size).
17985            let old = part_guard.take();
17986            let (co, cm) = old
17987                .as_ref()
17988                .map(|pp| (pp.0.len(), pp.1.len()))
17989                .unwrap_or((0, 0));
17990            if let Some(old) = old {
17991                self.fa_part_retired.lock().unwrap().push(old);
17992            }
17993            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
17994                eprintln!(
17995                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
17996                    co, o_len, cm, ml_len
17997                );
17998            }
17999            *part_guard = Some((
18000                self.alloc_uninit::<f32>(o_len.max(2 * co))?,
18001                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
18002                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
18003            ));
18004        }
18005        let pg = part_guard.as_mut().unwrap();
18006        self.gpu
18007            .stream()
18008            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
18009        self.gpu
18010            .stream()
18011            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
18012        self.gpu
18013            .stream()
18014            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
18015        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
18016        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
18017        let (nspm, spk) = (n_splits_max as i32, split_keys as i32);
18018        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
18019        let gqa = (n_head / n_head_kv).max(1) as u32;
18020        let f = self.func("fa_decode_vec_q_seqs_v4");
18021        // fa_v4_smem (11520B) + sV bf16 tile — the v4 eager arm's sizing on the default module.
18022        let shmem = (11520 + 32 * head_dim * 2) as u32;
18023        use cudarc::driver::sys::CUfunction_attribute_enum as A;
18024        f.set_attribute(
18025            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18026            shmem as i32,
18027        )?;
18028        let cfg = LaunchConfig {
18029            grid_dim: (n_head_kv as u32, n_splits_max as u32, b_n as u32),
18030            block_dim: (32, gqa, 1),
18031            shared_mem_bytes: shmem,
18032        };
18033        {
18034            let __s_b = self.gpu.stream();
18035            let mut b = __s_b.launch_builder(&f);
18036            b.arg(q)
18037                .arg(kv_ptrs)
18038                .arg(pos_seq)
18039                .arg(&mut *part_o)
18040                .arg(&mut *part_m)
18041                .arg(&mut *part_l)
18042                .arg(&hd)
18043                .arg(&nh)
18044                .arg(&nhkv)
18045                .arg(&scale)
18046                .arg(&nspm)
18047                .arg(&spk)
18048                .arg(&ktb)
18049                .arg(&vtb);
18050            unsafe {
18051                b.launch(cfg)?;
18052            }
18053        }
18054        let fc = self.func("fa_decode_combine_seqs");
18055        let cfg2 = LaunchConfig {
18056            grid_dim: (n_head as u32, b_n as u32, 1),
18057            block_dim: (head_dim as u32, 1, 1),
18058            shared_mem_bytes: 0,
18059        };
18060        let __s_b2 = self.gpu.stream();
18061        let mut b2 = __s_b2.launch_builder(&fc);
18062        b2.arg(&*part_o)
18063            .arg(&*part_m)
18064            .arg(&*part_l)
18065            .arg(o)
18066            .arg(&hd)
18067            .arg(&nh)
18068            .arg(pos_seq)
18069            .arg(&nspm)
18070            .arg(&spk);
18071        unsafe {
18072            b2.launch(cfg2)?;
18073        }
18074        Ok(())
18075    }
18076
18077    /// BATCHED-TICK increment 2: z-batched decode KV append — one launch appends this
18078    /// step's B rows, each into ITS OWN sequence cache at slot pos_seq[z], through the same
18079    /// [2B] interleaved pointer table the seqs FA reads. Each (block, z) warp executes the
18080    /// per-token appender's exact warp program on row z of the stacked [B, kv_dim] k/v —
18081    /// written cache bytes are BIT-IDENTICAL to the B per-seq calls it replaces
18082    /// (kernel-check pins the bytes). Default flash module only (callers exclude fp8-KV).
18083    #[allow(clippy::too_many_arguments)]
18084    pub fn append_kv_quantized_seqs(
18085        &self,
18086        k_rows: &CudaSlice<f32>,
18087        v_rows: &CudaSlice<f32>,
18088        kv_ptrs: &cudarc::driver::CudaView<u64>,
18089        pos_seq: &CudaSlice<i32>,
18090        b_n: usize,
18091        kv_dim_k: usize,
18092        kv_dim_v: usize,
18093        k_tok_bytes: usize,
18094        v_tok_bytes: usize,
18095    ) -> Result<(), Box<dyn std::error::Error>> {
18096        let f = self.func("append_quantize_kv_q8_0_q5_1_seqs");
18097        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
18098        let cfg = LaunchConfig {
18099            grid_dim: (nblk, b_n as u32, 1),
18100            block_dim: (32, 1, 1),
18101            shared_mem_bytes: 0,
18102        };
18103        let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
18104        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
18105        let __s_b = self.gpu.stream();
18106        let mut b = __s_b.launch_builder(&f);
18107        b.arg(k_rows)
18108            .arg(v_rows)
18109            .arg(kv_ptrs)
18110            .arg(pos_seq)
18111            .arg(&kdk)
18112            .arg(&kdv)
18113            .arg(&ktb)
18114            .arg(&vtb);
18115        unsafe {
18116            b.launch(cfg)?;
18117        }
18118        Ok(())
18119    }
18120
18121    /// True iff the MULTI-ROW verify FA (`fa_decode_rows`) is usable for a verify batch whose
18122    /// FIRST row attends `base_len + 1` keys: every row must take the SAME kernel eager decode
18123    /// would (the vec path) — mirrors fa_decode's gate exactly (MEMRA_NO_FA_VEC + FA_VEC_MIN_TKV +
18124    /// head_dim), evaluated at the MINIMUM row bound so no row could have picked scalar.
18125    /// MEMRA_FA_ROWS_OFF=1 is the A/B + fallback seam (per-row loop).
18126    pub fn fa_rows_eligible(&self, base_len: usize, head_dim: usize) -> bool {
18127        std::env::var("MEMRA_NO_FA_VEC").is_err()
18128            && std::env::var("MEMRA_FA_ROWS_OFF").is_err()
18129            && base_len + 1 >= fa_vec_min_tkv()
18130            && head_dim <= 256
18131            && head_dim % 32 == 0
18132    }
18133
18134    /// MULTI-ROW verify FA: run fa_decode_vec_q's EXACT per-row program for T causal query rows
18135    /// (row r attends keys [0..base_len+r+1)) in ONE kernel launch with grid.z = row, plus ONE
18136    /// row-batched combine. Replaces the T separate (fa_decode + combine) launches of the spec
18137    /// verify — same per-row split partition (n_splits_r = ceil(t_kv_r/split_keys), the
18138    /// fa_split_keys formula), same key-walk order, same reduce shapes => bit-identical outputs
18139    /// per row (kernel-check pins rows-vs-loop byte identity; run-spec is the end gate).
18140    /// Caller must have checked `fa_rows_eligible(base_len, head_dim)`.
18141    /// q is the verify's token-major [T, n_head, head_dim] stack; o is written [T, n_head, head_dim].
18142    #[allow(clippy::too_many_arguments)]
18143    pub fn fa_decode_rows(
18144        &self,
18145        q: &CudaSlice<f32>,
18146        k: &cudarc::driver::CudaView<u8>,
18147        v: &cudarc::driver::CudaView<u8>,
18148        o: &mut CudaSlice<f32>,
18149        head_dim: usize,
18150        n_head: usize,
18151        n_head_kv: usize,
18152        base_len: usize,
18153        t: usize,
18154        scale: f32,
18155        k_tok_bytes: usize,
18156        v_tok_bytes: usize,
18157        // hd512 dpl16 twin is DEVICE-LEN (graph arc): base_dev/plus feed the
18158        // kernel; host base_len keeps sizing the splits/partials. hd256 twins
18159        // keep the host arg. None is a bug for hd512 (asserted below).
18160        base_dev: Option<(&CudaSlice<i32>, i32)>,
18161        // K and V planes hold the same values (gemma globals, wv:=wk): pick
18162        // the _kv twin — V plane never read, value rides the q8_0 key dq.
18163        kv_shared: bool,
18164        // this layer's cache is e4m3 (gemma windowed under wkv): resolve the
18165        // hd256 rows kernel from the kf8vf8 module. PER-CALL — a global env
18166        // check here hijacked qwen/kernel-check hd256 rows (8 FAILs, 230ebbe).
18167        g: bool,
18168        // t=1 decode arm only: emit (int8, per-32 scales) from the dc combine
18169        // (hd512 path) — the standalone quantize launch folds away.
18170        mut q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
18171    ) -> Result<(), Box<dyn std::error::Error>> {
18172        debug_assert!(base_len + 1 >= fa_vec_min_tkv() && head_dim <= 512 && head_dim % 32 == 0);
18173        let t_kv_max = base_len + t; // LAST row's key bound
18174        let mut sp = fa_split_keys(t_kv_max, n_head_kv); // env/default — same value every row
18175        // hd512 split override (MEMRA_FA_SP512, 2026-07-11): gemma globals have n_head_kv=2 so
18176        // the grid is (2 x n_splits) — at depth ~29 splits = 58 blocks on 82 SMs (half idle,
18177        // rows_dpl16 8x off its byte floor). EVERY gemma hd512 caller shares THIS wrapper
18178        // (parity law), so the partition is freely tunable — verify and decode move together.
18179        if head_dim == 512 {
18180            static SP512: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
18181            // default 16 (2026-07-11 depth sweep, N=2: plain 155.4->156.5, depth spec
18182            // 236.9->250.4; 12/24/32 all worse). hd512 exists only on gemma globals.
18183            let v = *SP512.get_or_init(|| {
18184                std::env::var("MEMRA_FA_SP512")
18185                    .ok()
18186                    .and_then(|x| x.parse().ok())
18187                    .unwrap_or(0)
18188            });
18189            sp = if v >= 8 {
18190                v
18191            } else {
18192                FA_SP512_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
18193            };
18194        }
18195        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
18196        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
18197        let gqa = (n_head / n_head_kv).max(1) as u32;
18198        // LADDER-RUNG STRADDLE FIX (issue #10, 2026-07-13, g7e-proven): one sp for every row
18199        // diverges from eager decode when a split-ladder rung falls INSIDE the batch — row r's
18200        // eager twin used fa_split_keys(t_kv_r), the batch used fa_split_keys(t_kv_max), and
18201        // the different partition changes the combine's FP order (greedy tie flips at depth;
18202        // MEMRA_FA_SPLIT=64 pin -> PASS on the exact g7e failing config). Fix: group
18203        // consecutive rows by their OWN ladder value and launch once per group — each row then
18204        // executes the exact per-row program eager ran. Rungs land once per doubling, so this
18205        // is 1 launch in the common case and 2 on a crossing round. hd512 keeps one group (its
18206        // sp override is t_kv-independent by construction).
18207        let mut groups: Vec<(usize, usize, usize)> = Vec::new(); // (row0, t_g, sp_g)
18208        if head_dim == 512 || fa_split_keys(base_len + 1, n_head_kv) == sp {
18209            groups.push((0, t, sp));
18210        } else {
18211            let mut r0 = 0usize;
18212            while r0 < t {
18213                let sp_g = fa_split_keys(base_len + r0 + 1, n_head_kv);
18214                let mut r1 = r0 + 1;
18215                while r1 < t && fa_split_keys(base_len + r1 + 1, n_head_kv) == sp_g {
18216                    r1 += 1;
18217                }
18218                groups.push((r0, r1 - r0, sp_g));
18219                r0 = r1;
18220            }
18221        }
18222        // Deep-ctx smem twin for the VERIFY rows (2026-07-05): same threshold + rationale as
18223        // fa_decode's dispatch — at 40k the register path's GQA L2-reuse premise is dead and the
18224        // verify multiplies the 4x DRAM re-read by T rows. Bit-identical per (row,token,split).
18225        static SMEM_TKV_R: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
18226        let smem_tkv = *SMEM_TKV_R.get_or_init(|| {
18227            std::env::var("MEMRA_FA_SMEM_TKV")
18228                .ok()
18229                .and_then(|v| v.parse().ok())
18230                .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
18231        });
18232        let v4 = fa_v4_at(base_len + t) && head_dim == 256;
18233        let v3 = fa_v3_active(head_dim);
18234        let smem_rows =
18235            head_dim <= 256 && !v3 && !fa_v2_on() && smem_tkv > 0 && t_kv_max >= smem_tkv;
18236        // kv_shared twin RETIRED (2026-07-11 depth run-gen gate): the wv:=wk premise fails
18237        // POST-cache — cached K is k-normed+roped, cached V is not; the twin fed roped keys
18238        // in as values. Verify/decode/stream gates were blind (both sides shared the wrong
18239        // symbol — the parity law's blind spot); only prefill-vs-decode at depth caught it.
18240        let _ = kv_shared;
18241        // i2 twin: 2-key interleaved walk (MEMRA_FA_I2=0 reverts). i4 probed NEGATIVE
18242        // (157.3 vs 161.2 depth plain — register pressure past i2's sweet spot; jsonl).
18243        let i2 = head_dim == 512 && std::env::var("MEMRA_FA_I2").as_deref() != Ok("0");
18244        // v4-hd512 (MEMRA_FA_V512=1 opt-in, 2026-07-14): the v4 key-per-lane recipe on the
18245        // globals lane (depth profile: i2 ~4.6x off its byte floor — the v3-class
18246        // reduce-per-key latency signature). NEW NUMERIC CONFIG shared by every hd512
18247        // caller (decode+verify flip together); run-gen argmax + acceptance arbitrate.
18248        // T-BATCHED hd512 (DEFAULT ON 2026-07-14, MEMRA_FA_TB512=0 seam): one block per
18249        // (kv_head, split) stages its tile once and loops the rows over it — kills the
18250        // x t DRAM re-read of the full-ctx globals (depth cell +1.4%, plain flat, N=3
18251        // interleaved). FIXED absolute partition = NEW NUMERIC for the combine order,
18252        // shared by every hd512 caller through this wrapper (decode+verify flip together;
18253        // depth stream identical, acceptance unshifted, spec 256/256 x3 models).
18254        // Requires sp <= 32 (single staged tile; acc reused per row). The z-form v4_512
18255        // sibling (in-kernel dp4a port alone) probed FLAT — hd512 was DRAM-re-read-bound,
18256        // not unpack-bound; jsonl 2026-07-14.
18257        static TB512: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
18258        // gqa <= 16 = fa_v4_smem_512's q-array capacity; past it fall to the register twins.
18259        let tb512 = head_dim == 512
18260            && sp <= 32
18261            && n_head / n_head_kv.max(1) <= 16
18262            && *TB512.get_or_init(|| std::env::var("MEMRA_FA_TB512").as_deref() != Ok("0"));
18263        let fname = if tb512 {
18264            "fa_decode_vec_q_rows_v4_512_tb"
18265        } else if i2 {
18266            "fa_decode_vec_q_rows_dpl16_i2"
18267        } else if head_dim == 512 {
18268            "fa_decode_vec_q_rows_dpl16"
18269        }
18270        // gemma globals (parity law)
18271        else if v4 {
18272            "fa_decode_vec_q_rows_v4"
18273        } else if v3 {
18274            "fa_decode_vec_q_rows_v3"
18275        } else if fa_v2_on() {
18276            "fa_decode_vec_q_rows_v2"
18277        } else if smem_rows {
18278            "fa_decode_vec_q_rows_smem"
18279        } else {
18280            "fa_decode_vec_q_rows"
18281        };
18282        let f = if head_dim == 512 {
18283            self.fa_func(fname, head_dim)
18284        } else if g {
18285            // FP8-WINDOWED: hd256 rows over an e4m3 cache — kf8vf8 module, SAME symbol
18286            // choice as decode's kvmod dispatch (parity law: excluding v4 here paired
18287            // g-module rows against decode's g-module v4 — different programs, short-VG
18288            // maxdiff 2.0 / spec stream 0/128, 2026-07-12). rows_v4 is format-aware
18289            // since fda9790; only the smem twin stays excluded (V-stage q5_1-only).
18290            // hd128 (qwen fp8-KV) lands on the base/register rows via fname — the
18291            // dq macros are format-aware.
18292            self.func_g(if smem_rows {
18293                "fa_decode_vec_q_rows"
18294            } else {
18295                fname
18296            })
18297        } else {
18298            self.func(fname)
18299        };
18300        let shmem = if tb512 {
18301            // fa_v4_smem_512 (q 9KB gqa<=16 + k tile 18KB) + sV 32*512 (e4m3 module halves it)
18302            let gk = Self::gkv_on();
18303            let sh =
18304                (8192 + 1024 + 32 * 512 + 32 * 64 + 32 * head_dim * if gk { 1 } else { 2 }) as u32;
18305            use cudarc::driver::sys::CUfunction_attribute_enum as A;
18306            f.set_attribute(
18307                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18308                sh as i32,
18309            )?;
18310            sh
18311        } else if v4 || v3 || smem_rows || fa_v2_on() {
18312            // v4: fa_v4_smem (11.5KB) + sV; v3 stages sV only; v2/smem twins stage sK+sV.
18313            let sh = (if v4 {
18314                11520 + 32 * head_dim * if g { 1 } else { 2 }
18315            } else if v3 {
18316                32 * head_dim * 2
18317            } else {
18318                2 * 32 * head_dim * 2
18319            }) as u32;
18320            use cudarc::driver::sys::CUfunction_attribute_enum as A;
18321            f.set_attribute(
18322                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18323                sh as i32,
18324            )?;
18325            sh
18326        } else {
18327            0
18328        };
18329        // Per-GROUP launches (single group in the common case — identical to the pre-fix
18330        // single launch there): each group gets its own partials (the rows kernel indexes
18331        // partials by its LOCAL grid.z row) and q/o row-offset views.
18332        for &(r0, t_g, sp_g) in &groups {
18333            let n_splits_g = (base_len + r0 + t_g).div_ceil(sp_g);
18334            let (nspm, spk) = (n_splits_g as i32, sp_g as i32);
18335            let base_i = (base_len + r0) as i32;
18336            let o_len = t_g * n_head * n_splits_g * head_dim;
18337            let ml_len = t_g * n_head * n_splits_g;
18338            let mut part_guard = self.fa_part_pool.lock().unwrap();
18339            if part_guard
18340                .as_ref()
18341                .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
18342                .unwrap_or(true)
18343            {
18344                // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
18345                // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
18346                // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
18347                // later live allocations land at those addresses, and the next graph REPLAY writes
18348                // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
18349                // output corruption began the burst after the trunk's t_kv growth first realloc'd
18350                // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
18351                // the baked addresses alive (single-stream: eager writes the new buffers, replays
18352                // touch only the old — never concurrently). Doubling growth bounds retired VRAM
18353                // (total retired < final size).
18354                let old = part_guard.take();
18355                let (co, cm) = old
18356                    .as_ref()
18357                    .map(|pp| (pp.0.len(), pp.1.len()))
18358                    .unwrap_or((0, 0));
18359                if let Some(old) = old {
18360                    self.fa_part_retired.lock().unwrap().push(old);
18361                }
18362                if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
18363                    eprintln!(
18364                        "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
18365                        co, o_len, cm, ml_len
18366                    );
18367                }
18368                *part_guard = Some((
18369                    self.alloc_uninit::<f32>(o_len.max(2 * co))?,
18370                    self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
18371                    self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
18372                ));
18373            }
18374            let pg = part_guard.as_mut().unwrap();
18375            self.gpu
18376                .stream()
18377                .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
18378            self.gpu
18379                .stream()
18380                .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
18381            self.gpu
18382                .stream()
18383                .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
18384            let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
18385            let (part_o, part_m, part_l) = (&mut *part_o, &mut *part_m, &mut *part_l);
18386            let qv = self.view(q, t * n_head * head_dim);
18387            let q_g = qv.slice(r0 * n_head * head_dim..(r0 + t_g) * n_head * head_dim);
18388            let cfg = LaunchConfig {
18389                grid_dim: (n_head_kv as u32, n_splits_g as u32, t_g as u32),
18390                block_dim: (32, gqa, 1),
18391                shared_mem_bytes: shmem,
18392            };
18393            {
18394                let __s_b = self.gpu.stream();
18395                let mut b = __s_b.launch_builder(&f);
18396                if tb512 {
18397                    // rows-inner launch: grid.z dropped, the kernel loops n_rows itself.
18398                    let (bd, plus) =
18399                        base_dev.expect("hd512 rows twin requires a device base counter");
18400                    let plus_g = plus + r0 as i32;
18401                    let nr = t_g as i32;
18402                    if Self::pdl_on() && Self::pdl_wb_on() {
18403                        // wave-B2b: flavor mirrors fa_func(fname, 512) = gkv.
18404                        use cudarc::driver::{DevicePtr, DevicePtrMut};
18405                        let s = &self.gpu.stream();
18406                        let (pq, _b0) = q_g.device_ptr(s);
18407                        let (pk, _b1) = k.device_ptr(s);
18408                        let (pv, _b2) = v.device_ptr(s);
18409                        let (po, _b3) = part_o.device_ptr_mut(s);
18410                        let (pm, _b4) = part_m.device_ptr_mut(s);
18411                        let (pl, _b5) = part_l.device_ptr_mut(s);
18412                        let (pb, _b6) = bd.device_ptr(s);
18413                        let mut ps = [
18414                            &pq as *const _ as *mut std::ffi::c_void,
18415                            &pk as *const _ as *mut _,
18416                            &pv as *const _ as *mut _,
18417                            &po as *const _ as *mut _,
18418                            &pm as *const _ as *mut _,
18419                            &pl as *const _ as *mut _,
18420                            &hd as *const _ as *mut _,
18421                            &nh as *const _ as *mut _,
18422                            &nhkv as *const _ as *mut _,
18423                            &pb as *const _ as *mut _,
18424                            &plus_g as *const _ as *mut _,
18425                            &scale as *const _ as *mut _,
18426                            &nspm as *const _ as *mut _,
18427                            &spk as *const _ as *mut _,
18428                            &ktb as *const _ as *mut _,
18429                            &vtb as *const _ as *mut _,
18430                            &nr as *const _ as *mut _,
18431                        ];
18432                        unsafe {
18433                            self.launch_pdl_flash(
18434                                Self::gkv_on(),
18435                                "fa_decode_vec_q_rows_v4_512_tb",
18436                                (n_head_kv as u32, n_splits_g as u32, 1),
18437                                (32, gqa, 1),
18438                                shmem,
18439                                &mut ps,
18440                            )?;
18441                        }
18442                    } else {
18443                        let cfg_tb = LaunchConfig {
18444                            grid_dim: (n_head_kv as u32, n_splits_g as u32, 1),
18445                            block_dim: (32, gqa, 1),
18446                            shared_mem_bytes: shmem,
18447                        };
18448                        b.arg(&q_g)
18449                            .arg(k)
18450                            .arg(v)
18451                            .arg(&mut *part_o)
18452                            .arg(&mut *part_m)
18453                            .arg(&mut *part_l)
18454                            .arg(&hd)
18455                            .arg(&nh)
18456                            .arg(&nhkv)
18457                            .arg(bd)
18458                            .arg(&plus_g)
18459                            .arg(&scale)
18460                            .arg(&nspm)
18461                            .arg(&spk)
18462                            .arg(&ktb)
18463                            .arg(&vtb)
18464                            .arg(&nr);
18465                        unsafe {
18466                            b.launch(cfg_tb)?;
18467                        }
18468                    }
18469                } else if head_dim == 512 {
18470                    let (bd, plus) =
18471                        base_dev.expect("hd512 rows twin requires a device base counter");
18472                    let plus_g = plus + r0 as i32;
18473                    b.arg(&q_g)
18474                        .arg(k)
18475                        .arg(v)
18476                        .arg(&mut *part_o)
18477                        .arg(&mut *part_m)
18478                        .arg(&mut *part_l)
18479                        .arg(&hd)
18480                        .arg(&nh)
18481                        .arg(&nhkv)
18482                        .arg(bd)
18483                        .arg(&plus_g)
18484                        .arg(&scale)
18485                        .arg(&nspm)
18486                        .arg(&spk)
18487                        .arg(&ktb)
18488                        .arg(&vtb);
18489                    unsafe {
18490                        b.launch(cfg)?;
18491                    }
18492                } else {
18493                    b.arg(&q_g)
18494                        .arg(k)
18495                        .arg(v)
18496                        .arg(&mut *part_o)
18497                        .arg(&mut *part_m)
18498                        .arg(&mut *part_l)
18499                        .arg(&hd)
18500                        .arg(&nh)
18501                        .arg(&nhkv)
18502                        .arg(&base_i)
18503                        .arg(&scale)
18504                        .arg(&nspm)
18505                        .arg(&spk)
18506                        .arg(&ktb)
18507                        .arg(&vtb);
18508                    unsafe {
18509                        b.launch(cfg)?;
18510                    }
18511                }
18512            }
18513            let cfg2 = LaunchConfig {
18514                grid_dim: (n_head as u32, t_g as u32, 1),
18515                block_dim: (head_dim as u32, 1, 1),
18516                shared_mem_bytes: 0,
18517            };
18518            let mut o_g = o.slice_mut(r0 * n_head * head_dim..(r0 + t_g) * n_head * head_dim);
18519            if head_dim == 512 {
18520                // device-len combine (shared by verify/eager/graph — parity by symbol): the
18521                // per-row n_splits derives from the SAME counter the rows kernel read.
18522                let (bd, plus) = base_dev.unwrap();
18523                let plus_g = plus + r0 as i32;
18524                if let Some((oq, od)) = q8_out.as_mut() {
18525                    // wave-5b port (2026-07-23, t=1 decode only): q8-emitting dc combine.
18526                    debug_assert!(t == 1, "rows q8 emit is a t=1 decode arm");
18527                    if Self::pdl_on() && Self::pdl_wb_on() {
18528                        // wave-B2: flavor mirrors fa_func (hd512 + gkv → kf8vf8).
18529                        use cudarc::driver::{DevicePtr, DevicePtrMut};
18530                        let s = &self.gpu.stream();
18531                        let (po, _g0) = part_o.device_ptr(s);
18532                        let (pm, _g1) = part_m.device_ptr(s);
18533                        let (pl, _g2) = part_l.device_ptr(s);
18534                        let (pq, _g3) = oq.device_ptr_mut(s);
18535                        let (pd, _g4) = od.device_ptr_mut(s);
18536                        let (pb, _g5) = bd.device_ptr(s);
18537                        let mut ps = [
18538                            &po as *const _ as *mut std::ffi::c_void,
18539                            &pm as *const _ as *mut _,
18540                            &pl as *const _ as *mut _,
18541                            &pq as *const _ as *mut _,
18542                            &pd as *const _ as *mut _,
18543                            &hd as *const _ as *mut _,
18544                            &nh as *const _ as *mut _,
18545                            &pb as *const _ as *mut _,
18546                            &plus_g as *const _ as *mut _,
18547                            &nspm as *const _ as *mut _,
18548                            &spk as *const _ as *mut _,
18549                        ];
18550                        unsafe {
18551                            self.launch_pdl_flash(
18552                                Self::gkv_on(),
18553                                "fa_decode_combine_rows_dc_q8_1",
18554                                cfg2.grid_dim,
18555                                cfg2.block_dim,
18556                                0,
18557                                &mut ps,
18558                            )?;
18559                        }
18560                        continue;
18561                    }
18562                    let fc = self.fa_func("fa_decode_combine_rows_dc_q8_1", head_dim);
18563                    let __s_b2 = self.gpu.stream();
18564                    let mut b2 = __s_b2.launch_builder(&fc);
18565                    b2.arg(&*part_o)
18566                        .arg(&*part_m)
18567                        .arg(&*part_l)
18568                        .arg(&mut **oq)
18569                        .arg(&mut **od)
18570                        .arg(&hd)
18571                        .arg(&nh)
18572                        .arg(bd)
18573                        .arg(&plus_g)
18574                        .arg(&nspm)
18575                        .arg(&spk);
18576                    unsafe {
18577                        b2.launch(cfg2)?;
18578                    }
18579                    continue;
18580                }
18581                let fc = self.fa_func("fa_decode_combine_rows_dc", head_dim);
18582                let __s_b2 = self.gpu.stream();
18583                let mut b2 = __s_b2.launch_builder(&fc);
18584                b2.arg(&*part_o)
18585                    .arg(&*part_m)
18586                    .arg(&*part_l)
18587                    .arg(&mut o_g)
18588                    .arg(&hd)
18589                    .arg(&nh)
18590                    .arg(bd)
18591                    .arg(&plus_g)
18592                    .arg(&nspm)
18593                    .arg(&spk);
18594                unsafe {
18595                    b2.launch(cfg2)?;
18596                }
18597            } else {
18598                // q8 emit is wired for the hd512 dc-combine arm only — a Some here would
18599                // leave the caller's pair unwritten (consumer would read garbage).
18600                assert!(
18601                    q8_out.is_none(),
18602                    "rows q8 emit requires the hd512 dc combine"
18603                );
18604                let fc = self.func("fa_decode_combine_rows");
18605                let __s_b2 = self.gpu.stream();
18606                let mut b2 = __s_b2.launch_builder(&fc);
18607                b2.arg(&*part_o)
18608                    .arg(&*part_m)
18609                    .arg(&*part_l)
18610                    .arg(&mut o_g)
18611                    .arg(&hd)
18612                    .arg(&nh)
18613                    .arg(&base_i)
18614                    .arg(&nspm)
18615                    .arg(&spk);
18616                unsafe {
18617                    b2.launch(cfg2)?;
18618                }
18619            }
18620        }
18621        Ok(())
18622    }
18623
18624    /// WINDOWED verify rows (gemma R6 deep-ctx): every row attends exactly `window` keys —
18625    /// bit-identical per row to the T=1 decode's fa_decode over the window VIEW. Caller gates
18626    /// base_len + 1 >= window (no under-window rows) and head_dim == 256 (v4 stamp).
18627    #[allow(clippy::too_many_arguments)]
18628    pub fn fa_decode_rows_w(
18629        &self,
18630        q: &CudaSlice<f32>,
18631        k: &cudarc::driver::CudaView<u8>,
18632        v: &cudarc::driver::CudaView<u8>,
18633        o: &mut CudaSlice<f32>,
18634        head_dim: usize,
18635        n_head: usize,
18636        n_head_kv: usize,
18637        base_dev: &CudaSlice<i32>,
18638        base_plus: i32,
18639        t: usize,
18640        scale: f32,
18641        window: usize,
18642        k_tok_bytes: usize,
18643        v_tok_bytes: usize,
18644        q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
18645    ) -> Result<(), Box<dyn std::error::Error>> {
18646        // DEVICE-LEN (graph arc step 1, 2026-07-11): the causal base rides an i32 counter
18647        // (kernel T_kv = dev[0] + base_plus + r + 1) so depth graphs can replay with len
18648        // advancing on-device. dc paths pass kvl.len_d with plus=-1; verify/eager sync the
18649        // counter with one async set_i32_one first. Partials/splits size from `window` (host).
18650        debug_assert!(head_dim == 256);
18651        // windowed split (MEMRA_FA_SPW, default 32 — re-swept 2026-07-12 under the raw-e4m3 sV
18652        // occupancy ceiling (4 blocks/SM): t=1 decode is GRID-limited (win/sp splits x nkv
18653        // blocks), so smaller splits fill the ceiling — 1.7k 174.4/174.0 vs 48's 170.7/170.3,
18654        // 4.9k 159.8 vs 157.4 (N=2 interleaved, stable window). Spec serving prefers 64
18655        // (verify t=K+1 fills the grid via grid.z=t; depth K=7 281.3 vs 249.3 at 32) — set
18656        // MEMRA_FA_SPW=64 there, same config law as MEMRA_GEMMA_GKV=0. MUST be one value for
18657        // ALL widths: a t-keyed probe broke decode-vs-verify combine order (stream 9/128).
18658        let sp = {
18659            static SPW: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
18660            let v = *SPW.get_or_init(|| {
18661                std::env::var("MEMRA_FA_SPW")
18662                    .ok()
18663                    .and_then(|x| x.parse().ok())
18664                    .unwrap_or(0)
18665            });
18666            if v >= 8 {
18667                v
18668            } else {
18669                FA_SPW_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
18670            }
18671        };
18672        let n_splits_max = (window + sp - 1) / sp;
18673        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
18674        let (nspm, spk, wini) = (n_splits_max as i32, sp as i32, window as i32);
18675        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
18676        let gqa = (n_head / n_head_kv).max(1) as u32;
18677        let o_len = t * n_head * n_splits_max * head_dim;
18678        let ml_len = t * n_head * n_splits_max;
18679        let mut part_guard = self.fa_part_pool.lock().unwrap();
18680        if part_guard
18681            .as_ref()
18682            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
18683            .unwrap_or(true)
18684        {
18685            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
18686            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
18687            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
18688            // later live allocations land at those addresses, and the next graph REPLAY writes
18689            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
18690            // output corruption began the burst after the trunk's t_kv growth first realloc'd
18691            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
18692            // the baked addresses alive (single-stream: eager writes the new buffers, replays
18693            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
18694            // (total retired < final size).
18695            let old = part_guard.take();
18696            let (co, cm) = old
18697                .as_ref()
18698                .map(|pp| (pp.0.len(), pp.1.len()))
18699                .unwrap_or((0, 0));
18700            if let Some(old) = old {
18701                self.fa_part_retired.lock().unwrap().push(old);
18702            }
18703            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
18704                eprintln!(
18705                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
18706                    co, o_len, cm, ml_len
18707                );
18708            }
18709            *part_guard = Some((
18710                self.alloc_uninit::<f32>(o_len.max(2 * co))?,
18711                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
18712                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
18713            ));
18714        }
18715        let pg = part_guard.as_mut().unwrap();
18716        self.gpu
18717            .stream()
18718            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
18719        self.gpu
18720            .stream()
18721            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
18722        self.gpu
18723            .stream()
18724            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
18725        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
18726        // Lane pick: decode AND verify both land here in the windowed regime (parity law —
18727        // hybrid_forward verify_attn), so the pick only needs internal consistency, not
18728        // clone-of-decode bit fidelity (SASS-proven impossible for textually identical
18729        // kernels, jsonl 2026-07-10). v4 under the threshold; smem twin at/above the smem
18730        // floor (deep-ctx broadcast win); register twin between.
18731        static SMEM_TKV_W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
18732        let smem_tkv = *SMEM_TKV_W.get_or_init(|| {
18733            std::env::var("MEMRA_FA_SMEM_TKV")
18734                .ok()
18735                .and_then(|v| v.parse().ok())
18736                .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
18737        });
18738        // MULTI-ROW v4: resurrected 2026-07-14 (the '33 tok/s collapse' was a paired-map
18739        // partial-write bug, not the mechanism) and falsified HONESTLY at gqa 2: bit-exact
18740        // but −1.7% on the 31B depth cell — the sp helper warp already hides staging
18741        // in-block, and mr trades L2-cheap redundant bytes for serialized per-warp gqa
18742        // score/B3 chains. Arm deleted; jsonl row 2026-07-14 is the record.
18743        use cudarc::driver::sys::CUfunction_attribute_enum as A;
18744        // FP8-WINDOWED (wkv): the v4 family is format-aware (2026-07-12 KFMT/VFMT staging
18745        // arms) — wkv rides the SAME lane logic, resolved from the kf8vf8 module. One symbol
18746        // per (lane, format-module) keeps parity structural; the old register-i2 detour
18747        // (-33%) is retired.
18748        let wg = Self::wkv_on();
18749        // STAGING-PARALLEL v4 (MEMRA_FA_SPW2, default ON at gqa==1): warp 1 = staging helper
18750        // (v4 is 61% staging); score phases identical to v4_w. Same symbol all t.
18751        let sp2 =
18752            gqa <= 4 && fa_v4_at(window) && std::env::var("MEMRA_FA_SPW2").as_deref() != Ok("0");
18753        if sp2 {
18754            let sh = (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32;
18755            if Self::pdl_on() && Self::pdl_wb_on() {
18756                // wave-B2b: flavor mirrors wg.
18757                use cudarc::driver::{DevicePtr, DevicePtrMut};
18758                let s = &self.gpu.stream();
18759                let (pq, _b0) = q.device_ptr(s);
18760                let (pk, _b1) = k.device_ptr(s);
18761                let (pv, _b2) = v.device_ptr(s);
18762                let (po, _b3) = part_o.device_ptr_mut(s);
18763                let (pm, _b4) = part_m.device_ptr_mut(s);
18764                let (pl, _b5) = part_l.device_ptr_mut(s);
18765                let (pb, _b6) = base_dev.device_ptr(s);
18766                let mut ps = [
18767                    &pq as *const _ as *mut std::ffi::c_void,
18768                    &pk as *const _ as *mut _,
18769                    &pv as *const _ as *mut _,
18770                    &po as *const _ as *mut _,
18771                    &pm as *const _ as *mut _,
18772                    &pl as *const _ as *mut _,
18773                    &hd as *const _ as *mut _,
18774                    &nh as *const _ as *mut _,
18775                    &nhkv as *const _ as *mut _,
18776                    &pb as *const _ as *mut _,
18777                    &base_plus as *const _ as *mut _,
18778                    &scale as *const _ as *mut _,
18779                    &nspm as *const _ as *mut _,
18780                    &spk as *const _ as *mut _,
18781                    &ktb as *const _ as *mut _,
18782                    &vtb as *const _ as *mut _,
18783                    &wini as *const _ as *mut _,
18784                ];
18785                unsafe {
18786                    self.launch_pdl_flash(
18787                        wg,
18788                        "fa_decode_vec_q_rows_v4_w_sp",
18789                        (n_head_kv as u32, n_splits_max as u32, t as u32),
18790                        (32, gqa + 1, 1),
18791                        sh,
18792                        &mut ps,
18793                    )?;
18794                }
18795            } else {
18796                let f = if wg {
18797                    self.func_g("fa_decode_vec_q_rows_v4_w_sp")
18798                } else {
18799                    self.func("fa_decode_vec_q_rows_v4_w_sp")
18800                };
18801                f.set_attribute(
18802                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18803                    sh as i32,
18804                )?;
18805                let cfg = LaunchConfig {
18806                    grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
18807                    block_dim: (32, gqa + 1, 1),
18808                    shared_mem_bytes: sh,
18809                };
18810                let __s_b = self.gpu.stream();
18811                let mut b = __s_b.launch_builder(&f);
18812                b.arg(q)
18813                    .arg(k)
18814                    .arg(v)
18815                    .arg(&mut *part_o)
18816                    .arg(&mut *part_m)
18817                    .arg(&mut *part_l)
18818                    .arg(&hd)
18819                    .arg(&nh)
18820                    .arg(&nhkv)
18821                    .arg(base_dev)
18822                    .arg(&base_plus)
18823                    .arg(&scale)
18824                    .arg(&nspm)
18825                    .arg(&spk)
18826                    .arg(&ktb)
18827                    .arg(&vtb)
18828                    .arg(&wini);
18829                unsafe {
18830                    b.launch(cfg)?;
18831                }
18832            }
18833        } else {
18834            if fa_v4_at(window) && Self::pdl_on() && Self::pdl_wb_on() {
18835                // wave-B2b: the v4_w pick only (smem/reg twins stay builder-launched).
18836                let sh = (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32;
18837                use cudarc::driver::{DevicePtr, DevicePtrMut};
18838                let s = &self.gpu.stream();
18839                let (pq, _b0) = q.device_ptr(s);
18840                let (pk, _b1) = k.device_ptr(s);
18841                let (pv, _b2) = v.device_ptr(s);
18842                let (po, _b3) = part_o.device_ptr_mut(s);
18843                let (pm, _b4) = part_m.device_ptr_mut(s);
18844                let (pl, _b5) = part_l.device_ptr_mut(s);
18845                let (pb, _b6) = base_dev.device_ptr(s);
18846                let mut ps = [
18847                    &pq as *const _ as *mut std::ffi::c_void,
18848                    &pk as *const _ as *mut _,
18849                    &pv as *const _ as *mut _,
18850                    &po as *const _ as *mut _,
18851                    &pm as *const _ as *mut _,
18852                    &pl as *const _ as *mut _,
18853                    &hd as *const _ as *mut _,
18854                    &nh as *const _ as *mut _,
18855                    &nhkv as *const _ as *mut _,
18856                    &pb as *const _ as *mut _,
18857                    &base_plus as *const _ as *mut _,
18858                    &scale as *const _ as *mut _,
18859                    &nspm as *const _ as *mut _,
18860                    &spk as *const _ as *mut _,
18861                    &ktb as *const _ as *mut _,
18862                    &vtb as *const _ as *mut _,
18863                    &wini as *const _ as *mut _,
18864                ];
18865                unsafe {
18866                    self.launch_pdl_flash(
18867                        wg,
18868                        "fa_decode_vec_q_rows_v4_w",
18869                        (n_head_kv as u32, n_splits_max as u32, t as u32),
18870                        (32, gqa, 1),
18871                        sh,
18872                        &mut ps,
18873                    )?;
18874                }
18875            } else {
18876                let pick = |name: &str| {
18877                    if wg {
18878                        self.func_g(name)
18879                    } else {
18880                        self.func(name)
18881                    }
18882                };
18883                let (f, sh) = if fa_v4_at(window) {
18884                    let f = pick("fa_decode_vec_q_rows_v4_w");
18885                    (f, (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32)
18886                } else if smem_tkv > 0 && window >= smem_tkv {
18887                    // NOTE: the smem twin's V-stage is still q5_1-hardcoded — unreachable under wkv
18888                    // at the gemma window (v4 covers it); revisit if the smem floor ever drops.
18889                    (
18890                        pick("fa_decode_vec_q_rows_smem_w"),
18891                        (2 * 32 * head_dim * 2) as u32,
18892                    )
18893                } else {
18894                    (pick("fa_decode_vec_q_rows_reg_w"), 0u32)
18895                };
18896                f.set_attribute(
18897                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18898                    sh as i32,
18899                )?;
18900                let cfg = LaunchConfig {
18901                    grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
18902                    block_dim: (32, gqa, 1),
18903                    shared_mem_bytes: sh,
18904                };
18905                let __s_b = self.gpu.stream();
18906                let mut b = __s_b.launch_builder(&f);
18907                b.arg(q)
18908                    .arg(k)
18909                    .arg(v)
18910                    .arg(&mut *part_o)
18911                    .arg(&mut *part_m)
18912                    .arg(&mut *part_l)
18913                    .arg(&hd)
18914                    .arg(&nh)
18915                    .arg(&nhkv)
18916                    .arg(base_dev)
18917                    .arg(&base_plus)
18918                    .arg(&scale)
18919                    .arg(&nspm)
18920                    .arg(&spk)
18921                    .arg(&ktb)
18922                    .arg(&vtb)
18923                    .arg(&wini);
18924                unsafe {
18925                    b.launch(cfg)?;
18926                }
18927            }
18928        }
18929        let cfg2 = LaunchConfig {
18930            grid_dim: (n_head as u32, t as u32, 1),
18931            block_dim: (head_dim as u32, 1, 1),
18932            shared_mem_bytes: 0,
18933        };
18934        if let Some((oq, od)) = q8_out {
18935            // wave-5b port (2026-07-23): q8-emitting combine — the t=1 decode's wo matvec
18936            // consumes the pair directly; the standalone quantize launch folds away.
18937            if Self::pdl_on() && Self::pdl_wb_on() {
18938                // wave-B2: flavor mirrors the builder's wg choice.
18939                use cudarc::driver::{DevicePtr, DevicePtrMut};
18940                let s = &self.gpu.stream();
18941                let (po, _g0) = part_o.device_ptr(s);
18942                let (pm, _g1) = part_m.device_ptr(s);
18943                let (pl, _g2) = part_l.device_ptr(s);
18944                let (pq, _g3) = oq.device_ptr_mut(s);
18945                let (pd, _g4) = od.device_ptr_mut(s);
18946                let mut ps = [
18947                    &po as *const _ as *mut std::ffi::c_void,
18948                    &pm as *const _ as *mut _,
18949                    &pl as *const _ as *mut _,
18950                    &pq as *const _ as *mut _,
18951                    &pd as *const _ as *mut _,
18952                    &hd as *const _ as *mut _,
18953                    &nh as *const _ as *mut _,
18954                    &nspm as *const _ as *mut _,
18955                    &spk as *const _ as *mut _,
18956                    &wini as *const _ as *mut _,
18957                ];
18958                unsafe {
18959                    self.launch_pdl_flash(
18960                        wg,
18961                        "fa_decode_combine_rows_w_q8_1",
18962                        cfg2.grid_dim,
18963                        cfg2.block_dim,
18964                        0,
18965                        &mut ps,
18966                    )?;
18967                }
18968                return Ok(());
18969            }
18970            let fc = if wg {
18971                self.func_g("fa_decode_combine_rows_w_q8_1")
18972            } else {
18973                self.func("fa_decode_combine_rows_w_q8_1")
18974            };
18975            let __s_b2 = self.gpu.stream();
18976            let mut b2 = __s_b2.launch_builder(&fc);
18977            b2.arg(&*part_o)
18978                .arg(&*part_m)
18979                .arg(&*part_l)
18980                .arg(oq)
18981                .arg(od)
18982                .arg(&hd)
18983                .arg(&nh)
18984                .arg(&nspm)
18985                .arg(&spk)
18986                .arg(&wini);
18987            unsafe {
18988                b2.launch(cfg2)?;
18989            }
18990            return Ok(());
18991        }
18992        let fc = if wg {
18993            self.func_g("fa_decode_combine_rows_w")
18994        } else {
18995            self.func("fa_decode_combine_rows_w")
18996        };
18997        let __s_b2 = self.gpu.stream();
18998        let mut b2 = __s_b2.launch_builder(&fc);
18999        b2.arg(&*part_o)
19000            .arg(&*part_m)
19001            .arg(&*part_l)
19002            .arg(o)
19003            .arg(&hd)
19004            .arg(&nh)
19005            .arg(&nspm)
19006            .arg(&spk)
19007            .arg(&wini);
19008        unsafe {
19009            b2.launch(cfg2)?;
19010        }
19011        Ok(())
19012    }
19013
19014    /// ROUND-STREAM stage (c): fa rows with the causal base from a device counter. Two lanes:
19015    /// v3 (qwen stream, fa_v3_active) and v4 (gemma hd256 burst — rows_v4_dc, g-module aware);
19016    /// `t_kv_upper` sizes splits/partials — the same one-sp-for-all-rows approximation class
19017    /// the host rows path already uses (battery-arbitrated); actual per-row bounds derive
19018    /// in-kernel from the counter (+ base_plus, v4 lane only — v3's kernel has no plus arg).
19019    #[allow(clippy::too_many_arguments)]
19020    pub fn fa_decode_rows_dc(
19021        &self,
19022        q: &CudaSlice<f32>,
19023        k: &cudarc::driver::CudaView<u8>,
19024        v: &cudarc::driver::CudaView<u8>,
19025        o: &mut CudaSlice<f32>,
19026        head_dim: usize,
19027        n_head: usize,
19028        n_head_kv: usize,
19029        base_dev: &CudaSlice<i32>,
19030        t_kv_upper: usize,
19031        t: usize,
19032        scale: f32,
19033        k_tok_bytes: usize,
19034        v_tok_bytes: usize,
19035        base_plus: i32,
19036        g: bool,
19037    ) -> Result<(), Box<dyn std::error::Error>> {
19038        let v4 = head_dim == 256 && fa_v4_at(t_kv_upper);
19039        assert!(
19040            v4 || fa_v3_active(head_dim),
19041            "stream fa rows requires the v3 or v4 lane"
19042        );
19043        assert!(v4 || base_plus == 0, "v3_dc kernel takes no plus arg");
19044        if v4 {
19045            let sp = fa_split_keys(t_kv_upper, n_head_kv);
19046            let n_splits_max = (t_kv_upper + sp - 1) / sp;
19047            let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
19048            let (nspm, spk) = (n_splits_max as i32, sp as i32);
19049            let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
19050            let gqa = (n_head / n_head_kv).max(1) as u32;
19051            let o_len = t * n_head * n_splits_max * head_dim;
19052            let ml_len = t * n_head * n_splits_max;
19053            let mut part_guard = self.fa_part_pool.lock().unwrap();
19054            if part_guard
19055                .as_ref()
19056                .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
19057                .unwrap_or(true)
19058            {
19059                // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
19060                // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
19061                // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
19062                // later live allocations land at those addresses, and the next graph REPLAY writes
19063                // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
19064                // output corruption began the burst after the trunk's t_kv growth first realloc'd
19065                // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
19066                // the baked addresses alive (single-stream: eager writes the new buffers, replays
19067                // touch only the old — never concurrently). Doubling growth bounds retired VRAM
19068                // (total retired < final size).
19069                let old = part_guard.take();
19070                let (co, cm) = old
19071                    .as_ref()
19072                    .map(|pp| (pp.0.len(), pp.1.len()))
19073                    .unwrap_or((0, 0));
19074                if let Some(old) = old {
19075                    self.fa_part_retired.lock().unwrap().push(old);
19076                }
19077                if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
19078                    eprintln!(
19079                        "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
19080                        co, o_len, cm, ml_len
19081                    );
19082                }
19083                *part_guard = Some((
19084                    self.alloc_uninit::<f32>(o_len.max(2 * co))?,
19085                    self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
19086                    self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
19087                ));
19088            }
19089            let pg = part_guard.as_mut().unwrap();
19090            self.gpu
19091                .stream()
19092                .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
19093            self.gpu
19094                .stream()
19095                .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
19096            self.gpu
19097                .stream()
19098                .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
19099            let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
19100            let f = if g {
19101                self.func_g("fa_decode_vec_q_rows_v4_dc")
19102            } else {
19103                self.func("fa_decode_vec_q_rows_v4_dc")
19104            };
19105            let sh = (11520 + 32 * head_dim * if g { 1 } else { 2 }) as u32;
19106            use cudarc::driver::sys::CUfunction_attribute_enum as A;
19107            f.set_attribute(
19108                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19109                sh as i32,
19110            )?;
19111            let cfg = LaunchConfig {
19112                grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
19113                block_dim: (32, gqa, 1),
19114                shared_mem_bytes: sh,
19115            };
19116            let __s_b = self.gpu.stream();
19117            let mut b = __s_b.launch_builder(&f);
19118            b.arg(q)
19119                .arg(k)
19120                .arg(v)
19121                .arg(&mut *part_o)
19122                .arg(&mut *part_m)
19123                .arg(&mut *part_l)
19124                .arg(&hd)
19125                .arg(&nh)
19126                .arg(&nhkv)
19127                .arg(base_dev)
19128                .arg(&base_plus)
19129                .arg(&scale)
19130                .arg(&nspm)
19131                .arg(&spk)
19132                .arg(&ktb)
19133                .arg(&vtb);
19134            unsafe {
19135                b.launch(cfg)?;
19136            }
19137            let fc = self.func("fa_decode_combine_rows_dc");
19138            let cfg2 = LaunchConfig {
19139                grid_dim: (n_head as u32, t as u32, 1),
19140                block_dim: (head_dim as u32, 1, 1),
19141                shared_mem_bytes: 0,
19142            };
19143            let __s_b2 = self.gpu.stream();
19144            let mut b2 = __s_b2.launch_builder(&fc);
19145            b2.arg(&*part_o)
19146                .arg(&*part_m)
19147                .arg(&*part_l)
19148                .arg(o)
19149                .arg(&hd)
19150                .arg(&nh)
19151                .arg(base_dev)
19152                .arg(&base_plus)
19153                .arg(&nspm)
19154                .arg(&spk);
19155            unsafe {
19156                b2.launch(cfg2)?;
19157            }
19158            return Ok(());
19159        }
19160        let sp = fa_split_keys(t_kv_upper, n_head_kv);
19161        let n_splits_max = (t_kv_upper + sp - 1) / sp;
19162        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
19163        let (nspm, spk) = (n_splits_max as i32, sp as i32);
19164        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
19165        let gqa = (n_head / n_head_kv).max(1) as u32;
19166        let o_len = t * n_head * n_splits_max * head_dim;
19167        let ml_len = t * n_head * n_splits_max;
19168        let mut part_guard = self.fa_part_pool.lock().unwrap();
19169        if part_guard
19170            .as_ref()
19171            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
19172            .unwrap_or(true)
19173        {
19174            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
19175            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
19176            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
19177            // later live allocations land at those addresses, and the next graph REPLAY writes
19178            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
19179            // output corruption began the burst after the trunk's t_kv growth first realloc'd
19180            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
19181            // the baked addresses alive (single-stream: eager writes the new buffers, replays
19182            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
19183            // (total retired < final size).
19184            let old = part_guard.take();
19185            let (co, cm) = old
19186                .as_ref()
19187                .map(|pp| (pp.0.len(), pp.1.len()))
19188                .unwrap_or((0, 0));
19189            if let Some(old) = old {
19190                self.fa_part_retired.lock().unwrap().push(old);
19191            }
19192            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
19193                eprintln!(
19194                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
19195                    co, o_len, cm, ml_len
19196                );
19197            }
19198            *part_guard = Some((
19199                self.alloc_uninit::<f32>(o_len.max(2 * co))?,
19200                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
19201                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
19202            ));
19203        }
19204        let pg = part_guard.as_mut().unwrap();
19205        self.gpu
19206            .stream()
19207            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
19208        self.gpu
19209            .stream()
19210            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
19211        self.gpu
19212            .stream()
19213            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
19214        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
19215        let f = self.func("fa_decode_vec_q_rows_v3_dc");
19216        let sh = (32 * head_dim * 2) as u32;
19217        use cudarc::driver::sys::CUfunction_attribute_enum as A;
19218        f.set_attribute(
19219            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19220            sh as i32,
19221        )?;
19222        let cfg = LaunchConfig {
19223            grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
19224            block_dim: (32, gqa, 1),
19225            shared_mem_bytes: sh,
19226        };
19227        let __s_b = self.gpu.stream();
19228        let mut b = __s_b.launch_builder(&f);
19229        b.arg(q)
19230            .arg(k)
19231            .arg(v)
19232            .arg(&mut *part_o)
19233            .arg(&mut *part_m)
19234            .arg(&mut *part_l)
19235            .arg(&hd)
19236            .arg(&nh)
19237            .arg(&nhkv)
19238            .arg(base_dev)
19239            .arg(&scale)
19240            .arg(&nspm)
19241            .arg(&spk)
19242            .arg(&ktb)
19243            .arg(&vtb);
19244        unsafe {
19245            b.launch(cfg)?;
19246        }
19247        let fc = self.func("fa_decode_combine_rows_dc");
19248        let cfg2 = LaunchConfig {
19249            grid_dim: (n_head as u32, t as u32, 1),
19250            block_dim: (head_dim as u32, 1, 1),
19251            shared_mem_bytes: 0,
19252        };
19253        let plus0 = 0i32;
19254        let __s_b2 = self.gpu.stream();
19255        let mut b2 = __s_b2.launch_builder(&fc);
19256        b2.arg(&*part_o)
19257            .arg(&*part_m)
19258            .arg(&*part_l)
19259            .arg(o)
19260            .arg(&hd)
19261            .arg(&nh)
19262            .arg(base_dev)
19263            .arg(&plus0)
19264            .arg(&nspm)
19265            .arg(&spk);
19266        unsafe {
19267            b2.launch(cfg2)?;
19268        }
19269        Ok(())
19270    }
19271
19272    /// Device-counter variant of `fa_decode` (CUDA-GRAPH-PLAN Phase 2). The sequence length is read
19273    /// from `t_kv_dev[0]` (resident device i32[1]) for the attention loop bound + per-split key range;
19274    /// the GRID `n_splits` is sized for `bucket_max` (the bucket's max t_kv — baked at capture time).
19275    /// Empty splits (key range beyond the actual t_kv) write an empty partial (m=NEG_INF) so the
19276    /// shared combine skips them -> bit-correct for ANY actual t_kv <= bucket_max.
19277    ///
19278    /// BIT-IDENTITY (the gate): pass `bucket_max == actual_t_kv` and this reproduces `fa_decode`
19279    /// EXACTLY (same n_splits, same per, same split boundaries, same combine) while reading t_kv from
19280    /// device. Bucketing (bucket_max > t_kv) is for the future captured path and changes split
19281    /// grouping (different but mathematically-equal log-sum-exp merge).
19282    pub fn fa_decode_dc(
19283        &self,
19284        q: &CudaSlice<f32>,
19285        k: &cudarc::driver::CudaView<u8>,
19286        v: &cudarc::driver::CudaView<u8>,
19287        o: &mut CudaSlice<f32>,
19288        head_dim: usize,
19289        n_head: usize,
19290        n_head_kv: usize,
19291        t_kv_dev: &CudaSlice<i32>,
19292        bucket_max: usize,
19293        scale: f32,
19294        k_tok_bytes: usize,
19295        v_tok_bytes: usize,
19296        g: bool,
19297    ) -> Result<(), Box<dyn std::error::Error>> {
19298        self.fa_decode_dc_q8(
19299            q,
19300            k,
19301            v,
19302            o,
19303            head_dim,
19304            n_head,
19305            n_head_kv,
19306            t_kv_dev,
19307            bucket_max,
19308            scale,
19309            k_tok_bytes,
19310            v_tok_bytes,
19311            g,
19312            None,
19313        )
19314    }
19315
19316    /// `fa_decode_dc` with an optional q8_1 sink (wave 5b): when `q8_out` is given the
19317    /// combine emits (int8, per-32 scales) for the wo matmul_pre and skips the f32 O write.
19318    #[allow(clippy::too_many_arguments)]
19319    pub fn fa_decode_dc_q8(
19320        &self,
19321        q: &CudaSlice<f32>,
19322        k: &cudarc::driver::CudaView<u8>,
19323        v: &cudarc::driver::CudaView<u8>,
19324        o: &mut CudaSlice<f32>,
19325        head_dim: usize,
19326        n_head: usize,
19327        n_head_kv: usize,
19328        t_kv_dev: &CudaSlice<i32>,
19329        bucket_max: usize,
19330        scale: f32,
19331        k_tok_bytes: usize,
19332        v_tok_bytes: usize,
19333        g: bool,
19334        q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
19335    ) -> Result<(), Box<dyn std::error::Error>> {
19336        // The fa_vec gate + n_splits are sized from bucket_max (host, fixed at capture). The kernel
19337        // reads the ACTUAL t_kv from t_kv_dev for the per-split bound. DEFAULT-ON to MATCH the eager
19338        // `fa_decode` gate above — graph capture must mirror eager's kernel choice or the graph-vs-eager
19339        // bit-identity gate breaks. MEMRA_NO_FA_VEC forces scalar on BOTH paths in lockstep.
19340        // `g` = this layer's cache is e4m3 (gemma windowed under wkv) — every pick below must
19341        // mirror fa_decode_kvmod's g-routing or the graph diverges from eager (short/mid 1/96,
19342        // 2026-07-12).
19343        let mut fa_vec =
19344            std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
19345        if g && head_dim == 256 && !fa_v4_at(bucket_max) {
19346            fa_vec = false;
19347        } // mirror kvmod/geom
19348        let sp = fa_split_keys(bucket_max, n_head_kv);
19349        let n_splits = if fa_vec {
19350            ((bucket_max + sp - 1) / sp).max(1)
19351        } else {
19352            ((bucket_max + 255) / 256).max(1)
19353        };
19354        let o_len = n_head * n_splits * head_dim;
19355        let ml_len = n_head * n_splits;
19356        let mut part_guard = self.fa_part_pool.lock().unwrap();
19357        if part_guard
19358            .as_ref()
19359            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
19360            .unwrap_or(true)
19361        {
19362            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
19363            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
19364            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
19365            // later live allocations land at those addresses, and the next graph REPLAY writes
19366            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
19367            // output corruption began the burst after the trunk's t_kv growth first realloc'd
19368            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
19369            // the baked addresses alive (single-stream: eager writes the new buffers, replays
19370            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
19371            // (total retired < final size).
19372            let old = part_guard.take();
19373            let (co, cm) = old
19374                .as_ref()
19375                .map(|pp| (pp.0.len(), pp.1.len()))
19376                .unwrap_or((0, 0));
19377            if let Some(old) = old {
19378                self.fa_part_retired.lock().unwrap().push(old);
19379            }
19380            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
19381                eprintln!(
19382                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
19383                    co, o_len, cm, ml_len
19384                );
19385            }
19386            *part_guard = Some((
19387                self.alloc_uninit::<f32>(o_len.max(2 * co))?,
19388                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
19389                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
19390            ));
19391        }
19392        let pg = part_guard.as_mut().unwrap();
19393        self.gpu
19394            .stream()
19395            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
19396        self.gpu
19397            .stream()
19398            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
19399        self.gpu
19400            .stream()
19401            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
19402        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
19403        let (hd, nh, nhkv, nsp) = (
19404            head_dim as i32,
19405            n_head as i32,
19406            n_head_kv as i32,
19407            n_splits as i32,
19408        );
19409        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
19410        let fa_vec = fa_vec && head_dim <= 512 && head_dim % 32 == 0;
19411        // FA-DEEP pick keyed on bucket_max (the fa_v4_at precedent) — bit-identical twins,
19412        // so a threshold falling between t_kv and bucket_max cannot diverge eager-vs-graph.
19413        let deep = fa_vec
19414            && head_dim == 256
19415            && fa_v4_at(bucket_max)
19416            && !g
19417            && fa_deep_at(bucket_max)
19418            && !matches!(fa_v4_mode(), "noB3" | "stage");
19419        let (f, cfg) = if fa_vec
19420            && head_dim == 512
19421            && bucket_max >= {
19422                static FA512_MIN_DC: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
19423                *FA512_MIN_DC.get_or_init(|| {
19424                    std::env::var("MEMRA_FA512_MIN")
19425                        .ok()
19426                        .and_then(|v| v.parse().ok())
19427                        .unwrap_or(512)
19428                })
19429            } {
19430            // gemma globals dc twin (mirror the eager dpl16 pick incl the crossover floor).
19431            let gqa = (n_head / n_head_kv).max(1) as u32;
19432            (
19433                self.fa_func("fa_decode_vec_q_dpl16_dc", head_dim),
19434                LaunchConfig {
19435                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
19436                    block_dim: (32, gqa, 1),
19437                    shared_mem_bytes: 0,
19438                },
19439            )
19440        } else if fa_vec && head_dim == 512 {
19441            // under the 512 floor eager runs scalar — the SAME unified symbol, ctr non-null;
19442            // ns_eff in-kernel reproduces eager's ceil(t_kv/sp) partition for the LIVE len.
19443            let q_view = q.as_view();
19444            let mut o_view = o.as_view_mut();
19445            return self.fa_decode_scalar_unified(
19446                &q_view,
19447                k,
19448                v,
19449                &mut o_view,
19450                head_dim,
19451                n_head,
19452                n_head_kv,
19453                0,
19454                Some(t_kv_dev),
19455                scale,
19456                n_splits,
19457                sp,
19458                k_tok_bytes,
19459                v_tok_bytes,
19460                g,
19461                &mut *part_o,
19462                &mut *part_m,
19463                &mut *part_l,
19464                q8_out,
19465            );
19466        } else if fa_vec && head_dim == 256 && fa_v4_at(bucket_max) {
19467            // gemma/qwen v4 dc twin (eager default lane) — capture must mirror eager's pick,
19468            // incl the g-module route + raw-e4m3 sV sizing.
19469            let gqa = (n_head / n_head_kv).max(1) as u32;
19470            let fv = if g {
19471                self.func_g("fa_decode_vec_q_v4_dc")
19472            } else if deep {
19473                self.func("fa_decode_vec_q_v4_deep_dc")
19474            } else {
19475                self.func("fa_decode_vec_q_v4_dc")
19476            };
19477            let shmem =
19478                (if deep { 12160 } else { 11520 } + 32 * head_dim * if g { 1 } else { 2 }) as u32;
19479            use cudarc::driver::sys::CUfunction_attribute_enum as A;
19480            fv.set_attribute(
19481                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19482                shmem as i32,
19483            )?;
19484            (
19485                fv,
19486                LaunchConfig {
19487                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
19488                    block_dim: (32, gqa, 1),
19489                    shared_mem_bytes: shmem,
19490                },
19491            )
19492        } else if fa_vec && fa_v3_active(head_dim) {
19493            // FA v3 lane _dc twin: the captured graph must run the SAME walk body as eager
19494            // under MEMRA_FA_V3=1 (eager, rows-verify and graph switch together).
19495            let gqa = (n_head / n_head_kv).max(1) as u32;
19496            let fv = if g {
19497                self.func_g("fa_decode_vec_q_v3_dc")
19498            } else {
19499                self.func("fa_decode_vec_q_v3_dc")
19500            };
19501            let shmem = (32 * head_dim * 2) as u32; // sV bf16 [FA_DEC_TILE=32][hd]
19502            (
19503                fv,
19504                LaunchConfig {
19505                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
19506                    block_dim: (32, gqa, 1),
19507                    shared_mem_bytes: shmem,
19508                },
19509            )
19510        } else if fa_vec && fa_v2_on() {
19511            // FAVENDOR lane: v2 _dc twin — the captured graph must run the SAME walk body as
19512            // eager under MEMRA_FA_V2=1 or graph_decode_gate's bit-identity breaks (the flag is
19513            // a numeric config; eager, rows-verify and graph all switch together).
19514            let gqa = (n_head / n_head_kv).max(1) as u32;
19515            let fv = if g {
19516                self.func_g("fa_decode_vec_q_v2_dc")
19517            } else {
19518                self.func("fa_decode_vec_q_v2_dc")
19519            };
19520            let shmem = (2 * 32 * head_dim * 2) as u32; // sK+sV bf16 [FA_DEC_TILE=32][hd]
19521            (
19522                fv,
19523                LaunchConfig {
19524                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
19525                    block_dim: (32, gqa, 1),
19526                    shared_mem_bytes: shmem,
19527                },
19528            )
19529        } else if fa_vec {
19530            let gqa = (n_head / n_head_kv).max(1) as u32;
19531            // REGISTER-DEQUANT twin: zero dynamic smem (see fa_decode above).
19532            let fv = if g {
19533                self.func_g("fa_decode_vec_q_dc")
19534            } else {
19535                self.func("fa_decode_vec_q_dc")
19536            };
19537            (
19538                fv,
19539                LaunchConfig {
19540                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
19541                    block_dim: (32, gqa, 1),
19542                    shared_mem_bytes: 0,
19543                },
19544            )
19545        } else {
19546            let q_view = q.as_view();
19547            let mut o_view = o.as_view_mut();
19548            return self.fa_decode_scalar_unified(
19549                &q_view,
19550                k,
19551                v,
19552                &mut o_view,
19553                head_dim,
19554                n_head,
19555                n_head_kv,
19556                0,
19557                Some(t_kv_dev),
19558                scale,
19559                n_splits,
19560                if fa_vec { sp } else { 256 },
19561                k_tok_bytes,
19562                v_tok_bytes,
19563                g,
19564                &mut *part_o,
19565                &mut *part_m,
19566                &mut *part_l,
19567                q8_out,
19568            );
19569        };
19570        let ski = sp as i32; // one-partition law: the twins derive ns_eff from (T_kv, ski)
19571        let __s_b = self.gpu.stream();
19572        let mut b = __s_b.launch_builder(&f);
19573        b.arg(q)
19574            .arg(k)
19575            .arg(v)
19576            .arg(&mut *part_o)
19577            .arg(&mut *part_m)
19578            .arg(&mut *part_l)
19579            .arg(&hd)
19580            .arg(&nh)
19581            .arg(&nhkv)
19582            .arg(t_kv_dev)
19583            .arg(&scale)
19584            .arg(&nsp)
19585            .arg(&ski)
19586            .arg(&ktb)
19587            .arg(&vtb);
19588        unsafe {
19589            b.launch(cfg)?;
19590        }
19591        let cfg2 = LaunchConfig {
19592            grid_dim: (n_head as u32, 1, 1),
19593            block_dim: (head_dim as u32, 1, 1),
19594            shared_mem_bytes: 0,
19595        };
19596        if let Some((oq, od)) = q8_out {
19597            let fc = if g {
19598                self.func_g("fa_decode_combine_q8_1")
19599            } else {
19600                self.fa_func("fa_decode_combine_q8_1", head_dim)
19601            };
19602            let __s_b2 = self.gpu.stream();
19603            let mut b2 = __s_b2.launch_builder(&fc);
19604            b2.arg(&*part_o)
19605                .arg(&*part_m)
19606                .arg(&*part_l)
19607                .arg(oq)
19608                .arg(od)
19609                .arg(&hd)
19610                .arg(&nh)
19611                .arg(&nsp);
19612            unsafe {
19613                b2.launch(cfg2)?;
19614            }
19615            return Ok(());
19616        }
19617        let fc = if g {
19618            self.func_g("fa_decode_combine_f32")
19619        } else {
19620            self.fa_func("fa_decode_combine_f32", head_dim)
19621        };
19622        let __s_b2 = self.gpu.stream();
19623        let mut b2 = __s_b2.launch_builder(&fc);
19624        b2.arg(&*part_o)
19625            .arg(&*part_m)
19626            .arg(&*part_l)
19627            .arg(o)
19628            .arg(&hd)
19629            .arg(&nh)
19630            .arg(&nsp);
19631        unsafe {
19632            b2.launch(cfg2)?;
19633        }
19634        Ok(())
19635    }
19636
19637    /// EAGER fa_decode geometry for a given actual `t_kv` (CUDA-GRAPH-PLAN §3.3 bucketing). Returns
19638    /// `(fa_vec, n_splits)` EXACTLY as `fa_decode` computes them so the graph-capture path can key its
19639    /// bucket on the same `(kernel, n_splits)` pair and pass a `bucket_max` that reproduces eager's
19640    /// n_splits bit-for-bit. (Per = ceil(t_kv/n_splits) is then recomputed from the DEVICE t_kv inside
19641    /// the kernel and matches eager when n_splits matches — the bit-identity contract.)
19642    pub fn fa_geom_eager(
19643        &self,
19644        t_kv: usize,
19645        head_dim: usize,
19646        n_head_kv: usize,
19647        g: bool,
19648    ) -> (bool, usize) {
19649        // MUST mirror `fa_decode` / `fa_decode_dc` (default-ON 2026-06-28). This is the bucket-key
19650        // source: if it disagrees with the actual kernel pick, the graph captures the wrong path and
19651        // replay diverges from eager. All three sites read MEMRA_NO_FA_VEC in lockstep.
19652        let fa_ok = std::env::var("MEMRA_NO_FA_VEC").is_err() && t_kv >= fa_vec_min_tkv();
19653        // hd512 dpl16 vec lane (gemma globals, 2026-07-11 graph-arc fix): the original key
19654        // hardcoded vec = hd<=256, so for hd512 it bucketed by the SCALAR 256-key splits while
19655        // the dpl16/rows_dpl16 kernels split by the ladder — n_splits changed WITHIN a bucket
19656        // (mid-ctx graph mismatch at pos 19 + partials OOB at longer runs). Mirror the real
19657        // fa_decode dispatch: vec512 above the fa512 floor, vec256 as before.
19658        let vec512 = fa_ok && head_dim == 512 && t_kv >= fa512_min_tkv();
19659        let mut fa_vec = vec512 || (fa_ok && head_dim <= 256 && head_dim % 32 == 0);
19660        // g (fp8-windowed): mirror kvmod's clamp — only the v4 lane parses e4m3 in the vec
19661        // family; everything else falls to the g-module scalar.
19662        // hd256 under g: v4-or-scalar. hd128/other under g ride the REGISTER g-lane (the
19663        // dq_K_lane/dq_V_lane macros are format-aware; v3 is format-gated off, v2/smem
19664        // arms are excluded under g). Mirrored in kvmod / fa_decode_dc / fa_geom_eager.
19665        if g && head_dim == 256 && !fa_v4_at(t_kv) {
19666            fa_vec = false;
19667        }
19668        let sp = fa_split_keys(t_kv, n_head_kv);
19669        let n_splits = if fa_vec {
19670            ((t_kv + sp - 1) / sp).max(1)
19671        } else {
19672            ((t_kv + 255) / 256).max(1)
19673        };
19674        (fa_vec, n_splits)
19675    }
19676
19677    /// `bucket_max` (host t_kv to feed `fa_decode_dc` / `full_attn_decode_dc`) that makes the _dc
19678    /// kernel pick the SAME (fa_vec, n_splits) as eager would for actual `t_kv`. Because the dc
19679    /// launcher derives both from `bucket_max` via the same formulas, we just hand it `t_kv` itself:
19680    /// the n_splits is then identical, and the per-split boundaries (computed from the DEVICE t_kv in
19681    /// the kernel) match eager exactly. The bucket KEY (for the graph HashMap) is `(fa_vec, n_splits)`.
19682    pub fn fa_bucket_key(
19683        &self,
19684        t_kv: usize,
19685        head_dim: usize,
19686        n_head_kv: usize,
19687        g: bool,
19688    ) -> (bool, usize) {
19689        self.fa_geom_eager(t_kv, head_dim, n_head_kv, g)
19690    }
19691
19692    /// CUDA-graph capture wrapper (CUDA-GRAPH-PLAN §3.2, llama.cpp warmup pattern). Runs `step`
19693    /// inline TWICE (warmup — lets the caching allocator settle to stable pointers and any one-time
19694    /// kernel attribute/JIT happen outside capture), then captures a THIRD invocation on the Engine's
19695    /// decode stream (RELAXED mode) and instantiates it into a replayable `CudaGraph`. The closure
19696    /// must enqueue ONLY device work on `e.stream()` (no dtoh / no synchronize / no host branch on
19697    /// device data) — every per-step varying scalar must come from a device counter. Returns the
19698    /// instantiated graph; `CudaGraph::launch()` replays the whole step in one dispatch.
19699    /// `capture_graph` with CAPTURE-RETAIN: every Engine allocation made during the warmups
19700    /// and the capture is kept alive in the returned keeper — hold it as long as the graph
19701    /// replays (transients returning to the pool get reused by unrelated work and corrupt
19702    /// replays; the draft-graph root cause). Model-generic, next capture reuses it.
19703    pub fn capture_graph_retained<F>(
19704        &self,
19705        step: F,
19706    ) -> Result<
19707        (
19708            cudarc::driver::CudaGraph,
19709            Vec<Box<dyn std::any::Any + Send>>,
19710        ),
19711        Box<dyn std::error::Error>,
19712    >
19713    where
19714        F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
19715    {
19716        use cudarc::driver::sys::CUgraphInstantiate_flags;
19717        self.capture_graph_retained_flags(
19718            CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
19719            step,
19720        )
19721    }
19722
19723    /// Retained capture with an explicit instantiate flag. ALLOC-FREE captured graphs
19724    /// (zero mem nodes — the gemma slotted door) should pass UPLOAD instead of
19725    /// AUTO_FREE_ON_LAUNCH: the auto-free flag's launch-time mem-pool scan was measured at
19726    /// ~0.25us/node (205us on the 826-node step) even with nothing to free.
19727    pub fn capture_graph_retained_flags<F>(
19728        &self,
19729        flags: cudarc::driver::sys::CUgraphInstantiate_flags,
19730        mut step: F,
19731    ) -> Result<
19732        (
19733            cudarc::driver::CudaGraph,
19734            Vec<Box<dyn std::any::Any + Send>>,
19735        ),
19736        Box<dyn std::error::Error>,
19737    >
19738    where
19739        F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
19740    {
19741        use cudarc::driver::sys::CUstreamCaptureMode;
19742        // KEEP scope = WARMUPS ONLY (2026-07-13): keep_if_capturing retains via
19743        // CudaSlice::clone, which is a device ALLOC + D2D COPY on the stream — clones made
19744        // while the capture region is open become dead copy NODES replayed every launch
19745        // (E4B: 1440 copies = 0.74ms/token, the whole graph-vs-eager regression). The
19746        // warmup runs allocate the same transient sequence at the same pool addresses, so
19747        // retaining the warmup clones preserves the draft-graph fix without polluting the
19748        // captured graph.
19749        self.capture_keep.lock().unwrap().clear();
19750        let was_tracking = self.gpu.ctx.is_event_tracking();
19751        if was_tracking {
19752            unsafe {
19753                self.gpu.ctx.disable_event_tracking();
19754            }
19755        }
19756        let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
19757            self.capture_keep_on
19758                .store(true, std::sync::atomic::Ordering::Relaxed);
19759            let w = (|| {
19760                step(self)?;
19761                step(self)
19762            })();
19763            self.capture_keep_on
19764                .store(false, std::sync::atomic::Ordering::Relaxed);
19765            w?;
19766            self.gpu.stream().synchronize()?;
19767            self.gpu
19768                .stream()
19769                .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
19770            let r = step(self);
19771            let g = self.gpu.stream().end_capture(flags);
19772            r?;
19773            let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
19774            graph.upload()?;
19775            Ok(graph)
19776        };
19777        let result = run();
19778        self.capture_keep_on
19779            .store(false, std::sync::atomic::Ordering::Relaxed);
19780        if was_tracking {
19781            unsafe {
19782                self.gpu.ctx.enable_event_tracking();
19783            }
19784        }
19785        let keeper = std::mem::take(&mut *self.capture_keep.lock().unwrap());
19786        Ok((result?, keeper))
19787    }
19788
19789    pub fn capture_graph<F>(
19790        &self,
19791        mut step: F,
19792    ) -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>>
19793    where
19794        F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
19795    {
19796        use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
19797        // EVENT TRACKING OFF for capture. The Engine creates a 2nd stream (copy_stream) so cudarc is in
19798        // multi-stream mode and, by default, records a CudaEvent per CudaSlice alloc/use to serialize
19799        // cross-stream access. Those per-buffer event waits issue stream ops that are NOT permitted
19800        // inside a capture region (CUDA_ERROR_STREAM_CAPTURE_UNSUPPORTED). The captured decode step is
19801        // strictly SINGLE-STREAM (every kernel on gpu.stream), so this synchronization is unnecessary
19802        // here — disable it for the whole warmup+capture, re-enable after. SAFETY: the decode-dc path
19803        // touches only gpu.stream; no buffer crosses to copy_stream during capture.
19804        let was_tracking = self.gpu.ctx.is_event_tracking();
19805        if was_tracking {
19806            unsafe {
19807                self.gpu.ctx.disable_event_tracking();
19808            }
19809        }
19810        // Q1 PROBE (MEMRA_GRAPH_IFLAG): the generic capture body's cuMemAllocAsync nodes are
19811        // EXACTLY BALANCED by in-graph free nodes (measured census q27: 1589 ALLOC / 1589
19812        // FREE), so AUTO_FREE_ON_LAUNCH has nothing to reclaim at launch — it only pays its
19813        // per-node launch-time mem-pool scan. `upload` / `none` select the alternatives to
19814        // measure that scan's real cost on the generic path. Diagnostic door only; the
19815        // default stays AUTO_FREE until a measured A/B justifies moving it.
19816        let iflag = {
19817            static F: std::sync::OnceLock<CUgraphInstantiate_flags> = std::sync::OnceLock::new();
19818            *F.get_or_init(|| match std::env::var("MEMRA_GRAPH_IFLAG").as_deref() {
19819                // UPLOAD = the gemma slotted door's zero-mem-node choice; PRIORITY = the flag
19820                // hybrid_forward.rs:5935 actually ships (both drop the auto-free launch scan).
19821                Ok("upload") => CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_UPLOAD,
19822                Ok("priority") => {
19823                    CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY
19824                }
19825                _ => CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
19826            })
19827        };
19828        // MEMRA_GRAPH_CAPTIME=1 (Q1 lane): phase-resolved capture cost. Recapture is paid at
19829        // every kernel-class crossing, so it — not steady-state decode — is the quantity a
19830        // mem-node reduction could plausibly shrink. Only `instantiate` (cuStreamEndCapture +
19831        // cuGraphInstantiateWithFlags) and `upload` scale with node count; the warmups are
19832        // eager step executions and are node-count-invariant. Printing the split bounds the
19833        // refactor's ceiling instead of assuming it.
19834        let ct = {
19835            static T: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19836            *T.get_or_init(|| std::env::var("MEMRA_GRAPH_CAPTIME").as_deref() == Ok("1"))
19837        };
19838        // MEMRA_GRAPH_WARMUPS (Q1 lane; DEFAULT 1 since lane/graph-warmups 2026-08-05): the
19839        // phase split showed the eager warmups are 80% of recapture cost (q27 27.4 of 34.4 ms
19840        // pod / 42% of 52.6 ms 5090) — 3x larger than the ENTIRE mem-node ceiling the audit
19841        // chased, and node-count-invariant, so no capture-body refactor could touch it.
19842        // Warmup 2's theorized job was async-pool ADDRESS STABILITY: warmup 1's allocs may
19843        // grow/map the pool, warmup 2 re-walks the same sequence over the freed blocks so the
19844        // captured third run bakes settled addresses. That hazard is the #68 stale-baked-
19845        // address class — which the engine now guards STRUCTURALLY rather than by re-walking:
19846        // in-body transients are captured as BALANCED in-graph alloc/free node pairs (census
19847        // 1589/1589 — replays allocate for themselves; no baked transient pointers), every
19848        // externally-referenced buffer is stable-pointer by design (fa_part_pool retires-on-
19849        // grow and never frees, resident counters/scratch, cache set in place), and the
19850        // draft-graph path additionally rides capture_graph_retained (capture_keep holds all
19851        // warmup+capture allocs alive). One warmup therefore suffices for kernel-attr
19852        // settling and pool mapping. Arbitrated adversarially, not by taste:
19853        // graph-warmup-stress (pool-growth cycles large<->small x10, overlap arm, forced
19854        // recaptures over freed blocks — bit-identity vs eager + canary teeth) is GREEN at
19855        // warmups=1 on the deployment rig, plus graph-decode-gate 256-step bit-identity,
19856        // graph-session-gate, run-spec K=1..8 (receipts research/graph-warmups-5090-20260805/
19857        // + the pod's research/graph-allocfree-20260805/). Measured: recapture -38..-42% q27 /
19858        // -41% q9, decode +~1%, capture+prime -13ms. MEMRA_GRAPH_WARMUPS=2 = the rollback
19859        // seam; tools/graph-warmup-stress-gate.sh = the gate any regression re-runs.
19860        let warmups = {
19861            static W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
19862            *W.get_or_init(|| {
19863                std::env::var("MEMRA_GRAPH_WARMUPS")
19864                    .ok()
19865                    .and_then(|v| v.parse().ok())
19866                    .filter(|n| *n >= 1)
19867                    .unwrap_or(1)
19868            })
19869        };
19870        let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
19871            let t_w = std::time::Instant::now();
19872            // warmup: inline runs (no capture) so allocator pointers + kernel attrs are stable.
19873            for _ in 0..warmups {
19874                step(self)?;
19875            }
19876            self.gpu.stream().synchronize()?;
19877            let ms_warm = t_w.elapsed().as_secs_f64() * 1e3;
19878            // capture the third run.
19879            let t_c = std::time::Instant::now();
19880            self.gpu
19881                .stream()
19882                .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
19883            // If the body errors mid-capture, end the capture before propagating so the stream isn't
19884            // left in a capturing state.
19885            let r = step(self);
19886            let ms_body = t_c.elapsed().as_secs_f64() * 1e3;
19887            let t_i = std::time::Instant::now();
19888            let g = self.gpu.stream().end_capture(iflag);
19889            let ms_inst = t_i.elapsed().as_secs_f64() * 1e3;
19890            r?;
19891            let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
19892            let t_u = std::time::Instant::now();
19893            graph.upload()?;
19894            if ct {
19895                println!(
19896                    "[graph-captime] warmup2x {ms_warm:.2} ms  capture-body {ms_body:.2} ms  \
19897                          instantiate {ms_inst:.2} ms  upload {:.2} ms",
19898                    t_u.elapsed().as_secs_f64() * 1e3
19899                );
19900            }
19901            Ok(graph)
19902        };
19903        let result = run();
19904        if was_tracking {
19905            unsafe {
19906                self.gpu.ctx.enable_event_tracking();
19907            }
19908        }
19909        result
19910    }
19911
19912    /// gdn_scan variant where state_in/out are CudaViews (resident SSM state, in-place per step).
19913    pub fn gdn_scan_s128_view(
19914        &self,
19915        q: &CudaSlice<f32>,
19916        k: &CudaSlice<f32>,
19917        v: &CudaSlice<f32>,
19918        g: &CudaSlice<f32>,
19919        beta: &CudaSlice<f32>,
19920        state_in: &cudarc::driver::CudaView<f32>,
19921        state_out: &mut cudarc::driver::CudaViewMut<f32>,
19922        o: &mut CudaSlice<f32>,
19923        n_head: usize,
19924        t: usize,
19925        scale: f32,
19926    ) -> Result<(), Box<dyn std::error::Error>> {
19927        let f = self.func("gdn_scan_s128");
19928        const S_V: u32 = 128;
19929        const WARP: u32 = 32;
19930        const COLS: u32 = 4;
19931        let cfg = LaunchConfig {
19932            grid_dim: (n_head as u32, 1, S_V / COLS),
19933            block_dim: (WARP, COLS, 1),
19934            shared_mem_bytes: 0,
19935        };
19936        let (h, ti) = (n_head as i32, t as i32);
19937        let __s_b = self.gpu.stream();
19938        let mut b = __s_b.launch_builder(&f);
19939        b.arg(q)
19940            .arg(k)
19941            .arg(v)
19942            .arg(g)
19943            .arg(beta)
19944            .arg(state_in)
19945            .arg(state_out)
19946            .arg(o)
19947            .arg(&h)
19948            .arg(&ti)
19949            .arg(&scale);
19950        unsafe {
19951            b.launch(cfg)?;
19952        }
19953        Ok(())
19954    }
19955
19956    /// conv1d where the input is a CudaView (resident conv state assembled in place).
19957    pub fn ssm_conv1d_view(
19958        &self,
19959        x: &cudarc::driver::CudaView<f32>,
19960        w: &CudaSlice<f32>,
19961        y: &mut CudaSlice<f32>,
19962        conv_dim: usize,
19963        t: usize,
19964        d_conv: usize,
19965        silu: bool,
19966    ) -> Result<(), Box<dyn std::error::Error>> {
19967        let f = self.func("ssm_conv1d_silu_f32");
19968        // grid.x = channel, grid.y = T-tiles (block 256 strides over T) — parallel over both axes.
19969        let cfg = LaunchConfig {
19970            grid_dim: (conv_dim as u32, ((t as u32 + 255) / 256).max(1), 1),
19971            block_dim: (256, 1, 1),
19972            shared_mem_bytes: 0,
19973        };
19974        let (cd, ti, dc, s) = (conv_dim as i32, t as i32, d_conv as i32, silu as i32);
19975        let __s_b = self.gpu.stream();
19976        let mut b = __s_b.launch_builder(&f);
19977        b.arg(x).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc).arg(&s);
19978        unsafe {
19979            b.launch(cfg)?;
19980        }
19981        Ok(())
19982    }
19983
19984    /// Depthwise causal conv1d + optional SiLU.
19985    /// x:[conv_dim, T+d_conv-1] channel-major (first d_conv-1 cols = carried state),
19986    /// w:[d_conv, conv_dim] kernel-major, y:[conv_dim, T] channel-major.
19987    /// FUSED prefill conv (token-major input, zero left-state): replaces
19988    /// transpose + zeros + conv_left_pad + ssm_conv1d with ONE launch reading the matmul output
19989    /// directly. Output channel-major [conv_dim, T], SiLU applied. BIT-IDENTICAL accumulation.
19990    pub fn ssm_conv1d_tm(
19991        &self,
19992        qkv_tm: &CudaSlice<f32>,
19993        w: &CudaSlice<f32>,
19994        y: &mut CudaSlice<f32>,
19995        conv_dim: usize,
19996        t: usize,
19997        d_conv: usize,
19998    ) -> Result<(), Box<dyn std::error::Error>> {
19999        let f = self.func("ssm_conv1d_tm_f32");
20000        let cfg = LaunchConfig {
20001            grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
20002            block_dim: (256, 1, 1),
20003            shared_mem_bytes: 0,
20004        };
20005        let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
20006        let __s_b = self.gpu.stream();
20007        let mut b = __s_b.launch_builder(&f);
20008        b.arg(qkv_tm).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc);
20009        unsafe {
20010            b.launch(cfg)?;
20011        }
20012        Ok(())
20013    }
20014
20015    /// BATCHED verify conv (T>1, carried state): window reads the resident conv ring for
20016    /// negative rows; separate ring-update launch afterwards. BIT-IDENTICAL per value to the
20017    /// T=1 chain. T >= pad rides the pure input-column ring update (unchanged legacy path);
20018    /// T < pad (the MEMRA_SPEC_M2 t=2 verify arm) needs old-ring sources for the roll — the
20019    /// update kernel would race reading the ring it rewrites, so that arm clones the ring
20020    /// (dtod) and rolls via ssm_conv_ring_rebuild (PURE COPIES: the ring stores raw input
20021    /// columns; the final ring == what T sequential decode ring rolls leave).
20022    pub fn ssm_conv1d_tm_state(
20023        &self,
20024        qkv_tm: &CudaSlice<f32>,
20025        conv_state: &mut CudaSlice<f32>,
20026        w: &CudaSlice<f32>,
20027        y: &mut CudaSlice<f32>,
20028        conv_dim: usize,
20029        t: usize,
20030        d_conv: usize,
20031    ) -> Result<(), Box<dyn std::error::Error>> {
20032        self.ssm_conv1d_tm_state_pad(qkv_tm, conv_state, w, y, conv_dim, t, d_conv, None)
20033    }
20034
20035    /// task #14: `pad_len` = device true length for PADDED prime graphs — the ring update
20036    /// reads rows [len-pad, len) instead of the pad tail. None = the classic host-T path.
20037    #[allow(clippy::too_many_arguments)]
20038    pub fn ssm_conv1d_tm_state_pad(
20039        &self,
20040        qkv_tm: &CudaSlice<f32>,
20041        conv_state: &mut CudaSlice<f32>,
20042        w: &CudaSlice<f32>,
20043        y: &mut CudaSlice<f32>,
20044        conv_dim: usize,
20045        t: usize,
20046        d_conv: usize,
20047        pad_len: Option<&CudaSlice<i32>>,
20048    ) -> Result<(), Box<dyn std::error::Error>> {
20049        assert!(t >= 1, "ssm_conv1d_tm_state requires T >= 1");
20050        // clone BEFORE the window kernel is issued is not required (stream-ordered: the dtod and
20051        // the window kernel both read the pre-roll ring; the roll launches after both) — but
20052        // cloning first keeps the ordering trivially correct under any future stream split.
20053        let ring_old = if t < d_conv - 1 {
20054            Some(self.clone_dtod(conv_state)?)
20055        } else {
20056            None
20057        };
20058        {
20059            let f = self.func("ssm_conv1d_tm_state_f32");
20060            let cfg = LaunchConfig {
20061                grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
20062                block_dim: (256, 1, 1),
20063                shared_mem_bytes: 0,
20064            };
20065            let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
20066            let __s_b = self.gpu.stream();
20067            let mut b = __s_b.launch_builder(&f);
20068            b.arg(qkv_tm)
20069                .arg(&*conv_state)
20070                .arg(w)
20071                .arg(y)
20072                .arg(&cd)
20073                .arg(&ti)
20074                .arg(&dc);
20075            unsafe {
20076                b.launch(cfg)?;
20077            }
20078        }
20079        match (ring_old, pad_len) {
20080            (None, Some(len_d)) => {
20081                let f = self.func("ssm_conv_ring_update_dev_f32");
20082                let n = conv_dim * (d_conv - 1);
20083                let cfg = LaunchConfig::for_num_elems(n as u32);
20084                let (cd, dc) = (conv_dim as i32, d_conv as i32);
20085                let __s_b = self.gpu.stream();
20086                let mut b = __s_b.launch_builder(&f);
20087                b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
20088                unsafe {
20089                    b.launch(cfg)?;
20090                }
20091            }
20092            (None, None) => {
20093                let f = self.func("ssm_conv_ring_update_f32");
20094                let n = conv_dim * (d_conv - 1);
20095                let cfg = LaunchConfig::for_num_elems(n as u32);
20096                let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
20097                let __s_b = self.gpu.stream();
20098                let mut b = __s_b.launch_builder(&f);
20099                b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
20100                unsafe {
20101                    b.launch(cfg)?;
20102                }
20103            }
20104            (Some(old), _) => {
20105                self.ssm_conv_ring_rebuild(qkv_tm, &old, conv_state, conv_dim, t, d_conv)?
20106            }
20107        }
20108        Ok(())
20109    }
20110
20111    /// qkv-view twin (task #16): batched prime reads the concat GEMM output directly.
20112    pub fn ssm_conv1d_tm_state_pad_v(
20113        &self,
20114        qkv_tm: &cudarc::driver::CudaView<f32>,
20115        conv_state: &mut CudaSlice<f32>,
20116        w: &CudaSlice<f32>,
20117        y: &mut CudaSlice<f32>,
20118        conv_dim: usize,
20119        t: usize,
20120        d_conv: usize,
20121        pad_len: Option<&CudaSlice<i32>>,
20122    ) -> Result<(), Box<dyn std::error::Error>> {
20123        assert!(t >= 1, "ssm_conv1d_tm_state requires T >= 1");
20124        // clone BEFORE the window kernel is issued is not required (stream-ordered: the dtod and
20125        // the window kernel both read the pre-roll ring; the roll launches after both) — but
20126        // cloning first keeps the ordering trivially correct under any future stream split.
20127        let ring_old = if t < d_conv - 1 {
20128            Some(self.clone_dtod(conv_state)?)
20129        } else {
20130            None
20131        };
20132        {
20133            let f = self.func("ssm_conv1d_tm_state_f32");
20134            let cfg = LaunchConfig {
20135                grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
20136                block_dim: (256, 1, 1),
20137                shared_mem_bytes: 0,
20138            };
20139            let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
20140            let __s_b = self.gpu.stream();
20141            let mut b = __s_b.launch_builder(&f);
20142            b.arg(qkv_tm)
20143                .arg(&*conv_state)
20144                .arg(w)
20145                .arg(y)
20146                .arg(&cd)
20147                .arg(&ti)
20148                .arg(&dc);
20149            unsafe {
20150                b.launch(cfg)?;
20151            }
20152        }
20153        match (ring_old, pad_len) {
20154            (None, Some(len_d)) => {
20155                let f = self.func("ssm_conv_ring_update_dev_f32");
20156                let n = conv_dim * (d_conv - 1);
20157                let cfg = LaunchConfig::for_num_elems(n as u32);
20158                let (cd, dc) = (conv_dim as i32, d_conv as i32);
20159                let __s_b = self.gpu.stream();
20160                let mut b = __s_b.launch_builder(&f);
20161                b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
20162                unsafe {
20163                    b.launch(cfg)?;
20164                }
20165            }
20166            (None, None) => {
20167                let f = self.func("ssm_conv_ring_update_f32");
20168                let n = conv_dim * (d_conv - 1);
20169                let cfg = LaunchConfig::for_num_elems(n as u32);
20170                let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
20171                let __s_b = self.gpu.stream();
20172                let mut b = __s_b.launch_builder(&f);
20173                b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
20174                unsafe {
20175                    b.launch(cfg)?;
20176                }
20177            }
20178            (Some(_), _) => unreachable!(
20179                "ssm_conv1d_tm_state_pad_v: T < d_conv-1 has no view path (PRIME_MIN_T gates it)"
20180            ),
20181        }
20182        Ok(())
20183    }
20184
20185    /// PREFIX conv-ring rebuild (spec REPLAY-FREE partial accept): overwrite the resident ring
20186    /// with the state a T=1 chain holds after only the FIRST `tc` columns of `qkv_tm` — the last
20187    /// `pad` entries of [ring_old | cols 0..tc-1]. PURE COPIES (the ring stores raw inputs; no
20188    /// arithmetic, cannot perturb FP order). `ring_old` = the pre-round snapshot ring.
20189    pub fn ssm_conv_ring_rebuild(
20190        &self,
20191        qkv_tm: &CudaSlice<f32>,
20192        ring_old: &CudaSlice<f32>,
20193        conv_state: &mut CudaSlice<f32>,
20194        conv_dim: usize,
20195        tc: usize,
20196        d_conv: usize,
20197    ) -> Result<(), Box<dyn std::error::Error>> {
20198        let f = self.func("ssm_conv_ring_rebuild_f32");
20199        let n = conv_dim * (d_conv - 1);
20200        let cfg = LaunchConfig::for_num_elems(n as u32);
20201        let (cd, ti, dc) = (conv_dim as i32, tc as i32, d_conv as i32);
20202        let __s_b = self.gpu.stream();
20203        let mut b = __s_b.launch_builder(&f);
20204        b.arg(qkv_tm)
20205            .arg(ring_old)
20206            .arg(conv_state)
20207            .arg(&cd)
20208            .arg(&ti)
20209            .arg(&dc);
20210        unsafe {
20211            b.launch(cfg)?;
20212        }
20213        Ok(())
20214    }
20215
20216    /// FUSED decode GDN prep (T=1): repack + q/k L2-norm + beta sigmoid + g_log in one launch.
20217    /// Replaces 5 tiny serialized kernels on the decode critical path. L2 reduce runs as a 32-lane
20218    /// warp tree (vs l2_norm_f32's 256-thread two-level tree) — same math, different FP sum order;
20219    /// the argmax + run-spec gates are the authority.
20220    #[allow(clippy::too_many_arguments)]
20221    pub fn gdn_prep_decode(
20222        &self,
20223        conv_out: &CudaSlice<f32>,
20224        beta_raw: &CudaSlice<f32>,
20225        alpha: &CudaSlice<f32>,
20226        dt_bias: &CudaSlice<f32>,
20227        a: &CudaSlice<f32>,
20228        q_l2: &mut CudaSlice<f32>,
20229        k_l2: &mut CudaSlice<f32>,
20230        v_g: &mut CudaSlice<f32>,
20231        beta: &mut CudaSlice<f32>,
20232        g_log: &mut CudaSlice<f32>,
20233        d_state: usize,
20234        num_v: usize,
20235        num_k: usize,
20236        key_dim: usize,
20237        eps: f32,
20238    ) -> Result<(), Box<dyn std::error::Error>> {
20239        let f = self.func("gdn_prep_decode_f32");
20240        let cfg = LaunchConfig {
20241            grid_dim: (num_v as u32, 1, 1),
20242            block_dim: (32, 4, 1),
20243            shared_mem_bytes: 0,
20244        };
20245        let (ds, nv, nk, kd) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
20246        let __s_b = self.gpu.stream();
20247        let mut b = __s_b.launch_builder(&f);
20248        b.arg(conv_out)
20249            .arg(beta_raw)
20250            .arg(alpha)
20251            .arg(dt_bias)
20252            .arg(a)
20253            .arg(q_l2)
20254            .arg(k_l2)
20255            .arg(v_g)
20256            .arg(beta)
20257            .arg(g_log)
20258            .arg(&ds)
20259            .arg(&nv)
20260            .arg(&nk)
20261            .arg(&kd)
20262            .arg(&eps);
20263        unsafe {
20264            b.launch(cfg)?;
20265        }
20266        Ok(())
20267    }
20268
20269    /// FUSED prefill conv + GDN repack: token-major qkv -> q_g/k_g/v_g in ONE launch (no conv_out
20270    /// materialization, no qkv_to_gdn_repack pass). BIT-IDENTICAL values; scatter matches
20271    /// qkv_to_gdn_repack's modulo head-repeat mapping exactly.
20272    #[allow(clippy::too_many_arguments)]
20273    pub fn ssm_conv1d_gdn(
20274        &self,
20275        qkv_tm: &CudaSlice<f32>,
20276        w: &CudaSlice<f32>,
20277        q_g: &mut CudaSlice<f32>,
20278        k_g: &mut CudaSlice<f32>,
20279        v_g: &mut CudaSlice<f32>,
20280        conv_dim: usize,
20281        t: usize,
20282        d_conv: usize,
20283        d_state: usize,
20284        num_v: usize,
20285        num_k: usize,
20286        key_dim: usize,
20287    ) -> Result<(), Box<dyn std::error::Error>> {
20288        let f = self.func("ssm_conv1d_gdn_f32");
20289        let cfg = LaunchConfig {
20290            grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
20291            block_dim: (256, 1, 1),
20292            shared_mem_bytes: 0,
20293        };
20294        let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
20295        let (ds, nv, nk, kd) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
20296        let __s_b = self.gpu.stream();
20297        let mut b = __s_b.launch_builder(&f);
20298        b.arg(qkv_tm)
20299            .arg(w)
20300            .arg(q_g)
20301            .arg(k_g)
20302            .arg(v_g)
20303            .arg(&cd)
20304            .arg(&ti)
20305            .arg(&dc)
20306            .arg(&ds)
20307            .arg(&nv)
20308            .arg(&nk)
20309            .arg(&kd);
20310        unsafe {
20311            b.launch(cfg)?;
20312        }
20313        Ok(())
20314    }
20315
20316    pub fn ssm_conv1d(
20317        &self,
20318        x: &CudaSlice<f32>,
20319        w: &CudaSlice<f32>,
20320        y: &mut CudaSlice<f32>,
20321        conv_dim: usize,
20322        t: usize,
20323        d_conv: usize,
20324        silu: bool,
20325    ) -> Result<(), Box<dyn std::error::Error>> {
20326        let f = self.func("ssm_conv1d_silu_f32");
20327        let cfg = LaunchConfig {
20328            grid_dim: (conv_dim as u32, ((t as u32 + 255) / 256).max(1), 1),
20329            block_dim: (256, 1, 1),
20330            shared_mem_bytes: 0,
20331        };
20332        let (cd, ti, dc, s) = (conv_dim as i32, t as i32, d_conv as i32, silu as i32);
20333        let __s_b = self.gpu.stream();
20334        let mut b = __s_b.launch_builder(&f);
20335        b.arg(x).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc).arg(&s);
20336        unsafe {
20337            b.launch(cfg)?;
20338        }
20339        Ok(())
20340    }
20341
20342    /// Gated DeltaNet scan, S_v=128. q,k,v:[128,H,T]; g,beta:[H,T]; state:[128,128,H] transposed;
20343    /// o:[128,H,T]. Single sequence.
20344    pub fn gdn_scan_s128(
20345        &self,
20346        q: &CudaSlice<f32>,
20347        k: &CudaSlice<f32>,
20348        v: &CudaSlice<f32>,
20349        g: &CudaSlice<f32>,
20350        beta: &CudaSlice<f32>,
20351        state_in: &CudaSlice<f32>,
20352        state_out: &mut CudaSlice<f32>,
20353        o: &mut CudaSlice<f32>,
20354        n_head: usize,
20355        t: usize,
20356        scale: f32,
20357    ) -> Result<(), Box<dyn std::error::Error>> {
20358        let f = self.func("gdn_scan_s128");
20359        const S_V: u32 = 128;
20360        const WARP: u32 = 32;
20361        const COLS_PER_BLOCK: u32 = 4;
20362        let cfg = LaunchConfig {
20363            grid_dim: (n_head as u32, 1, S_V / COLS_PER_BLOCK),
20364            block_dim: (WARP, COLS_PER_BLOCK, 1),
20365            shared_mem_bytes: 0,
20366        };
20367        let (h, ti) = (n_head as i32, t as i32);
20368        let __s_b = self.gpu.stream();
20369        let mut b = __s_b.launch_builder(&f);
20370        b.arg(q)
20371            .arg(k)
20372            .arg(v)
20373            .arg(g)
20374            .arg(beta)
20375            .arg(state_in)
20376            .arg(state_out)
20377            .arg(o)
20378            .arg(&h)
20379            .arg(&ti)
20380            .arg(&scale);
20381        unsafe {
20382            b.launch(cfg)?;
20383        }
20384        Ok(())
20385    }
20386
20387    // ==== B2' batched decode state ops (decode_batch.rs) ====
20388    // Per-seq state pointers ride device u64 arrays (views into the per-step pointer table).
20389    // Bodies are the single-seq kernels per sequence — bit-identical per row.
20390
20391    #[allow(clippy::too_many_arguments)]
20392    pub fn ssm_conv1d_fused_decode_b(
20393        &self,
20394        qkv_cols: &CudaSlice<f32>,
20395        conv_state_ptrs: &cudarc::driver::CudaView<u64>,
20396        w: &CudaSlice<f32>,
20397        conv_outs: &mut CudaSlice<f32>,
20398        conv_dim: usize,
20399        d_conv: usize,
20400        b_n: usize,
20401    ) -> Result<(), Box<dyn std::error::Error>> {
20402        let f = self.func("ssm_conv1d_fused_decode_b_f32");
20403        let cfg = LaunchConfig {
20404            grid_dim: (((conv_dim + 255) / 256) as u32, 1, b_n as u32),
20405            block_dim: (256, 1, 1),
20406            shared_mem_bytes: 0,
20407        };
20408        let (cd, dc) = (conv_dim as i32, d_conv as i32);
20409        let __s_b = self.gpu.stream();
20410        let mut b = __s_b.launch_builder(&f);
20411        b.arg(qkv_cols)
20412            .arg(conv_state_ptrs)
20413            .arg(w)
20414            .arg(conv_outs)
20415            .arg(&cd)
20416            .arg(&dc);
20417        unsafe {
20418            b.launch(cfg)?;
20419        }
20420        Ok(())
20421    }
20422
20423    #[allow(clippy::too_many_arguments)]
20424    pub fn gdn_prep_decode_b(
20425        &self,
20426        conv_outs: &CudaSlice<f32>,
20427        beta_raws: &CudaSlice<f32>,
20428        alphas: &CudaSlice<f32>,
20429        dt_bias: &CudaSlice<f32>,
20430        a: &CudaSlice<f32>,
20431        q_l2: &mut CudaSlice<f32>,
20432        k_l2: &mut CudaSlice<f32>,
20433        v_g: &mut CudaSlice<f32>,
20434        beta: &mut CudaSlice<f32>,
20435        g_log: &mut CudaSlice<f32>,
20436        d_state: usize,
20437        num_v: usize,
20438        num_k: usize,
20439        key_dim: usize,
20440        eps: f32,
20441        conv_dim: usize,
20442        b_n: usize,
20443    ) -> Result<(), Box<dyn std::error::Error>> {
20444        let f = self.func("gdn_prep_decode_b_f32");
20445        let cfg = LaunchConfig {
20446            grid_dim: (num_v as u32, 1, b_n as u32),
20447            block_dim: (32, 4, 1),
20448            shared_mem_bytes: 0,
20449        };
20450        let (ds, nv, nk, kd, cd) = (
20451            d_state as i32,
20452            num_v as i32,
20453            num_k as i32,
20454            key_dim as i32,
20455            conv_dim as i32,
20456        );
20457        let __s_b = self.gpu.stream();
20458        let mut b = __s_b.launch_builder(&f);
20459        b.arg(conv_outs)
20460            .arg(beta_raws)
20461            .arg(alphas)
20462            .arg(dt_bias)
20463            .arg(a)
20464            .arg(q_l2)
20465            .arg(k_l2)
20466            .arg(v_g)
20467            .arg(beta)
20468            .arg(g_log)
20469            .arg(&ds)
20470            .arg(&nv)
20471            .arg(&nk)
20472            .arg(&kd)
20473            .arg(&eps)
20474            .arg(&cd);
20475        unsafe {
20476            b.launch(cfg)?;
20477        }
20478        Ok(())
20479    }
20480
20481    #[allow(clippy::too_many_arguments)]
20482    pub fn gdn_scan_s128_batched(
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_ptrs: &cudarc::driver::CudaView<u64>,
20490        state_out_ptrs: &cudarc::driver::CudaView<u64>,
20491        o: &mut CudaSlice<f32>,
20492        n_head: usize,
20493        b_n: usize,
20494        scale: f32,
20495    ) -> Result<(), Box<dyn std::error::Error>> {
20496        let f = self.func("gdn_scan_s128_b");
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, b_n as u32, S_V / COLS_PER_BLOCK),
20502            block_dim: (WARP, COLS_PER_BLOCK, 1),
20503            shared_mem_bytes: 0,
20504        };
20505        let h = n_head 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_ptrs)
20514            .arg(state_out_ptrs)
20515            .arg(o)
20516            .arg(&h)
20517            .arg(&scale);
20518        unsafe {
20519            b.launch(cfg)?;
20520        }
20521        Ok(())
20522    }
20523
20524    /// VIEW twins of the three GDN decode state wrappers (lane/verify-launchslim): identical
20525    /// launches, row args as views into the caller's packed [T, ...] buffers — the t-parallel
20526    /// verify's per-row loop passes slices instead of paying an arithmetic-free dtod per row
20527    /// per kernel (48 layers x T rows x 4 copies/round on the money path). Same kernels, same
20528    /// numeric class; only the pointer arithmetic moved host-side.
20529    #[allow(clippy::too_many_arguments)]
20530    pub fn ssm_conv1d_fused_decode_b_view(
20531        &self,
20532        qkv_cols: &cudarc::driver::CudaView<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_view(
20563        &self,
20564        conv_outs: &CudaSlice<f32>,
20565        beta_raws: &cudarc::driver::CudaView<f32>,
20566        alphas: &cudarc::driver::CudaView<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_view(
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 cudarc::driver::CudaViewMut<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    /// A4 seam: chunked WY GDN prefill. DEFAULT ON (`MEMRA_GDN_CHUNKED=0` = rollback to the
20663    /// sequential scan). Flipped 2026-07-04 with the full battery green: kernel-check ALL
20664    /// GREEN x {9B, 27B} incl the f64-truth chunk gates; run-gen argmax 82==82 both models
20665    /// on AND off (24/24 sweep runs); run-spec K={1,2,3,4,6,8} PASS x {9B synth, 9B text,
20666    /// 27B p2, 27B p3}; e2e first-16-token agreement 6/6 (full-256 drifts at index 47-125
20667    /// on 5/6 prompts — accepted cache-state-FP class, batched-prime precedent).
20668    /// PREFILL-ONLY: decode + spec verify never route here (decode==verify dispatch
20669    /// identity law); prime_cache/forward/forward_last are the only callers.
20670    pub fn gdn_chunked_enabled() -> bool {
20671        static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20672        *E.get_or_init(|| {
20673            std::env::var("MEMRA_GDN_CHUNKED")
20674                .map(|v| v != "0")
20675                .unwrap_or(true)
20676        })
20677    }
20678
20679    /// A4 chunk size (MEMRA_GDN_CHUNK, default 32 — the sweep winner: the O(T*C) chunk
20680    /// matrices grow with C while the sequential state pass is C-flat, so smaller chunks
20681    /// win; C=32/64 also get the register-history solve template). Clamped to multiples
20682    /// of 32 in [32, 128] (kernel row mappings require it).
20683    pub fn gdn_chunk_size() -> usize {
20684        static C: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
20685        *C.get_or_init(|| {
20686            let c: usize = std::env::var("MEMRA_GDN_CHUNK")
20687                .ok()
20688                .and_then(|v| v.parse().ok())
20689                .unwrap_or(32);
20690            c.clamp(32, 128) / 32 * 32
20691        })
20692    }
20693
20694    /// A4: chunked WY / blockwise-inverse GDN prefill (see cu/hybrid.cu K1-K5 header for the
20695    /// math). Same contract as `gdn_scan_s128` (layouts, state ping-pong) but chunk-parallel:
20696    /// NOT bit-identical to the sequential scan (chunked FP accumulation order); run-gen
20697    /// argmax + run-spec batteries are the accuracy authority. PREFILL callers only.
20698    #[allow(clippy::too_many_arguments)]
20699    /// task #18: K1-K3 of the chunked WY scan (shared by the per-seq path and the
20700    /// batched-prime varlen path). Returns (gcum, P, U, W); `A` is K3-internal.
20701    #[allow(clippy::too_many_arguments, clippy::type_complexity)]
20702    #[allow(clippy::too_many_arguments)]
20703    pub fn gdn_chunk_k123(
20704        &self,
20705        q: &CudaSlice<f32>,
20706        k: &CudaSlice<f32>,
20707        v: &CudaSlice<f32>,
20708        g: &CudaSlice<f32>,
20709        beta: &CudaSlice<f32>,
20710        wb16: Option<&mut CudaSlice<u8>>,
20711        n_head: usize,
20712        t: usize,
20713        c: usize,
20714        hk: usize,
20715        k2w: Option<(&CudaSlice<u8>, &CudaSlice<u8>, &mut CudaSlice<u8>)>,
20716    ) -> Result<
20717        (
20718            CudaSlice<f32>,
20719            CudaSlice<f32>,
20720            CudaSlice<f32>,
20721            CudaSlice<f32>,
20722        ),
20723        Box<dyn std::error::Error>,
20724    > {
20725        const D: usize = 128;
20726        let h = n_head;
20727        let nc = (t + c - 1) / c;
20728        let (hi, ti, ci) = (h as i32, t as i32, c as i32);
20729        let mut gcum = self.uninit(t * h)?;
20730        let mut a = self.uninit(nc * h * c * c)?;
20731        let mut p = self.uninit(nc * h * c * c)?;
20732        let mut u = self.uninit(nc * h * c * D)?;
20733        let mut w = self.uninit(nc * h * c * D)?;
20734        {
20735            // K1
20736            let f = self.func("gdn_chunk_cumgate_f32");
20737            let cfg = LaunchConfig {
20738                grid_dim: (nc as u32, h as u32, 1),
20739                block_dim: (32, 1, 1),
20740                shared_mem_bytes: 0,
20741            };
20742            let __s_b = self.gpu.stream();
20743            let mut b = __s_b.launch_builder(&f);
20744            b.arg(g).arg(&mut gcum).arg(&hi).arg(&ti).arg(&ci);
20745            unsafe {
20746                b.launch(cfg)?;
20747            }
20748        }
20749        if let Some((qb, kb, pb)) = k2w {
20750            // K2-wgmma (MEMRA_GDN_WGMMA path, c==32): A + pre-masked Pb16 in one kernel;
20751            // the P f32 buffer stays UNWRITTEN (its only wgmma-path consumer is Pb16).
20752            assert!(c == 32, "gdn_k2_wgmma is a C==32 tile");
20753            let f = self.func("gdn_k2_wgmma");
20754            let cfg = LaunchConfig {
20755                grid_dim: (nc as u32, h as u32, 1),
20756                block_dim: (128, 1, 1),
20757                shared_mem_bytes: 0,
20758            };
20759            let hki = hk as i32;
20760            let __s_b = self.gpu.stream();
20761            let mut b = __s_b.launch_builder(&f);
20762            b.arg(qb)
20763                .arg(kb)
20764                .arg(&gcum)
20765                .arg(beta)
20766                .arg(&mut a)
20767                .arg(&mut *pb)
20768                .arg(&hi)
20769                .arg(&ti)
20770                .arg(&ci)
20771                .arg(&hki);
20772            unsafe {
20773                b.launch(cfg)?;
20774            }
20775        } else if c <= 64 && !portable_mma_gated() {
20776            // K2 register-tiled (2x2 outputs/thread, whole-chunk smem k tile)
20777            let f = self.func("gdn_chunk_attn_f32");
20778            let jt = ((c + 31) / 32) as u32;
20779            let cfg = LaunchConfig {
20780                grid_dim: (nc as u32, h as u32, jt),
20781                block_dim: (256, 1, 1),
20782                shared_mem_bytes: 0,
20783            };
20784            let hki = hk as i32;
20785            let __s_b = self.gpu.stream();
20786            let mut b = __s_b.launch_builder(&f);
20787            b.arg(q)
20788                .arg(k)
20789                .arg(&gcum)
20790                .arg(beta)
20791                .arg(&mut a)
20792                .arg(&mut p)
20793                .arg(&hi)
20794                .arg(&ti)
20795                .arg(&ci)
20796                .arg(&hki);
20797            unsafe {
20798                b.launch(cfg)?;
20799            }
20800        } else {
20801            // K2 generic (C = 128, or the portable target's low-smem fallback)
20802            assert!(
20803                hk == h,
20804                "generic K2 is broadcast-only (de-broadcast rides C==32)"
20805            );
20806            let f = self.func("gdn_chunk_attn_g_f32");
20807            let cfg = LaunchConfig {
20808                grid_dim: (nc as u32, h as u32, 1),
20809                block_dim: (32, 8, 1),
20810                shared_mem_bytes: 0,
20811            };
20812            let __s_b = self.gpu.stream();
20813            let mut b = __s_b.launch_builder(&f);
20814            b.arg(q)
20815                .arg(k)
20816                .arg(&gcum)
20817                .arg(beta)
20818                .arg(&mut a)
20819                .arg(&mut p)
20820                .arg(&hi)
20821                .arg(&ti)
20822                .arg(&ci);
20823            unsafe {
20824                b.launch(cfg)?;
20825            }
20826        }
20827        {
20828            // K3 (register-history templates for C=32/64; local-memory generic otherwise)
20829            let cfg = LaunchConfig {
20830                grid_dim: (nc as u32, h as u32, 1),
20831                block_dim: (256, 1, 1),
20832                shared_mem_bytes: 0,
20833            };
20834            match c {
20835                32 | 64 => {
20836                    let f = self.func(if c == 32 {
20837                        "gdn_chunk_solve32_f32"
20838                    } else {
20839                        "gdn_chunk_solve64_f32"
20840                    });
20841                    // mirror-fold: W's bf16 twin emitted on store (0 = skip)
20842                    let wb: u64 = match wb16 {
20843                        Some(d) => self.addr_u8(d),
20844                        None => 0,
20845                    };
20846                    let hki = hk as i32;
20847                    let __s_b = self.gpu.stream();
20848                    let mut b = __s_b.launch_builder(&f);
20849                    b.arg(v)
20850                        .arg(k)
20851                        .arg(&a)
20852                        .arg(&gcum)
20853                        .arg(&mut u)
20854                        .arg(&mut w)
20855                        .arg(&wb)
20856                        .arg(&hi)
20857                        .arg(&ti)
20858                        .arg(&hki);
20859                    unsafe {
20860                        b.launch(cfg)?;
20861                    }
20862                }
20863                _ => {
20864                    assert!(hk == h, "generic K3 is broadcast-only");
20865                    let f = self.func("gdn_chunk_solve_f32");
20866                    let __s_b = self.gpu.stream();
20867                    let mut b = __s_b.launch_builder(&f);
20868                    b.arg(v)
20869                        .arg(k)
20870                        .arg(&a)
20871                        .arg(&gcum)
20872                        .arg(&mut u)
20873                        .arg(&mut w)
20874                        .arg(&hi)
20875                        .arg(&ti)
20876                        .arg(&ci);
20877                    unsafe {
20878                        b.launch(cfg)?;
20879                    }
20880                }
20881            }
20882        }
20883        Ok((gcum, p, u, w))
20884    }
20885
20886    /// task #21 de-broadcast seam: q/k stored at num_k distinct GQA heads instead of
20887    /// the num_v broadcast. MEMRA_GDN_DB=0 reverts. Only the chunked prefill path
20888    /// consumes the compact layout (hk plumbed; hk == H reproduces broadcast exactly).
20889    pub fn gdn_db_on() -> bool {
20890        std::env::var("MEMRA_GDN_DB").as_deref() != Ok("0")
20891    }
20892
20893    /// Whether the K4/K5 mma pair serves at chunk size `c` (mirrors gdn_scan_chunked's
20894    /// seam read — env re-read per call ON PURPOSE, kernel-check pins both configs).
20895    /// DEFAULT ON for sm_120a builds too (lane/moeprime-nvfp4-direct, 2026-08-21): the pair
20896    /// was qualified on 90a only and left env-opt-in elsewhere; measured on Blackwell it
20897    /// wins on BOTH rigs — one RTX PRO 6000 (ornith15 pp14715 12,036 -> 12,751/12,957,
20898    /// +6-8%, both orders) and the local 5090 (q38-27b pp6435 1,397/1,429 -> 1,427/1,446,
20899    /// both orders) — with kernel-check/run-gen/margin-gate/run-spec green under the flag.
20900    /// bf16 HMMA (m16n8k16) is sm_80-class PTX; only the wgmma nest stays Hopper-gated.
20901    pub fn gdn_mma_enabled(&self, c: usize) -> bool {
20902        !portable_mma_gated()
20903            && c == 32
20904            && match std::env::var("MEMRA_GDN_MMA").as_deref() {
20905                Ok("1") => true,
20906                Ok("0") => false,
20907                _ => gdn_mma_default_on(),
20908            }
20909    }
20910
20911    /// task #22: whether the fused K4+K5 (+K2) wgmma path serves (nested inside the
20912    /// mma config; same per-call env read discipline). HARD-gated to the Hopper build:
20913    /// the wgmma asm bodies exist only at __CUDA_ARCH__ == 900 (MEMRA_K45_REAL,
20914    /// wgmma_common.cuh) — on every other arch the kernel compiles EMPTY, so an env
20915    /// force would silently produce garbage. Required since the sm_120a mma default
20916    /// flip made MEMRA_GDN_WGMMA=1 alone reach this branch there.
20917    pub fn gdn_wgmma_on(&self, c: usize) -> bool {
20918        cfg!(memra_hopper_mma)
20919            && self.gdn_mma_enabled(c)
20920            && std::env::var("MEMRA_GDN_WGMMA").as_deref() != Ok("0")
20921    }
20922
20923    /// task #18 conv-fuse: carried-ring conv + SiLU + GDN repack in ONE pass (the
20924    /// conv_out intermediate and its transposed re-read disappear — 11.8ms of the
20925    /// T=2048 prime). Ring update stays the separate follow-up launch (pad-aware).
20926    /// BIT-IDENTICAL values to ssm_conv1d_tm_state_pad + qkv_to_gdn_repack.
20927    #[allow(clippy::too_many_arguments)]
20928    pub fn ssm_conv1d_gdn_state_pad(
20929        &self,
20930        qkv_tm: &cudarc::driver::CudaView<f32>,
20931        conv_state: &mut CudaSlice<f32>,
20932        w: &CudaSlice<f32>,
20933        q_g: &mut CudaSlice<f32>,
20934        k_g: &mut CudaSlice<f32>,
20935        v_g: &mut CudaSlice<f32>,
20936        conv_dim: usize,
20937        t: usize,
20938        d_conv: usize,
20939        d_state: usize,
20940        num_v: usize,
20941        num_k: usize,
20942        key_dim: usize,
20943        hk: usize,
20944        pad_len: Option<&CudaSlice<i32>>,
20945    ) -> Result<(), Box<dyn std::error::Error>> {
20946        assert!(
20947            t >= d_conv - 1,
20948            "fused state conv requires T >= pad (PRIME_MIN_T gates)"
20949        );
20950        {
20951            let f = self.func("ssm_conv1d_gdn_state_f32");
20952            let cfg = LaunchConfig {
20953                grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
20954                block_dim: (256, 1, 1),
20955                shared_mem_bytes: 0,
20956            };
20957            let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
20958            let (ds, nv, nk, kd, hki) = (
20959                d_state as i32,
20960                num_v as i32,
20961                num_k as i32,
20962                key_dim as i32,
20963                hk as i32,
20964            );
20965            let __s_b = self.gpu.stream();
20966            let mut b = __s_b.launch_builder(&f);
20967            b.arg(qkv_tm)
20968                .arg(&*conv_state)
20969                .arg(w)
20970                .arg(q_g)
20971                .arg(k_g)
20972                .arg(v_g)
20973                .arg(&cd)
20974                .arg(&ti)
20975                .arg(&dc)
20976                .arg(&ds)
20977                .arg(&nv)
20978                .arg(&nk)
20979                .arg(&kd)
20980                .arg(&hki);
20981            unsafe {
20982                b.launch(cfg)?;
20983            }
20984        }
20985        match pad_len {
20986            Some(len_d) => {
20987                let f = self.func("ssm_conv_ring_update_dev_f32");
20988                let n = conv_dim * (d_conv - 1);
20989                let cfg = LaunchConfig::for_num_elems(n as u32);
20990                let (cd, dc) = (conv_dim as i32, d_conv as i32);
20991                let __s_b = self.gpu.stream();
20992                let mut b = __s_b.launch_builder(&f);
20993                b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
20994                unsafe {
20995                    b.launch(cfg)?;
20996                }
20997            }
20998            None => {
20999                let f = self.func("ssm_conv_ring_update_f32");
21000                let n = conv_dim * (d_conv - 1);
21001                let cfg = LaunchConfig::for_num_elems(n as u32);
21002                let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
21003                let __s_b = self.gpu.stream();
21004                let mut b = __s_b.launch_builder(&f);
21005                b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
21006                unsafe {
21007                    b.launch(cfg)?;
21008                }
21009            }
21010        }
21011        Ok(())
21012    }
21013
21014    /// task #18 increment 2: allocate ONE sequence's chunk buffers (no launches) —
21015    /// K1-K5 all run varlen afterwards. `a`/`w` become struct members so the varlen
21016    /// K2/K3 can write them.
21017    pub fn gdn_chunk_alloc(
21018        &self,
21019        n_head: usize,
21020        t: usize,
21021        c: usize,
21022        hk: usize,
21023    ) -> Result<GdnChunkBufs, Box<dyn std::error::Error>> {
21024        const D: usize = 128;
21025        assert!(
21026            c == 32,
21027            "gdn_chunk_alloc: varlen chain is the C==32 mma pair"
21028        );
21029        let h = n_head;
21030        let nc = (t + c - 1) / c;
21031        Ok(GdnChunkBufs {
21032            gcum: self.uninit(t * h)?,
21033            a: self.uninit(nc * h * c * c)?,
21034            p: self.uninit(nc * h * c * c)?,
21035            u: self.uninit(nc * h * c * D)?,
21036            w: self.uninit(nc * h * c * D)?,
21037            kb16: self.alloc_u8_uninit(t * hk * D * 2)?,
21038            wb16: self.alloc_u8_uninit(nc * h * c * D * 2)?,
21039            y16: self.alloc_u8_uninit(nc * h * c * D * 2)?,
21040            ssnap16: self.alloc_u8_uninit(nc * h * D * D * 2)?,
21041            qb16: self.alloc_u8_uninit(t * hk * D * 2)?,
21042            pb16: self.alloc_u8_uninit(nc * h * c * c * 2)?,
21043            o: self.uninit(D * h * t)?,
21044            t,
21045            nc,
21046        })
21047    }
21048
21049    /// view-source twin of f32_to_bf16 (the batched FA3 v mirror reads a concat view).
21050    pub fn f32_to_bf16_v(
21051        &self,
21052        x: &cudarc::driver::CudaView<f32>,
21053        dst: &mut CudaSlice<u8>,
21054        n: usize,
21055    ) -> Result<(), Box<dyn std::error::Error>> {
21056        let f = self.func("f32_to_bf16_bulk");
21057        let ni = n as i64;
21058        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
21059        let __s_b = self.gpu.stream();
21060        let mut b = __s_b.launch_builder(&f);
21061        b.arg(x).arg(dst).arg(&ni);
21062        unsafe {
21063            b.launch(cfg)?;
21064        }
21065        Ok(())
21066    }
21067
21068    /// f32 -> bf16 bulk mirror into a caller buffer (the K4/K5 operand mirrors).
21069    pub fn f32_to_bf16_into(
21070        &self,
21071        x: &CudaSlice<f32>,
21072        dst: &mut CudaSlice<u8>,
21073        n: usize,
21074    ) -> Result<(), Box<dyn std::error::Error>> {
21075        let f = self.func("f32_to_bf16_bulk");
21076        let ni = n as i64;
21077        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
21078        let __s_b = self.gpu.stream();
21079        let mut b = __s_b.launch_builder(&f);
21080        b.arg(x).arg(dst).arg(&ni);
21081        unsafe {
21082            b.launch(cfg)?;
21083        }
21084        Ok(())
21085    }
21086
21087    /// task #18 increment 2: varlen K1+K2+K3 — three launches run every sequence's
21088    /// cumgate/attn/solve (per-block math identical to the per-seq kernels).
21089    pub fn gdn_chunk_k123_vl8(
21090        &self,
21091        seqs: &[GdnSeqVl],
21092        n_head: usize,
21093        hk: usize,
21094        wq: Option<&GdnWVl8>,
21095    ) -> Result<(), Box<dyn std::error::Error>> {
21096        let b = seqs.len();
21097        assert!(b >= 1 && b <= 8, "gdn_chunk_k123_vl8: 1..=8 sequences");
21098        let mut packed = [GdnSeqVl::default(); 8];
21099        packed[..b].copy_from_slice(seqs);
21100        let v = GdnVl8(packed);
21101        let (hi, ci) = (n_head as i32, 32i32);
21102        let max_nc = seqs.iter().map(|a| a.nc).max().unwrap() as u32;
21103        {
21104            let f = self.func("gdn_chunk_cumgate_vl");
21105            let cfg = LaunchConfig {
21106                grid_dim: (max_nc, n_head as u32, b as u32),
21107                block_dim: (32, 1, 1),
21108                shared_mem_bytes: 0,
21109            };
21110            let __s_lb = self.gpu.stream();
21111            let mut lb = __s_lb.launch_builder(&f);
21112            lb.arg(&v).arg(&hi).arg(&ci);
21113            unsafe {
21114                lb.launch(cfg)?;
21115            }
21116        }
21117        let hki = hk as i32;
21118        if let Some(w) = wq {
21119            // K2-wgmma vl twin (writes A + pre-masked Pb16)
21120            let f = self.func("gdn_k2_wgmma_vl");
21121            let cfg = LaunchConfig {
21122                grid_dim: (max_nc, n_head as u32, b as u32),
21123                block_dim: (128, 1, 1),
21124                shared_mem_bytes: 0,
21125            };
21126            let __s_lb = self.gpu.stream();
21127            let mut lb = __s_lb.launch_builder(&f);
21128            lb.arg(&v).arg(w).arg(&hi).arg(&ci).arg(&hki);
21129            unsafe {
21130                lb.launch(cfg)?;
21131            }
21132        } else {
21133            let f = self.func("gdn_chunk_attn_vl");
21134            let cfg = LaunchConfig {
21135                grid_dim: (max_nc, n_head as u32, b as u32),
21136                block_dim: (256, 1, 1),
21137                shared_mem_bytes: 0,
21138            };
21139            let __s_lb = self.gpu.stream();
21140            let mut lb = __s_lb.launch_builder(&f);
21141            lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
21142            unsafe {
21143                lb.launch(cfg)?;
21144            }
21145        }
21146        {
21147            let f = self.func("gdn_chunk_solve32_vl");
21148            let cfg = LaunchConfig {
21149                grid_dim: (max_nc, n_head as u32, b as u32),
21150                block_dim: (256, 1, 1),
21151                shared_mem_bytes: 0,
21152            };
21153            let __s_lb = self.gpu.stream();
21154            let mut lb = __s_lb.launch_builder(&f);
21155            lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
21156            unsafe {
21157                lb.launch(cfg)?;
21158            }
21159        }
21160        Ok(())
21161    }
21162
21163    /// task #18 increment 3: varlen PREP chain — conv(+ring) / repack / fused-l2 /
21164    /// fused gate-prep, 5 launches for every sequence (per-element math identical
21165    /// to the per-seq kernels; l2/gate fusions write disjoint outputs).
21166    #[allow(clippy::too_many_arguments)]
21167    pub fn gdn_prep_vl8(
21168        &self,
21169        seqs: &[GdnPrepVl],
21170        conv_w: &CudaSlice<f32>,
21171        dt_bias: &CudaSlice<f32>,
21172        a: &CudaSlice<f32>,
21173        conv_dim: usize,
21174        d_conv: usize,
21175        d_state: usize,
21176        num_v: usize,
21177        num_k: usize,
21178        key_dim: usize,
21179        hk: usize,
21180        eps: f32,
21181    ) -> Result<(), Box<dyn std::error::Error>> {
21182        let b = seqs.len();
21183        assert!(b >= 1 && b <= 8);
21184        let mut packed = [GdnPrepVl::default(); 8];
21185        packed[..b].copy_from_slice(seqs);
21186        let v = GdnPrepVl8(packed);
21187        let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
21188        let (cdi, dci) = (conv_dim as i32, d_conv as i32);
21189        let conv_fuse = std::env::var("MEMRA_CONV_FUSE").as_deref() != Ok("0");
21190        assert!(
21191            conv_fuse || hk == num_v,
21192            "de-broadcast requires the fused conv"
21193        );
21194        if conv_fuse {
21195            let f = self.func("ssm_conv1d_gdn_state_vl");
21196            let cfg = LaunchConfig {
21197                grid_dim: ((conv_dim as u32).div_ceil(256), max_t, b as u32),
21198                block_dim: (256, 1, 1),
21199                shared_mem_bytes: 0,
21200            };
21201            let (dsi, nvi, nki, kdi, hki) = (
21202                d_state as i32,
21203                num_v as i32,
21204                num_k as i32,
21205                key_dim as i32,
21206                hk as i32,
21207            );
21208            let __s_lb = self.gpu.stream();
21209            let mut lb = __s_lb.launch_builder(&f);
21210            lb.arg(&v)
21211                .arg(conv_w)
21212                .arg(&cdi)
21213                .arg(&dci)
21214                .arg(&dsi)
21215                .arg(&nvi)
21216                .arg(&nki)
21217                .arg(&kdi)
21218                .arg(&hki);
21219            unsafe {
21220                lb.launch(cfg)?;
21221            }
21222        } else {
21223            let f = self.func("ssm_conv1d_tm_state_vl");
21224            let cfg = LaunchConfig {
21225                grid_dim: ((conv_dim as u32).div_ceil(256), max_t, b as u32),
21226                block_dim: (256, 1, 1),
21227                shared_mem_bytes: 0,
21228            };
21229            let __s_lb = self.gpu.stream();
21230            let mut lb = __s_lb.launch_builder(&f);
21231            lb.arg(&v).arg(conv_w).arg(&cdi).arg(&dci);
21232            unsafe {
21233                lb.launch(cfg)?;
21234            }
21235        }
21236        {
21237            let f = self.func("ssm_conv_ring_update_vl");
21238            let n = (conv_dim * (d_conv - 1)) as u32;
21239            let cfg = LaunchConfig {
21240                grid_dim: (n.div_ceil(256), 1, b as u32),
21241                block_dim: (256, 1, 1),
21242                shared_mem_bytes: 0,
21243            };
21244            let __s_lb = self.gpu.stream();
21245            let mut lb = __s_lb.launch_builder(&f);
21246            lb.arg(&v).arg(&cdi).arg(&dci);
21247            unsafe {
21248                lb.launch(cfg)?;
21249            }
21250        }
21251        if !conv_fuse {
21252            let f = self.func("qkv_to_gdn_repack_vl");
21253            let n = max_t * (num_v * d_state) as u32;
21254            let cfg = LaunchConfig {
21255                grid_dim: (n.div_ceil(256), 1, b as u32),
21256                block_dim: (256, 1, 1),
21257                shared_mem_bytes: 0,
21258            };
21259            let (dsi, nvi, nki, kdi) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
21260            let __s_lb = self.gpu.stream();
21261            let mut lb = __s_lb.launch_builder(&f);
21262            lb.arg(&v).arg(&dsi).arg(&nvi).arg(&nki).arg(&kdi);
21263            unsafe {
21264                lb.launch(cfg)?;
21265            }
21266        }
21267        if Self::l2_v2_on(d_state) {
21268            let f = self.func("gdn_l2_v2_vl");
21269            let cfg = LaunchConfig {
21270                grid_dim: ((max_t * hk as u32).div_ceil(8), 2, b as u32),
21271                block_dim: (256, 1, 1),
21272                shared_mem_bytes: 0,
21273            };
21274            let (dsi, nvi) = (d_state as i32, hk as i32);
21275            let __s_lb = self.gpu.stream();
21276            let mut lb = __s_lb.launch_builder(&f);
21277            lb.arg(&v).arg(&dsi).arg(&nvi).arg(&eps);
21278            unsafe {
21279                lb.launch(cfg)?;
21280            }
21281        } else {
21282            let f = self.func("gdn_l2_vl");
21283            let cfg = LaunchConfig {
21284                grid_dim: (max_t * hk as u32, 2, b as u32),
21285                block_dim: (256, 1, 1),
21286                shared_mem_bytes: 0,
21287            };
21288            let (dsi, nvi) = (d_state as i32, hk as i32);
21289            let __s_lb = self.gpu.stream();
21290            let mut lb = __s_lb.launch_builder(&f);
21291            lb.arg(&v).arg(&dsi).arg(&nvi).arg(&eps);
21292            unsafe {
21293                lb.launch(cfg)?;
21294            }
21295        }
21296        {
21297            let f = self.func("gdn_gate_prep_vl");
21298            let n = max_t * num_v as u32;
21299            let cfg = LaunchConfig {
21300                grid_dim: (n.div_ceil(256), 1, b as u32),
21301                block_dim: (256, 1, 1),
21302                shared_mem_bytes: 0,
21303            };
21304            let nvi = num_v as i32;
21305            let __s_lb = self.gpu.stream();
21306            let mut lb = __s_lb.launch_builder(&f);
21307            lb.arg(&v).arg(dt_bias).arg(a).arg(&nvi);
21308            unsafe {
21309                lb.launch(cfg)?;
21310            }
21311        }
21312        Ok(())
21313    }
21314
21315    /// varlen bf16 mirrors over the gdnseq_t table (which: 0 = k_l2 -> kb16, 1 = w -> wb16).
21316    pub fn gdn_mirror_vl8(
21317        &self,
21318        seqs: &[GdnSeqVl],
21319        n_head: usize,
21320        which: i32,
21321        hk: usize,
21322    ) -> Result<(), Box<dyn std::error::Error>> {
21323        let b = seqs.len();
21324        assert!(b >= 1 && b <= 8);
21325        let mut packed = [GdnSeqVl::default(); 8];
21326        packed[..b].copy_from_slice(seqs);
21327        let v = GdnVl8(packed);
21328        let ept = (if which == 0 { hk } else { n_head } * 128) as i32;
21329        let max_n = seqs
21330            .iter()
21331            .map(|s| {
21332                if which == 0 {
21333                    s.t as i64 * ept as i64
21334                } else {
21335                    s.nc as i64 * ept as i64 * 32
21336                }
21337            })
21338            .max()
21339            .unwrap();
21340        let f = self.func("gdn_mirror_vl");
21341        let blocks = ((max_n as u32).div_ceil(4)).div_ceil(256);
21342        let cfg = LaunchConfig {
21343            grid_dim: (blocks, 1, b as u32),
21344            block_dim: (256, 1, 1),
21345            shared_mem_bytes: 0,
21346        };
21347        let __s_lb = self.gpu.stream();
21348        let mut lb = __s_lb.launch_builder(&f);
21349        lb.arg(&v).arg(&ept).arg(&which);
21350        unsafe {
21351            lb.launch(cfg)?;
21352        }
21353        Ok(())
21354    }
21355
21356    /// varlen gated-norm tail (+f16out) — one launch replaces B gated_rmsnorm calls.
21357    pub fn gdn_tail_vl8(
21358        &self,
21359        seqs: &[GdnPrepVl],
21360        norm_w: &CudaSlice<f32>,
21361        d_state: usize,
21362        num_v: usize,
21363        eps: f32,
21364    ) -> Result<(), Box<dyn std::error::Error>> {
21365        let b = seqs.len();
21366        assert!(b >= 1 && b <= 8);
21367        let mut packed = [GdnPrepVl::default(); 8];
21368        packed[..b].copy_from_slice(seqs);
21369        let v = GdnPrepVl8(packed);
21370        let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
21371        let f = self.func("gated_rmsnorm_f16out_vl");
21372        // block_dim MUST match gated_rmsnorm's (128): the reduction tree order pins the scale
21373        let cfg = LaunchConfig {
21374            grid_dim: (max_t * num_v as u32, 1, b as u32),
21375            block_dim: (128, 1, 1),
21376            shared_mem_bytes: 0,
21377        };
21378        let (dsi, nvi) = (d_state as i32, num_v as i32);
21379        let __s_lb = self.gpu.stream();
21380        let mut lb = __s_lb.launch_builder(&f);
21381        lb.arg(&v).arg(norm_w).arg(&dsi).arg(&nvi).arg(&eps);
21382        unsafe {
21383            lb.launch(cfg)?;
21384        }
21385        Ok(())
21386    }
21387
21388    /// Raw device address helpers for the varlen by-value arg struct (single-stream
21389    /// launches; every buffer outlives the call — the f16 FFI discipline).
21390    pub fn addr_f32(&self, x: &CudaSlice<f32>) -> u64 {
21391        use cudarc::driver::DevicePtr;
21392        let s = self.gpu.stream();
21393        let (p, _g) = x.device_ptr(&s);
21394        p as u64
21395    }
21396    pub fn addr_f32_mut(&self, x: &mut CudaSlice<f32>) -> u64 {
21397        use cudarc::driver::DevicePtrMut;
21398        let s = self.gpu.stream();
21399        let (p, _g) = x.device_ptr_mut(&s);
21400        p as u64
21401    }
21402    pub fn addr_f32v(&self, x: &cudarc::driver::CudaView<f32>) -> u64 {
21403        use cudarc::driver::DevicePtr;
21404        let s = self.gpu.stream();
21405        let (p, _g) = x.device_ptr(&s);
21406        p as u64
21407    }
21408    pub fn addr_u8(&self, x: &CudaSlice<u8>) -> u64 {
21409        use cudarc::driver::DevicePtr;
21410        let s = self.gpu.stream();
21411        let (p, _g) = x.device_ptr(&s);
21412        p as u64
21413    }
21414
21415    /// task #18: the varlen K4+K5 pair — TWO launches run every sequence's state pass
21416    /// and output pass (grid gains a seq dim; per-block math identical to the per-seq
21417    /// launches, so this is strictly bit-gateable against them).
21418    pub fn gdn_chunk_vl8(
21419        &self,
21420        seqs: &[GdnSeqVl],
21421        n_head: usize,
21422        scale: f32,
21423        hk: usize,
21424        wq: Option<&GdnWVl8>,
21425    ) -> Result<(), Box<dyn std::error::Error>> {
21426        const NSPLIT: u32 = 4;
21427        let b = seqs.len();
21428        assert!(b >= 1 && b <= 8, "gdn_chunk_vl8: 1..=8 sequences");
21429        let mut packed = [GdnSeqVl::default(); 8];
21430        packed[..b].copy_from_slice(seqs);
21431        let v = GdnVl8(packed);
21432        let (hi, ci) = (n_head as i32, 32i32);
21433        let max_nc = seqs.iter().map(|a| a.nc).max().unwrap() as u32;
21434        let hki = hk as i32;
21435        if let Some(w) = wq {
21436            // K4+K5 fused wgmma vl twin: one launch, Y/Ssnap never materialized.
21437            let f = self.func("gdn_k45_wgmma_vl");
21438            let cfg = LaunchConfig {
21439                grid_dim: (n_head as u32, NSPLIT, b as u32),
21440                block_dim: (256, 1, 1),
21441                shared_mem_bytes: 0,
21442            };
21443            let __s_lb = self.gpu.stream();
21444            let mut lb = __s_lb.launch_builder(&f);
21445            lb.arg(&v).arg(w).arg(&scale).arg(&hi).arg(&ci).arg(&hki);
21446            unsafe {
21447                lb.launch(cfg)?;
21448            }
21449            let _ = max_nc;
21450            return Ok(());
21451        }
21452        {
21453            let f = self.func("gdn_chunk_state_mma_vl");
21454            let cfg = LaunchConfig {
21455                grid_dim: (n_head as u32, NSPLIT, b as u32),
21456                block_dim: (256, 1, 1),
21457                shared_mem_bytes: 0,
21458            };
21459            let __s_lb = self.gpu.stream();
21460            let mut lb = __s_lb.launch_builder(&f);
21461            lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
21462            unsafe {
21463                lb.launch(cfg)?;
21464            }
21465        }
21466        {
21467            let f = self.func("gdn_chunk_output_mma_vl");
21468            let cfg = LaunchConfig {
21469                grid_dim: (max_nc, n_head as u32, b as u32),
21470                block_dim: (256, 1, 1),
21471                shared_mem_bytes: 0,
21472            };
21473            let __s_lb = self.gpu.stream();
21474            let mut lb = __s_lb.launch_builder(&f);
21475            lb.arg(&v).arg(&hi).arg(&ci).arg(&scale).arg(&hki);
21476            unsafe {
21477                lb.launch(cfg)?;
21478            }
21479        }
21480        Ok(())
21481    }
21482    pub fn gdn_scan_chunked(
21483        &self,
21484        q: &CudaSlice<f32>,
21485        k: &CudaSlice<f32>,
21486        v: &CudaSlice<f32>,
21487        g: &CudaSlice<f32>,
21488        beta: &CudaSlice<f32>,
21489        kb16_pre: Option<&CudaSlice<u8>>,
21490        qb16_pre: Option<&CudaSlice<u8>>,
21491        state_in: &CudaSlice<f32>,
21492        state_out: &mut CudaSlice<f32>,
21493        o: &mut CudaSlice<f32>,
21494        n_head: usize,
21495        t: usize,
21496        scale: f32,
21497        c: usize,
21498        hk: usize,
21499    ) -> Result<(), Box<dyn std::error::Error>> {
21500        const D: usize = 128;
21501        const NSPLIT: u32 = 4;
21502        assert!(c >= 1 && c <= 128, "gdn_scan_chunked: C must be in 1..=128");
21503        let h = n_head;
21504        let nc = (t + c - 1) / c;
21505        let (hi, ti, ci) = (h as i32, t as i32, c as i32);
21506        // mirror-fold (round 27): on the mma path W's bf16 twin is emitted by K3's store
21507        // (wb16 pre-allocated and threaded through k123) and k's by the producer l2 when
21508        // the caller hands `kb16_pre` — both standalone mirror passes disappear.
21509        // default mirrors gdn_mma_enabled (incl. the sm_120a-build ON, lane/moeprime-nvfp4-direct)
21510        let gdn_mma_pre = !portable_mma_gated()
21511            && c == 32
21512            && match std::env::var("MEMRA_GDN_MMA").as_deref() {
21513                Ok("1") => true,
21514                Ok("0") => false,
21515                _ => gdn_mma_default_on(),
21516            };
21517        let mut wb16_pre: Option<CudaSlice<u8>> = if gdn_mma_pre {
21518            Some(self.alloc_u8_uninit(nc * h * c * D * 2)?)
21519        } else {
21520            None
21521        };
21522        // K2-wgmma pre-work (MEMRA_GDN_WGMMA): the kb16/qb16 mirrors hoist ABOVE K123 so
21523        // K2 rides them via cp.async; K2 writes the pre-masked Pb16 directly (the
21524        // gdn_p_bf16_masked pass and the in-branch mirror builds disappear).
21525        // mirrors gdn_wgmma_on: hard Hopper-build gate (empty wgmma bodies elsewhere)
21526        let gdn_wgmma_pre = cfg!(memra_hopper_mma)
21527            && gdn_mma_pre
21528            && std::env::var("MEMRA_GDN_WGMMA").as_deref() != Ok("0");
21529        let nk = t * hk * D;
21530        let mut kb16_local: Option<CudaSlice<u8>> = None;
21531        if gdn_mma_pre && kb16_pre.is_none() {
21532            let mut kb = self.alloc_u8_uninit(nk * 2)?;
21533            let f = self.func("f32_to_bf16_bulk");
21534            let n2 = nk as i64;
21535            let cfg2 = LaunchConfig::for_num_elems((nk as u32).div_ceil(4));
21536            let __s_b = self.gpu.stream();
21537            let mut b = __s_b.launch_builder(&f);
21538            b.arg(k).arg(&mut kb).arg(&n2);
21539            unsafe {
21540                b.launch(cfg2)?;
21541            }
21542            kb16_local = Some(kb);
21543        }
21544        let kb16_ref0: Option<&CudaSlice<u8>> = kb16_local.as_ref().or(kb16_pre);
21545        if let Some(kb) = kb16_pre {
21546            assert!(kb.len() >= nk * 2, "kb16_pre too small");
21547        }
21548        let mut qb16: Option<CudaSlice<u8>> = None;
21549        let mut pb16: Option<CudaSlice<u8>> = None;
21550        if gdn_wgmma_pre {
21551            // mirror-fold (round 35): prep's l2 v2 emits qb16 in-epilogue (kb16 pattern);
21552            // the standalone bulk cvt only serves callers without the prep mirror.
21553            if qb16_pre.is_none() {
21554                let mut qb = self.alloc_u8_uninit(nk * 2)?;
21555                let f = self.func("f32_to_bf16_bulk");
21556                let n2 = nk as i64;
21557                let cfg2 = LaunchConfig::for_num_elems((nk as u32).div_ceil(4));
21558                let __s_b = self.gpu.stream();
21559                let mut b = __s_b.launch_builder(&f);
21560                b.arg(q).arg(&mut qb).arg(&n2);
21561                unsafe {
21562                    b.launch(cfg2)?;
21563                }
21564                qb16 = Some(qb);
21565            } else if let Some(qb) = qb16_pre {
21566                assert!(qb.len() >= nk * 2, "qb16_pre too small");
21567            }
21568            pb16 = Some(self.alloc_u8_uninit(nc * h * c * c * 2)?);
21569        }
21570        let qb16_ref0: Option<&CudaSlice<u8>> = qb16.as_ref().or(qb16_pre);
21571        let k2w = if gdn_wgmma_pre {
21572            Some((
21573                *qb16_ref0.as_ref().unwrap(),
21574                *kb16_ref0.as_ref().unwrap(),
21575                pb16.as_mut().unwrap(),
21576            ))
21577        } else {
21578            None
21579        };
21580        let (gcum, p, u, w) =
21581            self.gdn_chunk_k123(q, k, v, g, beta, wb16_pre.as_mut(), n_head, t, c, hk, k2w)?;
21582        let _ = &w;
21583        let mut y = self.uninit(nc * h * c * D)?;
21584        let mut ssnap = self.uninit(nc * h * D * D)?; // chunk-start state snapshots (K5 phase 1)
21585        // K4-MMA seam (MEMRA_GDN_MMA; harness verdict 1.75x — tools/bench_gdn_k4.cu, ledger
21586        // 2026-07-26): M in mma accumulator fragments, bf16 W/k mirrors through a cp.async
21587        // ring. C==32 only (the kernel's tile). PROMOTED default-ON on the Hopper lane
21588        // after the STATE-CARRY battery (2026-07-26): 2048-token prime (64 in-kernel state
21589        // carries) -> 256 greedy decode tokens IDENTICAL to f32 on 3 seeds, AND chunked-
21590        // continuation prime (MEMRA_PRIME_CHUNK=512, 4 cross-call carries via cache.recur)
21591        // IDENTICAL on 2 seeds; plus argmax MATCH, pp512 +3.5% (17286), oracle out
21592        // mean_rel ~1e-4. kernel-check pins BOTH configs (f32 tight band forced =0; mma
21593        // band 8e-2/8e-1 vs f64 truth). =0 reverts; portable stays f32. NOT read via
21594        // OnceLock ON PURPOSE: kernel-check toggles the env per call to pin both forms.
21595        // default mirrors gdn_mma_enabled (incl. the sm_120a-build ON) — all three read
21596        // sites must agree or the pre-work arms while the scan takes the scalar route.
21597        let gdn_mma = !portable_mma_gated()
21598            && c == 32
21599            && match std::env::var("MEMRA_GDN_MMA").as_deref() {
21600                Ok("1") => true,
21601                Ok("0") => false,
21602                _ => gdn_mma_default_on(),
21603            };
21604        if gdn_mma {
21605            let wb16 = wb16_pre
21606                .take()
21607                .expect("mma path pre-allocates wb16 (K3 store fold)");
21608            let kb16_ref: &CudaSlice<u8> = kb16_ref0.expect("mma path pre-builds kb16 above K123");
21609            // K4+K5 FUSED wgmma seam (MEMRA_GDN_WGMMA, task #22; harness verdict
21610            // tools/bench_gdn_wgmma.cu v5, ledger 1f08b997: in-band Y 1.07e-2 / state
21611            // 1.03e-2 / O 1.08e-2, 91.3us vs 70.4 K4-only at H=32 T=512). K5's output
21612            // pass runs inside the persistent-M kernel; Y and Ssnap are never
21613            // materialized. New numeric class (gk folds into k^T instead of ys) —
21614            // explicit opt-in until the state-carry battery promotes it. Env read per
21615            // call (kernel-check pins configs by toggling env, GDN_MMA precedent).
21616            // PROMOTED default-ON hopper (2026-07-27): full battery green — harness
21617            // in-band, argmax gate PASS, 3-seed greedy IDENTICAL after ~2k prime,
21618            // chunked-continuation IDENTICAL, kernel-check + decode-batch gates green,
21619            // official prefill lane +0.74% interleaved x5 (5/5 rounds). =0 reverts.
21620            if gdn_wgmma_pre {
21621                // qb16/pb16 pre-built above K123 (K2-wgmma wrote the masked Pb16).
21622                let qb16 = qb16_ref0.unwrap();
21623                let pb16 = pb16.as_ref().unwrap();
21624                {
21625                    let f = self.func("gdn_k45_wgmma");
21626                    let cfg = LaunchConfig {
21627                        grid_dim: (h as u32, 4, 1),
21628                        block_dim: (256, 1, 1),
21629                        shared_mem_bytes: 0,
21630                    };
21631                    let hki = hk as i32;
21632                    let __s_b = self.gpu.stream();
21633                    let mut b = __s_b.launch_builder(&f);
21634                    b.arg(kb16_ref)
21635                        .arg(&gcum)
21636                        .arg(beta)
21637                        .arg(&u)
21638                        .arg(&wb16)
21639                        .arg(qb16)
21640                        .arg(pb16)
21641                        .arg(o)
21642                        .arg(&scale)
21643                        .arg(state_in)
21644                        .arg(&mut *state_out)
21645                        .arg(&hi)
21646                        .arg(&ti)
21647                        .arg(&ci)
21648                        .arg(&hki);
21649                    unsafe {
21650                        b.launch(cfg)?;
21651                    }
21652                }
21653                return Ok(());
21654            }
21655            // COUPLED PAIR: K4-mma writes Y and Ssnap as bf16 (their only consumer is
21656            // K5-mma, which rounds to bf16 regardless — identical numerics, half the
21657            // traffic; harness K5 63.0 -> 35.3us). Fresh bf16 buffers replace the f32 ones.
21658            let mut y16 = self.alloc_u8_uninit(nc * h * c * D * 2)?;
21659            let mut ssnap16 = self.alloc_u8_uninit(nc * h * D * D * 2)?;
21660            {
21661                let f = self.func("gdn_chunk_state_mma");
21662                let cfg = LaunchConfig {
21663                    grid_dim: (h as u32, NSPLIT, 1),
21664                    block_dim: (256, 1, 1),
21665                    shared_mem_bytes: 0,
21666                };
21667                let hki = hk as i32;
21668                let __s_b = self.gpu.stream();
21669                let mut b = __s_b.launch_builder(&f);
21670                b.arg(kb16_ref)
21671                    .arg(&gcum)
21672                    .arg(beta)
21673                    .arg(&u)
21674                    .arg(&wb16)
21675                    .arg(&mut y16)
21676                    .arg(&mut ssnap16)
21677                    .arg(state_in)
21678                    .arg(&mut *state_out)
21679                    .arg(&hi)
21680                    .arg(&ti)
21681                    .arg(&ci)
21682                    .arg(&hki);
21683                unsafe {
21684                    b.launch(cfg)?;
21685                }
21686            }
21687            {
21688                // K5-mma (bf16 St/Y consumers)
21689                let f = self.func("gdn_chunk_output_mma");
21690                let jt = ((c + 31) / 32) as u32;
21691                let cfg = LaunchConfig {
21692                    grid_dim: (nc as u32, h as u32, jt),
21693                    block_dim: (256, 1, 1),
21694                    shared_mem_bytes: 0,
21695                };
21696                let hki = hk as i32;
21697                let __s_b = self.gpu.stream();
21698                let mut b = __s_b.launch_builder(&f);
21699                b.arg(q)
21700                    .arg(&gcum)
21701                    .arg(&p)
21702                    .arg(&y16)
21703                    .arg(&ssnap16)
21704                    .arg(o)
21705                    .arg(&hi)
21706                    .arg(&ti)
21707                    .arg(&ci)
21708                    .arg(&scale)
21709                    .arg(&hki);
21710                unsafe {
21711                    b.launch(cfg)?;
21712                }
21713            }
21714            return Ok(());
21715        }
21716        {
21717            // K4 (sequential over chunks inside; blocks col-partition the state)
21718            let f = self.func("gdn_chunk_state_f32");
21719            let cfg = LaunchConfig {
21720                grid_dim: (h as u32, NSPLIT, 1),
21721                block_dim: (256, 1, 1),
21722                shared_mem_bytes: 0,
21723            };
21724            let __s_b = self.gpu.stream();
21725            let mut b = __s_b.launch_builder(&f);
21726            b.arg(k)
21727                .arg(&gcum)
21728                .arg(beta)
21729                .arg(&u)
21730                .arg(&w)
21731                .arg(&mut y)
21732                .arg(&mut ssnap)
21733                .arg(state_in)
21734                .arg(&mut *state_out)
21735                .arg(&hi)
21736                .arg(&ti)
21737                .arg(&ci);
21738            unsafe {
21739                b.launch(cfg)?;
21740            }
21741        }
21742        {
21743            // K5 (j-blocked: grid.z = 32-row output blocks per chunk; writes o fully)
21744            let f = self.func("gdn_chunk_output_f32");
21745            let jt = ((c + 31) / 32) as u32;
21746            let cfg = LaunchConfig {
21747                grid_dim: (nc as u32, h as u32, jt),
21748                block_dim: (256, 1, 1),
21749                shared_mem_bytes: 0,
21750            };
21751            let __s_b = self.gpu.stream();
21752            let mut b = __s_b.launch_builder(&f);
21753            b.arg(q)
21754                .arg(&gcum)
21755                .arg(&p)
21756                .arg(&y)
21757                .arg(&ssnap)
21758                .arg(o)
21759                .arg(&hi)
21760                .arg(&ti)
21761                .arg(&ci)
21762                .arg(&scale);
21763            unsafe {
21764                b.launch(cfg)?;
21765            }
21766        }
21767        Ok(())
21768    }
21769
21770    /// PREFILL GDN scan dispatch (the A4 seam): chunked WY form when enabled and T is in the
21771    /// batched-prefill regime, else the sequential scan. Callers: hybrid_forward::linear_attn
21772    /// (forward/forward_last) + linear_attn_prime (prime_cache). Decode (T=1) and the spec
21773    /// verify call `gdn_scan_s128` DIRECTLY — the decode==verify dispatch identity is untouched.
21774    ///
21775    /// MEMRA_GDN_DIFF=1: numerical-oracle mode — runs BOTH forms on the same inputs, prints the
21776    /// per-call (== per-layer, in call order) output/state error distribution, and keeps the
21777    /// SEQUENTIAL results so the run stays on the shipped path (stage-1 prototype evidence).
21778    #[allow(clippy::too_many_arguments)]
21779    #[allow(clippy::too_many_arguments)]
21780    pub fn gdn_scan_prefill(
21781        &self,
21782        q: &CudaSlice<f32>,
21783        k: &CudaSlice<f32>,
21784        v: &CudaSlice<f32>,
21785        g: &CudaSlice<f32>,
21786        beta: &CudaSlice<f32>,
21787        kb16_pre: Option<&CudaSlice<u8>>,
21788        qb16_pre: Option<&CudaSlice<u8>>,
21789        state_in: &CudaSlice<f32>,
21790        state_out: &mut CudaSlice<f32>,
21791        o: &mut CudaSlice<f32>,
21792        n_head: usize,
21793        t: usize,
21794        scale: f32,
21795        hk: usize,
21796    ) -> Result<(), Box<dyn std::error::Error>> {
21797        if std::env::var("MEMRA_GDN_DIFF").is_ok() && t >= 16 {
21798            assert!(hk == n_head, "GDN_DIFF oracle is broadcast-only");
21799            return self.gdn_scan_diff(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale);
21800        }
21801        if Self::gdn_chunked_enabled() && t >= 16 {
21802            self.gdn_scan_chunked(
21803                q,
21804                k,
21805                v,
21806                g,
21807                beta,
21808                kb16_pre,
21809                qb16_pre,
21810                state_in,
21811                state_out,
21812                o,
21813                n_head,
21814                t,
21815                scale,
21816                Self::gdn_chunk_size(),
21817                hk,
21818            )
21819        } else {
21820            assert!(
21821                hk == n_head,
21822                "s128 scan is broadcast-only (prep guarantees by predicate)"
21823            );
21824            self.gdn_scan_s128(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale)
21825        }
21826    }
21827
21828    /// Stage-1 oracle: run sequential AND chunked, report per-call error stats, keep sequential.
21829    #[allow(clippy::too_many_arguments)]
21830    fn gdn_scan_diff(
21831        &self,
21832        q: &CudaSlice<f32>,
21833        k: &CudaSlice<f32>,
21834        v: &CudaSlice<f32>,
21835        g: &CudaSlice<f32>,
21836        beta: &CudaSlice<f32>,
21837        state_in: &CudaSlice<f32>,
21838        state_out: &mut CudaSlice<f32>,
21839        o: &mut CudaSlice<f32>,
21840        n_head: usize,
21841        t: usize,
21842        scale: f32,
21843    ) -> Result<(), Box<dyn std::error::Error>> {
21844        static CALL: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
21845        let call = CALL.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
21846        let mut o_c = self.uninit(o.len())?;
21847        let mut st_c = self.uninit(state_out.len())?;
21848        self.gdn_scan_chunked(
21849            q,
21850            k,
21851            v,
21852            g,
21853            beta,
21854            None,
21855            None,
21856            state_in,
21857            &mut st_c,
21858            &mut o_c,
21859            n_head,
21860            t,
21861            scale,
21862            Self::gdn_chunk_size(),
21863            n_head,
21864        )?;
21865        self.gdn_scan_s128(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale)?;
21866        let (oh_s, oh_c) = (self.dtoh(o)?, self.dtoh(&o_c)?);
21867        let (sh_s, sh_c) = (self.dtoh(state_out)?, self.dtoh(&st_c)?);
21868        let stats = |a: &[f32], b: &[f32]| -> (f32, f32, f64) {
21869            let mut max_abs = 0f32;
21870            let mut max_rel = 0f32;
21871            let mut sum_rel = 0f64;
21872            for (x, y) in a.iter().zip(b) {
21873                let ad = (x - y).abs();
21874                let rel = ad / x.abs().max(y.abs()).max(1e-3);
21875                if ad > max_abs {
21876                    max_abs = ad;
21877                }
21878                if rel > max_rel {
21879                    max_rel = rel;
21880                }
21881                sum_rel += rel as f64;
21882            }
21883            (max_abs, max_rel, sum_rel / a.len() as f64)
21884        };
21885        let (o_ma, o_mr, o_mean) = stats(&oh_s, &oh_c);
21886        let (s_ma, s_mr, s_mean) = stats(&sh_s, &sh_c);
21887        println!(
21888            "[gdn-diff call {call:3} T={t} C={}] out: max_abs={o_ma:.3e} max_rel={o_mr:.3e} mean_rel={o_mean:.3e} | \
21889                  state: max_abs={s_ma:.3e} max_rel={s_mr:.3e} mean_rel={s_mean:.3e}",
21890            Self::gdn_chunk_size()
21891        );
21892        Ok(())
21893    }
21894
21895    /// softplus-based g_log: g_log[h,t] = a[h] * softplus(alpha[h,t] + dt_bias[h]). a pre-negated.
21896    pub fn gdn_glog(
21897        &self,
21898        alpha: &CudaSlice<f32>,
21899        dt_bias: &CudaSlice<f32>,
21900        a: &CudaSlice<f32>,
21901        g_log: &mut CudaSlice<f32>,
21902        n_head: usize,
21903        t: usize,
21904    ) -> Result<(), Box<dyn std::error::Error>> {
21905        let f = self.func("gdn_glog_f32");
21906        let cfg = LaunchConfig::for_num_elems((n_head * t) as u32);
21907        let (h, ti) = (n_head as i32, t as i32);
21908        let __s_b = self.gpu.stream();
21909        let mut b = __s_b.launch_builder(&f);
21910        b.arg(alpha).arg(dt_bias).arg(a).arg(g_log).arg(&h).arg(&ti);
21911        unsafe {
21912            b.launch(cfg)?;
21913        }
21914        Ok(())
21915    }
21916
21917    /// view twins (task #16): the batched prime's GDN core reads the CONCAT projection
21918    /// buffers at row offsets (CudaView) — same kernels, same values, no split copies.
21919    pub fn sigmoid_v(
21920        &self,
21921        x: &cudarc::driver::CudaView<f32>,
21922        y: &mut CudaSlice<f32>,
21923        n: usize,
21924    ) -> Result<(), Box<dyn std::error::Error>> {
21925        let f = self.func("sigmoid_f32");
21926        let cfg = LaunchConfig::for_num_elems(n as u32);
21927        let ni = n as i32;
21928        let __s_b = self.gpu.stream();
21929        let mut b = __s_b.launch_builder(&f);
21930        b.arg(x).arg(y).arg(&ni);
21931        unsafe {
21932            b.launch(cfg)?;
21933        }
21934        Ok(())
21935    }
21936
21937    pub fn gdn_glog_v(
21938        &self,
21939        alpha: &cudarc::driver::CudaView<f32>,
21940        dt_bias: &CudaSlice<f32>,
21941        a: &CudaSlice<f32>,
21942        g_log: &mut CudaSlice<f32>,
21943        n_head: usize,
21944        t: usize,
21945    ) -> Result<(), Box<dyn std::error::Error>> {
21946        let f = self.func("gdn_glog_f32");
21947        let cfg = LaunchConfig::for_num_elems((n_head * t) as u32);
21948        let (h, ti) = (n_head as i32, t as i32);
21949        let __s_b = self.gpu.stream();
21950        let mut b = __s_b.launch_builder(&f);
21951        b.arg(alpha).arg(dt_bias).arg(a).arg(g_log).arg(&h).arg(&ti);
21952        unsafe {
21953            b.launch(cfg)?;
21954        }
21955        Ok(())
21956    }
21957
21958    pub fn sigmoid(
21959        &self,
21960        x: &CudaSlice<f32>,
21961        y: &mut CudaSlice<f32>,
21962        n: usize,
21963    ) -> Result<(), Box<dyn std::error::Error>> {
21964        let f = self.func("sigmoid_f32");
21965        let cfg = LaunchConfig::for_num_elems(n as u32);
21966        let ni = n as i32;
21967        let __s_b = self.gpu.stream();
21968        let mut b = __s_b.launch_builder(&f);
21969        b.arg(x).arg(y).arg(&ni);
21970        unsafe {
21971            b.launch(cfg)?;
21972        }
21973        Ok(())
21974    }
21975
21976    /// attn out-gate fused epilogue (task #17): dst = a * sigmoid(g) + fp16 twin, one launch
21977    /// (replaces sigmoid + mul + convert). Bit-identical class.
21978    pub fn sig_mul_f16out(
21979        &self,
21980        a: &CudaSlice<f32>,
21981        g: &CudaSlice<f32>,
21982        dst: &mut CudaSlice<f32>,
21983        dst16: &mut CudaSlice<u8>,
21984        n: usize,
21985    ) -> Result<(), Box<dyn std::error::Error>> {
21986        let f = self.func("sig_mul_f16out_f32");
21987        let cfg = LaunchConfig::for_num_elems(n as u32);
21988        let ni = n as i32;
21989        let __s_b = self.gpu.stream();
21990        let mut b = __s_b.launch_builder(&f);
21991        b.arg(a).arg(g).arg(dst).arg(dst16).arg(&ni);
21992        unsafe {
21993            b.launch(cfg)?;
21994        }
21995        Ok(())
21996    }
21997
21998    /// step35 (Step-3.7-Flash) SEPARATE head-wise attention gate: one scalar per query head,
21999    /// broadcast over head_dim. `dst = a * sigmoid(g)` where `a`/`dst` are `[head_dim, n_head, T]`
22000    /// (the `q_gate_split` layout) and `g` is the PRE-sigmoid `attn_gate` projection output in
22001    /// token-major `[T, n_head]`. `dst16` is the optional fp16 operand for wo (None -> skipped).
22002    ///
22003    /// NOT interchangeable with `sig_mul_f16out`, which gates FULL WIDTH (qwen35 packs one gate
22004    /// value per (head, dim) element inside wq). Using this for that, or that for this, silently
22005    /// applies the wrong number of distinct gate values.
22006    #[allow(clippy::too_many_arguments)]
22007    pub fn attn_head_gate(
22008        &self,
22009        a: &CudaSlice<f32>,
22010        g: &CudaSlice<f32>,
22011        dst: &mut CudaSlice<f32>,
22012        dst16: Option<&mut CudaSlice<u8>>,
22013        head_dim: usize,
22014        n_head: usize,
22015        t: usize,
22016    ) -> Result<(), Box<dyn std::error::Error>> {
22017        let f = self.func("attn_head_gate_f32");
22018        let cfg = LaunchConfig::for_num_elems((head_dim * n_head * t) as u32);
22019        let (hd, nh, ti) = (head_dim as i32, n_head as i32, t as i32);
22020        // nullable device pointer by value (0 = skip), same convention as `l2_norm_pp`.
22021        let d16: u64 = match dst16 {
22022            Some(d) => self.addr_u8(d),
22023            None => 0,
22024        };
22025        let __s_b = self.gpu.stream();
22026        let mut b = __s_b.launch_builder(&f);
22027        b.arg(a)
22028            .arg(g)
22029            .arg(dst)
22030            .arg(&d16)
22031            .arg(&hd)
22032            .arg(&nh)
22033            .arg(&ti);
22034        unsafe {
22035            b.launch(cfg)?;
22036        }
22037        Ok(())
22038    }
22039
22040    /// step35 CLAMPED SwiGLU: `dst = min(silu(gate*gs), limit) * clamp(up*us, +-limit)`.
22041    /// Verbatim from llama.cpp `llama-graph.cpp:2146-2165` (routed, `swiglu_clamp_exp`) and
22042    /// `:1751-1770` (shared, `swiglu_clamp_shexp`), non-DEEPSEEK4 branch.
22043    ///
22044    /// This is NOT `swigluoai_mul_scaled`: that one clamps the gate BEFORE swish and multiplies by
22045    /// `(1 + clamp(up))`. Caller MUST check `limit > 1e-6` (upstream's eps gate) and use the plain
22046    /// `silu_mul_scaled` path otherwise — at limit=0 this kernel would clamp every positive
22047    /// activation to zero. On Step-3.7-Flash only layers 43 (7.0) and 44 (16.0) have a live limit.
22048    #[allow(clippy::too_many_arguments)]
22049    pub fn swiglu_clamped_mul_scaled(
22050        &self,
22051        gate: &CudaSlice<f32>,
22052        up: &CudaSlice<f32>,
22053        gs: f32,
22054        us: f32,
22055        limit: f32,
22056        dst: &mut CudaSlice<f32>,
22057        n: usize,
22058    ) -> Result<(), Box<dyn std::error::Error>> {
22059        debug_assert!(
22060            limit > 1e-6,
22061            "swiglu_clamped needs a live limit; use silu_mul_scaled"
22062        );
22063        let f = self.func("swiglu_clamped_mul_scaled_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(gate)
22069            .arg(up)
22070            .arg(&gs)
22071            .arg(&us)
22072            .arg(&limit)
22073            .arg(dst)
22074            .arg(&ni);
22075        unsafe {
22076            b.launch(cfg)?;
22077        }
22078        Ok(())
22079    }
22080
22081    /// gated RMSNorm: dst = RMSNorm(o, w[ncols]) * silu(z), per row of ncols. nrows blocks.
22082    pub fn gated_rmsnorm(
22083        &self,
22084        o: &CudaSlice<f32>,
22085        w: &CudaSlice<f32>,
22086        z: &CudaSlice<f32>,
22087        dst: &mut CudaSlice<f32>,
22088        ncols: usize,
22089        nrows: usize,
22090        eps: f32,
22091    ) -> Result<(), Box<dyn std::error::Error>> {
22092        let f = self.func("gated_rmsnorm_f32");
22093        let cfg = LaunchConfig {
22094            grid_dim: (nrows as u32, 1, 1),
22095            block_dim: (128, 1, 1),
22096            shared_mem_bytes: 0,
22097        };
22098        let (nc, e) = (ncols as i32, eps);
22099        let __s_b = self.gpu.stream();
22100        let mut b = __s_b.launch_builder(&f);
22101        b.arg(o).arg(w).arg(z).arg(dst).arg(&nc).arg(&e);
22102        unsafe {
22103            b.launch(cfg)?;
22104        }
22105        Ok(())
22106    }
22107
22108    /// f16out twin of `gated_rmsnorm` (task #17): epilogue also emits the fp16 operand for
22109    /// the ssm_out GEMM. Bit-identical class (same floats + the cvt kernel's __float2half).
22110    pub fn gated_rmsnorm_f16out(
22111        &self,
22112        o: &CudaSlice<f32>,
22113        w: &CudaSlice<f32>,
22114        z: &CudaSlice<f32>,
22115        dst: &mut CudaSlice<f32>,
22116        dst16: &mut CudaSlice<u8>,
22117        ncols: usize,
22118        nrows: usize,
22119        eps: f32,
22120    ) -> Result<(), Box<dyn std::error::Error>> {
22121        let f = self.func("gated_rmsnorm_f16out_f32");
22122        // block_dim MUST match gated_rmsnorm's (128): the reduction tree order pins the scale
22123        let cfg = LaunchConfig {
22124            grid_dim: (nrows as u32, 1, 1),
22125            block_dim: (128, 1, 1),
22126            shared_mem_bytes: 0,
22127        };
22128        let (nc, e) = (ncols as i32, eps);
22129        let __s_b = self.gpu.stream();
22130        let mut b = __s_b.launch_builder(&f);
22131        b.arg(o).arg(w).arg(z).arg(dst).arg(dst16).arg(&nc).arg(&e);
22132        unsafe {
22133            b.launch(cfg)?;
22134        }
22135        Ok(())
22136    }
22137
22138    /// add+RMSNorm emitting the f32 normed row AND its q8_1 quantization in one launch (the MoE
22139    /// layer input: z feeds the router matmul as f32, the expert dp4a as q8_1). BIT-IDENTICAL to
22140    /// add_rms_norm + quantize_q8_1. Returns (q, d) alongside the caller-provided res/z buffers.
22141    #[allow(clippy::too_many_arguments)]
22142    pub fn add_rms_norm_zq8(
22143        &self,
22144        a: &CudaSlice<f32>,
22145        b_in: &CudaSlice<f32>,
22146        w: &CudaSlice<f32>,
22147        res: &mut CudaSlice<f32>,
22148        z: &mut CudaSlice<f32>,
22149        ncols: usize,
22150        nrows: usize,
22151        eps: f32,
22152    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
22153        assert!(ncols % 32 == 0);
22154        let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
22155        let mut d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
22156        let f = self.func("add_rms_norm_zq8");
22157        let cfg = LaunchConfig {
22158            grid_dim: (nrows as u32, 1, 1),
22159            block_dim: (1024, 1, 1),
22160            shared_mem_bytes: 0,
22161        };
22162        let (nc, ep) = (ncols as i32, eps);
22163        let __s_b = self.gpu.stream();
22164        let mut b = __s_b.launch_builder(&f);
22165        b.arg(a)
22166            .arg(b_in)
22167            .arg(w)
22168            .arg(res)
22169            .arg(z)
22170            .arg(&mut q)
22171            .arg(&mut d)
22172            .arg(&nc)
22173            .arg(&ep);
22174        unsafe {
22175            b.launch(cfg)?;
22176        }
22177        Ok((q, d))
22178    }
22179
22180    /// gated RMSNorm emitting q8_1 directly (fused quantize epilogue) — the ssm_out matvec input.
22181    /// BIT-IDENTICAL bytes to gated_rmsnorm + quantize_q8_1 (ncols % 32 == 0; blocks never straddle
22182    /// rows). Saves one launch per linear-attn layer (36/token on the 9B).
22183    /// z-view twins of gated_rmsnorm(+f16out) — task #16 batched-prime split removal.
22184    pub fn gated_rmsnorm_zv(
22185        &self,
22186        o: &CudaSlice<f32>,
22187        w: &CudaSlice<f32>,
22188        z: &cudarc::driver::CudaView<f32>,
22189        dst: &mut CudaSlice<f32>,
22190        ncols: usize,
22191        nrows: usize,
22192        eps: f32,
22193    ) -> Result<(), Box<dyn std::error::Error>> {
22194        let f = self.func("gated_rmsnorm_f32");
22195        let cfg = LaunchConfig {
22196            grid_dim: (nrows as u32, 1, 1),
22197            block_dim: (128, 1, 1),
22198            shared_mem_bytes: 0,
22199        };
22200        let (nc, e) = (ncols as i32, eps);
22201        let __s_b = self.gpu.stream();
22202        let mut b = __s_b.launch_builder(&f);
22203        b.arg(o).arg(w).arg(z).arg(dst).arg(&nc).arg(&e);
22204        unsafe {
22205            b.launch(cfg)?;
22206        }
22207        Ok(())
22208    }
22209
22210    pub fn gated_rmsnorm_f16out_zv(
22211        &self,
22212        o: &CudaSlice<f32>,
22213        w: &CudaSlice<f32>,
22214        z: &cudarc::driver::CudaView<f32>,
22215        dst: &mut CudaSlice<f32>,
22216        dst16: &mut CudaSlice<u8>,
22217        ncols: usize,
22218        nrows: usize,
22219        eps: f32,
22220    ) -> Result<(), Box<dyn std::error::Error>> {
22221        let f = self.func("gated_rmsnorm_f16out_f32");
22222        // block_dim MUST match gated_rmsnorm's (128): the reduction tree order pins the scale
22223        let cfg = LaunchConfig {
22224            grid_dim: (nrows as u32, 1, 1),
22225            block_dim: (128, 1, 1),
22226            shared_mem_bytes: 0,
22227        };
22228        let (nc, e) = (ncols as i32, eps);
22229        let __s_b = self.gpu.stream();
22230        let mut b = __s_b.launch_builder(&f);
22231        b.arg(o).arg(w).arg(z).arg(dst).arg(dst16).arg(&nc).arg(&e);
22232        unsafe {
22233            b.launch(cfg)?;
22234        }
22235        Ok(())
22236    }
22237
22238    pub fn gated_rmsnorm_q8_1(
22239        &self,
22240        o: &CudaSlice<f32>,
22241        w: &CudaSlice<f32>,
22242        z: &CudaSlice<f32>,
22243        ncols: usize,
22244        nrows: usize,
22245        eps: f32,
22246    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
22247        assert!(ncols % 32 == 0);
22248        let f = self.func("gated_rmsnorm_q8_1");
22249        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
22250        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
22251        let cfg = LaunchConfig {
22252            grid_dim: (nrows as u32, 1, 1),
22253            block_dim: (128, 1, 1),
22254            shared_mem_bytes: 0,
22255        };
22256        let (nc, ep) = (ncols as i32, eps);
22257        let __s_b = self.gpu.stream();
22258        let mut b = __s_b.launch_builder(&f);
22259        b.arg(o)
22260            .arg(w)
22261            .arg(z)
22262            .arg(&mut out_q)
22263            .arg(&mut out_d)
22264            .arg(&nc)
22265            .arg(&ep);
22266        unsafe {
22267            b.launch(cfg)?;
22268        }
22269        Ok((out_q, out_d))
22270    }
22271
22272    /// transpose [rows,cols] row-major -> [cols,rows] row-major.
22273    pub fn transpose(
22274        &self,
22275        inp: &CudaSlice<f32>,
22276        rows: usize,
22277        cols: usize,
22278    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22279        let f = self.func("transpose_f32");
22280        let mut out = self.zeros(rows * cols)?;
22281        let cfg = LaunchConfig::for_num_elems((rows * cols) as u32);
22282        let (r, c) = (rows as i32, cols as i32);
22283        let __s_b = self.gpu.stream();
22284        let mut b = __s_b.launch_builder(&f);
22285        b.arg(inp).arg(&mut out).arg(&r).arg(&c);
22286        unsafe {
22287            b.launch(cfg)?;
22288        }
22289        Ok(out)
22290    }
22291
22292    /// repeat-interleave heads: in[head_dim,n_in,T] -> out[head_dim,n_out,T].
22293    pub fn repeat_heads(
22294        &self,
22295        inp: &CudaSlice<f32>,
22296        out: &mut CudaSlice<f32>,
22297        head_dim: usize,
22298        n_in: usize,
22299        n_out: usize,
22300        t: usize,
22301    ) -> Result<(), Box<dyn std::error::Error>> {
22302        let f = self.func("repeat_heads_f32");
22303        let cfg = LaunchConfig::for_num_elems((head_dim * n_out * t) as u32);
22304        let (hd, ni, no, ti) = (head_dim as i32, n_in as i32, n_out as i32, t as i32);
22305        let __s_b = self.gpu.stream();
22306        let mut b = __s_b.launch_builder(&f);
22307        b.arg(inp).arg(out).arg(&hd).arg(&ni).arg(&no).arg(&ti);
22308        unsafe {
22309            b.launch(cfg)?;
22310        }
22311        Ok(())
22312    }
22313
22314    /// q|gate split (on-device). qf:[T, n_head*2*head_dim] -> q_out,gate_out:[head_dim,n_head,T].
22315    /// Replaces the dtoh->host-double-loop->htod in full_attn / full_attn_decode.
22316    ///
22317    /// Bounds-checked against `qf` before the launch: the kernel reads `2*head_dim*n_head*t`
22318    /// floats, and running it on a `wq` output that carries no fused gate reads 2x off the end
22319    /// (silently, on the device). A layout mismatch is a typed `FusedQGateExtent` here instead.
22320    pub fn q_gate_split(
22321        &self,
22322        qf: &CudaSlice<f32>,
22323        q_out: &mut CudaSlice<f32>,
22324        gate_out: &mut CudaSlice<f32>,
22325        head_dim: usize,
22326        n_head: usize,
22327        t: usize,
22328    ) -> Result<(), Box<dyn std::error::Error>> {
22329        memra_gguf::config::check_fused_q_gate_extent(qf.len(), head_dim, n_head, t)?;
22330        let out_need = head_dim * n_head * t;
22331        if q_out.len() < out_need || gate_out.len() < out_need {
22332            return Err(format!(
22333                "q_gate_split destinations too small: need {out_need} each, have q={} gate={}",
22334                q_out.len(),
22335                gate_out.len()
22336            )
22337            .into());
22338        }
22339        let f = self.func("q_gate_split_f32");
22340        let cfg = LaunchConfig::for_num_elems((head_dim * n_head * t) as u32);
22341        let (hd, nh, ti) = (head_dim as i32, n_head as i32, t as i32);
22342        let __s_b = self.gpu.stream();
22343        let mut b = __s_b.launch_builder(&f);
22344        b.arg(qf)
22345            .arg(q_out)
22346            .arg(gate_out)
22347            .arg(&hd)
22348            .arg(&nh)
22349            .arg(&ti);
22350        unsafe {
22351            b.launch(cfg)?;
22352        }
22353        Ok(())
22354    }
22355
22356    /// qkv->GDN repack (on-device). conv_out:[conv_dim,T] channel-major ->
22357    /// q_g/k_g/v_g:[d_state,num_v,T] with q/k head-repeat kh = vh % num_k (validated modulo mapping).
22358    /// Replaces the dtoh->host-q/k/v-repack->3x-htod in linear_attn / linear_attn_decode.
22359    pub fn qkv_to_gdn_repack(
22360        &self,
22361        conv_out: &CudaSlice<f32>,
22362        q_g: &mut CudaSlice<f32>,
22363        k_g: &mut CudaSlice<f32>,
22364        v_g: &mut CudaSlice<f32>,
22365        d_state: usize,
22366        num_v: usize,
22367        num_k: usize,
22368        key_dim: usize,
22369        t: usize,
22370    ) -> Result<(), Box<dyn std::error::Error>> {
22371        let f = self.func("qkv_to_gdn_repack_f32");
22372        let cfg = LaunchConfig::for_num_elems((d_state * num_v * t) as u32);
22373        let (ds, nv, nk, kd, ti) = (
22374            d_state as i32,
22375            num_v as i32,
22376            num_k as i32,
22377            key_dim as i32,
22378            t as i32,
22379        );
22380        let __s_b = self.gpu.stream();
22381        let mut b = __s_b.launch_builder(&f);
22382        b.arg(conv_out)
22383            .arg(q_g)
22384            .arg(k_g)
22385            .arg(v_g)
22386            .arg(&ds)
22387            .arg(&nv)
22388            .arg(&nk)
22389            .arg(&kd)
22390            .arg(&ti);
22391        unsafe {
22392            b.launch(cfg)?;
22393        }
22394        Ok(())
22395    }
22396
22397    /// conv left zero-pad (prefill from zero state). src:[conv_dim,T] -> dst:[conv_dim,T+pad],
22398    /// cols 0..pad = 0, cols pad..pad+T = src. `dst` MUST be pre-zeroed. No dtoh/host-loop/htod.
22399    pub fn conv_left_pad(
22400        &self,
22401        src: &CudaSlice<f32>,
22402        dst: &mut CudaSlice<f32>,
22403        conv_dim: usize,
22404        t: usize,
22405        pad: usize,
22406    ) -> Result<(), Box<dyn std::error::Error>> {
22407        let f = self.func("conv_left_pad_f32");
22408        let cfg = LaunchConfig::for_num_elems((conv_dim * t) as u32);
22409        let (cd, ti, p) = (conv_dim as i32, t as i32, pad as i32);
22410        let __s_b = self.gpu.stream();
22411        let mut b = __s_b.launch_builder(&f);
22412        b.arg(src).arg(dst).arg(&cd).arg(&ti).arg(&p);
22413        unsafe {
22414            b.launch(cfg)?;
22415        }
22416        Ok(())
22417    }
22418
22419    /// conv-state assemble + ring roll (decode T=1). conv_state:[conv_dim,pad] (resident),
22420    /// qkv_col:[conv_dim] -> conv_in:[conv_dim,pad+1]; AND rolls conv_state (keep last pad cols).
22421    /// Replaces the dtoh->host-conv-ring-assemble->ring-update->htod in linear_attn_decode.
22422    pub fn conv_assemble_and_roll(
22423        &self,
22424        qkv_col: &CudaSlice<f32>,
22425        conv_state: &mut CudaSlice<f32>,
22426        conv_in: &mut CudaSlice<f32>,
22427        conv_dim: usize,
22428        pad: usize,
22429    ) -> Result<(), Box<dyn std::error::Error>> {
22430        let f = self.func("conv_assemble_and_roll_f32");
22431        let cfg = LaunchConfig::for_num_elems(conv_dim as u32);
22432        let (cd, p) = (conv_dim as i32, pad as i32);
22433        let __s_b = self.gpu.stream();
22434        let mut b = __s_b.launch_builder(&f);
22435        b.arg(qkv_col).arg(conv_state).arg(conv_in).arg(&cd).arg(&p);
22436        unsafe {
22437            b.launch(cfg)?;
22438        }
22439        Ok(())
22440    }
22441
22442    /// RANK3 LEVER (conv fuse, T=1 DECODE): fused conv_assemble_and_roll + ssm_conv1d_silu in ONE
22443    /// launch. Assembles the conv window [conv_state | qkv_col] in registers, computes the depthwise
22444    /// causal conv + SiLU into `conv_out`, and rolls the ring — never materializing conv_in to HBM.
22445    /// Replaces e.conv_assemble_and_roll(...) + e.ssm_conv1d(...). BIT-IDENTICAL to that two-kernel
22446    /// sequence (same 8-wide accumulation order, same SiLU). `conv_out` is [conv_dim] (T=1).
22447    pub fn ssm_conv1d_fused_decode(
22448        &self,
22449        qkv_col: &CudaSlice<f32>,
22450        conv_state: &mut CudaSlice<f32>,
22451        w: &CudaSlice<f32>,
22452        conv_out: &mut CudaSlice<f32>,
22453        conv_dim: usize,
22454        d_conv: usize,
22455    ) -> Result<(), Box<dyn std::error::Error>> {
22456        let f = self.func("ssm_conv1d_fused_decode_f32");
22457        let cfg = LaunchConfig::for_num_elems(conv_dim as u32);
22458        let (cd, dc) = (conv_dim as i32, d_conv as i32);
22459        let __s_b = self.gpu.stream();
22460        let mut b = __s_b.launch_builder(&f);
22461        b.arg(qkv_col)
22462            .arg(conv_state)
22463            .arg(w)
22464            .arg(conv_out)
22465            .arg(&cd)
22466            .arg(&dc);
22467        unsafe {
22468            b.launch(cfg)?;
22469        }
22470        Ok(())
22471    }
22472
22473    /// Copy a contiguous range [start, start+len) out of src into a fresh slice (device→device via host).
22474    /// Used for qkv split views. Small/rare; not perf-critical in Stage 1.
22475    pub fn slice_range(
22476        &self,
22477        src: &CudaSlice<f32>,
22478        start: usize,
22479        len: usize,
22480    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22481        let host = self.gpu.stream().clone_dtoh(src)?;
22482        self.gpu.stream().synchronize()?;
22483        Ok(self.htod(&host[start..start + len])?)
22484    }
22485}
22486
22487#[cfg(test)]
22488mod target_dispatch_tests {
22489    use super::legacy_quant_gemm_allowed;
22490
22491    #[test]
22492    fn legacy_quant_gemm_arch_policy_honors_the_escape_hatch() {
22493        // sm_120a native lane
22494        assert!(legacy_quant_gemm_allowed(false, false, false));
22495        assert!(!legacy_quant_gemm_allowed(false, false, true));
22496        // pure portable lane (sm_89): gated
22497        assert!(!legacy_quant_gemm_allowed(true, false, false));
22498        assert!(!legacy_quant_gemm_allowed(true, false, true));
22499        // Hopper-MMA lane (sm_90a): portable build, int8-MMA GEMM re-admitted
22500        assert!(legacy_quant_gemm_allowed(true, true, false));
22501        assert!(!legacy_quant_gemm_allowed(true, true, true));
22502    }
22503
22504    #[cfg(all(memra_portable_cuda, not(memra_hopper_mma)))]
22505    #[test]
22506    fn portable_build_disables_legacy_quant_gemm_without_an_env_override() {
22507        assert!(!legacy_quant_gemm_allowed(
22508            cfg!(memra_portable_cuda),
22509            cfg!(memra_hopper_mma),
22510            false
22511        ));
22512    }
22513
22514    #[cfg(memra_hopper_mma)]
22515    #[test]
22516    fn hopper_mma_build_re_admits_legacy_quant_gemm() {
22517        assert!(legacy_quant_gemm_allowed(
22518            cfg!(memra_portable_cuda),
22519            cfg!(memra_hopper_mma),
22520            false
22521        ));
22522        assert!(super::portable_mma_gated() == false);
22523    }
22524}
22525
22526/// The memra-kv device seam (Phase D): the cache's 7 ops delegate to the engine's
22527/// inherent methods (inherent methods win name resolution, so no recursion).
22528impl memra_kv::KvDev for Engine {
22529    fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22530        Engine::zeros(self, n)
22531    }
22532    fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22533        Engine::uninit(self, n)
22534    }
22535    fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
22536        Engine::alloc_u8(self, n)
22537    }
22538    fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
22539        Engine::htod_i32(self, v)
22540    }
22541    fn clone_dtod(
22542        &self,
22543        src: &CudaSlice<f32>,
22544    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22545        Engine::clone_dtod(self, src)
22546    }
22547    fn copy_into(
22548        &self,
22549        dst: &mut CudaSlice<f32>,
22550        off: usize,
22551        src: &CudaSlice<f32>,
22552        len: usize,
22553    ) -> Result<(), Box<dyn std::error::Error>> {
22554        Engine::copy_into(self, dst, off, src, len)
22555    }
22556    fn set_i32_one(
22557        &self,
22558        d: &mut CudaSlice<i32>,
22559        v: i32,
22560    ) -> Result<(), Box<dyn std::error::Error>> {
22561        Engine::set_i32_one(self, d, v)
22562    }
22563}
22564
22565#[cfg(test)]
22566mod fused_gate_bounds_tests {
22567    use super::*;
22568
22569    /// The fused `[q|gate]` split's read-site guard, on the device.
22570    ///
22571    /// `q_gate_split_f32` reads `2*head_dim*n_head*T` floats out of `qf`. A checkpoint whose gate
22572    /// is a SEPARATE tensor produces a `wq` output of exactly half that, so before 2026-08-19 the
22573    /// kernel launched and read 2x past the end of the allocation — an out-of-bounds DEVICE read:
22574    /// no panic, no error, just whatever memory follows. The guard turns it into a typed
22575    /// `FusedQGateExtent` before the launch.
22576    ///
22577    /// Catch demonstration for this test (guard temporarily removed, then restored):
22578    /// `compute-sanitizer --tool memcheck` on the half-width case reported invalid `__global__`
22579    /// reads of size 4 in `q_gate_split_f32`; with the guard in place the same run is clean and
22580    /// the call returns `Err`. Receipt in the lane report.
22581    #[test]
22582    #[ignore = "requires a CUDA GPU"]
22583    fn q_gate_split_refuses_a_separate_gate_wq_instead_of_reading_past_it() {
22584        let e = Engine::new(0).unwrap();
22585        let (head_dim, n_head, t) = (8usize, 4usize, 2usize);
22586        let fused = 2 * head_dim * n_head * t;
22587        let out_n = head_dim * n_head * t;
22588
22589        // half-width `qf` = the separate-gate / ungated layout. MUST be refused.
22590        let narrow = e.htod(&vec![1.0f32; out_n]).unwrap();
22591        let mut q = e.uninit(out_n).unwrap();
22592        let mut gate = e.uninit(out_n).unwrap();
22593        let err = e
22594            .q_gate_split(&narrow, &mut q, &mut gate, head_dim, n_head, t)
22595            .expect_err("half-width wq must be refused, not read past")
22596            .to_string();
22597        assert!(err.contains("NO fused gate"), "{err}");
22598        assert!(err.contains(&format!("{fused}")), "{err}");
22599
22600        // full-width `qf` = a real qwen3.5 fused layout. MUST still run, and split correctly:
22601        // per head hh the block is [q(head_dim) | gate(head_dim)] at stride 2*head_dim.
22602        let host: Vec<f32> = (0..fused).map(|i| i as f32).collect();
22603        let wide = e.htod(&host).unwrap();
22604        e.q_gate_split(&wide, &mut q, &mut gate, head_dim, n_head, t)
22605            .expect("full-width wq splits");
22606        let (qh, gh) = (e.dtoh(&q).unwrap(), e.dtoh(&gate).unwrap());
22607        for tok in 0..t {
22608            for hh in 0..n_head {
22609                for d in 0..head_dim {
22610                    let base = tok * (n_head * 2 * head_dim) + hh * (2 * head_dim);
22611                    let idx = tok * (n_head * head_dim) + hh * head_dim + d;
22612                    assert_eq!(qh[idx], host[base + d], "q t{tok} h{hh} d{d}");
22613                    assert_eq!(gh[idx], host[base + head_dim + d], "gate t{tok} h{hh} d{d}");
22614                }
22615            }
22616        }
22617
22618        // undersized destinations are refused too (the other half of the extent contract)
22619        let mut small = e.uninit(out_n - 1).unwrap();
22620        assert!(
22621            e.q_gate_split(&wide, &mut small, &mut gate, head_dim, n_head, t)
22622                .is_err()
22623        );
22624    }
22625}
22626
22627/// FULL-WIDTH-ROPE CONTRACT on the fused rms_norm+qkv+rope kernels
22628/// (lane/graph-s-key-exactness-20260819, probe O-10 follow-up). CPU-only: the guard runs before
22629/// any launch, so the refusal is testable without a device.
22630#[cfg(test)]
22631mod fused_rope_width_tests {
22632    use super::Engine;
22633
22634    /// gemma-4: rotary width == head width on both classes (GGUF 256/256 and 512/512, and the
22635    /// safetensors route derives the same), which is why the fusion is legal there today.
22636    #[test]
22637    fn full_width_is_accepted() {
22638        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 256, 256).is_ok());
22639        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_cat", 512, 512).is_ok());
22640        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append", 128, 128).is_ok());
22641    }
22642
22643    /// The widths the gemma-4 31B OFFICIAL artifact declares, read from its own GGUF header
22644    /// (`gemma-4-31B-it-official-Q8_0-MTP.gguf`, box3, 2026-08-19):
22645    ///
22646    /// ```text
22647    /// attention.key_length     512   rope.dimension_count     512   (global class)
22648    /// attention.key_length_swa 256   rope.dimension_count_swa 256   (SWA class)
22649    /// ```
22650    ///
22651    /// Both classes satisfy `n_rot == head_dim`, which is why the fusion is legal for gemma and
22652    /// why `HybridModel::gemma4_rope_dims` can feed this guard without refusing what we serve.
22653    /// An artifact that ever declares otherwise gets a loud refusal at the first fused launch
22654    /// instead of a silently over-rotated head.
22655    #[test]
22656    fn gemma4_official_artifact_widths_pass() {
22657        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 512, 512).is_ok());
22658        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append_dc", 256, 256).is_ok());
22659    }
22660
22661    /// The bug this guard exists to make impossible: a partial-rotary arch fused onto a kernel
22662    /// with no `n_dims`, silently rotating the pass-through band.
22663    #[test]
22664    fn partial_rotary_is_refused_with_the_geometry_named() {
22665        // qwen3.5: n_rot 64 of head_dim 256 (the shape probe O-10 pinned in the split path).
22666        let err = Engine::full_width_rope_only("rms_norm_qkv_rope", 64, 256)
22667            .expect_err("partial rotary must refuse");
22668        let msg = err.to_string();
22669        assert!(msg.contains("PARTIAL ROTARY REFUSED"), "{msg}");
22670        assert!(msg.contains("n_rot 64"), "{msg}");
22671        assert!(msg.contains("head_dim 256"), "{msg}");
22672        assert!(
22673            msg.contains("64..256"),
22674            "names the band it would corrupt: {msg}"
22675        );
22676        // step35 full-attn: 64 of 128 (upstream halves n_rot_full).
22677        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append_dc", 64, 128).is_err());
22678        // and the reverse mismatch (a wider rope than the head) is not "close enough" either.
22679        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 256, 128).is_err());
22680    }
22681}