Skip to main content

memra_engine/
lib.rs

1//! memra engine: Stage-1 correctness-first forward-pass kernels + ops, on sm_120 via cudarc.
2
3use cudarc::driver::{
4    CudaContext, CudaFunction, CudaModule, CudaSlice, CudaStream, LaunchConfig, PushKernelArg,
5};
6use cudarc::nvrtc::Ptx;
7use std::sync::{Arc, Mutex};
8
9#[cfg(debug_assertions)]
10pub(crate) fn debug_assert_tensor_stream_device<T>(
11    tensor: &CudaSlice<T>,
12    stream: &CudaStream,
13    site: &str,
14) {
15    let tensor_dev = tensor.ordinal();
16    let stream_dev = stream.context().ordinal();
17    assert_eq!(
18        tensor_dev, stream_dev,
19        "PP cross-device tensor read at {site}: tensor on dev{tensor_dev}, stream on dev{stream_dev}"
20    );
21}
22
23pub use memra_gguf;
24pub use memra_runtime;
25
26pub mod forward;
27pub mod hybrid;
28pub mod hybrid_forward;
29pub mod model;
30pub mod sigrouter_contract;
31pub mod vision;
32pub mod vision_gemma;
33pub mod vision_pre;
34/// The dual cache lives in the shared `memra-kv` crate (Phase D extraction); this
35/// re-export keeps every `crate::cache::` / `memra_engine::cache::` path unchanged.
36pub mod cache {
37    pub use memra_kv::*;
38}
39pub mod decode;
40pub mod decode_batch;
41pub mod dflash;
42pub mod eagle;
43pub mod gemma_spec;
44pub mod graph_update;
45/// MLA (multi-head latent attention) CPU f32 reference — GLM-5.2 bring-up lane increment 1.
46/// Naive vs absorbed decode forms + NORM/NEOX rope permutation, unit-tested; the permanent
47/// oracle for the MLA kernel family (`research/mla-bringup-20260801/DESIGN.md`). No CUDA deps.
48pub mod mla;
49pub mod moesd;
50pub mod parallel;
51pub mod pp;
52pub mod round_stream;
53pub mod spec;
54pub use memra_sampling as sampler;
55
56/// In-house MoE router GEMV on the spec-verify small-t path (DEFAULT ON since 2026-07-10:
57/// battery green on 35B p2/p3 K=1..8, acceptance bit-identical, +2-4% spec e2e — replaces
58/// ~240 per-column cuBLAS gemv launches/round). MEMRA_ROUTER_KERNEL=0 is the rollback seam.
59/// MoE grouped f16 GEMM door (experimental until gated), f16-mirror numeric class:
60/// per-layer expert dequant to f16 + one grouped f16 GEMM over the CSR groups.
61///   MEMRA_MOE_F16G=1  cublasGemmGroupedBatchedEx (round 46 arc 2). The grouped API issues
62///                     through cublas-internal streams NOT ordered with ours — v1 pays a full
63///                     stream sync per projection (round-47 ledgered defect).
64///   MEMRA_MOE_F16G=2  single-kernel grouped GEMM on the engine stream (round 49): ordered by
65///                     construction, zero syncs, f32 C with the act row-scale folded in.
66/// DEFAULT (2026-08-01, round 49 promotion): mode 1 on the Hopper lane — with the 41/41
67/// dequant coverage fix the q35 board-2048 prime measured 5490 (MMQ) / 8380 (mode 1,
68/// +53%) / 7990 (mode 2) x3 interleaved on the H100, argmax MATCH — the last board loss
69/// flips. The 5090 measured FLAT (858GB/s makes the dequant-workspace traffic cancel the
70/// GEMM win) — but that verdict is for expert banks the int8-MMA MMQ arm can take
71/// (IQ3_S/IQ4_XS/Q4_0). MEMRA_MOE_F16G=0 kills anywhere.
72///
73/// HOPPER RE-VERDICT (2026-08-02, lane/h100-flip-full): mode 2 with full direct coverage
74/// (Q4_K/Q6_K/IQ4_XS/IQ3_S tile loaders, lane/iq-direct-loaders) + the deep tail
75/// (lane/sk-tail-form) FLIPS past cublas mode 1 on the H100 — q35 board-2048 prime
76/// 13163.6 (mode 2, cross=32) vs 8626.5 (mode 1) vs 8073.4 (round-51 sk form), +52.6%,
77/// interleaved x5 zero overlap, argmax MATCH 30/30. The round-54 NO-FLIP (8547 vs 8112)
78/// was coverage-priced at 5.2% direct; ~100% coverage kills the workspace pass and the
79/// verdict inverts. Hopper naked default -> mode 2 (this arm); the gemma (gelu) site
80/// stays env-explicit-only via moe_f16g_gemma_on (Err => closed, unaffected by this arm).
81///
82/// MODE-2 DEFAULT (sm_120a naked, 2026-08-02, lane/f16g-default-rearb): with the direct
83/// tile loaders covering Q4_K/Q6_K/IQ4_XS/IQ3_S, the sk visitor beats the int8-MMA MMQ
84/// tiles on the IQ-bank models too (q35 board-2048 +33.9%, KAT pp512 +46.7% / pp2048
85/// +30.6% — research/iq-direct-loaders-20260802 §3-5, confirmed + full battery in
86/// research/f16g-default-rearb-20260802/), so every f16g-admitted expert layer rides
87/// mode 2 naked. Decode/verify stay on dp4a (t >= 16 floor). f16-mirror numeric class
88/// for naked q35/KAT prefill+prime — new token-sha anchors stamped in the rearb lane.
89///
90/// AUTO-KQUANT (mode 3, 2026-08-02, lane/q4k-expert-prefill): the previous sm_120a
91/// default, kept reachable via MEMRA_MOE_F16G=3. The mode-2 sk form is admitted ONLY for
92/// layers the MMA MMQ arm rejects (k-quant expert projections — Q3_K/Q4_K/Q6_K), i.e.
93/// exactly where the baseline is the per-pair moe_pairs_matvec_q8_em fallback with zero
94/// token reuse (Ornith-35B Q4_K_M board-2048 1098.2 -> 3453.7, 3.14x,
95/// research/q4k-expert-prefill-20260802/). Its "IQ banks keep their measured-faster MMQ
96/// tiles" ruling was priced BEFORE the IQ direct loaders and is refuted on the 5090 —
97/// the k-quant-only admission survives as the rollback seam, not the default.
98/// The gemma (gelu) site stays env-explicit-only (moe_f16g_gemma_on).
99pub fn moe_f16g_mode() -> u8 {
100    static M: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
101    *M.get_or_init(|| match std::env::var("MEMRA_MOE_F16G").as_deref() {
102        Ok("0") => 0,
103        Ok("2") => 2,
104        Ok("3") => 3,
105        Ok(_) => 1,
106        // Both arches independently re-arbitrated to mode 2 on 2026-08-02
107        // (5090: lane/f16g-default-rearb; H100: lane/h100-flip-full) — unset = 2 everywhere.
108        Err(_) => 2,
109    })
110}
111/// Mode-2 sk kernel form policy (round 51, lane/sk-bm128): the single-kernel grouped GEMM runs
112/// as a persistent problem-visitor over the real CSR tiles with two tile forms. Returns
113/// (shape_sel, cross) for the FFI:
114///   MEMRA_F16G_SK=0    -> (-1, _): the round-49 grid-scan kernel (rollback seam).
115///   MEMRA_F16G_SK=32   -> all groups on the 32x64x32 2-stage form (cross = i32::MAX).
116///   MEMRA_F16G_SK=128  -> all groups on the 128x64x64 3-stage form (cross = 1; groups fall
117///                         back to 32x64 in-launcher when the device/in_f can't take it).
118///   unset              -> hybrid split: groups with m_e >= MEMRA_F16G_SK_CROSS ride the 128
119///                         form. Default cross = 64 (5090 sweep 2026-08-01, receipts
120///                         research/sk-bm128-20260801/; H100 re-swept on the direct+tail
121///                         form 2026-08-02, lane/h100-flip-full: {16,32,64} ->
122///                         12868/13192/13225 — 64 wins there too, the pre-direct 32
123///                         verdict was stale).
124pub fn moe_f16g_sk_params() -> (i32, i32) {
125    static P: std::sync::OnceLock<(i32, i32)> = std::sync::OnceLock::new();
126    *P.get_or_init(|| match std::env::var("MEMRA_F16G_SK").as_deref() {
127        Ok("0") => (-1, 0),
128        Ok("32") => (0, i32::MAX),
129        Ok("128") => (0, 1),
130        _ => {
131            let cross = std::env::var("MEMRA_F16G_SK_CROSS")
132                .ok()
133                .and_then(|v| v.parse().ok())
134                .unwrap_or(64);
135            (0, cross)
136        }
137    })
138}
139/// DIRECT-FROM-QUANT sk tile loaders (lane/kquant-tile-loaders, 2026-08-02; IQ classes added
140/// by lane/iq-direct-loaders): Q4_K/Q6_K/IQ4_XS/IQ3_S expert projections on the mode-2/3 sk
141/// visitor forms dequant their weight tiles in-register from the quant superblocks instead of
142/// running the per-(layer,projection) dequant pass into an f16 workspace (41.8% of Ornith-35B
143/// t=512 kernel time — the pp512 wall, research/q4k-expert-prefill-20260802 §5; the IQ classes
144/// are 94.8% of q35's bank bytes — the h100-sk-direct coverage pricing). Bit-identical to the
145/// workspace path by construction (kernel-check "f16g-kq-direct" gates it bitwise) — a
146/// data-movement change, not a numeric-class change. Default ON; MEMRA_F16G_DIRECT=0 reverts
147/// to the workspace path everywhere; MEMRA_F16G_DIRECT=kq keeps the k-quant loaders and
148/// reverts only the IQ classes (the iq-direct-loaders A/B seam — the pre-lane shipped config).
149pub fn moe_f16g_direct_on(qtype: i32) -> bool {
150    static M: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
151    let m = *M.get_or_init(|| match std::env::var("MEMRA_F16G_DIRECT").as_deref() {
152        Ok("0") => 0,
153        Ok("kq") => 1,
154        _ => 2,
155    });
156    match m {
157        0 => false,
158        1 => qtype == QT_Q4_K || qtype == QT_Q6_K,
159        _ => true,
160    }
161}
162/// DEEP-TAIL sk form (lane/sk-tail-form, 2026-08-02): groups below the visitor crossover ride
163/// a 32x64x64 3-STAGE cp.async tile instead of the round-51 32x64x32 2-stage — the same 32-row
164/// tile (zero extra padding), 2 k-blocks in flight instead of 1 and half the syncs per k. The
165/// H100 ncu pricing (research/sk-bm128-20260801) put the 2-stage tail at 31% of the sk GEMM
166/// stage under q35's routing skew. Bit-identical to every other sk form by construction
167/// (kernel-check "f16g-sk" gates all tail arms maxdiff==0); exists in both the workspace-f16
168/// and direct-from-quant variants. Default ON; MEMRA_F16G_TAIL=0 = rollback to the 2-stage
169/// tail. in_f % 64 != 0 falls back in-launcher.
170pub fn moe_f16g_tail_on() -> bool {
171    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
172    *ON.get_or_init(|| std::env::var("MEMRA_F16G_TAIL").as_deref() != Ok("0"))
173}
174
175/// Per-model door for the gemma-MoE (gelu) grouped path: round 49's Hopper default
176/// REGRESSED g26 board-2048 prefill -8.3% interleaved x5 on-box (def median 10380,
177/// wild 8.9k-11.7k spread; off 11317, ±0.13%) — the +6-15% probe verdict didn't
178/// survive the board workload (stale-verdict law, round 50). The silu/qwen class
179/// keeps the round-49 default (q35 +53% board-2048). Explicit MEMRA_MOE_F16G=1/2
180/// still opens this door for A/B.
181pub fn moe_f16g_gemma_on() -> bool {
182    static M: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
183    *M.get_or_init(|| !matches!(std::env::var("MEMRA_MOE_F16G").as_deref(), Ok("0") | Err(_)))
184}
185
186/// Fused act-epilogue (silu/gelu-mul + q8_1_mmq quantize in one launch) for the MoE prefill
187/// MMA arms. Byte-identical to the two-pass path (kernel-check gated) — default ON.
188/// MEMRA_MOE_FUSE_ACTQ=0 is the rollback/A-B seam.
189pub fn moe_fuse_actq_on() -> bool {
190    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
191    *ON.get_or_init(|| std::env::var("MEMRA_MOE_FUSE_ACTQ").as_deref() != Ok("0"))
192}
193
194/// PREFILL router m-invariance (lane/concat-prime-exact, 2026-08-02). The batched cuBLASLt
195/// router GEMM changes a row's logits when OTHER rows join the call (probed: first change at
196/// m=65 on the Ornith-35B router, 3.9e-3 — while the MMQ/f16 trunk GEMMs are bit-identical
197/// across m). Feeding a top-k discontinuity, that made a served request's expert selection a
198/// function of its CO-ARRIVALS under cross-request prime batching. The in-house router GEMV
199/// is m-invariant, so prefill uses it too and routing depends on a session's own tokens only.
200/// DEFAULT ON: it is the serving isolation contract, and it is the same kernel decode and spec
201/// verify already use (dispatch parity, one router kernel for every t).
202/// MEMRA_ROUTER_PREFILL_EXACT=0 reverts to the batched GEMM.
203pub fn router_prefill_exact_on() -> bool {
204    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
205    *ON.get_or_init(|| std::env::var("MEMRA_ROUTER_PREFILL_EXACT").as_deref() != Ok("0"))
206}
207
208pub fn router_kernel_on() -> bool {
209    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
210    *ON.get_or_init(|| {
211        let on = std::env::var("MEMRA_ROUTER_KERNEL").as_deref() != Ok("0");
212        if !on {
213            eprintln!("[memra] router kernel OFF (rollback: per-column cuBLAS gemv)");
214        }
215        on
216    })
217}
218
219/// FAST-ROUTER batch twin (lane/fast-router, 2026-08-02). The concat-prime exactness fix
220/// (router_prefill_exact_on) routes prefill through router_gemv — m-invariant, but a
221/// per-(expert,token) GEMV program with zero operand reuse, so q35 board-2048 prefill paid
222/// -10% on the 5090. router_gemv_f32_w8_batch register-tiles (8x8 expert-x-token) the same
223/// per-row FP chains (BIT-IDENTICAL per row — kernel-check sweeps m=1..2048 on real router
224/// weights), so the t crossover below is pure perf, not a numeric config. Swept on-box
225/// (research/fast-router-20260802/crossover-router*.jsonl): plain wins t<=4, batch +7-9%
226/// at t=8, 1.9x at t=16 rising to 3.45x at t=2048 — MIN_T=8. Decode t=1 and spec verify
227/// t<8 keep the plain w8 form. MEMRA_ROUTER_BATCH=0 forces plain at every t (rollback
228/// seam, perf-only: bits are equal by the kernel-check gate).
229/// Killed arms (same sweep, JSONL is the record): the 8x16 tile lost to 8x8 at every t
230/// (128-accumulator register pressure beats the halved w-traffic), and the same-shape
231/// sigmoid_dot_rows twin (out_f=1) measured 0.62-0.89x at every prefill t
232/// (launch-latency-bound, ~7us/layer at m=2048) — both bit-identity-PASSED before dying.
233pub const ROUTER_BATCH_MIN_T: usize = 8;
234pub fn router_batch_on() -> bool {
235    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
236    *ON.get_or_init(|| std::env::var("MEMRA_ROUTER_BATCH").as_deref() != Ok("0"))
237}
238mod cpu_experts;
239#[cfg(memra_cutlass)]
240pub mod cutlass_ffi;
241pub mod dsv4_ffi;
242pub mod dsv4_gpu;
243pub mod f16_ffi;
244pub mod fp8_ffi;
245pub mod mmq_ffi;
246pub mod moe_cache;
247pub mod prime_graph;
248pub mod spill;
249mod spill_pread;
250
251// Fatbins are EMBEDDED (crates-release lane, 2026-08-04): build.rs still writes them to
252// OUT_DIR, but the bytes ship inside the binary via include_bytes! and load through
253// cuModuleLoadData. Distribution contract: a prebuilt or cargo-installed binary must be
254// self-contained — the old baked OUT_DIR *paths* pointed at the builder's temp dir and
255// broke every machine that wasn't the build machine. Same bytes, same module image;
256// the runtime MEMRA_GEMM_FATBIN tune-seam override below is preserved.
257const FATBIN: &[u8] = include_bytes!(env!("MEMRA_ENGINE_FATBIN"));
258const HYBRID_FATBIN: &[u8] = include_bytes!(env!("MEMRA_HYBRID_FATBIN"));
259const QMATVEC_FATBIN: &[u8] = include_bytes!(env!("MEMRA_QMATVEC_FATBIN"));
260const FLASH_FATBIN: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN"));
261const GEMM_FATBIN: &[u8] = include_bytes!(env!("MEMRA_GEMM_FATBIN"));
262const ROUTER_FATBIN: &[u8] = include_bytes!(env!("MEMRA_ROUTER_FATBIN"));
263/// spec_sample.cu: sampled-spec primitives (Philox Gumbel-max / softmax gather / residual sampler).
264const SAMPLE_FATBIN: &[u8] = include_bytes!(env!("MEMRA_SAMPLE_FATBIN"));
265
266/// TUNE SEAM (tools/sweep): a RUNTIME `MEMRA_GEMM_FATBIN=<path>` overrides the baked-in
267/// qmatvec_gemm.cu fatbin path (build.rs bakes the same name at COMPILE time via
268/// cargo:rustc-env — that constant is the default). Lets the sweep harness swap in a
269/// `-D`-tuned fatbin per process with NO rust rebuild. Unset at runtime => the
270/// compile-time default (zero behavior change).
271fn gemm_fatbin_bytes() -> std::borrow::Cow<'static, [u8]> {
272    assert!(
273        !(portable_mma_gated() && std::env::var_os("MEMRA_GEMM_FATBIN").is_some()),
274        "MEMRA_GEMM_FATBIN overrides are not allowed in the portable CUDA lane"
275    );
276    match std::env::var("MEMRA_GEMM_FATBIN") {
277        Ok(path) => std::borrow::Cow::Owned(
278            std::fs::read(&path).unwrap_or_else(|e| panic!("MEMRA_GEMM_FATBIN read {path}: {e}")),
279        ),
280        Err(_) => std::borrow::Cow::Borrowed(GEMM_FATBIN),
281    }
282}
283
284/// Phase A (ARCHITECTURE-H100.md): sm_90a re-enables the portable-PTX tensor-core paths
285/// (int8 mma.m16n8k32/k16.s8, bf16 m16n8k16, ldmatrix, cp.async — all sm_80-class, native
286/// on Hopper) that the portable boot lane gates off. Dispatch guards that used to test
287/// `cfg!(memra_portable_cuda)` test this instead; sm_89 keeps the pure-portable behavior.
288/// The sm_120a/sm_100a-only MMA kinds (mxf4nvf4, kind::f8f6f4) are NOT covered — their
289/// launchers stay fail-closed stubs on 90a and their dispatch arms stay arch-gated.
290pub(crate) const fn portable_mma_gated() -> bool {
291    cfg!(memra_portable_cuda) && !cfg!(memra_hopper_mma)
292}
293
294/// The GDN K4/K5 mma pair's UNSET-env default — ONE definition for the three read sites
295/// (gdn_mma_enabled, the k123 pre-work, gdn_scan_chunked's dispatch). They read the env
296/// per call ON PURPOSE (kernel-check toggles it to pin both configs), so the shared part
297/// is this compile-time constant: ON for Hopper-MMA builds (the original 90a promotion)
298/// and for sm_120a builds (lane/moeprime-nvfp4-direct, 2026-08-21 — measured on one RTX
299/// PRO 6000 ornith15 pp14715 +6-8% and the local 5090 q38-27b +1-2%, both orders both
300/// rigs). A site defaulting differently from its peers arms the mma pre-work while the
301/// scan takes the scalar route — measured as a 0.8% LOSS, the drift this helper kills.
302pub(crate) const fn gdn_mma_default_on() -> bool {
303    cfg!(memra_hopper_mma) || konst_eq(env!("MEMRA_BUILT_CUDA_ARCH"), "120a")
304}
305
306/// const str-eq (std `==` on &str is not const-stable on this toolchain floor).
307const fn konst_eq(a: &str, b: &str) -> bool {
308    let (a, b) = (a.as_bytes(), b.as_bytes());
309    if a.len() != b.len() {
310        return false;
311    }
312    let mut i = 0;
313    while i < a.len() {
314        if a[i] != b[i] {
315            return false;
316        }
317        i += 1;
318    }
319    true
320}
321
322/// The legacy quantized prefill GEMMs are tuned and validated for sm_120a; sm_90a re-admits
323/// them through the Hopper-MMA lane (int8 m16n8k32.s8 is sm_80-class PTX).  Keep the policy
324/// in a pure helper so the dispatch guard can be regression-tested without constructing an
325/// Engine or allocating a GPU tensor.
326const fn legacy_quant_gemm_allowed(portable_cuda: bool, hopper_mma: bool, no_gemm: bool) -> bool {
327    (!portable_cuda || hopper_mma) && !no_gemm
328}
329
330// ---- KV-cache format selection (kvbytes lane, 2026-07-08; default OFF = daily config) ----
331// `MEMRA_KV_K` = q8_0 (default, 34 B/32elem) | fp8 (raw e4m3, 32 B — the -6% K-bytes arm)
332// `MEMRA_KV_V` = q5_1 (default, 24 B/32elem) | q4_0 (18 B, -25% V bytes) | fp8 (32 B, +33%)
333// A non-default format is a NEW NUMERIC CONFIG: its own run-gen argmax baseline is legal,
334// but the gate battery (kernel-check, run-spec self-consistency) must pass WITHIN it and
335// the choice is explicit env, never silent. flash_attn.cu is compiled once per format pair
336// (build.rs); the kernels keep their names — Engine::new just loads the matching fatbin.
337const FLASH_FATBIN_VQ4: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_VQ4"));
338const FLASH_FATBIN_VF8: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_VF8"));
339const FLASH_FATBIN_KF8: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_KF8"));
340const FLASH_FATBIN_KF8VQ4: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_KF8VQ4"));
341const FLASH_FATBIN_KF8VF8: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_KF8VF8"));
342
343/// KV format policy moved to the shared `memra-kv` crate (Phase D); re-exported so the
344/// fatbin router below and every existing `crate::kv_blk_bytes()` call site is unchanged.
345pub use memra_kv::{kv_blk_bytes, kv_cache_formats};
346
347/// The flash_attn fatbin matching the selected KV formats.
348fn flash_fatbin_bytes() -> &'static [u8] {
349    match kv_cache_formats() {
350        ("q8_0", "q5_1") => FLASH_FATBIN,
351        ("q8_0", "q4_0") => FLASH_FATBIN_VQ4,
352        ("q8_0", "fp8") => FLASH_FATBIN_VF8,
353        ("fp8", "q5_1") => FLASH_FATBIN_KF8,
354        ("fp8", "q4_0") => FLASH_FATBIN_KF8VQ4,
355        ("fp8", "fp8") => FLASH_FATBIN_KF8VF8,
356        other => unreachable!("kv_cache_formats returned {other:?}"),
357    }
358}
359
360/// TUNE SEAM (tools/sweep): kernel1 (Q8_0/Q4_K/Q5_K) launch-tile override,
361/// `MEMRA_GEMM_K1_LAUNCH="BM,BN,NWARP"`. MUST match the `-D K1_BM/K1_BN/NWARP` the swept
362/// fatbin was compiled with (the .cu tile and the host launch grid/block have to agree —
363/// the hardcoded (128,128,8) in qmatvec_gemm/qmatvec_gemm_raw is the shipped default).
364/// Kernel2 (Q6_K/NVFP4) launch is untouched. Unset or malformed => None => shipped
365/// defaults (zero behavior change).
366fn k1_launch_override() -> Option<(u32, u32, u32)> {
367    static K1: std::sync::OnceLock<Option<(u32, u32, u32)>> = std::sync::OnceLock::new();
368    *K1.get_or_init(|| {
369        let v = std::env::var("MEMRA_GEMM_K1_LAUNCH").ok()?;
370        let p: Vec<u32> = v.split(',').filter_map(|s| s.trim().parse().ok()).collect();
371        match p.as_slice() {
372            [bm, bn, w] => Some((*bm, *bn, *w)),
373            _ => None,
374        }
375    })
376}
377
378/// H100 wgmma prefill-GEMM seam (task 8, ARCHITECTURE-H100.md): OPT-IN (MEMRA_WGMMA=1).
379/// v0 verdict (2026-07-26, N=5 pp512 9B-Q8_0): wgmma 3845 tok/s vs MMQ 8692 — the
380/// standalone harness's "688us MMQ ref" was a pp2048-shape figure, so v0 (unpipelined,
381/// 64x64 tile, wait_group<0> every 32-K step) is ~3x SLOWER per launch at m=512 model
382/// shapes. Default stays MMQ until the pipelined version beats it N=5 (repo law).
383/// Correctness stays pinned regardless: kernel-check's wgmma case is cfg-gated, not env-gated.
384pub(crate) fn wgmma_gemm_enabled() -> bool {
385    static V: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
386    *V.get_or_init(|| std::env::var("MEMRA_WGMMA").as_deref() == Ok("1"))
387}
388
389/// TUNE SEAM: keys per FA-decode split (`MEMRA_FA_SPLIT` forces a fixed size; default 64). Smaller
390/// splits raise grid.y so grid = n_head_kv * n_splits fills the 82 SMs at short/mid ctx (vec path
391/// launches only n_head_kv=8 CTAs per split). Swept clock-locked 2026-07-03 (graph tg128): 32 beat
392/// 64 at ctx 128/512 (+0.5/+1.2%) and lost at 2048 (-3%) — BUT the adaptive 32/64 default BROKE the
393/// MTP spec-decode exact-match gate (run-spec K=1/2 self-consistency FAIL with 32; PASS with 64):
394/// the split count changes the combine's FP summation order, and the spec verify's batched forward
395/// only argmax-matches single-step decode under the 64-split order on real prompts. Spec exactness
396/// (the bigger lever) outranks a <=1.2% decode win -> default stays FIXED 64; sweeps use the env.
397/// Takes t_kv so eager, _dc capture, and fa_geom_eager stay signature-compatible for future
398/// adaptive retries (any retry MUST pass run-spec self-consistency first).
399/// Minimum t_kv for the warp-per-token vec FA path (below it the scalar path's 4x-more-blocks
400/// hides latency better — measured crossover, see `fa_decode`). Shared by fa_decode / fa_decode_dc /
401/// fa_geom_eager / fa_decode_rows-eligibility (spec verify) so the kernel pick NEVER diverges
402/// between eager decode and the verify (the spec-exactness law).
403pub const FA_VEC_MIN_TKV: usize = 96;
404/// Env-overridable crossover (MEMRA_FA_VEC_MIN, default FA_VEC_MIN_TKV). The 96 floor was
405/// measured on the qwen geometry (nkv=2); gemma4 SWA layers run nkv=8 = 4x the vec grid,
406/// which moves the crossover — sweep per model, adopt per the battery.
407pub fn fa_vec_min_tkv() -> usize {
408    static V: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
409    *V.get_or_init(|| {
410        std::env::var("MEMRA_FA_VEC_MIN")
411            .ok()
412            .and_then(|v| v.parse().ok())
413            .unwrap_or_else(|| FA_VEC_MIN_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
414    })
415}
416
417/// f16-P/V class (DEFAULT since 2026-07-23 stamp v4; MEMRA_FA_F16PV=0 = f32-class rollback):
418/// llama-fa=1-style f16 P + f16 P@V accumulation on the hd512/SWA prefill stamps
419/// (KQ/softmax/normalize stay f32). Laptop stamp: 12B 1.045x, 31B 0.979x vs llama.
420///
421/// SPEC-SERVING FLIP (2026-07-26, the wkv acceptance-law pattern): with MEMRA_DRAFT set the
422/// default is OFF. f16 P/V shifts the PRIME's hidden states/KV in the sub-argmax logit
423/// space the drafter feeds on — argmax gates stay MATCH while depth acceptance falls off a
424/// cliff (26B d1736 0.883 -> 0.405, -40% e2e; f16pv-off alone restores 0.846/314 tok/s —
425/// the perf-ci acceptance battery is the only gate that sees this class). Explicit
426/// MEMRA_FA_F16PV always wins; plain serving keeps the f16 prefill win.
427pub fn fa_f16pv_on() -> bool {
428    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
429    *ON.get_or_init(|| {
430        std::env::var("MEMRA_FA_F16PV")
431            .map(|v| v != "0")
432            .unwrap_or_else(|_| std::env::var("MEMRA_DRAFT").is_err())
433    })
434}
435
436/// hd512 head-pair arm (DEFAULT since stamp v4; MEMRA_FA512_HP=0 reverts to sp16): GQA
437/// ncols2=2 — 2 heads per CTA share each staged K/V tile, Q register-resident. Engages
438/// when n_head is even and the GQA group (n_head/n_head_kv) is even.
439pub fn fa512_hp_on() -> bool {
440    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
441    *ON.get_or_init(|| std::env::var("MEMRA_FA512_HP").as_deref() != Ok("0"))
442}
443
444/// SWA head-pair arm (DEFAULT since stamp v4; MEMRA_FAW_HP=0 reverts to p1): llama-class
445/// windowed geometry — 32 q-rows x 2 heads per CTA sharing staged K/V, f16 P@V
446/// accumulation. Even n_head and even GQA group required (guarded per call).
447pub fn faw_hp_on() -> bool {
448    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
449    *ON.get_or_init(|| std::env::var("MEMRA_FAW_HP").as_deref() != Ok("0"))
450}
451
452/// 4-warp sp16 experiment arm (MEMRA_FA512_W4=1, requires the f16pv door): GEMM0 split-K
453/// 4-way + GEMM1 4x128 O-dims. Own partial-sum order — oracle-band gated. Returns warp
454/// count (2 = base sp16). 8-warp arm measured NEGATIVE 2026-07-23 (jsonl) and removed.
455pub fn fa512_wide_warps() -> usize {
456    static N: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
457    *N.get_or_init(|| match std::env::var("MEMRA_FA512_W4").as_deref() {
458        Ok("1") => 4,
459        _ => 2,
460    })
461}
462
463/// hd-512 vec crossover floor (MEMRA_FA512_MIN, default 512) — shared by fa_decode dispatch
464/// and the gemma global-layer rows/parity call sites.
465pub fn fa512_min_tkv() -> usize {
466    static FA512_MIN: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
467    *FA512_MIN.get_or_init(|| {
468        std::env::var("MEMRA_FA512_MIN")
469            .ok()
470            .and_then(|v| v.parse().ok())
471            .unwrap_or(512)
472    })
473}
474/// Per-model crossover default, set at model load BEFORE the first decode (per-model
475/// numeric-config adoption law). qwen keeps the measured 96; gemma4 (nkv=8 SWA) measured
476/// vec-always fastest: 119.9 (96) / 130.0 (48) / 133.2 (1) tok/s tg128-regime, 2026-07-10.
477pub static FA_VEC_MIN_DEFAULT: std::sync::atomic::AtomicUsize =
478    std::sync::atomic::AtomicUsize::new(FA_VEC_MIN_TKV);
479/// Per-model windowed-split default (MEMRA_FA_SPW overrides): gemma MoE (26B, nkv=8) measured
480/// 32 (grid-limited t=1 under the raw-e4m3 sV ceiling, 2026-07-12); dense gemma (31B)
481/// measured 64 (37.13/37.12 vs 36.87/36.86 at 1.7k, N=2 — different attention geometry).
482pub static FA_SPW_DEFAULT: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(32);
483/// Per-model hd512 (gemma globals) split default (MEMRA_FA_SP512 overrides): 26B measured 16
484/// (2026-07-11 N=2), dense 31B measured 32 (36.86/36.93 vs 36.73/36.73 at 1.7k, 2026-07-12).
485/// fused t=1 q4_0 pair/triple row mapping: true = mr1 (one row/warp). Per-model default
486/// (dense gemma wins +1.1% short / +0.6% depth on the 31B; MoE 26B REGRESSES −1.2% —
487/// its shared-expert fused2 shapes lose to the finer grid). MEMRA_Q40_MR env still wins.
488pub static FUSED_MR1_DEFAULT: std::sync::atomic::AtomicBool =
489    std::sync::atomic::AtomicBool::new(false);
490/// Per-model router-GEMV form (2026-07-31): the 8-warp twin is +8.8% on the H100 q35
491/// decode step (router was 14.8% of it) with argmax + spec self-consistency green on
492/// qwen-class MoE both rigs. The gemma-4 26B knife-edge block (2026-07-31, single
493/// synthetic prompt) was RE-ARBITRATED 2026-08-01 on 6 real prompts — gate outcomes
494/// identical to the lone-warp arm, +13% g26 decode — so gemma4 rides the default too
495/// (research/g26-decode-20260801/). MEMRA_ROUTER_V2 env overrides either way.
496pub static ROUTER_W8_DEFAULT: std::sync::atomic::AtomicBool =
497    std::sync::atomic::AtomicBool::new(true);
498pub static FA_SP512_DEFAULT: std::sync::atomic::AtomicUsize =
499    std::sync::atomic::AtomicUsize::new(16);
500/// Per-model rms_norm block size (per-model numeric-config law: the per-thread partial-sum
501/// split changes with blockDim -> different FP order -> battery-arbitrated per model).
502/// qwen keeps the shipped 256; gemma4 adopts 1024 (single-row 2816-col norms are one-block
503/// latency-bound at 256 threads — 7us/launch measured).
504pub static RMS_BLOCK_DEFAULT: std::sync::atomic::AtomicU32 = std::sync::atomic::AtomicU32::new(256);
505/// gemma4 fa split ladder switch (set at model load; see fa_split_keys).
506pub static FA_SP_GEMMA: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
507/// Per-model stream-k entry override for SPEC serving (-1 = unset → env/default;
508/// 0 = force tiling; 1 = admit the deterministic form selector). The former timing
509/// selector made identical boots choose different fold orders; `MEMRA_MMQ_SK_FORM` is the
510/// explicit numerical-form seam. mmq_ffi reads this before the env.
511pub static MMQ_SK_FORCE: std::sync::atomic::AtomicI8 = std::sync::atomic::AtomicI8::new(-1);
512/// Per-model FP8-KV door — lives in memra-kv next to the format policy it drives
513/// (re-export keeps `crate::KV_FP8_FORCE` setters in model.rs/hybrid.rs working).
514pub use memra_kv::KV_FP8_FORCE;
515pub(crate) fn rms_block() -> u32 {
516    static V: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
517    *V.get_or_init(|| {
518        std::env::var("MEMRA_RMS_BLOCK")
519            .ok()
520            .and_then(|v| v.parse().ok())
521            .unwrap_or_else(|| RMS_BLOCK_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
522    })
523}
524
525pub(crate) fn fa_split_keys(t_kv: usize, n_head_kv: usize) -> usize {
526    static S: std::sync::OnceLock<Option<usize>> = std::sync::OnceLock::new();
527    if let Some(forced) = *S.get_or_init(|| {
528        std::env::var("MEMRA_FA_SPLIT")
529            .ok()
530            .and_then(|v| v.parse().ok())
531            .filter(|&s: &usize| s >= 8 && s % 8 == 0)
532    }) {
533        return forced;
534    }
535    // CTX-ADAPTIVE default (2026-07-05 40k sweep: sp32 24.5 vs sp128 26.0 tok/s = +5.8% — at
536    // deep ctx the n_splits count explodes (40k/32 = 1265 splits x 8 kv-heads) and the combine
537    // + partial-buffer cost dominates; at short ctx small splits fill the SMs). Exactness: split
538    // size only changes the PARTITION of keys; the rows/combine order per split is fixed and the
539    // gate battery (kernel-check + run-spec K=1..8) arbitrates every default change.
540    //
541    // SM-AWARE SHORT-CTX RUNG (2026-07-06 g7e): the 32-key rung was tuned on the 82-SM 5090.
542    // On 188 SMs the vec grid (n_head_kv x n_splits CTAs) starves at short ctx — the 35B has
543    // n_head_kv=2, so ctx128/split32 = 8 CTAs on 188 SMs. Measured on g7e (N=1 sweep + N=3
544    // interleaved confirm): 35B ctx128 sp16 179 vs sp32 161 (+11%), ctx512 178 vs 158, ctx2048
545    // flat, ctx>=4096 sp64 edges sp16 by ~3%; 27B ctx128 70.9 vs 66.3 (+7%); 9B 177 vs 163
546    // (+9%). Rigs <=100 SMs keep the validated 5090 ladder EXACTLY (default unchanged there —
547    // rig-divergence law: this branch is measured on 188 SMs only).
548    // gemma4 all-16 ladder probe REVERTED (2026-07-10): +1.3 plain at d1736 (157.5 vs 156.2)
549    // but depth VERIFY collapsed (spec 203.5 -> 169 — the windowed rows' per-row combine over
550    // 64 splits). The mixed default (swa nkv=8 -> 32, globals nkv=2 -> 8-ladder) stays; a
551    // caller-split policy would break row-vs-decode split parity. FA_SP_GEMMA kept as a seam.
552    if FA_SP_GEMMA.load(std::sync::atomic::Ordering::Relaxed)
553        && std::env::var("MEMRA_FA_SP16").as_deref() == Ok("1")
554    {
555        return if t_kv <= 8192 {
556            16
557        } else if t_kv <= 16384 {
558            64
559        } else {
560            128
561        };
562    }
563    let big_rig = fa_sm_count() >= 128;
564    if big_rig {
565        let _ = n_head_kv;
566        if t_kv <= 2048 {
567            16
568        } else if t_kv <= 16384 {
569            64
570        } else {
571            128
572        }
573    } else if n_head_kv <= 4 {
574        // KV-HEAD-AWARE RUNG (2026-07-08, 5090): the 8192->32 rung was validated on kv=8 models
575        // (27B/9B: 8 heads x n_splits fills 82 SMs). The 35B has n_head_kv=2 — at ctx512/sp32
576        // the vec grid is 2 x 20 = 40 CTAs on 82 SMs (half idle). Measured (35B, run-gen 128tok
577        // N=1 sweep + N=3 confirm): sp8 162.1 / sp16 161.3 / sp32 159.4 at short ctx.
578        // DEPTH TAPER (same day, the deep-ctx lesson re-learned on this rung): sp8 at d6257 =
579        // 782 splits -> combine + partial-buffer cost dominates (141.2 tok/s); the d6257 sweep
580        // says sp64 = 153.0 (sp16/32 147, sp96 147.6, sp128 141). Few-kv-head models need the
581        // taper EARLIER than kv=8 (per-split grid 4x thinner, same per-split combine cost).
582        // Crossover hunt: sp8 vs sp64 = 156.7/155.9 at d3072, 151.7/155.6 at d4096 -> boundary 3072.
583        // RUNG RE-SWEPT UNDER THE DEEP KERNEL (2026-08-02, lane/ladder-3072 — the stale-verdict
584        // law: the 3072 boundary was calibrated on the conflicted v4 core; the deep rewrite cut
585        // vec cost ~1.2-1.4x while combine scales with n_splits, so sp8's combine bill
586        // dominates far earlier). Kernel receipts (quiet-rig nsys, deep vec + combine us):
587        // d1024 sp8 17.1 vs sp64 10.6; d2048 31.0 vs 12.2; d3072 44.0 vs 18.3. e2e run-gen
588        // tg128 N=3 interleaved (KAT + q35, research/ladder-3072-20260802/): sp8 loses at
589        // EVERY depth >= 1024 (KAT d2048 182.6 vs 188.0 = -2.9%, d3072 175.9 vs 186.4 =
590        // -5.6%; q35 d4096 169.2 vs 182.6 = -7.4%); d512 flat (+-0.2%, inside noise). sp32
591        // ties sp64 within noise in the mid band and loses at d4096 -> no extra rung.
592        // Boundary 3072 -> 512: sp8 keeps only the short-ctx band it was validated on
593        // (ctx128-512); sp64 takes over where the deep kernel made combine the bill.
594        if t_kv <= 512 {
595            8
596        } else if t_kv <= 16384 {
597            64
598        } else {
599            128
600        }
601    } else {
602        if t_kv <= 8192 {
603            32
604        } else if t_kv <= 16384 {
605            64
606        } else {
607            128
608        }
609    }
610}
611
612/// SM count of device 0, cached (used by fa_split_keys' rig-size rung; primary-context query,
613/// same attribute Engine::batched_variant reads).
614fn fa_sm_count() -> i32 {
615    static N: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
616    *N.get_or_init(|| {
617        cudarc::driver::result::init().ok();
618        cudarc::driver::result::device::get(0)
619            .and_then(|d| unsafe { cudarc::driver::result::device::get_attribute(
620                d, cudarc::driver::sys::CUdevice_attribute_enum::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT) })
621            .unwrap_or(82)
622    })
623}
624
625/// FA-prefill kernel-name suffix for a head_dim (the template-stamped twins in flash_attn.cu):
626/// 256 = the original names (qwen35 class, dispatch unchanged), 128 = `_hd128` (MiniMax-M3).
627/// Any other dim errors — callers gate to sdpa_naive before dispatching FA.
628fn fa_hd_suffix(head_dim: usize) -> Result<&'static str, Box<dyn std::error::Error>> {
629    match head_dim {
630        256 => Ok(""),
631        128 => Ok("_hd128"),
632        d => Err(format!(
633            "fa_prefill: no kernel stamped for head_dim={d} (only 256/128); \
634                          callers must gate to sdpa_naive"
635        )
636        .into()),
637    }
638}
639
640/// Quant type codes matching qmatvec.cu QType enum.
641pub const QT_Q8_0: i32 = 0;
642pub const QT_Q4_K: i32 = 1;
643pub const QT_Q6_K: i32 = 2;
644pub const QT_Q5_K: i32 = 3;
645pub const QT_Q3_K: i32 = 4;
646pub const QT_IQ4_XS: i32 = 5;
647pub const QT_IQ3_S: i32 = 6;
648pub const QT_NVFP4: i32 = 7;
649/// Checkpoint-native FP8-E4M3 (MEMRA_ST_E4M3, lane e4m3dec): raw safetensors e4m3 weight bytes
650/// [out_f, in_f] row-major (row_bytes == in_f), per-tensor f32 weight_scale in GpuTensor `scale`
651/// (fused at the mmvq write / post-matmul scale_inplace). Decode = qmatvec_e4m3_mmvq (+ _b2/_b4/_b8
652/// batched twins); prefill (m>=16) = the cuBLASLt FP8 GEMM on the SAME resident bytes (fp8_ffi.rs)
653/// — ONE weight copy total, no Q8_0 re-encode duplicate.
654pub const QT_F8_E4M3: i32 = 10;
655/// Device-side tag for the A6 SPLIT-PLANE repacked NVFP4 layout (Stage-A generic kernel only;
656/// GpuTensor keeps qtype=QT_NVFP4 + an `rp` flag — this tag never lives in a GpuTensor).
657pub const QT_NVFP4_RP: i32 = 9;
658/// Unquantized f32 weight (safetensors MoE Path A: experts dequantized to f32 host-resident).
659pub const QT_F32: i32 = 8;
660pub const QT_BF16: i32 = 11;
661pub const QT_Q4_0: i32 = 12; // gemma-4 QAT GGUF weight format (18B/32: fp16 d + nibbles)
662/// GGUF Q2_K. Appended after the existing Q4_0 code so kernel ABI values do not move.
663/// Mixed-expert artifacts use the generic f32-dequant staged kernel until a target-rig-gated
664/// dp4a/MMQ implementation exists.
665pub const QT_Q2_K: i32 = 13;
666/// Checkpoint-native FP8-E4M3 with a BLOCK-128 weight-scale GRID (lane/fp8-blk128-decode,
667/// 2026-08-05) — the Qwen-official FP8 / DeepSeek-V3 scale class. Same raw e4m3 bytes as
668/// `QT_F8_E4M3` ([out_f, in_f] row-major, row_bytes == in_f), but the dequant scale is
669/// `GpuTensor::Quant.blk` (`Fp8BlockScales`, [ceil(out_f/128), ceil(in_f/128)] f32) and the
670/// scalar `scale` field is 1.0 by the layout contract.
671///
672/// WHY A DISTINCT CODE rather than `QT_F8_E4M3` + a `blk` flag: every existing QT_F8_E4M3
673/// consumer (qmatvec_e4m3_mmvq and its batched/fused twins, e4m3_fused_params,
674/// matmul_pre_dual_noscale's F8 arm, try_fp8_gemm) threads exactly ONE scalar weight scale. Under
675/// a shared code, any consumer that was not taught the grid would still MATCH and would dequant
676/// every tile at scale 1.0 — a silent numeric corruption. Under a distinct code every untaught
677/// consumer refuses loudly instead (`mmvq_supports`/`gemm_supports`/`mmq_supports` return false;
678/// the mmvq name match panics), so a missed dispatch site is a crash or a refusal receipt, never
679/// wrong numbers. Decode = `qmatvec_e4m3_blk_mmvq`; prefill (m>=16) = the per-block FP8 MMQ tile
680/// on the SAME resident bytes+grid (fp8_ffi::try_fp8_blk_mmq) — ONE weight copy total.
681pub const QT_F8_E4M3_BLK: i32 = 14;
682
683/// Engine device context: CUDA context, stream, loaded kernel modules, cuBLASLt (via runtime::Gpu).
684pub struct Engine {
685    pub gpu: memra_runtime::Gpu,
686    module: Arc<CudaModule>,
687    hybrid: Arc<CudaModule>,
688    qmatvec: Arc<CudaModule>,
689    flash: Arc<CudaModule>,
690    /// FP8-GLOBALS module (2026-07-11): the kf8vf8 fatbin loaded ALONGSIDE the default —
691    /// gemma GLOBAL layers (hd512) append + attend in e4m3 (dequant-latency arc, HANDOVER).
692    /// Lazy: loaded on first global-format use; None until then.
693    flash_g: std::sync::OnceLock<Arc<CudaModule>>,
694    gemm: Arc<CudaModule>,
695    router: Arc<CudaModule>,
696    /// Sampled-spec kernels (research/sampled-spec-impl-map.md piece A).
697    sample: Arc<CudaModule>,
698    /// EDGE-1 §B: one shared SLRU expert-residency cache, lazily built on first MoE dispatch under
699    /// MEMRA_MOE_CACHE. `Mutex` makes it multi-agent safe (§E.2); the lock covers only lookup/admit/
700    /// memcpy-issue (µs), NOT the GEMM, so streams still overlap. `None` => cache disabled.
701    moe_cache: Mutex<Option<crate::moe_cache::MoeSlotCache>>,
702    /// Exact retained expert-block lengths collected after model load. Mixed-layout models use
703    /// this inventory to preallocate fixed-address size classes instead of sizing every slot to
704    /// the single largest block. The cache still owns every address for its full lifetime.
705    moe_cache_layout: Mutex<Option<Vec<usize>>>,
706    /// CAPTURE-RETAIN mode (graph arc, 2026-07-12): while a graph capture (and its allocator
707    /// warmups) runs, every Engine allocation is ALSO kept alive here — a captured graph's
708    /// transient buffers must never return to the pool, or later allocations (e.g. the spec
709    /// verify between replays) reuse their addresses and the replay reads/writes live memory
710    /// (the draft-graph corruption root cause). Fast-path cost when off: one relaxed atomic.
711    capture_keep_on: std::sync::atomic::AtomicBool,
712    /// VERIFY-EXACT scope (dflash lane, 2026-07-13): when set, matmul/matmul_pre skip the
713    /// m>=16 prefill-GEMM branches so a t>=16 batched VERIFY rides the decode-exact b-tier
714    /// class (the parity law). The t=16 dflash verify tripped the GEMM threshold — 770us/
715    /// matmul (54% of the round) AND a different FP order than decode (issue-10 landmine).
716    verify_exact: std::sync::atomic::AtomicBool,
717    capture_keep: Mutex<Vec<Box<dyn std::any::Any + Send>>>,
718    /// EDGE-1 §C.2: dedicated H2D copy stream for async prefetch (event-synced to the compute stream).
719    pub copy_stream: Arc<CudaStream>,
720    /// Resident CUTLASS NVFP4 prefill scratch (workspace + a_packed + sfa_linear + sfa_sw + y + alpha),
721    /// allocated ONCE and grown to the largest prefill GEMM shape, then reused per-call. Removes the
722    /// 6 fresh allocations + alpha htod that `cutlass_fp4_gemm` did every prefill matmul (~200/prefill).
723    /// Safe as a single shared buffer because all GPU compute serializes on the one `gpu.stream` worker
724    /// thread (the server runs one GPU worker; no concurrent CUTLASS GEMMs share this scratch). `None`
725    /// until the first CUTLASS FP4 GEMM. Mutex guards lazy build/grow only (matches `moe_cache`).
726    #[cfg(memra_cutlass)]
727    cutlass_scratch: Mutex<Option<crate::cutlass_ffi::CutlassScratch>>,
728    /// FP8-ACT PREFILL scratch (MEMRA_PP_FP8): quantized-activation buffer + scale block + cuBLASLt
729    /// workspace, allocated once and grown to the largest prefill m*k (see fp8_ffi.rs). `None`
730    /// until the first FP8 prefill GEMM; Mutex guards lazy build/grow only (matches cutlass_scratch).
731    fp8_scratch: Mutex<Option<crate::fp8_ffi::Fp8Scratch>>,
732    /// f16-P/V door: pooled V re-encode buffer (bf16->f16) for the hd512 _pre path. Lazy-grow;
733    /// per-call cudaMalloc was a laptop-regression suspect (VRAM pressure, 31B nkv=4 = 4x bytes).
734    fa_vf16_scratch: Mutex<Option<CudaSlice<u8>>>,
735    /// Pooled fa-decode split partials (part_o, part_m, part_l): per-call zeros() was 3
736    /// alloc+memset pairs per fa launch (~144 mem nodes per decode token — the graph door's
737    /// residual launch tax) — lazy-grow, memset-prefix per use, stream-ordered reuse.
738    fa_part_pool: Mutex<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>>,
739    /// Retired fa-part pool generations (#68): old buffers whose addresses captured graphs may
740    /// have baked — kept alive for the Engine's lifetime instead of returning to the async pool
741    /// (see the RETIRE-ON-GROW comment at the realloc sites). Doubling growth bounds the total.
742    fa_part_retired: Mutex<Vec<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>>,
743    /// name -> resolved CudaFunction (capture-safe lookups; see `func`).
744    fn_cache: Mutex<std::collections::HashMap<String, CudaFunction>>,
745    f16_scratch: Mutex<Option<crate::f16_ffi::F16Scratch>>,
746    /// RANK1 LEVER (parallel argmax): resident pass-1 partials scratch (part_v[NB] f32, part_i[NB] i32),
747    /// allocated ONCE on first parallel-argmax call and reused. Stable pointers so the 2-pass argmax
748    /// is CUDA-graph-capturable (the buffer is referenced by both captured passes; lazy-allocated
749    /// before capture under the generate_graph tracking-off window so it carries no events).
750    argmax_partials: Mutex<Option<(CudaSlice<f32>, CudaSlice<i32>)>>,
751    /// ARC B (chunk-prime dequant-once): resident bf16 K/V workspace for `fa_prefill_view_ws`
752    /// ((K bytes, V bytes) u8 buffers holding [t_kv, kv_dim] bf16). Grown lazily to the largest
753    /// (t_kv, kv_dim) seen, REUSED across layers/chunks/calls (contents rewritten per launch —
754    /// safe because all compute serializes on the one gpu.stream). ~82MB at 40k ctx on the 27B.
755    prime_deqw_ws: Mutex<Option<(CudaSlice<u8>, CudaSlice<u8>)>>,
756    /// LAUNCH-STRUCTURE STAGE 1: persistent PINNED (cacheable, flags=0) host staging buffer for the
757    /// fused-router sel/w readback — one async DtoH pair + ONE sync instead of two synced dtohs.
758    /// Grown lazily; reused every MoE layer (single-threaded decode serializes on the sync).
759    router_stage: Mutex<Option<PinnedStage>>,
760}
761
762/// FAVENDOR lane env gate (2026-07-08): MEMRA_FA_V2=1 dispatches the llama-fattn-vec-mechanism
763/// decode kernels (fa_decode_vec_q_v2 / fa_decode_vec_q_rows_v2 / fa_decode_vec_q_v2_dc):
764/// tile-batched online softmax (one alpha rescale per 32-key tile instead of per key) + wide-load
765/// block dequant in the staging phase. NOTE rev2: llama's register streaming (no smem) was ALSO
766/// tried and measured 2x WORSE at depth in our gqa-warps frame — the smem KV-tile broadcast stays
767/// (see the kernel comment). NEW NUMERIC CONFIG (tile-level softmax regrouping changes FP order vs
768/// the per-key twins) — own argmax baseline; eager decode, the spec-verify rows path AND the
769/// graph _dc path switch TOGETHER (the spec-exactness law). Default OFF. Read per call (not
770/// OnceLock) so the gate battery can A/B within one process, matching the MEMRA_NO_FA_VEC pattern.
771fn fa_v2_on() -> bool {
772    // DEFAULT ON since 2026-07-08 (MEMRA_FA_V2=0 reverts): tile-batched online softmax, e2e
773    // measured across every model x depth — 35B 168.7->173.4 (d512) / 153.1->158.5 (d6257),
774    // 9B 131.2->132.7 / 108.4->124.5 (+15% — the engine-wide depth-slope fix), 27B 47.2->47.7 /
775    // 42.2->44.9. One-time numeric-config change; kernel-check + argmax + spec self-consistency
776    // + graph bit-identity green on all three models.
777    std::env::var("MEMRA_FA_V2")
778        .map(|v| v != "0")
779        .unwrap_or(true)
780}
781
782/// FA v3 gate (default ON since 2026-07-09; MEMRA_FA_V3=0 reverts to v2 — research/fa/fa_v3_design.md):
783/// HYBRID decode twins (fa_decode_vec_q_v3 / _rows_v3 / _v3_dc): llama's int8-dp4a K.Q with
784/// register-quantized Q (no K dequant, no K smem) + OUR CTA-shared staged bf16 V tile + OUR
785/// split partition/combine. NEW NUMERIC CONFIG (int8 Q quantization changes the K.Q accumulation
786/// vs the bf16-roundtrip FMA chain) — own argmax baseline; eager decode, the spec-verify rows
787/// path AND the graph _dc path switch TOGETHER (the spec-exactness law). Read per call so the
788/// gate battery can A/B within one process (the MEMRA_FA_V2 pattern).
789fn fa_v3_on() -> bool {
790    // DEFAULT ON since 2026-07-09 (MEMRA_FA_V3=0 reverts to v2): dp4a-K hybrid FA decode —
791    // fa kernel -21-23% at depth (micro), 35B spec p3 +5% (190->200, the last spec cell),
792    // d6257 +1.7%. Own numeric config; full battery green on 35B+9B incl graph bit-identity.
793    std::env::var("MEMRA_FA_V3")
794        .map(|v| v != "0")
795        .unwrap_or(true)
796}
797
798/// The v3 dp4a K path reads RAW q8_0 bytes (34B blocks) and stages q5_1 V verbatim — it is only
799/// correct on the DEFAULT KV formats — and needs dpl % 4 == 0 consecutive quants per lane
800/// (head_dim % 128 == 0; both daily models are hd256). All three dispatch sites share this
801/// predicate so the twins can never diverge.
802fn fa_v4_mode() -> &'static str {
803    static M: std::sync::OnceLock<String> = std::sync::OnceLock::new();
804    M.get_or_init(|| std::env::var("MEMRA_FA_V4").unwrap_or_default())
805}
806fn fa_v4_on() -> bool {
807    fa_v4_mode() != "0"
808} // DEFAULT ON 2026-07-10 (MEMRA_FA_V4=0 rollback)
809/// t_kv-conditional v4 pick (gemma depth lesson 2026-07-10: v4's key-per-lane pipeline starves
810/// at the 1024-window with short splits — MEMRA_FA_V4=0 measured depth plain 158.0 vs 156.7).
811/// Threshold MEMRA_FA_V4_MAX (default usize::MAX = unchanged behavior; gemma sets 1024 at load
812/// via FA_V4_MAX_DEFAULT). Applied at EVERY dispatch site (eager, rows, rows_w, dc) so verify
813/// stays kernel-family-identical to decode at the same t_kv.
814/// Per-model deep-ctx smem floor default (MEMRA_FA_SMEM_TKV env overrides): gemma pushes it
815/// above the 1024 window so the windowed decode + verify rows share the REGISTER family.
816pub static FA_SMEM_TKV_DEFAULT: std::sync::atomic::AtomicUsize =
817    std::sync::atomic::AtomicUsize::new(1024);
818pub static FA_V4_MAX_DEFAULT: std::sync::atomic::AtomicUsize =
819    std::sync::atomic::AtomicUsize::new(usize::MAX);
820pub fn fa_v4_at_pub(t_kv: usize) -> bool {
821    fa_v4_at(t_kv)
822}
823fn fa_v4_at(t_kv: usize) -> bool {
824    static M: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
825    let mx = *M.get_or_init(|| {
826        std::env::var("MEMRA_FA_V4_MAX")
827            .ok()
828            .and_then(|v| v.parse().ok())
829            .unwrap_or_else(|| FA_V4_MAX_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
830    });
831    fa_v4_on() && t_kv < mx
832}
833/// FA-DEEP gate (2026-08-02, lane fa-decode-deep): deep-ctx v4 twins
834/// (fa_decode_vec_q_v4_deep / _deep_dc) — the depth-decode lane's priced fix. Unlike
835/// v2/v3/v4 this is NOT a numeric config: the deep twins run the v4 program VERBATIM
836/// (same split partition, same softmax/accumulation order, same partials/combine) and only
837/// move the smem physical layout (bank de-conflict row pads) + the load schedule (next-tile
838/// L2 prefetch) — kernel-check pins bitdiff==0 vs the v4 twins across depths, so eager /
839/// rows-verify / graph / seqs stay mutually bit-identical wherever the threshold falls.
840/// Engages at t_kv >= MEMRA_FA_DEEP_MIN. The swept floor is 0 = ALWAYS ON where v4 ran
841/// (fa-deep-bench fine grid 96..6144, 2026-08-02: deep flat-or-better at EVERY depth,
842/// 1.01-1.26x, no losing cell — so there is no engagement boundary and no new
843/// capture-recapture edge; the env stays as a sweep/diagnostic seam only).
844/// MEMRA_FA_DEEP=0 is the rollback seam. Read per call so the battery + bench can A/B
845/// within one process (the v2/v3 pattern).
846pub const FA_DEEP_MIN_DEFAULT: usize = 0;
847fn fa_deep_at(t_kv: usize) -> bool {
848    if std::env::var("MEMRA_FA_DEEP").as_deref() == Ok("0") {
849        return false;
850    }
851    let min = std::env::var("MEMRA_FA_DEEP_MIN")
852        .ok()
853        .and_then(|v| v.parse().ok())
854        .unwrap_or(FA_DEEP_MIN_DEFAULT);
855    t_kv >= min
856}
857/// Public twin (kernel-check builds the deep-vs-v4 bit pin; bench sweeps the floor).
858pub fn fa_deep_at_pub(t_kv: usize) -> bool {
859    fa_deep_at(t_kv)
860}
861
862fn fa_v3_active(head_dim: usize) -> bool {
863    // v3's dp4a-K walk reads raw q8_0 K bytes — no e4m3 arm; the fp8-KV arm (MEMRA_KV_FP8)
864    // must fall back like any non-default KV format (the rows_dc stream path asserts on it).
865    fa_v3_on()
866        && head_dim % 128 == 0
867        && kv_cache_formats() == ("q8_0", "q5_1")
868        && !Engine::kv_fp8_on()
869}
870
871/// BATCHED-TICK increment 2 (2026-08-01): true iff a row at this t_kv would take the v4
872/// eager arm in `fa_decode_kvmod`'s dispatch — the exact precondition for the z-batched
873/// `fa_decode_vec_q_seqs_v4` twin to reproduce its per-seq program bit-identically.
874/// Mirrors the kvmod predicates: vec on + above the vec floor + hd256 + inside the v4
875/// window + the PRODUCTION v4 body (the noB3/stage phase probes are wrong-output) + the
876/// default flash module (no fp8-KV g-module). Callers must ALSO group rows on one
877/// `fa_split_keys` rung (the rows-twins' straddle law) before batching.
878pub fn fa_seqs_eligible(t_kv: usize, head_dim: usize) -> bool {
879    std::env::var("MEMRA_NO_FA_VEC").is_err()
880        && t_kv >= fa_vec_min_tkv()
881        && head_dim == 256
882        && fa_v4_at(t_kv)
883        && !matches!(fa_v4_mode(), "noB3" | "stage")
884        && !Engine::kv_fp8_on()
885}
886/// Public twin of the crate-private split ladder (kernel-check builds the seqs-vs-loop pin).
887pub fn fa_split_keys_pub(t_kv: usize, n_head_kv: usize) -> usize {
888    fa_split_keys(t_kv, n_head_kv)
889}
890
891/// A raw pinned (page-locked, CACHEABLE — flags=0, not write-combined) host allocation for
892/// DtoH staging. cudarc's `alloc_pinned` uses CU_MEMHOSTALLOC_WRITECOMBINED, which is right for
893/// HtoD streams but pathologically slow for host READS — the router readback is host-read-heavy,
894/// so we allocate through `result::malloc_host` with flags=0 directly.
895struct PinnedStage {
896    ptr: *mut u8,
897    cap: usize,
898}
899unsafe impl Send for PinnedStage {}
900impl PinnedStage {
901    fn new(cap: usize) -> Result<Self, Box<dyn std::error::Error>> {
902        let ptr = unsafe { cudarc::driver::result::malloc_host(cap, 0)? } as *mut u8;
903        Ok(PinnedStage { ptr, cap })
904    }
905}
906impl Drop for PinnedStage {
907    fn drop(&mut self) {
908        let _ = unsafe { cudarc::driver::result::free_host(self.ptr as _) };
909    }
910}
911
912/// Number of pass-1 blocks for the parallel argmax (fan-out across SMs to saturate HBM). 256 blocks
913/// x 256 threads = 65536 threads covering the 248K-vocab scan in ~4 strided loads/thread.
914pub const ARGMAX_NB: usize = 256;
915
916/// crate-visible alias for the batched FA3 shim entry (hybrid_forward's batch arm).
917pub(crate) use memra_fa3_vl as fa3_vl_raw;
918
919unsafe extern "C" {
920    /// FA3 v10 shim (cu/fa3_prefill.cu): TMA-swizzled wgmma FA, fresh causal hd256.
921    fn memra_fa3_prefill(
922        q16: *const core::ffi::c_void,
923        k16: *const core::ffi::c_void,
924        v16: *const core::ffi::c_void,
925        o: *mut f32,
926        t: i32,
927        h: i32,
928        hkv: i32,
929        d: i32,
930        scale: f32,
931        stream: *mut core::ffi::c_void,
932    ) -> i32;
933    /// batched varlen twin: host arrays of device pointers per seq (B <= 8).
934    pub(crate) fn memra_fa3_vl(
935        q16s: *const *const core::ffi::c_void,
936        k16s: *const *const core::ffi::c_void,
937        v16s: *const *const core::ffi::c_void,
938        os: *const *mut f32,
939        ts: *const i32,
940        b: i32,
941        h: i32,
942        hkv: i32,
943        d: i32,
944        scale: f32,
945        stream: *mut core::ffi::c_void,
946    ) -> i32;
947}
948
949/// STAGE-2 GROUPED DECODE: 8 expert weight-block device pointers passed BY VALUE as one kernel
950/// param (matches the CUDA `wptr8_t` struct: 8x 64-bit pointers, `#[repr(C)]` => identical
951/// layout). The pointers are SLRU cache-slot base addresses — fixed for the engine's lifetime
952/// (slots are never re-allocated), so passing raw values is stable across the launch.
953#[repr(C)]
954#[derive(Clone, Copy)]
955pub struct WPtr8(pub [u64; 8]);
956unsafe impl cudarc::driver::DeviceRepr for WPtr8 {}
957
958/// task #18 varlen GDN: per-seq args for gdn_chunk_{state,output}_mma_vl — one launch
959/// runs all B<=8 sequences' K4/K5 (CUDA `gdnseq_t`/`gdnvl_t`, layout-identical repr(C)).
960/// Raw addresses are valid for the launch: every referenced buffer outlives the call and
961/// all work is on the single compute stream (same discipline as the f16 GEMM FFI).
962#[repr(C)]
963#[derive(Clone, Copy, Default)]
964pub struct GdnSeqVl {
965    pub kb16: u64,
966    pub gcum: u64,
967    pub beta: u64,
968    pub u: u64,
969    pub wb16: u64,
970    pub y: u64,
971    pub ssnap: u64,
972    pub state_in: u64,
973    pub state_out: u64,
974    pub q: u64,
975    pub p: u64,
976    pub o: u64,
977    pub k: u64,
978    pub v: u64,
979    pub g: u64,
980    pub a: u64,
981    pub w: u64,
982    pub t: i32,
983    pub nc: i32,
984}
985unsafe impl cudarc::driver::DeviceRepr for GdnSeqVl {}
986#[repr(C)]
987#[derive(Clone, Copy)]
988pub struct GdnVl8(pub [GdnSeqVl; 8]);
989unsafe impl cudarc::driver::DeviceRepr for GdnVl8 {}
990
991/// task #22: per-seq wgmma-fused extras (CUDA `gdnw_t`/`gdnwvl_t`) — qb16 mirror +
992/// pre-masked Pb16, riding NEXT TO GdnSeqVl so the base struct stays untouched.
993#[repr(C)]
994#[derive(Clone, Copy, Default)]
995pub struct GdnWVl {
996    pub qb16: u64,
997    pub pb16: u64,
998}
999unsafe impl cudarc::driver::DeviceRepr for GdnWVl {}
1000#[repr(C)]
1001#[derive(Clone, Copy)]
1002pub struct GdnWVl8(pub [GdnWVl; 8]);
1003unsafe impl cudarc::driver::DeviceRepr for GdnWVl8 {}
1004
1005/// task #18 increment 3: per-seq PREP/TAIL args (CUDA `gdnprep_t`/`gdnprepvl_t`).
1006#[repr(C)]
1007#[derive(Clone, Copy, Default)]
1008pub struct GdnPrepVl {
1009    pub qkv: u64,
1010    pub conv_state: u64,
1011    pub conv_out: u64,
1012    pub q_g: u64,
1013    pub k_g: u64,
1014    pub v_g: u64,
1015    pub q_l2: u64,
1016    pub k_l2: u64,
1017    pub beta_raw: u64,
1018    pub alpha: u64,
1019    pub beta: u64,
1020    pub g_log: u64,
1021    pub o: u64,
1022    pub z: u64,
1023    pub gn: u64,
1024    pub gn16: u64,
1025    pub kb16: u64,
1026    pub qb16: u64,
1027    pub t: i32,
1028    pub pad: i32,
1029}
1030unsafe impl cudarc::driver::DeviceRepr for GdnPrepVl {}
1031#[repr(C)]
1032#[derive(Clone, Copy)]
1033pub struct GdnPrepVl8(pub [GdnPrepVl; 8]);
1034unsafe impl cudarc::driver::DeviceRepr for GdnPrepVl8 {}
1035
1036/// task #18 (attn side): per-seq varlen FA args (CUDA `faseq_t`/`favl_t`).
1037#[repr(C)]
1038#[derive(Clone, Copy, Default)]
1039pub struct FaSeqVl {
1040    pub q: u64,
1041    pub k16: u64,
1042    pub v16: u64,
1043    pub o: u64,
1044    pub kf: u64,
1045    pub vf: u64,
1046    pub t: i32,
1047    pub pad: i32,
1048}
1049unsafe impl cudarc::driver::DeviceRepr for FaSeqVl {}
1050#[repr(C)]
1051#[derive(Clone, Copy)]
1052pub struct FaVl8(pub [FaSeqVl; 8]);
1053unsafe impl cudarc::driver::DeviceRepr for FaVl8 {}
1054
1055/// task #18 (attn pre-FA): per-seq split/norm/rope/append args (CUDA `attnpre_t`).
1056#[repr(C)]
1057#[derive(Clone, Copy, Default)]
1058pub struct AttnPreVl {
1059    pub qf: u64,
1060    pub kf: u64,
1061    pub vf: u64,
1062    pub q: u64,
1063    pub gate: u64,
1064    pub qn: u64,
1065    pub kn: u64,
1066    pub kc: u64,
1067    pub vc: u64,
1068    pub t: i32,
1069    pub pad: i32,
1070}
1071unsafe impl cudarc::driver::DeviceRepr for AttnPreVl {}
1072#[repr(C)]
1073#[derive(Clone, Copy)]
1074pub struct AttnPreVl8(pub [AttnPreVl; 8]);
1075unsafe impl cudarc::driver::DeviceRepr for AttnPreVl8 {}
1076
1077/// task #18 increment 2: one sequence's FULL chunk-buffer set (alloc-only; the
1078/// varlen K1-K5 chain fills them).
1079pub struct GdnChunkBufs {
1080    pub gcum: CudaSlice<f32>,
1081    pub a: CudaSlice<f32>,
1082    pub p: CudaSlice<f32>,
1083    pub u: CudaSlice<f32>,
1084    pub w: CudaSlice<f32>,
1085    pub kb16: CudaSlice<u8>,
1086    pub wb16: CudaSlice<u8>,
1087    pub y16: CudaSlice<u8>,
1088    pub ssnap16: CudaSlice<u8>,
1089    pub qb16: CudaSlice<u8>,
1090    pub pb16: CudaSlice<u8>,
1091    pub o: CudaSlice<f32>,
1092    pub t: usize,
1093    pub nc: usize,
1094}
1095
1096/// STAGE-2 GROUPED DECODE: the 8 routed-expert weights by value (CUDA `f32x8_t`).
1097#[repr(C)]
1098#[derive(Clone, Copy)]
1099pub struct F32x8(pub [f32; 8]);
1100unsafe impl cudarc::driver::DeviceRepr for F32x8 {}
1101
1102/// Harness timing contract: wall nanos of the LAST generate/generate_spec prompt prime on this
1103/// process. Bench binaries read it right after the call to print gen-only throughput without the
1104/// prime-subtraction hack (which amplifies prime jitter into the gen number at long prompts).
1105pub static PRIME_NANOS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1106
1107impl Engine {
1108    pub fn new(ordinal: usize) -> Result<Self, Box<dyn std::error::Error>> {
1109        let gpu = memra_runtime::Gpu::new(ordinal)?;
1110        // ARCH GUARD (unified dual-arch engine): the fatbins carry single-arch SASS, so a
1111        // binary/device mismatch otherwise dies at first module load with an opaque CUDA
1112        // error. Fail early with the rebuild hint instead. MEMRA_ARCH_CHECK=0 skips.
1113        if std::env::var("MEMRA_ARCH_CHECK").as_deref() != Ok("0") {
1114            use cudarc::driver::sys::CUdevice_attribute_enum as A;
1115            let (maj, min) = cudarc::driver::result::device::get(ordinal as i32)
1116                .and_then(|d| unsafe {
1117                    Ok((
1118                        cudarc::driver::result::device::get_attribute(
1119                            d,
1120                            A::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR,
1121                        )?,
1122                        cudarc::driver::result::device::get_attribute(
1123                            d,
1124                            A::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR,
1125                        )?,
1126                    ))
1127                })
1128                .unwrap_or((0, 0));
1129            let built = env!("MEMRA_BUILT_CUDA_ARCH");
1130            let ok = matches!(
1131                (built, maj, min),
1132                ("120a", 12, 0) | ("120a", 12, 1) | ("100a", 10, 0) | ("90a", 9, 0) | ("89", 8, 9)
1133            );
1134            if !ok {
1135                return Err(format!(
1136                    "memra was built for sm_{built} but device {ordinal} reports compute \
1137                     capability {maj}.{min}. Rebuild on this machine (MEMRA_CUDA_ARCH \
1138                     auto-detects the GPU) or set MEMRA_ARCH_CHECK=0 to bypass."
1139                )
1140                .into());
1141            }
1142        }
1143        // Default async-pool RELEASE_THRESHOLD is 0: freed blocks return to the OS at every
1144        // sync, so cuMemAllocAsync NODES inside captured graphs re-map memory on EVERY
1145        // cuGraphLaunch (measured 226us/launch on the gemma graph door, 2026-07-23 osrt).
1146        // Pinning the threshold keeps the pool cached -> alloc nodes become pointer bumps.
1147        unsafe {
1148            use cudarc::driver::sys;
1149            let dev: sys::CUdevice = ordinal as sys::CUdevice;
1150            let mut pool: sys::CUmemoryPool = std::ptr::null_mut();
1151            if sys::cuDeviceGetDefaultMemPool(&mut pool, dev) == sys::CUresult::CUDA_SUCCESS {
1152                let mut thresh: u64 = u64::MAX;
1153                let _ = sys::cuMemPoolSetAttribute(
1154                    pool,
1155                    sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_RELEASE_THRESHOLD,
1156                    &mut thresh as *mut u64 as *mut core::ffi::c_void,
1157                );
1158            }
1159        }
1160        let module = gpu.ctx.load_module(Ptx::from_binary(FATBIN.to_vec()))?;
1161        let hybrid = gpu
1162            .ctx
1163            .load_module(Ptx::from_binary(HYBRID_FATBIN.to_vec()))?;
1164        let qmatvec = gpu
1165            .ctx
1166            .load_module(Ptx::from_binary(QMATVEC_FATBIN.to_vec()))?;
1167        let flash = gpu
1168            .ctx
1169            .load_module(Ptx::from_binary(flash_fatbin_bytes().to_vec()))?;
1170        let gemm = gpu
1171            .ctx
1172            .load_module(Ptx::from_binary(gemm_fatbin_bytes().into_owned()))?;
1173        let router = gpu
1174            .ctx
1175            .load_module(Ptx::from_binary(ROUTER_FATBIN.to_vec()))?;
1176        let sample = gpu
1177            .ctx
1178            .load_module(Ptx::from_binary(SAMPLE_FATBIN.to_vec()))?;
1179        let copy_stream = gpu.ctx.new_stream()?;
1180        // DECODE EVENT-TRACKING ELISION — DEFAULT ON (2026-07-05; MEMRA_EVT=1 = escape hatch).
1181        // cudarc is in multi-stream mode (main stream +
1182        // copy_stream are both created streams), so with tracking on EVERY launch arg records a
1183        // read/write CudaEvent and inserts cuStreamWaitEvent on prior events. On the 35B MoE decode
1184        // that is ~19k cuStreamWaitEvent + ~9k cuEventRecord + ~6k event create/destroy per token
1185        // (~7 ms/tok host time, measured nsys 2026-07-04 g7e), and +4.6% measured on 27B decode —
1186        // protecting NOTHING: every hot-path kernel/memcpy runs on the ONE gpu.stream.
1187        // CROSS-STREAM HAZARD AUDIT: MoeSlotCache in-memory prefetch uses copy_stream. Every
1188        // overwrite explicitly records the prior compute point and makes copy_stream wait; every
1189        // consumer explicitly waits for the copy completion event. The opt-in positioned-read
1190        // proof stays on gpu.stream and retains an explicit event solely to guard pinned-source
1191        // reuse. Graph-capture sites use only gpu.stream, so these handoffs never rely on cudarc's
1192        // implicit event tracking.
1193        // SAFETY: single-stream ordering is total; the runtime mem-pool is configured with
1194        // internal-dependency reuse (memra-runtime), so alloc reuse is stream-ordered too.
1195        if std::env::var("MEMRA_EVT")
1196            .map(|v| v == "1")
1197            .unwrap_or(false)
1198        {
1199            // escape hatch: keep cudarc's implicit cross-stream event tracking.
1200        } else {
1201            unsafe {
1202                gpu.ctx.disable_event_tracking();
1203            }
1204        }
1205        Ok(Self {
1206            gpu,
1207            module,
1208            hybrid,
1209            qmatvec,
1210            flash,
1211            flash_g: std::sync::OnceLock::new(),
1212            gemm,
1213            router,
1214            sample,
1215            moe_cache: Mutex::new(None),
1216            moe_cache_layout: Mutex::new(None),
1217            copy_stream,
1218            capture_keep_on: std::sync::atomic::AtomicBool::new(false),
1219            verify_exact: std::sync::atomic::AtomicBool::new(false),
1220            capture_keep: Mutex::new(Vec::new()),
1221            argmax_partials: Mutex::new(None),
1222            prime_deqw_ws: Mutex::new(None),
1223            router_stage: Mutex::new(None),
1224            fp8_scratch: Mutex::new(None),
1225            fa_vf16_scratch: Mutex::new(None),
1226            fa_part_pool: Mutex::new(None),
1227            fa_part_retired: Mutex::new(Vec::new()),
1228            fn_cache: Mutex::new(Default::default()),
1229            f16_scratch: Mutex::new(None),
1230            #[cfg(memra_cutlass)]
1231            cutlass_scratch: Mutex::new(None),
1232        })
1233    }
1234
1235    pub fn ctx(&self) -> &Arc<CudaContext> {
1236        &self.gpu.ctx
1237    }
1238
1239    /// Bytes the async pool holds MAPPED but NOT LIVE (reserved - used), i.e. freed blocks
1240    /// parked in the pool because `Engine::new` pins RELEASE_THRESHOLD to u64::MAX above.
1241    ///
1242    /// Why this is a public engine surface: `mem_get_info`'s `free` DOES NOT SEE these bytes —
1243    /// they are mapped to this process, so `free` counts them as gone, yet the very next
1244    /// `alloc_u8` is satisfied from them without touching `free` at all. Any admission or
1245    /// budget decision that reads `free` alone therefore under-counts real headroom by exactly
1246    /// this amount. Effective allocatable headroom is `free + pool_cached_bytes()`.
1247    ///
1248    /// MEASURED SIZE (c=64 serve burst, 9B NVFP4 + draft, 24GB card, 2026-08-06): 34-89 MB
1249    /// during the burst — SMALL. The admission gate adds it because a term that can only ever
1250    /// under-count headroom does not belong in a gate that queues real work, but the honest
1251    /// reading of this number is that pool caching is NOT where a long-running server's VRAM
1252    /// hides on this path: reserved ~= used throughout, so the memory the driver reports as
1253    /// gone is genuinely LIVE (see `pool_reserved_used` for the diagnostic pair).
1254    ///
1255    /// Returns 0 if the pool cannot be queried (never a false-positive headroom claim).
1256    pub fn pool_cached_bytes(&self) -> usize {
1257        let (reserved, used) = self.pool_reserved_used();
1258        reserved.saturating_sub(used)
1259    }
1260
1261    /// Raw async-pool occupancy: (RESERVED_MEM_CURRENT, USED_MEM_CURRENT) in bytes. Reserved is
1262    /// what the pool has mapped from the driver; used is what is live inside it. Exposed for
1263    /// admission/VRAM diagnostics — the pair distinguishes "memory is parked in the pool and
1264    /// `free` cannot see it" (reserved >> used) from "memory is genuinely held live by some
1265    /// owner" (reserved ~= used), which are opposite bugs with opposite fixes.
1266    /// (0, 0) if the pool cannot be queried.
1267    pub fn pool_reserved_used(&self) -> (usize, usize) {
1268        use cudarc::driver::sys;
1269        unsafe {
1270            let mut pool: sys::CUmemoryPool = std::ptr::null_mut();
1271            if sys::cuDeviceGetDefaultMemPool(&mut pool, self.gpu.ctx.ordinal() as sys::CUdevice)
1272                != sys::CUresult::CUDA_SUCCESS
1273            {
1274                return (0, 0);
1275            }
1276            let (mut reserved, mut used) = (0u64, 0u64);
1277            if sys::cuMemPoolGetAttribute(
1278                pool,
1279                sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_RESERVED_MEM_CURRENT,
1280                &mut reserved as *mut u64 as *mut core::ffi::c_void,
1281            ) != sys::CUresult::CUDA_SUCCESS
1282            {
1283                return (0, 0);
1284            }
1285            if sys::cuMemPoolGetAttribute(
1286                pool,
1287                sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_USED_MEM_CURRENT,
1288                &mut used as *mut u64 as *mut core::ffi::c_void,
1289            ) != sys::CUresult::CUDA_SUCCESS
1290            {
1291                return (0, 0);
1292            }
1293            (reserved as usize, used as usize)
1294        }
1295    }
1296
1297    /// Ambient stream (by value since M1-PP2 increment 2): the thread's pp2 stage stream
1298    /// when a stage scope is active, else the main compute stream — see `Gpu::stream`.
1299    pub fn stream(&self) -> Arc<CudaStream> {
1300        self.gpu.stream()
1301    }
1302    /// FP8-GLOBALS switch (MEMRA_GEMMA_GKV, default ON): gemma global (hd512) layers keep
1303    /// their KV in e4m3 — the dequant-latency arc (HANDOVER). Windowed layers stay q8_0/q5_1.
1304    pub fn gkv_on() -> bool {
1305        memra_kv::gkv_on()
1306    }
1307
1308    /// FP8-WINDOWED switch (MEMRA_GEMMA_WKV — measured 2026-07-12 in a validity-gated
1309    /// window: 1.7k 174.1-174.4 vs 168.6-169.4 default (+3%), 4.9k 158.7-160.4; vs llama
1310    /// same-window 159.5-160.2 / 140.6 = 1.09x / 1.13x): gemma windowed (hd256 SWA)
1311    /// layers hold e4m3 KV and ride the format-aware v4 lane from the kf8vf8 module.
1312    /// SERVING-MODE DEFAULT (2026-07-12, the 31B spec unlock): fp8-windowed KV GUTS the
1313    /// MTP drafter's acceptance — its single swa attention reads the windowed cache and
1314    /// e4m3 noise flips its argmaxes (31B short accept .758 -> 1.000 with q8/q5, spec 88
1315    /// -> 122.7 vs llama-mtp 112; depth .59 -> .78; 26B depth .57 -> .89). So the default
1316    /// keys on serving intent: SPEC serving (MEMRA_DRAFT set) -> OFF, plain -> ON (its
1317    /// depth-plain +3% stands). Explicit MEMRA_GEMMA_WKV always wins. GKV (globals) stays
1318    /// ON for both — no acceptance cost measured.
1319    pub fn wkv_on() -> bool {
1320        memra_kv::wkv_on()
1321    }
1322
1323    /// QWEN FP8-KV switch (MEMRA_KV_FP8 explicit; else the per-model KV_FP8_FORCE door set
1324    /// at model load; else OFF). Non-gemma full-attn layers hold e4m3 K/V via the kf8vf8
1325    /// module. Per-model verdict 2026-07-12: 9B +0.7-4% scaling with depth, 27B flat,
1326    /// 35B −2% (fp8 format-gates its v3 dp4a lane) — so the 9B class defaults ON
1327    /// (adopted 2026-07-28 with the deferred acceptance battery), others stay OFF.
1328    pub fn kv_fp8_on() -> bool {
1329        memra_kv::kv_fp8_on()
1330    }
1331
1332    /// fa kernel routed by head_dim: hd512 (gemma globals) resolves from the kf8vf8 module
1333    /// when the fp8-globals arm is on; everything else from the default flash module.
1334    fn fa_func(&self, name: &str, head_dim: usize) -> CudaFunction {
1335        if head_dim == 512 && Self::gkv_on() {
1336            self.func_g(name)
1337        } else {
1338            self.func(name)
1339        }
1340    }
1341
1342    /// Kernel from the FP8-GLOBALS (kf8vf8) flash module — gemma global-layer arm only.
1343    /// Format-AGNOSTIC kernels (e.g. fa_decode_combine_f32) are not compiled into the
1344    /// per-format fatbins; fall back to the base modules for those.
1345    fn func_g(&self, name: &str) -> CudaFunction {
1346        let m = self.flash_g.get_or_init(|| {
1347            self.gpu
1348                .ctx
1349                .load_module(cudarc::nvrtc::Ptx::from_binary(
1350                    FLASH_FATBIN_KF8VF8.to_vec(),
1351                ))
1352                .expect("load kf8vf8 flash fatbin (fp8-globals arm)")
1353        });
1354        let key = format!("g:{name}");
1355        if let Some(f) = self.fn_cache.lock().unwrap().get(&key) {
1356            return f.clone();
1357        }
1358        let f = match m.load_function(name) {
1359            Ok(f) => f,
1360            Err(_) => self.func(name),
1361        };
1362        self.fn_cache.lock().unwrap().insert(key, f.clone());
1363        f
1364    }
1365
1366    fn func(&self, name: &str) -> CudaFunction {
1367        // Resolution cache: cuModuleGetFunction fails inside a CUDA-graph capture region,
1368        // so capture-time lookups MUST be host-memory hits (warmups populate the cache).
1369        if let Some(f) = self.fn_cache.lock().unwrap().get(name) {
1370            return f.clone();
1371        }
1372        let f = self
1373            .module
1374            .load_function(name)
1375            .or_else(|_| self.hybrid.load_function(name))
1376            .or_else(|_| self.qmatvec.load_function(name))
1377            .or_else(|_| self.flash.load_function(name))
1378            .or_else(|_| self.gemm.load_function(name))
1379            .or_else(|_| self.router.load_function(name))
1380            .or_else(|_| self.sample.load_function(name))
1381            .unwrap_or_else(|_| panic!("kernel {name} not in any fatbin"));
1382        self.fn_cache
1383            .lock()
1384            .unwrap()
1385            .insert(name.to_string(), f.clone());
1386        f
1387    }
1388
1389    /// Scatter trimmed draft logits into full-vocab space: dst = -inf everywhere, then
1390    /// dst[d2t[i]] = src[i]. Two launches (fill, scatter) — no grid-wide sync needed.
1391    pub fn scatter_trim_logits(
1392        &self,
1393        src: &CudaSlice<f32>,
1394        d2t: &CudaSlice<u32>,
1395        dst: &mut CudaSlice<f32>,
1396        d_vocab: usize,
1397        n_vocab: usize,
1398    ) -> Result<(), Box<dyn std::error::Error>> {
1399        let f1 = self.func("scatter_trim_logits_f32");
1400        let f2 = self.func("scatter_trim_logits_pass2_f32");
1401        let (dv, nv) = (d_vocab as i32, n_vocab as i32);
1402        let cfg1 = LaunchConfig {
1403            grid_dim: (256, 1, 1),
1404            block_dim: (256, 1, 1),
1405            shared_mem_bytes: 0,
1406        };
1407        let __s_b1 = self.gpu.stream();
1408        let mut b1 = __s_b1.launch_builder(&f1);
1409        b1.arg(src).arg(d2t).arg(&mut *dst).arg(&dv).arg(&nv);
1410        unsafe {
1411            b1.launch(cfg1)?;
1412        }
1413        let cfg2 = LaunchConfig {
1414            grid_dim: (d_vocab.div_ceil(256) as u32, 1, 1),
1415            block_dim: (256, 1, 1),
1416            shared_mem_bytes: 0,
1417        };
1418        let __s_b2 = self.gpu.stream();
1419        let mut b2 = __s_b2.launch_builder(&f2);
1420        b2.arg(src).arg(d2t).arg(&mut *dst).arg(&dv);
1421        unsafe {
1422            b2.launch(cfg2)?;
1423        }
1424        Ok(())
1425    }
1426
1427    // ---- FILTERED-SPEC (feat/filtered-spec): top-k/p/min-p transforms applied symmetrically
1428    // to p and q — rejection sampling stays distribution-exact for the filtered target. ----
1429
1430    /// Per-row filtered-softmax stats: out[r] = (threshold_e, renorm_mass_e, row_max) for the
1431    /// filter (top_k, top_p, min_p) at `temp`. Rows index into x with row_stride f32s.
1432    #[allow(clippy::too_many_arguments)]
1433    pub fn filter_stats(
1434        &self,
1435        x: &CudaSlice<f32>,
1436        row_stride: usize,
1437        rows: &CudaSlice<i32>,
1438        out_th: &mut CudaSlice<f32>,
1439        out_z: &mut CudaSlice<f32>,
1440        out_max: &mut CudaSlice<f32>,
1441        n: usize,
1442        nrow: usize,
1443        temp: f32,
1444        top_k: i32,
1445        top_p: f32,
1446        min_p: f32,
1447    ) -> Result<(), Box<dyn std::error::Error>> {
1448        // A top-K-selection form of this kernel (3 vocab passes vs the search's ~51) was
1449        // implemented and REFUTED on 2026-08-21 (lane/moebatch-q35moe): the 248k-vocab row is
1450        // L2-resident, so the extra passes are near-free while the per-thread selection list
1451        // spills to local memory — B=8 tick 12.8/11.2 ms (cap 64/32) vs 10.4 ms for this
1452        // kernel, and serve c8 agg ~648 vs ~666. The receipts row is the record; the real
1453        // filtered-sampling win was batching the per-row launches (decode_batch.rs).
1454        //
1455        // COOPERATIVE MULTI-BLOCK FORM (lane/samplat, 2026-08-21): the surviving inefficiency
1456        // was WIDTH, not passes — one block per row left ~94% of the device idle for ~620us
1457        // per B=8 serve tick (5.9% of the tick, box4 nsys receipt). filter_stats_coop_f32
1458        // splits each row across 16 blocks with grid-synced bisection totals — same algorithm,
1459        // slice-partial f32 sums (accepted device-sampling class; sample-check arbitrates).
1460        // Admission: cooperative grid must co-reside (16*nrow blocks vs SM count).
1461        // MEMRA_FILTER_COOP=0 is the rollback seam to the single-block form.
1462        static COOP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1463        let coop_on =
1464            *COOP_ON.get_or_init(|| std::env::var("MEMRA_FILTER_COOP").as_deref() != Ok("0"));
1465        let (ni, nr, rs) = (n as i32, nrow as i32, row_stride as i64);
1466        if coop_on && 16 * nrow <= self.sm_count() as usize {
1467            let f = self.func("filter_stats_coop_f32");
1468            let mut ws = self.alloc_uninit::<f32>(nrow * (2 * 16 + 2))?;
1469            let cfg = LaunchConfig {
1470                grid_dim: (16, nrow as u32, 1),
1471                block_dim: (512, 1, 1),
1472                shared_mem_bytes: 0,
1473            };
1474            let __s_b = self.gpu.stream();
1475            let mut b = __s_b.launch_builder(&f);
1476            b.arg(x)
1477                .arg(&rs)
1478                .arg(rows)
1479                .arg(&mut *out_th)
1480                .arg(&mut *out_z)
1481                .arg(&mut *out_max)
1482                .arg(&mut ws)
1483                .arg(&ni)
1484                .arg(&nr)
1485                .arg(&temp)
1486                .arg(&top_k)
1487                .arg(&top_p)
1488                .arg(&min_p);
1489            unsafe {
1490                b.launch_cooperative(cfg)?;
1491            }
1492            return Ok(());
1493        }
1494        let f = self.func("filter_stats_f32");
1495        let cfg = LaunchConfig {
1496            grid_dim: (nrow as u32, 1, 1),
1497            block_dim: (1024, 1, 1),
1498            shared_mem_bytes: 0,
1499        };
1500        let __s_b = self.gpu.stream();
1501        let mut b = __s_b.launch_builder(&f);
1502        b.arg(x)
1503            .arg(&rs)
1504            .arg(rows)
1505            .arg(&mut *out_th)
1506            .arg(&mut *out_z)
1507            .arg(&mut *out_max)
1508            .arg(&ni)
1509            .arg(&nr)
1510            .arg(&temp)
1511            .arg(&top_k)
1512            .arg(&top_p)
1513            .arg(&min_p);
1514        unsafe {
1515            b.launch(cfg)?;
1516        }
1517        Ok(())
1518    }
1519
1520    /// out[pair] = filtered-softmax prob of ids[pair] in row rows[pair] (th/z per PAIR).
1521    #[allow(clippy::too_many_arguments)]
1522    pub fn softmax_gather_filtered(
1523        &self,
1524        x: &CudaSlice<f32>,
1525        row_stride: usize,
1526        ids: &CudaSlice<u32>,
1527        rows: &CudaSlice<i32>,
1528        th: &CudaSlice<f32>,
1529        z: &CudaSlice<f32>,
1530        out: &mut CudaSlice<f32>,
1531        n: usize,
1532        npair: usize,
1533        temp: f32,
1534    ) -> Result<(), Box<dyn std::error::Error>> {
1535        let f = self.func("softmax_gather_filtered_f32");
1536        let (ni, np, rs) = (n as i32, npair as i32, row_stride as i64);
1537        let cfg = LaunchConfig {
1538            grid_dim: (npair as u32, 1, 1),
1539            block_dim: (256, 1, 1),
1540            shared_mem_bytes: 0,
1541        };
1542        let __s_b = self.gpu.stream();
1543        let mut b = __s_b.launch_builder(&f);
1544        b.arg(x)
1545            .arg(&rs)
1546            .arg(ids)
1547            .arg(rows)
1548            .arg(th)
1549            .arg(z)
1550            .arg(&mut *out)
1551            .arg(&ni)
1552            .arg(&np)
1553            .arg(&temp);
1554        unsafe {
1555            b.launch(cfg)?;
1556        }
1557        Ok(())
1558    }
1559
1560    /// Filtered residual sample: token ~ norm(max(0, fp - fq)) with fp/fq the filtered softmaxes.
1561    #[allow(clippy::too_many_arguments)]
1562    pub fn residual_sample_filtered(
1563        &self,
1564        p: &CudaSlice<f32>,
1565        q: Option<&CudaSlice<f32>>,
1566        n: usize,
1567        temp: f32,
1568        seed: u64,
1569        stream_pos: u32,
1570        p_stats: (f32, f32, f32),
1571        q_stats: (f32, f32, f32),
1572        out_tok: &mut CudaSlice<u32>,
1573    ) -> Result<(), Box<dyn std::error::Error>> {
1574        let f = self.func("residual_sample_filtered_f32");
1575        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
1576        let has_q: i32 = q.is_some() as i32;
1577        let qbuf = q.unwrap_or(p);
1578        let (pm, pth, pz) = p_stats;
1579        let (qm, qth, qz) = q_stats;
1580        let cfg = LaunchConfig {
1581            grid_dim: (1, 1, 1),
1582            block_dim: (1024, 1, 1),
1583            shared_mem_bytes: 0,
1584        };
1585        let __s_b = self.gpu.stream();
1586        let mut b = __s_b.launch_builder(&f);
1587        b.arg(p)
1588            .arg(qbuf)
1589            .arg(&has_q)
1590            .arg(&ni)
1591            .arg(&temp)
1592            .arg(&slo)
1593            .arg(&shi)
1594            .arg(&stream_pos)
1595            .arg(&pm)
1596            .arg(&pth)
1597            .arg(&pz)
1598            .arg(&qm)
1599            .arg(&qth)
1600            .arg(&qz)
1601            .arg(&mut *out_tok);
1602        unsafe {
1603            b.launch(cfg)?;
1604        }
1605        Ok(())
1606    }
1607
1608    /// Sparse-q filtered residual sample: token ~ norm(max(0, fp - q)) where fp is the
1609    /// FILTERED softmax of `p` (stats from `filter_stats`) and q is a PROBABILITY vector
1610    /// supported on `cand_ids` (<=32 ids — the DFlash2 selector's candidate-set proposal;
1611    /// lane/dspark-sampled-admission-20260820). Same event semantics/Philox tag as
1612    /// `residual_sample_filtered` — one uniform per (seed, stream_pos).
1613    #[allow(clippy::too_many_arguments)]
1614    pub fn residual_sample_sparse_q(
1615        &self,
1616        p: &CudaSlice<f32>,
1617        cand_ids: &CudaSlice<u32>,
1618        q_probs: &CudaSlice<f32>,
1619        n_cand: usize,
1620        n: usize,
1621        temp: f32,
1622        seed: u64,
1623        stream_pos: u32,
1624        p_stats: (f32, f32, f32),
1625        out_tok: &mut CudaSlice<u32>,
1626    ) -> Result<(), Box<dyn std::error::Error>> {
1627        assert!(
1628            n_cand >= 1 && n_cand <= 32,
1629            "residual_sample_sparse_q supports 1..=32 candidates, got {n_cand}"
1630        );
1631        let f = self.func("residual_sample_sparse_q_f32");
1632        let (ni, nc) = (n as i32, n_cand as i32);
1633        let (slo, shi) = ((seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
1634        let (pm, pth, pz) = p_stats;
1635        let cfg = LaunchConfig {
1636            grid_dim: (1, 1, 1),
1637            block_dim: (1024, 1, 1),
1638            shared_mem_bytes: 0,
1639        };
1640        let __s_b = self.gpu.stream();
1641        let mut b = __s_b.launch_builder(&f);
1642        b.arg(p)
1643            .arg(cand_ids)
1644            .arg(q_probs)
1645            .arg(&nc)
1646            .arg(&ni)
1647            .arg(&temp)
1648            .arg(&slo)
1649            .arg(&shi)
1650            .arg(&stream_pos)
1651            .arg(&pm)
1652            .arg(&pth)
1653            .arg(&pz)
1654            .arg(&mut *out_tok);
1655        unsafe {
1656            b.launch(cfg)?;
1657        }
1658        Ok(())
1659    }
1660
1661    /// Gumbel-max draw from the FILTERED distribution (masked perturb; argmax after).
1662    #[allow(clippy::too_many_arguments)]
1663    pub fn gumbel_perturb_filtered(
1664        &self,
1665        x: &CudaSlice<f32>,
1666        y: &mut CudaSlice<f32>,
1667        n: usize,
1668        seed: u64,
1669        stream_pos: u32,
1670        temp: f32,
1671        row_max: f32,
1672        th: f32,
1673    ) -> Result<(), Box<dyn std::error::Error>> {
1674        let f = self.func("gumbel_perturb_filtered_f32");
1675        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
1676        let cfg = LaunchConfig {
1677            grid_dim: (n.div_ceil(256) as u32, 1, 1),
1678            block_dim: (256, 1, 1),
1679            shared_mem_bytes: 0,
1680        };
1681        let __s_b = self.gpu.stream();
1682        let mut b = __s_b.launch_builder(&f);
1683        b.arg(x)
1684            .arg(&mut *y)
1685            .arg(&ni)
1686            .arg(&slo)
1687            .arg(&shi)
1688            .arg(&stream_pos)
1689            .arg(&temp)
1690            .arg(&row_max)
1691            .arg(&th);
1692        unsafe {
1693            b.launch(cfg)?;
1694        }
1695        Ok(())
1696    }
1697
1698    /// Keskar penalties applied IN PLACE to a logits buffer: history token ids get
1699    /// rep-divided/multiplied + freq*count + presence subtracted. Symmetric p/q usage keeps
1700    /// filtered rejection sampling exact for the penalized target.
1701    #[allow(clippy::too_many_arguments)]
1702    pub fn penalize_logits(
1703        &self,
1704        x: &mut CudaSlice<f32>,
1705        hist: &CudaSlice<u32>,
1706        n_hist: usize,
1707        rep: f32,
1708        freq: f32,
1709        present: f32,
1710        n: usize,
1711    ) -> Result<(), Box<dyn std::error::Error>> {
1712        if n_hist == 0 {
1713            return Ok(());
1714        }
1715        let f = self.func("penalize_logits_f32");
1716        let (nh, ni) = (n_hist as i32, n as i32);
1717        let cfg = LaunchConfig {
1718            grid_dim: (n_hist.div_ceil(128) as u32, 1, 1),
1719            block_dim: (128, 1, 1),
1720            shared_mem_bytes: 0,
1721        };
1722        let __s_b = self.gpu.stream();
1723        let mut b = __s_b.launch_builder(&f);
1724        b.arg(&mut *x)
1725            .arg(hist)
1726            .arg(&nh)
1727            .arg(&rep)
1728            .arg(&freq)
1729            .arg(&present)
1730            .arg(&ni);
1731        unsafe {
1732            b.launch(cfg)?;
1733        }
1734        Ok(())
1735    }
1736
1737    /// Rows variant: penalize `nrow` contiguous rows of length n in one launch.
1738    #[allow(clippy::too_many_arguments)]
1739    pub fn penalize_logits_rows(
1740        &self,
1741        x: &mut CudaSlice<f32>,
1742        hist: &CudaSlice<u32>,
1743        n_hist: usize,
1744        rep: f32,
1745        freq: f32,
1746        present: f32,
1747        n: usize,
1748        nrow: usize,
1749    ) -> Result<(), Box<dyn std::error::Error>> {
1750        if n_hist == 0 || nrow == 0 {
1751            return Ok(());
1752        }
1753        let f = self.func("penalize_logits_rows_f32");
1754        let (nh, ni, nr) = (n_hist as i32, n as i32, nrow as i32);
1755        let cfg = LaunchConfig {
1756            grid_dim: (n_hist.div_ceil(128) as u32, nrow as u32, 1),
1757            block_dim: (128, 1, 1),
1758            shared_mem_bytes: 0,
1759        };
1760        let __s_b = self.gpu.stream();
1761        let mut b = __s_b.launch_builder(&f);
1762        b.arg(&mut *x)
1763            .arg(hist)
1764            .arg(&nh)
1765            .arg(&rep)
1766            .arg(&freq)
1767            .arg(&present)
1768            .arg(&ni)
1769            .arg(&nr);
1770        unsafe {
1771            b.launch(cfg)?;
1772        }
1773        Ok(())
1774    }
1775
1776    /// ROW-INCREMENTAL penalties (dspark penalized-sampled admission): row r of `x`
1777    /// penalizes over the last `min(win, n_hist0 + r)` entries of `hist[..n_hist0 + r]`,
1778    /// where `hist` = [session window (n_hist0) ++ per-row drafted tokens (nrow-1)]. This
1779    /// is the within-round evolving penalty state block drafting needs: verify row r's
1780    /// target is penalized by every token committed before it INCLUDING same-round
1781    /// accepts — `penalize_logits_rows` (one shared window) is the frozen-window
1782    /// approximation this exists to replace on the dspark route.
1783    #[allow(clippy::too_many_arguments)]
1784    pub fn penalize_logits_rows_inc(
1785        &self,
1786        x: &mut CudaSlice<f32>,
1787        hist: &CudaSlice<u32>,
1788        n_hist0: usize,
1789        rep: f32,
1790        freq: f32,
1791        present: f32,
1792        n: usize,
1793        nrow: usize,
1794        win: usize,
1795    ) -> Result<(), Box<dyn std::error::Error>> {
1796        if nrow == 0 || win == 0 || (n_hist0 == 0 && nrow == 1) {
1797            return Ok(());
1798        }
1799        debug_assert!(
1800            hist.len() >= n_hist0 + nrow - 1,
1801            "rows-inc hist must carry n_hist0 + nrow - 1 ids"
1802        );
1803        let f = self.func("penalize_logits_rows_inc_f32");
1804        let max_len = win.min(n_hist0 + nrow - 1).max(1);
1805        let (nh, ni, nr, wi) = (n_hist0 as i32, n as i32, nrow as i32, win as i32);
1806        let cfg = LaunchConfig {
1807            grid_dim: (max_len.div_ceil(128) as u32, nrow as u32, 1),
1808            block_dim: (128, 1, 1),
1809            shared_mem_bytes: 0,
1810        };
1811        let __s_b = self.gpu.stream();
1812        let mut b = __s_b.launch_builder(&f);
1813        b.arg(&mut *x)
1814            .arg(hist)
1815            .arg(&nh)
1816            .arg(&rep)
1817            .arg(&freq)
1818            .arg(&present)
1819            .arg(&ni)
1820            .arg(&nr)
1821            .arg(&wi);
1822        unsafe {
1823            b.launch(cfg)?;
1824        }
1825        Ok(())
1826    }
1827
1828    /// WEIGHT PREFETCH (SOTA item 3, 2026-07-13, DEFAULT ON): during a bandwidth-idle
1829    /// window (the fa launch reads KV, not weights) prefetch the NEXT matvec's
1830    /// decode-plane bytes into L2 so it reads L2-warm. Value-free scheduling op — same
1831    /// class as prefetch_l2 (numerics untouched by construction). Wired only where it
1832    /// measured positive: the E4B dc attn arm (+0.65%). 26B (flat — MoE ffn dominates),
1833    /// 31B (−0.2% — decode at the DRAM wall) and the ffn gate/up cascade (−1% — 29MB/layer
1834    /// floods the fill path) all probed and NOT wired. MEMRA_WPF=0 rollback seam.
1835    pub fn wpf_level() -> u32 {
1836        static ON: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
1837        *ON.get_or_init(|| {
1838            std::env::var("MEMRA_WPF")
1839                .ok()
1840                .and_then(|v| v.parse().ok())
1841                .unwrap_or(1)
1842        })
1843    }
1844
1845    /// PDL launch arm (SOTA item 2, 2026-07-13, DEFAULT ON): the six MEMRA_PDL_ENTRY glue
1846    /// kernels launch through cuLaunchKernelEx with PROGRAMMATIC_STREAM_SERIALIZATION — the
1847    /// grid launches while the predecessor drains (~120ns/kernel back, pdl_probe), the
1848    /// kernels' entry grid-dep sync restores read order (SASS-audited: ACQBULK precedes
1849    /// every LDG in all six). Valid windows: E4B +1.0-1.2% (128 AND 384-tok gens);
1850    /// 26B/31B/qwen flat no-harm. Battery: kernel-check GREEN, run-gen tokens IDENTICAL x3
1851    /// gemma, spec 64/64 E4B K=1/4/8 + 26B/31B K=4 + qwen PASS. Works eager AND under
1852    /// capture (capture encodes native programmatic edges — the post-capture edge-REWRITE
1853    /// arm died: engine graphs hold cuMemAllocAsync alloc nodes, edge edits on those return
1854    /// CUDA_ERROR_NOT_SUPPORTED). MEMRA_PDL=0 rollback seam.
1855    /// See the `verify_exact` field. Scoped by the dflash round around its t=16 verify.
1856    pub fn set_verify_exact(&self, on: bool) {
1857        self.verify_exact
1858            .store(on, std::sync::atomic::Ordering::Relaxed);
1859    }
1860    pub(crate) fn verify_exact_on(&self) -> bool {
1861        self.verify_exact.load(std::sync::atomic::Ordering::Relaxed)
1862    }
1863
1864    /// m=1 norm+rope+append fold seam (2026-07-23): MEMRA_QKV_APPEND=0 reverts to the
1865    /// fused-norm-rope + standalone-append pair (the exact-oracle bisect arm).
1866    pub fn qkv_append_on() -> bool {
1867        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1868        *ON.get_or_init(|| {
1869            std::env::var("MEMRA_QKV_APPEND")
1870                .map(|v| v != "0")
1871                .unwrap_or(true)
1872        })
1873    }
1874
1875    /// PDL wave-B1a seam: the four dense-glue kernels (rms_norm_f32, add_rms_norm_f32,
1876    /// add_scale_rms_norm_q8_1, quantize_q8_1). MEMRA_PDL_WB=0 reverts alone.
1877    pub fn pdl_wb_on() -> bool {
1878        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1879        *ON.get_or_init(|| {
1880            std::env::var("MEMRA_PDL_WB")
1881                .map(|v| v != "0")
1882                .unwrap_or(true)
1883        })
1884    }
1885
1886    /// Trunk-kernels norm ILP seam (lane/dspark-trunk-kernels-20260820): the T-row verify
1887    /// norms (rms_norm_f32 / add_rms_norm_f32 at grid=T, block=256) are serial-latency
1888    /// chains — 20 strided scalar load->fma rounds measured 11.8-12.2us/inst (nsys-B verify
1889    /// scope: 130 inst/rd = 1.51 ms/rd). The `_v2` twins unroll the element loop 4-deep
1890    /// (independent loads in flight; SAME per-thread element order into ONE accumulator,
1891    /// reduce VERBATIM) — BIT-IDENTICAL per row at every (ncols, blockDim).
1892    /// MEMRA_NORM_ILP=0 reverts to the v1 kernels alone.
1893    pub fn norm_ilp_on() -> bool {
1894        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1895        *ON.get_or_init(|| {
1896            std::env::var("MEMRA_NORM_ILP")
1897                .map(|v| v != "0")
1898                .unwrap_or(true)
1899        })
1900    }
1901
1902    /// Trunk-kernels FFN dual seam (lane/dspark-trunk-kernels-20260820): the qwen35
1903    /// t-parallel verify FFN pair rides the PROVEN dual gate+up doors
1904    /// (`matmul_decode_exact_dual_pre` + `silu_mul_scaled_q8_1`, the q27 verify shape —
1905    /// bit-identical per (tensor,token,row), kernel-check-pinned, MEMRA_SPEC_DUAL_T
1906    /// receipts) instead of two singles + silu_mul + a standalone quantize. The doors
1907    /// existed but the qwen35 body never called them (nsys-B verify scope: gate+up singles
1908    /// = 107 launches/rd at grid 4352). MEMRA_TK_FFN_DUAL=0 reverts to the singles chain.
1909    pub fn tk_ffn_dual_on() -> bool {
1910        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1911        *ON.get_or_init(|| {
1912            std::env::var("MEMRA_TK_FFN_DUAL")
1913                .map(|v| v != "0")
1914                .unwrap_or(true)
1915        })
1916    }
1917
1918    /// PDL wave-A seam: the mmvq matvec PDL launches only (the six glue kernels keep
1919    /// their own MEMRA_PDL master seam). MEMRA_PDL_MMVQ=0 reverts wave-A alone — the
1920    /// per-model no-harm bisect knob.
1921    pub fn pdl_mmvq_on() -> bool {
1922        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1923        *ON.get_or_init(|| {
1924            std::env::var("MEMRA_PDL_MMVQ")
1925                .map(|v| v != "0")
1926                .unwrap_or(true)
1927        })
1928    }
1929
1930    pub fn pdl_on() -> bool {
1931        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1932        *ON.get_or_init(|| std::env::var("MEMRA_PDL").map(|v| v != "0").unwrap_or(true))
1933    }
1934
1935    /// PDL wave-B seam (gap-diagnosis arc, GAP-DIAGNOSIS.md verdict 8): the gemma
1936    /// NVFP4mix decode chain's hot kernels — nvfp4 fused2/mr2 and the q8_0 `_rp`
1937    /// singles — join the wave-A launch class. Scheduling-only (the entry macro waits
1938    /// on the producer before any read), bit-identical by construction.
1939    /// MEMRA_PDL_NVFP4=0 reverts wave-B alone.
1940    pub fn pdl_nvfp4q8_on() -> bool {
1941        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1942        *ON.get_or_init(|| {
1943            std::env::var("MEMRA_PDL_NVFP4")
1944                .map(|v| v != "0")
1945                .unwrap_or(true)
1946        })
1947    }
1948
1949    /// Raw CUfunction for a PDL-attributed launch: the SAME kernels.fatbin loaded once more
1950    /// through the raw driver API (cudarc hides its CUfunction handles; a duplicate module
1951    /// of tiny glue kernels is free). Resolved lazily per name, cached process-wide.
1952    /// Fused t=1 q4_0 mr policy: env MEMRA_Q40_MR wins (1/2); else the per-model
1953    /// FUSED_MR1_DEFAULT (dense gemma = mr1, MoE = mr2 — see the static's doc).
1954    fn q40_mr1_on() -> bool {
1955        static Q40MR: std::sync::OnceLock<Option<u32>> = std::sync::OnceLock::new();
1956        match *Q40MR.get_or_init(|| {
1957            std::env::var("MEMRA_Q40_MR")
1958                .ok()
1959                .and_then(|v| v.parse().ok())
1960        }) {
1961            Some(v) => v == 1,
1962            None => crate::FUSED_MR1_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
1963        }
1964    }
1965
1966    /// PDL wave-B2: flash-module PDL functions. `g` selects the kf8vf8 flavor — the
1967    /// caller MUST pass the SAME flavor its builder launch would resolve (fa_func/func_g
1968    /// mirror); the flavors differ semantically (KV byte formats), a wrong-module launch
1969    /// writes wrong bytes silently.
1970    fn pdl_func_flash(
1971        &self,
1972        g: bool,
1973        name: &'static str,
1974    ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
1975        use cudarc::driver::sys as cu;
1976        // PER-CONTEXT caches (M1-PP2 cross-device fix, 8x box 2026-08-02): CUmodule and
1977        // CUfunction handles are CONTEXT-scoped, and a remote-stage Engine
1978        // (MEMRA_PP_DEVICES=a,b) lives in the other device's primary context. The old
1979        // process-wide OnceLock cache handed stage 1 the dev-a handles, so every stage-1
1980        // launch_pdl* died CUDA_ERROR_INVALID_HANDLE. Key module + function caches by
1981        // this engine's CUcontext; single-context runs behave exactly as before.
1982        static MODS: std::sync::Mutex<Option<std::collections::HashMap<(usize, bool), usize>>> =
1983            std::sync::Mutex::new(None);
1984        static FNS: std::sync::Mutex<
1985            Option<std::collections::HashMap<(usize, bool, &'static str), usize>>,
1986        > = std::sync::Mutex::new(None);
1987        let ctx_key = self.ctx().cu_ctx() as usize;
1988        if let Some(&f) = FNS
1989            .lock()
1990            .unwrap()
1991            .get_or_insert_with(Default::default)
1992            .get(&(ctx_key, g, name))
1993        {
1994            return Ok(f as cu::CUfunction);
1995        }
1996        let module = {
1997            let mut mods = MODS.lock().unwrap();
1998            let map = mods.get_or_insert_with(Default::default);
1999            match map.get(&(ctx_key, g)) {
2000                Some(&m) => m,
2001                None => {
2002                    let m = self.pdl_load_module_in_ctx(if g {
2003                        FLASH_FATBIN_KF8VF8
2004                    } else {
2005                        FLASH_FATBIN
2006                    })?;
2007                    map.insert((ctx_key, g), m);
2008                    m
2009                }
2010            }
2011        };
2012        let cname = std::ffi::CString::new(name)?;
2013        let mut f: cu::CUfunction = std::ptr::null_mut();
2014        let r = unsafe { cu::cuModuleGetFunction(&mut f, module as cu::CUmodule, cname.as_ptr()) };
2015        if r != cu::CUresult::CUDA_SUCCESS {
2016            return Err(format!("pdl_func_flash {name} (g={g}): {r:?}").into());
2017        }
2018        FNS.lock()
2019            .unwrap()
2020            .get_or_insert_with(Default::default)
2021            .insert((ctx_key, g, name), f as usize);
2022        Ok(f)
2023    }
2024
2025    /// Load a fatbin as a raw CUmodule IN THIS ENGINE'S CONTEXT. `cuModuleLoadData` binds
2026    /// the module to the thread's CURRENT context — a remote-stage engine must not
2027    /// inherit the primary's (the INVALID_HANDLE class above). Restores the caller's
2028    /// current context before returning.
2029    fn pdl_load_module_in_ctx(&self, bytes: &[u8]) -> Result<usize, Box<dyn std::error::Error>> {
2030        use cudarc::driver::sys as cu;
2031        let mut prev: cu::CUcontext = std::ptr::null_mut();
2032        unsafe {
2033            cu::cuCtxGetCurrent(&mut prev).result()?;
2034        }
2035        self.ctx().bind_to_thread()?;
2036        let mut m: cu::CUmodule = std::ptr::null_mut();
2037        let r = unsafe { cu::cuModuleLoadData(&mut m, bytes.as_ptr() as *const std::ffi::c_void) };
2038        let restore = if prev.is_null() {
2039            cu::CUresult::CUDA_SUCCESS
2040        } else {
2041            unsafe { cu::cuCtxSetCurrent(prev) }
2042        };
2043        if r != cu::CUresult::CUDA_SUCCESS {
2044            return Err(format!("pdl module load: {r:?}").into());
2045        }
2046        if restore != cu::CUresult::CUDA_SUCCESS {
2047            return Err(format!("pdl module load: ctx restore {restore:?}").into());
2048        }
2049        Ok(m as usize)
2050    }
2051
2052    fn pdl_func(
2053        &self,
2054        name: &'static str,
2055    ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
2056        use cudarc::driver::sys as cu;
2057        // PER-CONTEXT caches — same M1-PP2 cross-device fix as pdl_func_flash (handles
2058        // are context-scoped; key everything by this engine's CUcontext).
2059        static MODULES: std::sync::Mutex<Option<std::collections::HashMap<usize, usize>>> =
2060            std::sync::Mutex::new(None);
2061        // PDL wave-A: the mmvq kernels live in the qmatvec fatbin, not kernels.cu — second
2062        // duplicate module, loaded lazily on the first kernels-module miss.
2063        static QMODULES: std::sync::Mutex<Option<std::collections::HashMap<usize, usize>>> =
2064            std::sync::Mutex::new(None);
2065        static FNS: std::sync::Mutex<
2066            Option<std::collections::HashMap<(usize, &'static str), usize>>,
2067        > = std::sync::Mutex::new(None);
2068        let ctx_key = self.ctx().cu_ctx() as usize;
2069        if let Some(&f) = FNS
2070            .lock()
2071            .unwrap()
2072            .get_or_insert_with(Default::default)
2073            .get(&(ctx_key, name))
2074        {
2075            return Ok(f as cu::CUfunction);
2076        }
2077        let module = {
2078            let mut mods = MODULES.lock().unwrap();
2079            let map = mods.get_or_insert_with(Default::default);
2080            match map.get(&ctx_key) {
2081                Some(&m) => m,
2082                None => {
2083                    let m = self.pdl_load_module_in_ctx(FATBIN)?;
2084                    map.insert(ctx_key, m);
2085                    m
2086                }
2087            }
2088        };
2089        let cname = std::ffi::CString::new(name)?;
2090        let mut f: cu::CUfunction = std::ptr::null_mut();
2091        let mut r =
2092            unsafe { cu::cuModuleGetFunction(&mut f, module as cu::CUmodule, cname.as_ptr()) };
2093        if r == cu::CUresult::CUDA_ERROR_NOT_FOUND {
2094            let qmodule = {
2095                let mut mods = QMODULES.lock().unwrap();
2096                let map = mods.get_or_insert_with(Default::default);
2097                match map.get(&ctx_key) {
2098                    Some(&m) => m,
2099                    None => {
2100                        let m = self.pdl_load_module_in_ctx(QMATVEC_FATBIN)?;
2101                        map.insert(ctx_key, m);
2102                        m
2103                    }
2104                }
2105            };
2106            r = unsafe { cu::cuModuleGetFunction(&mut f, qmodule as cu::CUmodule, cname.as_ptr()) };
2107        }
2108        if r != cu::CUresult::CUDA_SUCCESS {
2109            return Err(format!("pdl_func {name}: {r:?}").into());
2110        }
2111        FNS.lock()
2112            .unwrap()
2113            .get_or_insert_with(Default::default)
2114            .insert((ctx_key, name), f as usize);
2115        Ok(f)
2116    }
2117
2118    /// cuLaunchKernelEx with CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION on the
2119    /// compute stream. ONLY legal for kernels whose entry carries MEMRA_PDL_ENTRY.
2120    ///
2121    /// # Safety
2122    /// `params` must match the kernel's exact parameter list (order, types, count) —
2123    /// a mismatch corrupts the launch silently.
2124    /// Flash-module twin of `launch_pdl` — `g` picks the kf8vf8 flavor (must mirror the
2125    /// builder path's fa_func/func_g choice exactly).
2126    ///
2127    /// # Safety
2128    /// Same contract as `launch_pdl`.
2129    unsafe fn launch_pdl_flash(
2130        &self,
2131        g: bool,
2132        name: &'static str,
2133        grid: (u32, u32, u32),
2134        block: (u32, u32, u32),
2135        smem: u32,
2136        params: &mut [*mut std::ffi::c_void],
2137    ) -> Result<(), Box<dyn std::error::Error>> {
2138        use cudarc::driver::sys as cu;
2139        let f = self.pdl_func_flash(g, name)?;
2140        if smem > 0 {
2141            // mirror the builder path's opt-in ceiling (idempotent host-side set).
2142            let r =
2143                unsafe {
2144                    cu::cuFuncSetAttribute(f,
2145                cu::CUfunction_attribute_enum::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
2146                smem as i32)
2147                };
2148            if r != cu::CUresult::CUDA_SUCCESS {
2149                return Err(format!("pdl smem attr {name}: {r:?}").into());
2150            }
2151        }
2152        let mut attr = cu::CUlaunchAttribute {
2153            id: cu::CUlaunchAttributeID::CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION,
2154            pad: [0; 4],
2155            value: cu::CUlaunchAttributeValue {
2156                programmaticStreamSerializationAllowed: 1,
2157            },
2158        };
2159        let cfg = cu::CUlaunchConfig {
2160            gridDimX: grid.0,
2161            gridDimY: grid.1,
2162            gridDimZ: grid.2,
2163            blockDimX: block.0,
2164            blockDimY: block.1,
2165            blockDimZ: block.2,
2166            sharedMemBytes: smem,
2167            hStream: self.gpu.stream().cu_stream(),
2168            attrs: &mut attr,
2169            numAttrs: 1,
2170        };
2171        let r = unsafe { cu::cuLaunchKernelEx(&cfg, f, params.as_mut_ptr(), std::ptr::null_mut()) };
2172        if r != cu::CUresult::CUDA_SUCCESS {
2173            return Err(format!("launch_pdl_flash {name}: {r:?}").into());
2174        }
2175        Ok(())
2176    }
2177
2178    unsafe fn launch_pdl(
2179        &self,
2180        name: &'static str,
2181        grid: (u32, u32, u32),
2182        block: (u32, u32, u32),
2183        params: &mut [*mut std::ffi::c_void],
2184    ) -> Result<(), Box<dyn std::error::Error>> {
2185        use cudarc::driver::sys as cu;
2186        let f = self.pdl_func(name)?;
2187        let mut attr = cu::CUlaunchAttribute {
2188            id: cu::CUlaunchAttributeID::CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION,
2189            pad: [0; 4],
2190            value: cu::CUlaunchAttributeValue {
2191                programmaticStreamSerializationAllowed: 1,
2192            },
2193        };
2194        let cfg = cu::CUlaunchConfig {
2195            gridDimX: grid.0,
2196            gridDimY: grid.1,
2197            gridDimZ: grid.2,
2198            blockDimX: block.0,
2199            blockDimY: block.1,
2200            blockDimZ: block.2,
2201            sharedMemBytes: 0,
2202            hStream: self.gpu.stream().cu_stream(),
2203            attrs: &mut attr,
2204            numAttrs: 1,
2205        };
2206        let r = unsafe { cu::cuLaunchKernelEx(&cfg, f, params.as_mut_ptr(), std::ptr::null_mut()) };
2207        if r != cu::CUresult::CUDA_SUCCESS {
2208            return Err(format!("launch_pdl {name}: {r:?}").into());
2209        }
2210        Ok(())
2211    }
2212
2213    /// L2-prefetch a quant weight's DECODE plane (the rp4 split-plane mirror when present —
2214    /// that is what the m<=8 dispatch reads — else the raw block bytes). No-op on float arms.
2215    pub fn prefetch_weight_l2(
2216        &self,
2217        w: &crate::model::GpuTensor,
2218    ) -> Result<(), Box<dyn std::error::Error>> {
2219        if let crate::model::GpuTensor::Quant { bytes, rp4, .. } = w {
2220            let p = rp4.as_ref().unwrap_or(bytes);
2221            self.prefetch_l2(p, p.len())?;
2222        }
2223        Ok(())
2224    }
2225
2226    /// DSpark markov chain ops (dflash lane): gather one bf16 row of a [V, rank] table
2227    /// by the DEVICE token id at tok[idx] into f32.
2228    pub fn gather_row_bf16(
2229        &self,
2230        table: &CudaSlice<u8>,
2231        tok: &CudaSlice<u32>,
2232        idx: usize,
2233        dst: &mut CudaSlice<f32>,
2234        ncols: usize,
2235    ) -> Result<(), Box<dyn std::error::Error>> {
2236        let f = self.func("gather_row_bf16_f32");
2237        let cfg = LaunchConfig {
2238            grid_dim: (ncols.div_ceil(256) as u32, 1, 1),
2239            block_dim: (256, 1, 1),
2240            shared_mem_bytes: 0,
2241        };
2242        let (nc, ix) = (ncols as i32, idx as i32);
2243        let __s_b = self.gpu.stream();
2244        let mut b = __s_b.launch_builder(&f);
2245        b.arg(table).arg(tok).arg(&ix).arg(dst).arg(&nc);
2246        unsafe {
2247            b.launch(cfg)?;
2248        }
2249        Ok(())
2250    }
2251
2252    /// DFlash2 grouped dynamic causal conv (dflash lane, DFLASH2-EVAL-20260820.md):
2253    /// out[p,c] = sum_{o<ksize, o<=p} (base[half][o][c] + dyn[p][half][o][group(c)])
2254    /// * x[p-o][c]. `dyn_` is the kernel_projection GEMM output [rows, 2*ksize*groups];
2255    /// `base` is base_kernel [2, ksize, hidden] flattened; `half` picks prepare(0) /
2256    /// finish(1).
2257    #[allow(clippy::too_many_arguments)]
2258    pub fn dflash2_dynconv(
2259        &self,
2260        x: &CudaSlice<f32>,
2261        dyn_: &CudaSlice<f32>,
2262        base: &CudaSlice<f32>,
2263        out: &mut CudaSlice<f32>,
2264        rows: usize,
2265        hidden: usize,
2266        group_size: usize,
2267        ksize: usize,
2268        half: usize,
2269    ) -> Result<(), Box<dyn std::error::Error>> {
2270        assert_eq!(hidden % group_size, 0, "hidden % group_size != 0");
2271        let f = self.func("dflash2_dynconv_f32");
2272        let n = rows * hidden;
2273        let cfg = LaunchConfig {
2274            grid_dim: (n.div_ceil(256) as u32, 1, 1),
2275            block_dim: (256, 1, 1),
2276            shared_mem_bytes: 0,
2277        };
2278        let (ri, hi, gi, ki, hf) = (
2279            rows as i32,
2280            hidden as i32,
2281            group_size as i32,
2282            ksize as i32,
2283            half as i32,
2284        );
2285        let __s_b = self.gpu.stream();
2286        let mut b = __s_b.launch_builder(&f);
2287        b.arg(x)
2288            .arg(dyn_)
2289            .arg(base)
2290            .arg(out)
2291            .arg(&ri)
2292            .arg(&hi)
2293            .arg(&gi)
2294            .arg(&ki)
2295            .arg(&hf);
2296        unsafe {
2297            b.launch(cfg)?;
2298        }
2299        Ok(())
2300    }
2301
2302    /// Per-row top-k (k <= 32) over a [n_rows, n_cols] logits matrix (DFlash2
2303    /// candidate selector). Returns (values [n_rows, k], column indices [n_rows, k]),
2304    /// value-descending, ties to the lower index.
2305    pub fn topk_rows(
2306        &self,
2307        logits: &CudaSlice<f32>,
2308        n_rows: usize,
2309        n_cols: usize,
2310        k: usize,
2311    ) -> Result<(CudaSlice<f32>, CudaSlice<u32>), Box<dyn std::error::Error>> {
2312        assert!(k <= 32 && k >= 1, "topk_rows supports 1..=32, got {k}");
2313        assert!(k <= n_cols, "topk_rows: k {k} > n_cols {n_cols}");
2314        let f = self.func("topk_rows_f32");
2315        let nth = 256usize;
2316        let mut vals = self.uninit(n_rows * k)?;
2317        let mut idxs = self.gpu.stream().alloc_zeros::<u32>(n_rows * k)?;
2318        let cfg = LaunchConfig {
2319            grid_dim: (n_rows as u32, 1, 1),
2320            block_dim: (nth as u32, 1, 1),
2321            shared_mem_bytes: (nth * k * 8) as u32,
2322        };
2323        let (nr, nc, ki) = (n_rows as i32, n_cols as i32, k as i32);
2324        let __s_b = self.gpu.stream();
2325        let mut b = __s_b.launch_builder(&f);
2326        b.arg(logits)
2327            .arg(&nr)
2328            .arg(&nc)
2329            .arg(&ki)
2330            .arg(&mut vals)
2331            .arg(&mut idxs);
2332        unsafe {
2333            b.launch(cfg)?;
2334        }
2335        Ok((vals, idxs))
2336    }
2337
2338    /// logits[row_off .. row_off+n] += bias[0..n] (in place, one row).
2339    pub fn add_row_inplace(
2340        &self,
2341        logits: &mut CudaSlice<f32>,
2342        bias: &CudaSlice<f32>,
2343        n: usize,
2344        row_off: usize,
2345    ) -> Result<(), Box<dyn std::error::Error>> {
2346        let f = self.func("add_row_inplace_f32");
2347        let cfg = LaunchConfig {
2348            grid_dim: (n.div_ceil(256) as u32, 1, 1),
2349            block_dim: (256, 1, 1),
2350            shared_mem_bytes: 0,
2351        };
2352        let (ni, off) = (n as i32, row_off as i64);
2353        let __s_b = self.gpu.stream();
2354        let mut b = __s_b.launch_builder(&f);
2355        b.arg(logits).arg(bias).arg(&ni).arg(&off);
2356        unsafe {
2357            b.launch(cfg)?;
2358        }
2359        Ok(())
2360    }
2361
2362    /// L2 prefetch of a device byte range (latency-hiding arc; value-free scheduling op).
2363    pub fn prefetch_l2(
2364        &self,
2365        p: &CudaSlice<u8>,
2366        n: usize,
2367    ) -> Result<(), Box<dyn std::error::Error>> {
2368        let f = self.func("prefetch_l2_bytes");
2369        let lines = n.div_ceil(128);
2370        let ni = n as i64;
2371        let cfg = LaunchConfig {
2372            grid_dim: (lines.div_ceil(256) as u32, 1, 1),
2373            block_dim: (256, 1, 1),
2374            shared_mem_bytes: 0,
2375        };
2376        let __s_b = self.gpu.stream();
2377        let mut b = __s_b.launch_builder(&f);
2378        b.arg(p).arg(&ni);
2379        unsafe {
2380            b.launch(cfg)?;
2381        }
2382        Ok(())
2383    }
2384
2385    /// MoE router GEMV (MEMRA_ROUTER_KERNEL): deterministic warp-per-(expert,token) f32 dot.
2386    /// Different FP order than the cuBLAS path it replaces — battery-gated numeric config.
2387    pub fn router_gemv(
2388        &self,
2389        w: &CudaSlice<f32>,
2390        x: &CudaSlice<f32>,
2391        n_embd: usize,
2392        n_experts: usize,
2393        t: usize,
2394    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2395        // float4 v2 probed 2026-07-14: +0.25% but flips near-tie routing (new FP order,
2396        // stream differs) — too small to justify a numeric config change; deleted.
2397        // w8 twin (2026-07-31): on the 132-SM H100 the lone-warp form is 14.8% of the q35
2398        // decode step (latency-bound) — the calculus flipped. MEMRA_ROUTER_V2=0 reverts to
2399        // the warp form (rollback seam; new FP order, battery-arbitrated per model).
2400        let w8 = match std::env::var("MEMRA_ROUTER_V2").as_deref() {
2401            Ok("0") => false,
2402            Ok(_) => true,
2403            Err(_) => ROUTER_W8_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
2404        };
2405        // FAST-ROUTER batch twin (lane/fast-router, 2026-08-02): at prefill m the per-(e,tok)
2406        // w8 form re-streams both operand rows per output (GEMV program at GEMM shape — the
2407        // concat-prime exactness fix paid -10% q35 board-2048 prefill through it). The batch
2408        // twin (8x8 expert-x-token register tile) is BIT-IDENTICAL per row (same k order,
2409        // same tree, same fold — kernel-check sweeps m=1..2048 on real router weights), so
2410        // the crossover is pure perf, not a numeric config. MIN_T from the on-box sweep
2411        // (research/fast-router-20260802/crossover-router*.jsonl); decode t=1 and small-t
2412        // spec verify keep the plain w8 form. MEMRA_ROUTER_BATCH=0: rollback seam
2413        // (perf-only, bits equal).
2414        let batch = w8 && t >= ROUTER_BATCH_MIN_T && router_batch_on();
2415        self.router_gemv_form(w, x, n_embd, n_experts, t, w8, batch)
2416    }
2417
2418    /// Form-explicit router GEMV launch (kernel-check bit-identity gate + crossover bench
2419    /// force both forms; `batch` requires `w8`).
2420    pub fn router_gemv_form(
2421        &self,
2422        w: &CudaSlice<f32>,
2423        x: &CudaSlice<f32>,
2424        n_embd: usize,
2425        n_experts: usize,
2426        t: usize,
2427        w8: bool,
2428        batch: bool,
2429    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2430        debug_assert!(!batch || w8, "batch twin exists for the w8 form only");
2431        let mut y = self.alloc_uninit::<f32>(t * n_experts)?;
2432        let f = if batch {
2433            self.func("router_gemv_f32_w8_batch")
2434        } else if w8 {
2435            self.func("router_gemv_f32_w8")
2436        } else {
2437            self.func("router_gemv_f32")
2438        };
2439        let (ne, nx, ti) = (n_embd as i32, n_experts as i32, t as i32);
2440        let cfg = if batch {
2441            LaunchConfig {
2442                grid_dim: (n_experts.div_ceil(8) as u32, t.div_ceil(8) as u32, 1),
2443                block_dim: (32, 8, 1),
2444                shared_mem_bytes: 0,
2445            }
2446        } else {
2447            LaunchConfig {
2448                grid_dim: (n_experts as u32, t as u32, 1),
2449                block_dim: (32, if w8 { 8 } else { 1 }, 1),
2450                shared_mem_bytes: 0,
2451            }
2452        };
2453        let __s_b = self.gpu.stream();
2454        let mut b = __s_b.launch_builder(&f);
2455        b.arg(w).arg(x).arg(&mut y).arg(&ne).arg(&nx).arg(&ti);
2456        unsafe {
2457            b.launch(cfg)?;
2458        }
2459        Ok(y)
2460    }
2461
2462    /// f32 row permute: dst[idx[i], :] = src[i, :] (grouped-GEMM CSR -> pair-id reorder).
2463    pub fn rows_permute(
2464        &self,
2465        src: &CudaSlice<f32>,
2466        idx: &CudaSlice<i32>,
2467        nrows: usize,
2468        ncols: usize,
2469    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2470        let mut dst = self.alloc_uninit::<f32>(nrows * ncols)?;
2471        let f = self.func("rows_permute_f32");
2472        let (nc, nr) = (ncols as i32, nrows as i32);
2473        let cfg = LaunchConfig {
2474            grid_dim: (nrows as u32, 1, 1),
2475            block_dim: (256, 1, 1),
2476            shared_mem_bytes: 0,
2477        };
2478        let __s_b = self.gpu.stream();
2479        let mut b = __s_b.launch_builder(&f);
2480        b.arg(src).arg(idx).arg(&mut dst).arg(&nc).arg(&nr);
2481        unsafe {
2482            b.launch(cfg)?;
2483        }
2484        Ok(dst)
2485    }
2486
2487    /// shexp gate fused dot: g[tok] = sigmoid(dot(x[tok,:], w)) — replaces the per-layer
2488    /// cuBLASLt m=1 GEMM + separate sigmoid launch on the qwen35moe decode path (the
2489    /// splitKreduce x40/step dig, 2026-07-31). One fold order for every t, so the t=1
2490    /// decode chain and the small-t spec-verify chain match per row by construction.
2491    pub fn sigmoid_dot_rows(
2492        &self,
2493        x: &CudaSlice<f32>,
2494        w: &CudaSlice<f32>,
2495        n_embd: usize,
2496        t: usize,
2497    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2498        // MEMRA_SHEXP_DOT=0: rollback seam to the cuBLASLt linear + sigmoid pair (numeric
2499        // config; same class as MEMRA_ROUTER_V2).
2500        static OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2501        if *OFF.get_or_init(|| std::env::var("MEMRA_SHEXP_DOT").as_deref() == Ok("0")) {
2502            let gs = self.linear(x, w, t, n_embd, 1)?;
2503            let mut g = self.uninit(t)?;
2504            self.sigmoid(&gs, &mut g, t)?;
2505            return Ok(g);
2506        }
2507        // FAST-ROUTER lane note (2026-08-02): a register-tiled 8-token batch twin of this
2508        // kernel was built, proven bit-identical, and measured SLOWER at every prefill t on
2509        // the 5090 (0.62-0.89x — launch-latency-bound op, ~7us/layer at m=2048;
2510        // research/fast-router-20260802/crossover-router.jsonl). Dispatch arm killed per
2511        // flags doctrine; this per-token form serves every t.
2512        let mut g = self.alloc_uninit::<f32>(t)?;
2513        let f = self.func("sigmoid_dot_rows_f32");
2514        let (ne, ti) = (n_embd as i32, t as i32);
2515        let cfg = LaunchConfig {
2516            grid_dim: (t as u32, 1, 1),
2517            block_dim: (32, 8, 1),
2518            shared_mem_bytes: 0,
2519        };
2520        let __s_b = self.gpu.stream();
2521        let mut b = __s_b.launch_builder(&f);
2522        b.arg(x).arg(w).arg(&mut g).arg(&ne).arg(&ti);
2523        unsafe {
2524            b.launch(cfg)?;
2525        }
2526        Ok(g)
2527    }
2528
2529    /// ROUND-STREAM stream rollback: all counters <- pos_start + base + n_acc.
2530    pub fn spec_rollback_stream(
2531        &self,
2532        len_ptrs: &CudaSlice<u64>,
2533        pos_start: &CudaSlice<i32>,
2534        acc: &CudaSlice<u32>,
2535        base: usize,
2536        n_rows: usize,
2537    ) -> Result<(), Box<dyn std::error::Error>> {
2538        let f = self.func("spec_rollback_stream");
2539        let (b, nr) = (base as i32, n_rows as i32);
2540        let cfg = LaunchConfig {
2541            grid_dim: (n_rows.div_ceil(64) as u32, 1, 1),
2542            block_dim: (64, 1, 1),
2543            shared_mem_bytes: 0,
2544        };
2545        let __s_bl = self.gpu.stream();
2546        let mut bl = __s_bl.launch_builder(&f);
2547        bl.arg(len_ptrs).arg(pos_start).arg(acc).arg(&b).arg(&nr);
2548        unsafe {
2549            bl.launch(cfg)?;
2550        }
2551        Ok(())
2552    }
2553
2554    /// PLAIN-DECODE GRAPH ring store: ring[(pos_start - base) % cap] = vam[0].
2555    pub fn plain_tok_ring(
2556        &self,
2557        vam: &CudaSlice<u32>,
2558        pos_start: &CudaSlice<i32>,
2559        base: usize,
2560        ring: &mut CudaSlice<u32>,
2561    ) -> Result<(), Box<dyn std::error::Error>> {
2562        let f = self.func("plain_tok_ring");
2563        let (b, cap) = (base as i32, ring.len() as i32);
2564        let cfg = LaunchConfig {
2565            grid_dim: (1, 1, 1),
2566            block_dim: (32, 1, 1),
2567            shared_mem_bytes: 0,
2568        };
2569        let __s_bl = self.gpu.stream();
2570        let mut bl = __s_bl.launch_builder(&f);
2571        bl.arg(vam).arg(pos_start).arg(&b).arg(&mut *ring).arg(&cap);
2572        unsafe {
2573            bl.launch(cfg)?;
2574        }
2575        Ok(())
2576    }
2577
2578    /// ROUND-STREAM stage (c) 4 epilogue: ring commit + tiny counter copies.
2579    pub fn spec_ring_commit(
2580        &self,
2581        vtok: &CudaSlice<u32>,
2582        acc: &CudaSlice<u32>,
2583        brk: &CudaSlice<u32>,
2584        ring: &mut CudaSlice<u32>,
2585        pend: &mut CudaSlice<u32>,
2586    ) -> Result<(), Box<dyn std::error::Error>> {
2587        let f = self.func("spec_ring_commit");
2588        let cfg = LaunchConfig {
2589            grid_dim: (1, 1, 1),
2590            block_dim: (32, 1, 1),
2591            shared_mem_bytes: 0,
2592        };
2593        let __s_b = self.gpu.stream();
2594        let mut b = __s_b.launch_builder(&f);
2595        b.arg(vtok).arg(acc).arg(brk).arg(ring).arg(pend);
2596        unsafe {
2597            b.launch(cfg)?;
2598        }
2599        Ok(())
2600    }
2601    pub fn i32_copy_add(
2602        &self,
2603        src: &CudaSlice<i32>,
2604        dst: &mut CudaSlice<i32>,
2605        delta: i32,
2606    ) -> Result<(), Box<dyn std::error::Error>> {
2607        let f = self.func("i32_copy_add");
2608        let cfg = LaunchConfig {
2609            grid_dim: (1, 1, 1),
2610            block_dim: (32, 1, 1),
2611            shared_mem_bytes: 0,
2612        };
2613        let __s_b = self.gpu.stream();
2614        let mut b = __s_b.launch_builder(&f);
2615        b.arg(src).arg(dst).arg(&delta);
2616        unsafe {
2617            b.launch(cfg)?;
2618        }
2619        Ok(())
2620    }
2621    pub fn u32_copy(
2622        &self,
2623        src: &CudaSlice<u32>,
2624        dst: &mut CudaSlice<u32>,
2625    ) -> Result<(), Box<dyn std::error::Error>> {
2626        let f = self.func("u32_copy");
2627        let cfg = LaunchConfig {
2628            grid_dim: (1, 1, 1),
2629            block_dim: (32, 1, 1),
2630            shared_mem_bytes: 0,
2631        };
2632        let __s_b = self.gpu.stream();
2633        let mut b = __s_b.launch_builder(&f);
2634        b.arg(src).arg(dst);
2635        unsafe {
2636            b.launch(cfg)?;
2637        }
2638        Ok(())
2639    }
2640
2641    /// ROUND-GRAPH adaptive depth: brk[0] <- clamp(acc[0] + 1, floor, cap) — the host
2642    /// adaptive policy as a captured device op (policy-identical: the accept walk depth
2643    /// caps acceptance exactly like drafting fewer tokens).
2644    pub fn spec_adapt_k(
2645        &self,
2646        acc: &CudaSlice<u32>,
2647        brk: &mut CudaSlice<u32>,
2648        floor: usize,
2649        cap: usize,
2650    ) -> Result<(), Box<dyn std::error::Error>> {
2651        let f = self.func("spec_adapt_k");
2652        let (fl, cp) = (floor as i32, cap as i32);
2653        let cfg = LaunchConfig {
2654            grid_dim: (1, 1, 1),
2655            block_dim: (32, 1, 1),
2656            shared_mem_bytes: 0,
2657        };
2658        let __s_b = self.gpu.stream();
2659        let mut b = __s_b.launch_builder(&f);
2660        b.arg(acc).arg(brk).arg(&fl).arg(&cp);
2661        unsafe {
2662            b.launch(cfg)?;
2663        }
2664        Ok(())
2665    }
2666
2667    /// ROUND-STREAM stage (c) 3: accept walk fully device-driven (brk + assembled vtok).
2668    pub fn spec_accept_greedy_dc(
2669        &self,
2670        preds: &CudaSlice<u32>,
2671        vtok: &CudaSlice<u32>,
2672        last_pred: &CudaSlice<u32>,
2673        brk: &CudaSlice<u32>,
2674        out: &mut CudaSlice<u32>,
2675    ) -> Result<(), Box<dyn std::error::Error>> {
2676        let f = self.func("spec_accept_greedy_dc");
2677        let cfg = LaunchConfig {
2678            grid_dim: (1, 1, 1),
2679            block_dim: (32, 1, 1),
2680            shared_mem_bytes: 0,
2681        };
2682        let __s_b = self.gpu.stream();
2683        let mut b = __s_b.launch_builder(&f);
2684        b.arg(preds).arg(vtok).arg(last_pred).arg(brk).arg(out);
2685        unsafe {
2686            b.launch(cfg)?;
2687        }
2688        Ok(())
2689    }
2690
2691    /// ROUND-STREAM stage (c) 2: verify-chain device-pos entries.
2692    pub fn pos_iota(
2693        &self,
2694        pos0: &CudaSlice<i32>,
2695        out: &mut CudaSlice<i32>,
2696        t: usize,
2697    ) -> Result<(), Box<dyn std::error::Error>> {
2698        let f = self.func("pos_iota_i32");
2699        let ti = t as i32;
2700        let cfg = LaunchConfig {
2701            grid_dim: (1, 1, 1),
2702            block_dim: (t.max(1) as u32, 1, 1),
2703            shared_mem_bytes: 0,
2704        };
2705        let __s_b = self.gpu.stream();
2706        let mut b = __s_b.launch_builder(&f);
2707        b.arg(pos0).arg(out).arg(&ti);
2708        unsafe {
2709            b.launch(cfg)?;
2710        }
2711        Ok(())
2712    }
2713    #[allow(clippy::too_many_arguments)]
2714    pub fn append_kv_quantized_rows_dc(
2715        &self,
2716        k_rows: &CudaSlice<f32>,
2717        v_rows: &CudaSlice<f32>,
2718        kc: &mut CudaSlice<u8>,
2719        vc: &mut CudaSlice<u8>,
2720        t0_dev: &CudaSlice<i32>,
2721        t: usize,
2722        kv_dim_k: usize,
2723        kv_dim_v: usize,
2724        k_tok_bytes: usize,
2725        v_tok_bytes: usize,
2726        g: bool,
2727    ) -> Result<(), Box<dyn std::error::Error>> {
2728        let f = if g {
2729            self.func_g("append_quantize_kv_q8_0_q5_1_rows_dc")
2730        } else {
2731            self.func("append_quantize_kv_q8_0_q5_1_rows_dc")
2732        };
2733        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
2734        let cfg = LaunchConfig {
2735            grid_dim: (nblk, t as u32, 1),
2736            block_dim: (32, 1, 1),
2737            shared_mem_bytes: 0,
2738        };
2739        let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
2740        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
2741        let __s_b = self.gpu.stream();
2742        let mut b = __s_b.launch_builder(&f);
2743        b.arg(k_rows)
2744            .arg(v_rows)
2745            .arg(kc)
2746            .arg(vc)
2747            .arg(t0_dev)
2748            .arg(&kdk)
2749            .arg(&kdv)
2750            .arg(&ktb)
2751            .arg(&vtb);
2752        unsafe {
2753            b.launch(cfg)?;
2754        }
2755        Ok(())
2756    }
2757
2758    /// t=1 dc append with a FUSED len_d increment (wave 5c) — one launch replaces
2759    /// append_rows_dc + inc_seqlen. Single block (read-before-inc ordering).
2760    #[allow(clippy::too_many_arguments)]
2761    pub fn append_kv_quantized_row_dc_inc(
2762        &self,
2763        k_row: &CudaSlice<f32>,
2764        v_row: &CudaSlice<f32>,
2765        kc: &mut CudaSlice<u8>,
2766        vc: &mut CudaSlice<u8>,
2767        t0_dev: &mut CudaSlice<i32>,
2768        kv_dim_k: usize,
2769        kv_dim_v: usize,
2770        k_tok_bytes: usize,
2771        v_tok_bytes: usize,
2772        g: bool,
2773    ) -> Result<(), Box<dyn std::error::Error>> {
2774        let f = if g {
2775            self.func_g("append_quantize_kv_q8_0_q5_1_dc_inc")
2776        } else {
2777            self.func("append_quantize_kv_q8_0_q5_1_dc_inc")
2778        };
2779        let nthreads = ((kv_dim_k.max(kv_dim_v) / 32) * 32).min(1024) as u32;
2780        let cfg = LaunchConfig {
2781            grid_dim: (1, 1, 1),
2782            block_dim: (nthreads, 1, 1),
2783            shared_mem_bytes: 0,
2784        };
2785        let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
2786        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
2787        let __s_b = self.gpu.stream();
2788        let mut b = __s_b.launch_builder(&f);
2789        b.arg(k_row)
2790            .arg(v_row)
2791            .arg(kc)
2792            .arg(vc)
2793            .arg(t0_dev)
2794            .arg(&kdk)
2795            .arg(&kdv)
2796            .arg(&ktb)
2797            .arg(&vtb);
2798        unsafe {
2799            b.launch(cfg)?;
2800        }
2801        Ok(())
2802    }
2803
2804    /// ROUND-STREAM: draft-chain pack + in-graph d2t remap (see kernels.cu headers).
2805    pub fn pack_tok_p(
2806        &self,
2807        tok: &CudaSlice<u32>,
2808        p: &CudaSlice<f32>,
2809        out: &mut CudaSlice<u32>,
2810        slot: usize,
2811    ) -> Result<(), Box<dyn std::error::Error>> {
2812        let f = self.func("pack_tok_p");
2813        let sl = slot as i32;
2814        let cfg = LaunchConfig {
2815            grid_dim: (1, 1, 1),
2816            block_dim: (32, 1, 1),
2817            shared_mem_bytes: 0,
2818        };
2819        let __s_b = self.gpu.stream();
2820        let mut b = __s_b.launch_builder(&f);
2821        b.arg(tok).arg(p).arg(out).arg(&sl);
2822        unsafe {
2823            b.launch(cfg)?;
2824        }
2825        Ok(())
2826    }
2827    pub fn tok_map_u32(
2828        &self,
2829        tok: &mut CudaSlice<u32>,
2830        map: &CudaSlice<u32>,
2831    ) -> Result<(), Box<dyn std::error::Error>> {
2832        let f = self.func("tok_map_u32");
2833        let cfg = LaunchConfig {
2834            grid_dim: (1, 1, 1),
2835            block_dim: (32, 1, 1),
2836            shared_mem_bytes: 0,
2837        };
2838        let __s_b = self.gpu.stream();
2839        let mut b = __s_b.launch_builder(&f);
2840        b.arg(tok).arg(map);
2841        unsafe {
2842            b.launch(cfg)?;
2843        }
2844        Ok(())
2845    }
2846
2847    /// ROUND-STREAM stage (c) 1: device verify-token assembly + p-min break derivation.
2848    #[allow(clippy::too_many_arguments)]
2849    pub fn spec_assemble_verify(
2850        &self,
2851        tokp: &CudaSlice<u32>,
2852        pend: &CudaSlice<u32>,
2853        d2t: Option<&CudaSlice<u32>>,
2854        vtok: &mut CudaSlice<u32>,
2855        brk: &mut CudaSlice<u32>,
2856        p_min: f32,
2857        k: usize,
2858        pmin0: bool,
2859    ) -> Result<(), Box<dyn std::error::Error>> {
2860        let f = self.func("spec_assemble_verify");
2861        let (ki, pm) = (k as i32, if pmin0 { 1i32 } else { 0i32 });
2862        let cfg = LaunchConfig {
2863            grid_dim: (1, 1, 1),
2864            block_dim: (32, 1, 1),
2865            shared_mem_bytes: 0,
2866        };
2867        let __s_b = self.gpu.stream();
2868        let mut b = __s_b.launch_builder(&f);
2869        match d2t {
2870            Some(m) => {
2871                b.arg(tokp)
2872                    .arg(pend)
2873                    .arg(m)
2874                    .arg(vtok)
2875                    .arg(brk)
2876                    .arg(&p_min)
2877                    .arg(&ki)
2878                    .arg(&pm);
2879                unsafe {
2880                    b.launch(cfg)?;
2881                }
2882            }
2883            None => {
2884                let null: u64 = 0;
2885                b.arg(tokp)
2886                    .arg(pend)
2887                    .arg(&null)
2888                    .arg(vtok)
2889                    .arg(brk)
2890                    .arg(&p_min)
2891                    .arg(&ki)
2892                    .arg(&pm);
2893                unsafe {
2894                    b.launch(cfg)?;
2895                }
2896            }
2897        }
2898        Ok(())
2899    }
2900
2901    /// ROUND-STREAM stage (b) 3b: recur-restore twins with device-j (see hybrid.cu headers).
2902    #[allow(clippy::too_many_arguments)]
2903    pub fn ssm_conv_ring_rebuild_dc(
2904        &self,
2905        qkv_tm: &CudaSlice<f32>,
2906        ring_old: &CudaSlice<f32>,
2907        conv_state: &mut CudaSlice<f32>,
2908        conv_dim: usize,
2909        acc: &CudaSlice<u32>,
2910        base: usize,
2911        t_v: usize,
2912        d_conv: usize,
2913    ) -> Result<(), Box<dyn std::error::Error>> {
2914        let f = self.func("ssm_conv_ring_rebuild_f32_dc");
2915        let n = conv_dim * (d_conv - 1);
2916        let cfg = LaunchConfig::for_num_elems(n as u32);
2917        let (cd, b0, tv, dc) = (conv_dim as i32, base as i32, t_v as i32, d_conv as i32);
2918        let __s_b = self.gpu.stream();
2919        let mut b = __s_b.launch_builder(&f);
2920        b.arg(qkv_tm)
2921            .arg(ring_old)
2922            .arg(conv_state)
2923            .arg(&cd)
2924            .arg(acc)
2925            .arg(&b0)
2926            .arg(&tv)
2927            .arg(&dc);
2928        unsafe {
2929            b.launch(cfg)?;
2930        }
2931        Ok(())
2932    }
2933    #[allow(clippy::too_many_arguments)]
2934    pub fn gdn_scan_s128_dc(
2935        &self,
2936        q: &CudaSlice<f32>,
2937        k: &CudaSlice<f32>,
2938        v: &CudaSlice<f32>,
2939        g: &CudaSlice<f32>,
2940        beta: &CudaSlice<f32>,
2941        state_in: &CudaSlice<f32>,
2942        state_out: &mut CudaSlice<f32>,
2943        o: &mut CudaSlice<f32>,
2944        n_head: usize,
2945        acc: &CudaSlice<u32>,
2946        base: usize,
2947        t_v: usize,
2948        scale: f32,
2949    ) -> Result<(), Box<dyn std::error::Error>> {
2950        let f = self.func("gdn_scan_s128_dc");
2951        const S_V: u32 = 128;
2952        const WARP: u32 = 32;
2953        const COLS_PER_BLOCK: u32 = 4;
2954        let cfg = LaunchConfig {
2955            grid_dim: (n_head as u32, 1, S_V / COLS_PER_BLOCK),
2956            block_dim: (WARP, COLS_PER_BLOCK, 1),
2957            shared_mem_bytes: 0,
2958        };
2959        let (h, b0, tv) = (n_head as i32, base as i32, t_v as i32);
2960        let __s_b = self.gpu.stream();
2961        let mut b = __s_b.launch_builder(&f);
2962        b.arg(q)
2963            .arg(k)
2964            .arg(v)
2965            .arg(g)
2966            .arg(beta)
2967            .arg(state_in)
2968            .arg(state_out)
2969            .arg(o)
2970            .arg(&h)
2971            .arg(acc)
2972            .arg(&b0)
2973            .arg(&tv)
2974            .arg(&scale);
2975        unsafe {
2976            b.launch(cfg)?;
2977        }
2978        Ok(())
2979    }
2980
2981    /// ROUND-STREAM stage (b) 3a: device per-layer KV-len rollback (see spec_rollback_kv).
2982    pub fn spec_rollback_kv(
2983        &self,
2984        len_ptrs: &CudaSlice<u64>,
2985        saved: &CudaSlice<i32>,
2986        acc: &CudaSlice<u32>,
2987        base: usize,
2988        n_layer: usize,
2989    ) -> Result<(), Box<dyn std::error::Error>> {
2990        let f = self.func("spec_rollback_kv");
2991        let (b, nl) = (base as i32, n_layer as i32);
2992        let cfg = LaunchConfig {
2993            grid_dim: (n_layer.div_ceil(64) as u32, 1, 1),
2994            block_dim: (64, 1, 1),
2995            shared_mem_bytes: 0,
2996        };
2997        let __s_bl = self.gpu.stream();
2998        let mut bl = __s_bl.launch_builder(&f);
2999        bl.arg(len_ptrs).arg(saved).arg(acc).arg(&b).arg(&nl);
3000        unsafe {
3001            bl.launch(cfg)?;
3002        }
3003        Ok(())
3004    }
3005
3006    /// OPTIPIPE increment 1: derive the K=1 successor-valid bit on device.
3007    pub fn spec_fork_valid(
3008        &self,
3009        acc: &CudaSlice<u32>,
3010        optimistic_pending: u32,
3011        valid: &mut CudaSlice<u32>,
3012    ) -> Result<(), Box<dyn std::error::Error>> {
3013        let f = self.func("spec_fork_valid");
3014        let cfg = LaunchConfig {
3015            grid_dim: (1, 1, 1),
3016            block_dim: (1, 1, 1),
3017            shared_mem_bytes: 0,
3018        };
3019        let __s_bl = self.gpu.stream();
3020        let mut bl = __s_bl.launch_builder(&f);
3021        bl.arg(acc).arg(&optimistic_pending).arg(valid);
3022        unsafe {
3023            bl.launch(cfg)?;
3024        }
3025        Ok(())
3026    }
3027
3028    /// OPTIPIPE increment 1: leave stage-local KV lengths on hit, restore them on miss.
3029    pub fn spec_fork_reconcile_kv(
3030        &self,
3031        len_ptrs: &CudaSlice<u64>,
3032        saved: &CudaSlice<i32>,
3033        acc: &CudaSlice<u32>,
3034        valid: &CudaSlice<u32>,
3035        base: usize,
3036        n_layer: usize,
3037    ) -> Result<(), Box<dyn std::error::Error>> {
3038        let f = self.func("spec_fork_reconcile_kv");
3039        let (b, nl) = (base as i32, n_layer as i32);
3040        let cfg = LaunchConfig {
3041            grid_dim: (n_layer.div_ceil(64) as u32, 1, 1),
3042            block_dim: (64, 1, 1),
3043            shared_mem_bytes: 0,
3044        };
3045        let __s_bl = self.gpu.stream();
3046        let mut bl = __s_bl.launch_builder(&f);
3047        bl.arg(len_ptrs)
3048            .arg(saved)
3049            .arg(acc)
3050            .arg(valid)
3051            .arg(&b)
3052            .arg(&nl);
3053        unsafe {
3054            bl.launch(cfg)?;
3055        }
3056        Ok(())
3057    }
3058
3059    /// OPTIPIPE increment 1: conditionally restore one stage-owned recurrent-state buffer.
3060    pub fn spec_fork_restore_f32(
3061        &self,
3062        snapshot: &CudaSlice<f32>,
3063        state: &mut CudaSlice<f32>,
3064        valid: &CudaSlice<u32>,
3065    ) -> Result<(), Box<dyn std::error::Error>> {
3066        assert_eq!(
3067            snapshot.len(),
3068            state.len(),
3069            "fork recurrent snapshot shape mismatch"
3070        );
3071        let f = self.func("spec_fork_restore_f32");
3072        let n = state.len() as i32;
3073        let blocks = state.len().div_ceil(256).min(65535).max(1) as u32;
3074        let cfg = LaunchConfig {
3075            grid_dim: (blocks, 1, 1),
3076            block_dim: (256, 1, 1),
3077            shared_mem_bytes: 0,
3078        };
3079        let __s_bl = self.gpu.stream();
3080        let mut bl = __s_bl.launch_builder(&f);
3081        bl.arg(snapshot).arg(state).arg(valid).arg(&n);
3082        unsafe {
3083            bl.launch(cfg)?;
3084        }
3085        Ok(())
3086    }
3087
3088    /// ROUND-STREAM stage (b): device next-round seed gather (see spec_seed_gather header).
3089    /// Caller D2Ds h_seed into fill_prev after (both slots carry the same value in every arm).
3090    pub fn spec_seed_gather(
3091        &self,
3092        vx: &CudaSlice<f32>,
3093        fill_prev: &CudaSlice<f32>,
3094        acc: &CudaSlice<u32>,
3095        h_seed: &mut CudaSlice<f32>,
3096        base: usize,
3097        n_embd: usize,
3098    ) -> Result<(), Box<dyn std::error::Error>> {
3099        let f = self.func("spec_seed_gather");
3100        let (b, ne) = (base as i32, n_embd as i32);
3101        let cfg = LaunchConfig {
3102            grid_dim: (n_embd.div_ceil(256) as u32, 1, 1),
3103            block_dim: (256, 1, 1),
3104            shared_mem_bytes: 0,
3105        };
3106        let __s_bl = self.gpu.stream();
3107        let mut bl = __s_bl.launch_builder(&f);
3108        bl.arg(vx)
3109            .arg(fill_prev)
3110            .arg(acc)
3111            .arg(h_seed)
3112            .arg(&b)
3113            .arg(&ne);
3114        unsafe {
3115            bl.launch(cfg)?;
3116        }
3117        Ok(())
3118    }
3119
3120    /// ROUND-STREAM stage (a): device greedy accept walk (see spec_accept_greedy header).
3121    pub fn spec_accept_greedy(
3122        &self,
3123        preds: &CudaSlice<u32>,
3124        draft: &CudaSlice<u32>,
3125        last_pred: u32,
3126        base: usize,
3127        k_round: usize,
3128        out: &mut CudaSlice<u32>,
3129    ) -> Result<(), Box<dyn std::error::Error>> {
3130        let f = self.func("spec_accept_greedy");
3131        let (b, k) = (base as i32, k_round as i32);
3132        let cfg = LaunchConfig {
3133            grid_dim: (1, 1, 1),
3134            block_dim: (32, 1, 1),
3135            shared_mem_bytes: 0,
3136        };
3137        let __s_bl = self.gpu.stream();
3138        let mut bl = __s_bl.launch_builder(&f);
3139        bl.arg(preds)
3140            .arg(draft)
3141            .arg(&last_pred)
3142            .arg(&b)
3143            .arg(&k)
3144            .arg(out);
3145        unsafe {
3146            bl.launch(cfg)?;
3147        }
3148        Ok(())
3149    }
3150
3151    // ================= SAMPLED-SPEC PRIMITIVES (spec_sample.cu, piece A) =================
3152    // Counter-based randomness: every call takes (seed, stream_pos) — the caller owns the
3153    // event counter (one per sampled token). temp <= 0 arms are exact greedy limits.
3154
3155    /// y = x/temp + Gumbel(Philox(seed, stream_pos)) over n logits (then run device argmax on y
3156    /// = one categorical sample at temperature `temp`). temp<=0: y = x (pure copy).
3157    pub fn gumbel_perturb(
3158        &self,
3159        x: &CudaSlice<f32>,
3160        y: &mut CudaSlice<f32>,
3161        n: usize,
3162        seed: u64,
3163        stream_pos: u32,
3164        temp: f32,
3165    ) -> Result<(), Box<dyn std::error::Error>> {
3166        let f = self.func("gumbel_perturb_f32");
3167        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3168        let cfg = LaunchConfig {
3169            grid_dim: (n.div_ceil(256) as u32, 1, 1),
3170            block_dim: (256, 1, 1),
3171            shared_mem_bytes: 0,
3172        };
3173        let __s_b = self.gpu.stream();
3174        let mut b = __s_b.launch_builder(&f);
3175        b.arg(x)
3176            .arg(&mut *y)
3177            .arg(&ni)
3178            .arg(&slo)
3179            .arg(&shi)
3180            .arg(&stream_pos)
3181            .arg(&temp);
3182        unsafe {
3183            b.launch(cfg)?;
3184        }
3185        Ok(())
3186    }
3187
3188    /// GRAMMAR TOKEN MASK (constrained decoding, lane/constrained-full): ban every vocab id
3189    /// whose bit is unset in the packed llguidance bitset, IN PLACE on row `col` of a stacked
3190    /// [B, n_vocab] logits buffer. `mask` = the SimpleVob u32 words H2D'd verbatim
3191    /// (~n_vocab/8 bytes/step — trivial on PCIe); ids >= 32*mask_words (padded lm_head tail)
3192    /// are banned too, the device twin of constrained::apply_mask. Banned value -FLT_MAX ==
3193    /// the argmax/gumbel kernels' init sentinel, so a fully-banned tail can never win and
3194    /// ordering matches the host -inf mask bit-for-bit for every finite logit.
3195    pub fn mask_logits_col(
3196        &self,
3197        logits: &mut CudaSlice<f32>,
3198        mask: &CudaSlice<u32>,
3199        col: usize,
3200        n: usize,
3201        mask_words: usize,
3202    ) -> Result<(), Box<dyn std::error::Error>> {
3203        let f = self.func("mask_logits_f32");
3204        let (ci, ni, mw) = (col as i32, n as i32, mask_words as i32);
3205        let cfg = LaunchConfig {
3206            grid_dim: (n.div_ceil(256).min(1024) as u32, 1, 1),
3207            block_dim: (256, 1, 1),
3208            shared_mem_bytes: 0,
3209        };
3210        let __s_b = self.gpu.stream();
3211        let mut b = __s_b.launch_builder(&f);
3212        b.arg(&mut *logits).arg(mask).arg(&ci).arg(&ni).arg(&mw);
3213        unsafe {
3214            b.launch(cfg)?;
3215        }
3216        Ok(())
3217    }
3218
3219    /// Column-`col` twin of `gumbel_perturb` over stacked logits [B, n_vocab] (the batched
3220    /// serving tick's device sampler): y = x[col]/temp + gumbel(seed, stream_pos, lane).
3221    /// SAME kernel/Philox mapping as `gumbel_perturb` — bit-identical perturbation for the
3222    /// same (seed, stream_pos, temp) regardless of which batch column the row sits in
3223    /// (the lane index is the in-row position; `col` only moves the input pointer). That
3224    /// pointer-invariance IS the serving isolation contract for sampled rows.
3225    pub fn gumbel_perturb_col(
3226        &self,
3227        x: &CudaSlice<f32>,
3228        col: usize,
3229        y: &mut CudaSlice<f32>,
3230        n: usize,
3231        seed: u64,
3232        stream_pos: u32,
3233        temp: f32,
3234    ) -> Result<(), Box<dyn std::error::Error>> {
3235        let f = self.func("gumbel_perturb_f32");
3236        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3237        let col_view = x.slice(col * n..(col + 1) * n);
3238        let cfg = LaunchConfig {
3239            grid_dim: (n.div_ceil(256) as u32, 1, 1),
3240            block_dim: (256, 1, 1),
3241            shared_mem_bytes: 0,
3242        };
3243        let __s_b = self.gpu.stream();
3244        let mut b = __s_b.launch_builder(&f);
3245        b.arg(&col_view)
3246            .arg(&mut *y)
3247            .arg(&ni)
3248            .arg(&slo)
3249            .arg(&shi)
3250            .arg(&stream_pos)
3251            .arg(&temp);
3252        unsafe {
3253            b.launch(cfg)?;
3254        }
3255        Ok(())
3256    }
3257
3258    /// Filtered twin of `gumbel_perturb_col`: the per-row (row_max, th) floor comes from
3259    /// DEVICE buffers (`filter_stats` output slots at `stat_idx`) — one filtered draw from
3260    /// the top-k/top-p/min-p-truncated softmax with no stat D2H and no row copy. Same
3261    /// Philox mapping as every gumbel kernel (pointer-invariant across batch columns —
3262    /// the serving isolation contract for sampled rows).
3263    #[allow(clippy::too_many_arguments)]
3264    pub fn gumbel_perturb_filtered_col(
3265        &self,
3266        x: &CudaSlice<f32>,
3267        col: usize,
3268        y: &mut CudaSlice<f32>,
3269        n: usize,
3270        seed: u64,
3271        stream_pos: u32,
3272        temp: f32,
3273        stat_max: &CudaSlice<f32>,
3274        stat_th: &CudaSlice<f32>,
3275        stat_idx: usize,
3276    ) -> Result<(), Box<dyn std::error::Error>> {
3277        let f = self.func("gumbel_perturb_filtered_col_f32");
3278        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3279        let (ci, si) = (col as i32, stat_idx as i32);
3280        let cfg = LaunchConfig {
3281            grid_dim: (n.div_ceil(256) as u32, 1, 1),
3282            block_dim: (256, 1, 1),
3283            shared_mem_bytes: 0,
3284        };
3285        let __s_b = self.gpu.stream();
3286        let mut b = __s_b.launch_builder(&f);
3287        b.arg(x)
3288            .arg(&ci)
3289            .arg(&mut *y)
3290            .arg(&ni)
3291            .arg(&slo)
3292            .arg(&shi)
3293            .arg(&stream_pos)
3294            .arg(&temp)
3295            .arg(stat_max)
3296            .arg(stat_th)
3297            .arg(&si);
3298        unsafe {
3299            b.launch(cfg)?;
3300        }
3301        Ok(())
3302    }
3303
3304    /// In-graph sampling-event counter bump (spec_sample.cu kernel 5): ctr[0] += 1. The sampled
3305    /// graph-draft chain replays with FIXED kernel args, so the Philox event counter must be
3306    /// DEVICE data — the host seeds it once per round; every replay bumps it before the perturb
3307    /// reads it (counter is data, not state — graph-replay-safe).
3308    pub fn sctr_inc(&self, ctr: &mut CudaSlice<u32>) -> Result<(), Box<dyn std::error::Error>> {
3309        let f = self.func("memra_sctr_inc");
3310        let cfg = LaunchConfig {
3311            grid_dim: (1, 1, 1),
3312            block_dim: (1, 1, 1),
3313            shared_mem_bytes: 0,
3314        };
3315        let __s_b = self.gpu.stream();
3316        let mut b = __s_b.launch_builder(&f);
3317        b.arg(&mut *ctr);
3318        unsafe {
3319            b.launch(cfg)?;
3320        }
3321        Ok(())
3322    }
3323
3324    /// Graph-capturable `gumbel_perturb`: the sampling-event counter comes from DEVICE memory
3325    /// (`ctr[0]`) instead of a host scalar. Identical math to `gumbel_perturb` at
3326    /// stream_pos == ctr[0] (same Philox call, same lane mapping) — the eager and graph sampled
3327    /// chains produce bit-identical perturbations for the same (seed, counter, temp).
3328    pub fn gumbel_perturb_ctr(
3329        &self,
3330        x: &CudaSlice<f32>,
3331        y: &mut CudaSlice<f32>,
3332        n: usize,
3333        seed: u64,
3334        ctr: &CudaSlice<u32>,
3335        temp: f32,
3336    ) -> Result<(), Box<dyn std::error::Error>> {
3337        let f = self.func("gumbel_perturb_ctr_f32");
3338        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3339        let cfg = LaunchConfig {
3340            grid_dim: (n.div_ceil(256) as u32, 1, 1),
3341            block_dim: (256, 1, 1),
3342            shared_mem_bytes: 0,
3343        };
3344        let __s_b = self.gpu.stream();
3345        let mut b = __s_b.launch_builder(&f);
3346        b.arg(x)
3347            .arg(&mut *y)
3348            .arg(&ni)
3349            .arg(&slo)
3350            .arg(&shi)
3351            .arg(ctr)
3352            .arg(&temp);
3353        unsafe {
3354            b.launch(cfg)?;
3355        }
3356        Ok(())
3357    }
3358
3359    /// out[pair] = softmax_temp(x[rows[pair]])[ids[pair]] for npair (row, id) pairs; rows index
3360    /// into x with `row_stride` f32s per row. temp<=0: out = 1.0 iff id is the row argmax
3361    /// (smallest-index tie-break — matches the argmax-gate contract).
3362    pub fn softmax_gather(
3363        &self,
3364        x: &CudaSlice<f32>,
3365        row_stride: usize,
3366        ids: &CudaSlice<u32>,
3367        rows: &CudaSlice<i32>,
3368        out: &mut CudaSlice<f32>,
3369        n: usize,
3370        npair: usize,
3371        temp: f32,
3372    ) -> Result<(), Box<dyn std::error::Error>> {
3373        let f = self.func("softmax_gather_f32");
3374        let (ni, rs) = (n as i32, row_stride as i64);
3375        let np = npair as i32;
3376        let cfg = LaunchConfig {
3377            grid_dim: (npair as u32, 1, 1),
3378            block_dim: (256, 1, 1),
3379            shared_mem_bytes: 0,
3380        };
3381        let __s_b = self.gpu.stream();
3382        let mut b = __s_b.launch_builder(&f);
3383        b.arg(x)
3384            .arg(&rs)
3385            .arg(ids)
3386            .arg(rows)
3387            .arg(&mut *out)
3388            .arg(&ni)
3389            .arg(&np)
3390            .arg(&temp);
3391        unsafe {
3392            b.launch(cfg)?;
3393        }
3394        Ok(())
3395    }
3396
3397    /// Sample token from norm(max(0, softmax_temp(p) - softmax_temp(q))) (q = None -> plain
3398    /// categorical from softmax_temp(p)). Row stats (max, sumexp at temp) must be precomputed
3399    /// (softmax_gather's pass-1 values; see spec.rs caller). Deterministic fixed-order CDF walk.
3400    pub fn residual_sample(
3401        &self,
3402        p: &CudaSlice<f32>,
3403        q: Option<&CudaSlice<f32>>,
3404        n: usize,
3405        temp: f32,
3406        seed: u64,
3407        stream_pos: u32,
3408        out_tok: &mut CudaSlice<u32>,
3409    ) -> Result<(), Box<dyn std::error::Error>> {
3410        let f = self.func("residual_sample_f32");
3411        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3412        let nth = 1024u32;
3413        let cfg = LaunchConfig {
3414            grid_dim: (1, 1, 1),
3415            block_dim: (nth, 1, 1),
3416            shared_mem_bytes: 0,
3417        };
3418        let has_q: i32 = q.is_some() as i32;
3419        let qbuf = q.unwrap_or(p); // dummy when absent; kernel gates on has_q
3420        let __s_b = self.gpu.stream();
3421        let mut b = __s_b.launch_builder(&f);
3422        b.arg(p)
3423            .arg(qbuf)
3424            .arg(&has_q)
3425            .arg(&ni)
3426            .arg(&temp)
3427            .arg(&slo)
3428            .arg(&shi)
3429            .arg(&stream_pos)
3430            .arg(&mut *out_tok);
3431        unsafe {
3432            b.launch(cfg)?;
3433        }
3434        Ok(())
3435    }
3436
3437    /// Access the shared MoE residency cache (EDGE-1 §B), building it on first use under
3438    /// MEMRA_MOE_CACHE. The closure runs while the lock is held — keep it to lookup/admit/issue, not
3439    /// the GEMM. `max_block_bytes` sizes the slots (largest of gate/up/down). Returns the closure's
3440    /// result. If MEMRA_MOE_CACHE is unset this is never called (the caller checks the env first).
3441    pub fn with_moe_cache<R>(
3442        &self,
3443        max_block_bytes: usize,
3444        f: impl FnOnce(
3445            &mut crate::moe_cache::MoeSlotCache,
3446            &Engine,
3447        ) -> Result<R, Box<dyn std::error::Error>>,
3448    ) -> Result<R, Box<dyn std::error::Error>> {
3449        let mut guard = self.moe_cache.lock().unwrap();
3450        if guard.is_none() {
3451            *guard = Some(crate::moe_cache::MoeSlotCache::new(self, max_block_bytes)?);
3452        }
3453        let cache = guard.as_mut().unwrap();
3454        f(cache, self)
3455    }
3456
3457    /// Freeze the already-built MoE residency set. This never constructs a cache: callers use it
3458    /// only after a real prefill has populated the machine-specific CPU/GPU working set.
3459    pub fn freeze_moe_cache(&self) {
3460        if let Some(cache) = self.moe_cache.lock().unwrap().as_mut() {
3461            cache.freeze();
3462        }
3463    }
3464
3465    /// The current residency set as (layer, proj, ex) triples, or None if no cache was built.
3466    /// Never constructs a cache.
3467    pub fn export_moe_residency(&self) -> Option<Vec<(u16, u8, u16)>> {
3468        self.moe_cache
3469            .lock()
3470            .unwrap()
3471            .as_ref()
3472            .map(crate::moe_cache::MoeSlotCache::export_residency)
3473    }
3474
3475    pub(crate) fn moe_cache_frozen(&self) -> bool {
3476        self.moe_cache
3477            .lock()
3478            .unwrap()
3479            .as_ref()
3480            .is_some_and(crate::moe_cache::MoeSlotCache::is_frozen)
3481    }
3482
3483    /// A frozen heterogeneous CPU/GPU expert split cannot use Hy3's ordinary batched prefill
3484    /// efficiently: T>=PRIME_MIN_T bypasses the CPU backend and transiently rereads every missing
3485    /// expert through the GPU spill path. Replay the short prompt through decode after freezing,
3486    /// while leaving the profiling warmup's established batched behavior untouched.
3487    /// (`pub`: run-gen's #46 batched-prime gate skips itself when generation will take the
3488    /// tokenwise arm anyway.)
3489    pub fn frozen_cpu_experts_prefer_tokenwise_prime(&self) -> bool {
3490        crate::cpu_experts::configured()
3491            && self.moe_cache_frozen()
3492            && std::env::var("MEMRA_CPU_EXPERT_BATCHED_PRIME").as_deref() != Ok("1")
3493    }
3494
3495    /// Install the loaded model's exact retained expert-block inventory before lazy cache build.
3496    pub(crate) fn configure_moe_cache_layout(&self, block_bytes: Vec<usize>) {
3497        assert!(
3498            self.moe_cache.lock().unwrap().is_none(),
3499            "MoE cache layout configured after cache construction"
3500        );
3501        *self.moe_cache_layout.lock().unwrap() = Some(block_bytes);
3502    }
3503
3504    pub(crate) fn moe_cache_layout(&self) -> Option<Vec<usize>> {
3505        self.moe_cache_layout.lock().unwrap().clone()
3506    }
3507
3508    /// True if the MoE residency cache is enabled (MEMRA_MOE_CACHE set).
3509    pub fn moe_cache_enabled() -> bool {
3510        std::env::var("MEMRA_MOE_CACHE").as_deref() != Ok("0")
3511    }
3512
3513    /// Snapshot the MoE cache counters (hits, misses, staged_bytes, n_slots) for the §D.4 PCIe gate.
3514    /// Returns None if the cache was never built (disabled or no MoE forward ran).
3515    pub fn moe_cache_stats(&self) -> Option<(u64, u64, u64, usize)> {
3516        let guard = self.moe_cache.lock().unwrap();
3517        guard
3518            .as_ref()
3519            .map(|c| (c.hits, c.misses, c.staged_bytes, c.n_slots()))
3520    }
3521
3522    /// Experimental CPU expert backend counters: completed layer calls, experts served, and the
3523    /// sum of backend wall nanoseconds. The timer includes explicit disk->RAM fills on cache misses;
3524    /// callers compare a before/after snapshot around a decode window.
3525    pub fn cpu_expert_stats(
3526        &self,
3527    ) -> Option<(u64, u64, u64, u64, u64, u64, u64, u64, u64, u64, u64)> {
3528        crate::cpu_experts::configured().then(crate::cpu_experts::stats)
3529    }
3530
3531    /// Caller-blocked nanoseconds at CPU expert joins. Compare before/after snapshots to measure
3532    /// the backend tail that resident-GPU expert work did not hide.
3533    pub fn cpu_expert_predictor_stats(&self) -> (u64, u64) {
3534        crate::cpu_experts::predictor_stats()
3535    }
3536
3537    pub fn cpu_expert_exposed_wait_ns(&self) -> Option<u64> {
3538        crate::cpu_experts::configured().then(crate::cpu_experts::exposed_wait_ns)
3539    }
3540
3541    /// CPU-routed expert selections grouped by how many of their three projections were already
3542    /// resident in HBM. This makes otherwise-stranded partial residency visible to tuning runs.
3543    pub fn cpu_expert_gpu_residency_stats(&self) -> Option<(u64, u64, u64)> {
3544        crate::cpu_experts::configured().then(crate::cpu_experts::incomplete_gpu_residency_stats)
3545    }
3546
3547    /// Positioned-read proof-backend counters:
3548    /// `(reads, bytes, read_errors, short_reads, mmap_fallbacks, buffer_waits, ring_full)`.
3549    pub fn moe_pread_stats(&self) -> Option<(u64, u64, u64, u64, u64, u64, u64)> {
3550        let guard = self.moe_cache.lock().unwrap();
3551        guard
3552            .as_ref()
3553            .and_then(|cache| cache.pread_stats())
3554            .map(|stats| {
3555                (
3556                    stats.reads,
3557                    stats.bytes,
3558                    stats.read_errors,
3559                    stats.short_reads,
3560                    stats.fallbacks,
3561                    stats.buffer_waits,
3562                    stats.ring_full,
3563                )
3564            })
3565    }
3566
3567    /// Spill configuration values that warned and substituted their documented defaults.
3568    pub fn spill_config_fallbacks(&self) -> u64 {
3569        crate::spill_pread::config_fallbacks()
3570    }
3571
3572    /// Reset the MoE cache perf counters (to separate warmup from steady-state windows).
3573    pub fn moe_cache_reset_counters(&self) {
3574        if let Some(c) = self.moe_cache.lock().unwrap().as_mut() {
3575            c.reset_counters();
3576        }
3577    }
3578
3579    pub fn htod_bytes(&self, v: &[u8]) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
3580        Ok(self.gpu.stream().clone_htod(v)?)
3581    }
3582
3583    /// `htod_bytes` with a mapped (uninit) tail pad: the wide-load expert dots read up to 6B
3584    /// past the final q4_0 block through their aligned window — the bytes never reach a
3585    /// result (funnelshift discards them) but must be mapped memory.
3586    pub fn htod_bytes_padded(
3587        &self,
3588        v: &[u8],
3589        pad: usize,
3590    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
3591        let mut d = self.alloc_u8_uninit(v.len() + pad)?;
3592        {
3593            let mut view = d.slice_mut(0..v.len());
3594            self.gpu.stream().memcpy_htod(v, &mut view)?;
3595        }
3596        Ok(d)
3597    }
3598
3599    /// Device-to-device copy of `src` into `dst[off..off+len]` (f32). For in-place KV append.
3600    pub fn copy_into(
3601        &self,
3602        dst: &mut CudaSlice<f32>,
3603        off: usize,
3604        src: &CudaSlice<f32>,
3605        len: usize,
3606    ) -> Result<(), Box<dyn std::error::Error>> {
3607        let mut view = dst.slice_mut(off..off + len);
3608        self.gpu
3609            .stream()
3610            .memcpy_dtod(&src.slice(0..len), &mut view)?;
3611        Ok(())
3612    }
3613
3614    /// View a sub-range of a device buffer (for attending over [0..len) of a KV cache).
3615    /// u8 twin of copy_into (D2D byte-range copy at an offset).
3616    pub fn copy_u8_into(
3617        &self,
3618        dst: &mut CudaSlice<u8>,
3619        off: usize,
3620        src: &CudaSlice<u8>,
3621        len: usize,
3622    ) -> Result<(), Box<dyn std::error::Error>> {
3623        let mut view = dst.slice_mut(off..off + len);
3624        self.gpu
3625            .stream()
3626            .memcpy_dtod(&src.slice(0..len), &mut view)?;
3627        Ok(())
3628    }
3629
3630    /// D2D byte-range copy with explicit source and destination offsets.
3631    pub fn copy_u8_range_into(
3632        &self,
3633        dst: &mut CudaSlice<u8>,
3634        dst_off: usize,
3635        src: &CudaSlice<u8>,
3636        src_off: usize,
3637        len: usize,
3638    ) -> Result<(), Box<dyn std::error::Error>> {
3639        let mut dst_view = dst.slice_mut(dst_off..dst_off + len);
3640        self.gpu
3641            .stream()
3642            .memcpy_dtod(&src.slice(src_off..src_off + len), &mut dst_view)?;
3643        Ok(())
3644    }
3645
3646    /// Resolve an absolute append slot to the Step35 SWA layer's physical rows. At wrap, copy
3647    /// only the aligned live prefix through temporary device storage and rebase it at row zero,
3648    /// keeping the audited attention range contiguous without changing its absolute start.
3649    pub fn prepare_kv_append(
3650        &self,
3651        kv: &mut crate::cache::KvLayer,
3652        retain_from: usize,
3653        append_rows: usize,
3654    ) -> Result<usize, Box<dyn std::error::Error>> {
3655        let Some(plan) = kv
3656            .ring
3657            .as_ref()
3658            .map(|ring| ring.append_plan(kv.len, retain_from, append_rows))
3659            .transpose()?
3660        else {
3661            return Ok(kv.len);
3662        };
3663        match plan {
3664            crate::cache::KvRingAppend::Contiguous { write_row } => Ok(write_row),
3665            crate::cache::KvRingAppend::Rebase {
3666                src_row,
3667                keep_rows,
3668                new_base,
3669                write_row,
3670            } => {
3671                if keep_rows > 0 {
3672                    let k_len = keep_rows * kv.k_tok_bytes;
3673                    let v_len = keep_rows * kv.v_tok_bytes;
3674                    let mut k_tmp = self.alloc_u8_uninit(k_len)?;
3675                    let mut v_tmp = self.alloc_u8_uninit(v_len)?;
3676                    self.copy_u8_range_into(&mut k_tmp, 0, &kv.k, src_row * kv.k_tok_bytes, k_len)?;
3677                    self.copy_u8_range_into(&mut v_tmp, 0, &kv.v, src_row * kv.v_tok_bytes, v_len)?;
3678                    self.copy_u8_into(&mut kv.k, 0, &k_tmp, k_len)?;
3679                    self.copy_u8_into(&mut kv.v, 0, &v_tmp, v_len)?;
3680                }
3681                kv.ring.as_mut().unwrap().apply_rebase(new_base);
3682                Ok(write_row)
3683            }
3684        }
3685    }
3686
3687    /// H2D write of `src` into `dst[off..off+src.len()]` (u8). In-place row updates for the
3688    /// adaptive trim head: no realloc, so captured graphs keep their baked addresses.
3689    pub fn htod_u8_into(
3690        &self,
3691        dst: &mut CudaSlice<u8>,
3692        off: usize,
3693        src: &[u8],
3694    ) -> Result<(), Box<dyn std::error::Error>> {
3695        let mut view = dst.slice_mut(off..off + src.len());
3696        self.gpu.stream().memcpy_htod(src, &mut view)?;
3697        Ok(())
3698    }
3699
3700    pub fn view<'a>(&self, b: &'a CudaSlice<f32>, len: usize) -> cudarc::driver::CudaView<'a, f32> {
3701        b.slice(0..len)
3702    }
3703
3704    /// View the first `len` BYTES of a u8 device buffer (quantized KV cache: [0..t_kv*tok_bytes)).
3705    /// Byte-range view (gemma4 R6 window offset into the quantized KV stream).
3706    pub fn view_u8_range<'a>(
3707        &self,
3708        b: &'a CudaSlice<u8>,
3709        start: usize,
3710        end: usize,
3711    ) -> cudarc::driver::CudaView<'a, u8> {
3712        b.slice(start..end)
3713    }
3714    pub fn view_u8<'a>(
3715        &self,
3716        b: &'a CudaSlice<u8>,
3717        len: usize,
3718    ) -> cudarc::driver::CudaView<'a, u8> {
3719        b.slice(0..len)
3720    }
3721
3722    /// Append-quantize ONE token's post-RoPE K (q8_0) and V (q5_1) into the resident byte caches at
3723    /// token index `t` (KVQUANT-PLAN §C). One CTA (one warp) per 32-element block; the kernel writes
3724    /// the f16 scale(s) + packed quants for K and V. k_row/v_row are f32 [kv_dim_k]/[kv_dim_v].
3725    pub fn append_kv_quantized(
3726        &self,
3727        k_row: &CudaSlice<f32>,
3728        v_row: &CudaSlice<f32>,
3729        kc: &mut CudaSlice<u8>,
3730        vc: &mut CudaSlice<u8>,
3731        t: usize,
3732        kv_dim_k: usize,
3733        kv_dim_v: usize,
3734        k_tok_bytes: usize,
3735        v_tok_bytes: usize,
3736        g: bool,
3737    ) -> Result<(), Box<dyn std::error::Error>> {
3738        let f = if g {
3739            self.func_g("append_quantize_kv_q8_0_q5_1")
3740        } else {
3741            self.func("append_quantize_kv_q8_0_q5_1")
3742        };
3743        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
3744        let cfg = LaunchConfig {
3745            grid_dim: (nblk, 1, 1),
3746            block_dim: (32, 1, 1),
3747            shared_mem_bytes: 0,
3748        };
3749        let (ti, kdk, kdv) = (t as i32, kv_dim_k as i32, kv_dim_v as i32);
3750        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3751        let __s_b = self.gpu.stream();
3752        let mut b = __s_b.launch_builder(&f);
3753        b.arg(k_row)
3754            .arg(v_row)
3755            .arg(kc)
3756            .arg(vc)
3757            .arg(&ti)
3758            .arg(&kdk)
3759            .arg(&kdv)
3760            .arg(&ktb)
3761            .arg(&vtb);
3762        unsafe {
3763            b.launch(cfg)?;
3764        }
3765        Ok(())
3766    }
3767
3768    /// Device-counter variant of `append_kv_quantized` (CUDA-GRAPH-PLAN Phase 2): the write slot
3769    /// `t` is read from `t_dev[0]` (a resident device i32[1]) instead of a host int arg, so the
3770    /// launch args are FIXED across decode steps (graph-capturable). Identical quant math.
3771    pub fn append_kv_quantized_dc(
3772        &self,
3773        k_row: &CudaSlice<f32>,
3774        v_row: &CudaSlice<f32>,
3775        kc: &mut CudaSlice<u8>,
3776        vc: &mut CudaSlice<u8>,
3777        t_dev: &CudaSlice<i32>,
3778        kv_dim_k: usize,
3779        kv_dim_v: usize,
3780        k_tok_bytes: usize,
3781        v_tok_bytes: usize,
3782        g: bool,
3783    ) -> Result<(), Box<dyn std::error::Error>> {
3784        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
3785        let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
3786        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3787        // PDL wave-B2: flash-module flavor mirrors the builder path's g flag exactly.
3788        if Self::pdl_on() && Self::pdl_wb_on() {
3789            use cudarc::driver::{DevicePtr, DevicePtrMut};
3790            let s = &self.gpu.stream();
3791            let (pk, _g0) = k_row.device_ptr(s);
3792            let (pv, _g1) = v_row.device_ptr(s);
3793            let (pkc, _g2) = kc.device_ptr_mut(s);
3794            let (pvc, _g3) = vc.device_ptr_mut(s);
3795            let (pt, _g4) = t_dev.device_ptr(s);
3796            let mut ps = [
3797                &pk as *const _ as *mut std::ffi::c_void,
3798                &pv as *const _ as *mut _,
3799                &pkc as *const _ as *mut _,
3800                &pvc as *const _ as *mut _,
3801                &pt as *const _ as *mut _,
3802                &kdk as *const _ as *mut _,
3803                &kdv as *const _ as *mut _,
3804                &ktb as *const _ as *mut _,
3805                &vtb as *const _ as *mut _,
3806            ];
3807            unsafe {
3808                self.launch_pdl_flash(
3809                    g,
3810                    "append_quantize_kv_q8_0_q5_1_dc",
3811                    (nblk, 1, 1),
3812                    (32, 1, 1),
3813                    0,
3814                    &mut ps,
3815                )?;
3816            }
3817            return Ok(());
3818        }
3819        let f = if g {
3820            self.func_g("append_quantize_kv_q8_0_q5_1_dc")
3821        } else {
3822            self.func("append_quantize_kv_q8_0_q5_1_dc")
3823        };
3824        let cfg = LaunchConfig {
3825            grid_dim: (nblk, 1, 1),
3826            block_dim: (32, 1, 1),
3827            shared_mem_bytes: 0,
3828        };
3829        let __s_b = self.gpu.stream();
3830        let mut b = __s_b.launch_builder(&f);
3831        b.arg(k_row)
3832            .arg(v_row)
3833            .arg(kc)
3834            .arg(vc)
3835            .arg(t_dev)
3836            .arg(&kdk)
3837            .arg(&kdv)
3838            .arg(&ktb)
3839            .arg(&vtb);
3840        unsafe {
3841            b.launch(cfg)?;
3842        }
3843        Ok(())
3844    }
3845
3846    /// Append-quantize T token rows in one shot (BATCHED PROMPT PRIME). k_rows/v_rows are
3847    /// token-major [T, kv_dim] post-RoPE f32; rows land at cache slots t0..t0+T. Default = the
3848    /// batched `_rows` kernel: one (nblk, T) launch whose per-(block,token) warp program is the
3849    /// per-token append kernel verbatim -> every written row is BIT-IDENTICAL to T sequential
3850    /// `append_kv_quantized_view` calls (kernel_check pins the bytes). MEMRA_PRIME_APPEND_LOOP=1
3851    /// forces the T-launch per-row loop (the A/B seam that measured the launch overhead).
3852    #[allow(clippy::too_many_arguments)]
3853    pub fn append_kv_quantized_rows(
3854        &self,
3855        k_rows: &CudaSlice<f32>,
3856        v_rows: &CudaSlice<f32>,
3857        kc: &mut CudaSlice<u8>,
3858        vc: &mut CudaSlice<u8>,
3859        t0: usize,
3860        t: usize,
3861        kv_dim_k: usize,
3862        kv_dim_v: usize,
3863        k_tok_bytes: usize,
3864        v_tok_bytes: usize,
3865        g: bool,
3866    ) -> Result<(), Box<dyn std::error::Error>> {
3867        if std::env::var("MEMRA_PRIME_APPEND_LOOP").is_ok() {
3868            for i in 0..t {
3869                let k_row = k_rows.slice(i * kv_dim_k..(i + 1) * kv_dim_k);
3870                let v_row = v_rows.slice(i * kv_dim_v..(i + 1) * kv_dim_v);
3871                self.append_kv_quantized_view(
3872                    &k_row,
3873                    &v_row,
3874                    kc,
3875                    vc,
3876                    t0 + i,
3877                    kv_dim_k,
3878                    kv_dim_v,
3879                    k_tok_bytes,
3880                    v_tok_bytes,
3881                    g,
3882                )?;
3883            }
3884            return Ok(());
3885        }
3886        let f = if g {
3887            self.func_g("append_quantize_kv_q8_0_q5_1_rows")
3888        } else {
3889            self.func("append_quantize_kv_q8_0_q5_1_rows")
3890        };
3891        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
3892        let cfg = LaunchConfig {
3893            grid_dim: (nblk, t as u32, 1),
3894            block_dim: (32, 1, 1),
3895            shared_mem_bytes: 0,
3896        };
3897        let (t0i, kdk, kdv) = (t0 as i32, kv_dim_k as i32, kv_dim_v as i32);
3898        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3899        let __s_b = self.gpu.stream();
3900        let mut b = __s_b.launch_builder(&f);
3901        b.arg(k_rows)
3902            .arg(v_rows)
3903            .arg(kc)
3904            .arg(vc)
3905            .arg(&t0i)
3906            .arg(&kdk)
3907            .arg(&kdv)
3908            .arg(&ktb)
3909            .arg(&vtb);
3910        unsafe {
3911            b.launch(cfg)?;
3912        }
3913        Ok(())
3914    }
3915
3916    /// Increment a device i32[1] counter in place (p[0] += 1) via the resident `inc_i32` kernel.
3917    /// Used to advance the device-resident seqlen/pos counters inside the decode-dc path (and,
3918    /// later, inside a captured graph) without a host round-trip.
3919    pub fn inc_seqlen(&self, p: &mut CudaSlice<i32>) -> Result<(), Box<dyn std::error::Error>> {
3920        let f = self.func("inc_i32");
3921        let cfg = LaunchConfig {
3922            grid_dim: (1, 1, 1),
3923            block_dim: (1, 1, 1),
3924            shared_mem_bytes: 0,
3925        };
3926        let __s_b = self.gpu.stream();
3927        let mut b = __s_b.launch_builder(&f);
3928        b.arg(p);
3929        unsafe {
3930            b.launch(cfg)?;
3931        }
3932        Ok(())
3933    }
3934
3935    /// Like `append_kv_quantized` but k_row/v_row are CudaViews (one token's row sliced out of a
3936    /// token-major [T, kv_dim] activation buffer — the MTP verify path appends T tokens).
3937    pub fn append_kv_quantized_view(
3938        &self,
3939        k_row: &cudarc::driver::CudaView<f32>,
3940        v_row: &cudarc::driver::CudaView<f32>,
3941        kc: &mut CudaSlice<u8>,
3942        vc: &mut CudaSlice<u8>,
3943        t: usize,
3944        kv_dim_k: usize,
3945        kv_dim_v: usize,
3946        k_tok_bytes: usize,
3947        v_tok_bytes: usize,
3948        g: bool,
3949    ) -> Result<(), Box<dyn std::error::Error>> {
3950        let f = if g {
3951            self.func_g("append_quantize_kv_q8_0_q5_1")
3952        } else {
3953            self.func("append_quantize_kv_q8_0_q5_1")
3954        };
3955        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
3956        let cfg = LaunchConfig {
3957            grid_dim: (nblk, 1, 1),
3958            block_dim: (32, 1, 1),
3959            shared_mem_bytes: 0,
3960        };
3961        let (ti, kdk, kdv) = (t as i32, kv_dim_k as i32, kv_dim_v as i32);
3962        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3963        let __s_b = self.gpu.stream();
3964        let mut b = __s_b.launch_builder(&f);
3965        b.arg(k_row)
3966            .arg(v_row)
3967            .arg(kc)
3968            .arg(vc)
3969            .arg(&ti)
3970            .arg(&kdk)
3971            .arg(&kdv)
3972            .arg(&ktb)
3973            .arg(&vtb);
3974        unsafe {
3975            b.launch(cfg)?;
3976        }
3977        Ok(())
3978    }
3979
3980    /// Device-to-device copy of a CudaView `src` into `dst[off..off+len]` (f32). Like `copy_into`
3981    /// but the source is a sub-view (e.g. one column of a token-major activation buffer).
3982    pub fn copy_view_into(
3983        &self,
3984        dst: &mut CudaSlice<f32>,
3985        off: usize,
3986        src: &cudarc::driver::CudaView<f32>,
3987        len: usize,
3988    ) -> Result<(), Box<dyn std::error::Error>> {
3989        let mut view = dst.slice_mut(off..off + len);
3990        self.gpu
3991            .stream()
3992            .memcpy_dtod(&src.slice(0..len), &mut view)?;
3993        Ok(())
3994    }
3995
3996    /// Real device-to-device COPY of `src` into a freshly allocated buffer (NOT an Arc clone).
3997    /// Used for cache snapshots (MTP-PLAN §D.4): `CudaSlice::clone()` only bumps a refcount and
3998    /// would alias the live buffer; this allocs new device memory and memcpy_dtod's the contents.
3999    pub fn clone_dtod(
4000        &self,
4001        src: &CudaSlice<f32>,
4002    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4003        let mut dst = self.gpu.stream().alloc_zeros::<f32>(src.len())?;
4004        self.gpu.stream().memcpy_dtod(src, &mut dst)?;
4005        Ok(dst)
4006    }
4007
4008    /// D2D row extraction: copy a view (e.g. one row of a [B, n] batch buffer) into `dst`.
4009    /// Stream-ordered, async — decode_batch's per-sequence row plumbing.
4010    pub fn dtod_copy_view(
4011        &self,
4012        src: &cudarc::driver::CudaView<f32>,
4013        dst: &mut CudaSlice<f32>,
4014    ) -> Result<(), Box<dyn std::error::Error>> {
4015        self.gpu.stream().memcpy_dtod(src, dst)?;
4016        Ok(())
4017    }
4018
4019    /// D2D i8 twin of `dtod_copy_view` (q8_1 activation rows).
4020    pub fn dtod_copy_view_i8(
4021        &self,
4022        src: &cudarc::driver::CudaView<i8>,
4023        dst: &mut CudaSlice<i8>,
4024    ) -> Result<(), Box<dyn std::error::Error>> {
4025        self.gpu.stream().memcpy_dtod(src, dst)?;
4026        Ok(())
4027    }
4028
4029    /// D2D row placement: copy `src` into `dst[offset .. offset+src.len()]`.
4030    pub fn dtod_copy_into(
4031        &self,
4032        src: &CudaSlice<f32>,
4033        dst: &mut CudaSlice<f32>,
4034        offset: usize,
4035    ) -> Result<(), Box<dyn std::error::Error>> {
4036        let n = src.len();
4037        let mut dv = dst.slice_mut(offset..offset + n);
4038        self.gpu.stream().memcpy_dtod(src, &mut dv)?;
4039        Ok(())
4040    }
4041
4042    /// Batched uniform-size D2D copy (engine-bundle slice 1, DSF-ROUNDCOST-20260820 §1.1):
4043    /// `n` disjoint regions of `words` f32 each; `table` = [src_0..src_{n-1}, dst_0..dst_{n-1}]
4044    /// raw device pointers. ONE kernel launch replaces `n` memcpy_dtod dispatches — the dspark
4045    /// round's snap/commit copy dribble (~0.9 ms/round of dispatch serialization measured).
4046    /// Bytes and stream order are identical to the memcpy sequence it replaces.
4047    pub fn copy_batch_uniform_f32(
4048        &self,
4049        table: &CudaSlice<u64>,
4050        n: usize,
4051        words: usize,
4052    ) -> Result<(), Box<dyn std::error::Error>> {
4053        if n == 0 || words == 0 {
4054            return Ok(());
4055        }
4056        debug_assert!(
4057            table.len() >= 2 * n,
4058            "pointer table must hold n srcs + n dsts"
4059        );
4060        let f = self.func("copy_batch_uniform_f32");
4061        // Enough blocks to stream a multi-MB region, few enough that (chunks x n) stays a
4062        // sane grid: 512K-word ssm regions get 48 grid-striding blocks each.
4063        let chunks = (words / 4).max(1).div_ceil(256).min(48) as u32;
4064        let (ni, wi) = (n as i32, words as i32);
4065        let cfg = LaunchConfig {
4066            grid_dim: (chunks, n as u32, 1),
4067            block_dim: (256, 1, 1),
4068            shared_mem_bytes: 0,
4069        };
4070        let __s = self.gpu.stream();
4071        let mut b = __s.launch_builder(&f);
4072        b.arg(table).arg(&ni).arg(&wi);
4073        unsafe {
4074            b.launch(cfg)?;
4075        }
4076        Ok(())
4077    }
4078
4079    /// H2D refresh of an EXISTING u64 pointer table IN PLACE (stable pointer — the batched
4080    /// state-copy tables are refreshed per round because the GDN ssm handles ping-pong).
4081    pub fn htod_u64_into(
4082        &self,
4083        v: &[u64],
4084        dst: &mut CudaSlice<u64>,
4085    ) -> Result<(), Box<dyn std::error::Error>> {
4086        let mut view = dst.slice_mut(0..v.len());
4087        self.gpu.stream().memcpy_htod(v, &mut view)?;
4088        Ok(())
4089    }
4090
4091    /// Indirect-source copy (engine-bundle slice 3): the src ADDRESS is loaded from a
4092    /// device pointer-table entry at run time, so a captured graph follows the gdn
4093    /// ping-pong through the same table its scan kernels read — a baked memcpy node
4094    /// would keep the capture-time physical buffer. `dst_off`/`words` in f32 elements.
4095    pub fn copy_indirect_src_f32(
4096        &self,
4097        src_entry: &cudarc::driver::CudaView<u64>,
4098        dst: &mut CudaSlice<f32>,
4099        dst_off: usize,
4100        words: usize,
4101    ) -> Result<(), Box<dyn std::error::Error>> {
4102        let f = self.func("copy_indirect_src_f32");
4103        let chunks = (words / 4).max(1).div_ceil(256).min(48) as u32;
4104        let wi = words as i32;
4105        let cfg = LaunchConfig {
4106            grid_dim: (chunks, 1, 1),
4107            block_dim: (256, 1, 1),
4108            shared_mem_bytes: 0,
4109        };
4110        let mut dv = dst.slice_mut(dst_off..dst_off + words);
4111        let __s = self.gpu.stream();
4112        let mut b = __s.launch_builder(&f);
4113        b.arg(src_entry).arg(&mut dv).arg(&wi);
4114        unsafe {
4115            b.launch(cfg)?;
4116        }
4117        Ok(())
4118    }
4119
4120    /// Uninitialized i8 device buffer (decode_batch q8_1 row scratch).
4121    pub fn uninit_i8(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
4122        self.alloc_uninit::<i8>(n)
4123    }
4124
4125    /// Resident-quantized linear (Stage-A: f32 dequant-in-kernel). y[m,out]=x[m,in]@W[out,in]^T.
4126    pub fn qmatvec(
4127        &self,
4128        w: &CudaSlice<u8>,
4129        x: &CudaSlice<f32>,
4130        m: usize,
4131        in_f: usize,
4132        out_f: usize,
4133        qtype: i32,
4134        row_bytes: usize,
4135    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4136        let f = self.func("qmatvec_f32");
4137        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
4138        let cfg = LaunchConfig {
4139            grid_dim: (out_f as u32, m as u32, 1),
4140            block_dim: (256, 1, 1),
4141            shared_mem_bytes: 0,
4142        };
4143        let (inf, outf, mi, qt, rb) =
4144            (in_f as i32, out_f as i32, m as i32, qtype, row_bytes as i64);
4145        let __s_b = self.gpu.stream();
4146        let mut b = __s_b.launch_builder(&f);
4147        b.arg(w)
4148            .arg(x)
4149            .arg(&mut y)
4150            .arg(&inf)
4151            .arg(&outf)
4152            .arg(&mi)
4153            .arg(&qt)
4154            .arg(&rb);
4155        unsafe {
4156            b.launch(cfg)?;
4157        }
4158        Ok(y)
4159    }
4160
4161    /// Allocate a reusable u8 GPU scratch buffer (for staged expert weights).
4162    pub fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
4163        let s = self.gpu.stream().alloc_zeros::<u8>(n)?;
4164        self.keep_if_capturing(&s);
4165        Ok(s)
4166    }
4167
4168    /// Uninitialized u8 scratch — skips alloc_zeros' memset. ONLY for staging buffers whose read
4169    /// range is fully overwritten by a stage_expert H2D before any kernel reads it (LAUNCH-STRUCTURE
4170    /// STAGE 2: the per-layer MoE scratch trio was 3 dead ~1MB memsets per layer per decode token).
4171    pub fn alloc_u8_uninit(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
4172        let s = unsafe { self.gpu.stream().alloc::<u8>(n)? };
4173        self.keep_if_capturing(&s);
4174        Ok(s)
4175    }
4176
4177    /// Zero a SUB-RANGE of an f32 buffer (CudaViewMut) — the row-sized memset the moe_out
4178    /// memset-elision uses for tokens that fall off the gdec fast path (LAUNCH-STRUCTURE STAGE 2).
4179    pub fn memset_zeros_view(
4180        &self,
4181        dst: &mut cudarc::driver::CudaViewMut<f32>,
4182    ) -> Result<(), Box<dyn std::error::Error>> {
4183        self.gpu.stream().memset_zeros(dst)?;
4184        Ok(())
4185    }
4186
4187    /// EDGE-1 staging: copy `host_bytes` (a sub-slice of a HostExps buffer) into `scratch`
4188    /// at byte offset `off` (async H2D on the default stream). Length is host_bytes.len().
4189    /// The qmatvec_view that reads `scratch[off..]` is enqueued on the SAME stream after this,
4190    /// so ordering is guaranteed without an explicit sync (Stage-1; Stage-2 prefetch on a 2nd
4191    /// stream would require an event).
4192    pub fn stage_expert(
4193        &self,
4194        host_bytes: &[u8],
4195        scratch: &mut CudaSlice<u8>,
4196        off: usize,
4197    ) -> Result<(), Box<dyn std::error::Error>> {
4198        let mut dst = scratch.slice_mut(off..off + host_bytes.len()); // CudaViewMut<u8>
4199        self.gpu.stream().memcpy_htod(host_bytes, &mut dst)?; // accepts &[u8] HostSlice src
4200        Ok(())
4201    }
4202
4203    /// EDGE-1 §A: fused MoE router. `logits` is the router output [t, n_expert] (device, f32, the
4204    /// `gate_inp @ z` result). Returns (sel_idx [t, n_used] i32, sel_w [t, n_used] f32): the top-k
4205    /// expert ids (DESC by prob, ascending-index tiebreak) and renormalized weights. Replaces the
4206    /// host dtoh + softmax-256 + stable DESC top-8 sort + renorm (hybrid_forward.rs ~281-298).
4207    /// One CTA per token row, 256 threads (one per expert).
4208    pub fn moe_router_topk(
4209        &self,
4210        logits: &CudaSlice<f32>,
4211        t: usize,
4212        n_expert: usize,
4213        n_used: usize,
4214    ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4215        let f = self.func("moe_router_topk_f32");
4216        let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?; // kernel fully overwrites
4217        let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?; // kernel fully overwrites
4218        let cfg = LaunchConfig {
4219            grid_dim: (t as u32, 1, 1),
4220            block_dim: (n_expert as u32, 1, 1),
4221            shared_mem_bytes: 0,
4222        };
4223        let (ne, nu) = (n_expert as i32, n_used as i32);
4224        let __s_b = self.gpu.stream();
4225        let mut b = __s_b.launch_builder(&f);
4226        b.arg(logits)
4227            .arg(&mut sel_idx)
4228            .arg(&mut sel_w)
4229            .arg(&ne)
4230            .arg(&nu);
4231        unsafe {
4232            b.launch(cfg)?;
4233        }
4234        Ok((sel_idx, sel_w))
4235    }
4236
4237    /// gemma4 twin: per-expert output scale folded into the topk renorm write (replaces the
4238    /// separate moe_w_exscale launch; value chain identical: (w/ws) * s[sel]).
4239    pub fn moe_router_topk_scaled(
4240        &self,
4241        logits: &CudaSlice<f32>,
4242        t: usize,
4243        n_expert: usize,
4244        n_used: usize,
4245        ex_scale: &CudaSlice<f32>,
4246    ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4247        // barrier-lean v2 twin (per-warp top-k + one-warp merge) FALSIFIED 2026-07-14:
4248        // bit-identical streams but −1.4% (26B plain N=3 interleaved) — at t=1 the grid is
4249        // ONE block, so the 6.6us is launch/dependency overhead, not the barrier chain;
4250        // fewer barriers bought nothing and the merge structure cost. jsonl is the record.
4251        let f = self.func("moe_router_topk_scaled_f32");
4252        let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?;
4253        let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?;
4254        let cfg = LaunchConfig {
4255            grid_dim: (t as u32, 1, 1),
4256            block_dim: (n_expert as u32, 1, 1),
4257            shared_mem_bytes: 0,
4258        };
4259        let (ne, nu) = (n_expert as i32, n_used as i32);
4260        let __s_b = self.gpu.stream();
4261        let mut b = __s_b.launch_builder(&f);
4262        b.arg(logits)
4263            .arg(&mut sel_idx)
4264            .arg(&mut sel_w)
4265            .arg(&ne)
4266            .arg(&nu)
4267            .arg(ex_scale);
4268        unsafe {
4269            b.launch(cfg)?;
4270        }
4271        Ok((sel_idx, sel_w))
4272    }
4273
4274    /// LAUNCH-STRUCTURE STAGE 1 (2026-07-05): fused router + SINGLE-SYNC host readback. The old
4275    /// MEMRA_FUSED_ROUTER path lost 2% at t=1 because it paid TWO full stream syncs (dtoh_i32 then
4276    /// dtoh, each = clone_dtoh + synchronize) + two alloc_zeros memsets per MoE layer, where the
4277    /// host route pays ONE sync on the 1KB logits dtoh. This variant: uninit outputs (kernel fully
4278    /// overwrites), both DtoH copies issued ASYNC into a persistent PINNED host staging buffer
4279    /// (flags=0 — cacheable, NOT cudarc's WRITECOMBINED default, so the host-side reads of sel/w
4280    /// stay cached), then ONE synchronize. Numerics identical to `moe_router_topk` (same kernel).
4281    pub fn moe_router_topk_host(
4282        &self,
4283        logits: &CudaSlice<f32>,
4284        t: usize,
4285        n_expert: usize,
4286        n_used: usize,
4287    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
4288        let f = self.func("moe_router_topk_f32");
4289        let n = t * n_used;
4290        let mut sel_idx = self.alloc_uninit::<i32>(n)?;
4291        let mut sel_w = self.alloc_uninit::<f32>(n)?;
4292        let cfg = LaunchConfig {
4293            grid_dim: (t as u32, 1, 1),
4294            block_dim: (n_expert as u32, 1, 1),
4295            shared_mem_bytes: 0,
4296        };
4297        let (ne, nu) = (n_expert as i32, n_used as i32);
4298        let __s_b = self.gpu.stream();
4299        let mut b = __s_b.launch_builder(&f);
4300        b.arg(logits)
4301            .arg(&mut sel_idx)
4302            .arg(&mut sel_w)
4303            .arg(&ne)
4304            .arg(&nu);
4305        unsafe {
4306            b.launch(cfg)?;
4307        }
4308        // single-sync readback: sel (i32) at offset 0, w (f32) at offset n*4 of the pinned stage.
4309        let bytes = n * 8;
4310        let mut guard = self.router_stage.lock().unwrap();
4311        if guard.as_ref().map(|p| p.cap < bytes).unwrap_or(true) {
4312            *guard = Some(PinnedStage::new(bytes.max(4096))?);
4313        }
4314        let stage = guard.as_mut().unwrap();
4315        let (si, sw) = unsafe {
4316            (
4317                std::slice::from_raw_parts_mut(stage.ptr as *mut i32, n),
4318                std::slice::from_raw_parts_mut(stage.ptr.add(n * 4) as *mut f32, n),
4319            )
4320        };
4321        self.gpu.stream().memcpy_dtoh(&sel_idx, si)?; // async (pinned dst)
4322        self.gpu.stream().memcpy_dtoh(&sel_w, sw)?; // async (pinned dst)
4323        self.gpu.stream().synchronize()?; // ONE sync for both
4324        Ok((si.iter().map(|&i| i as u32).collect(), sw.to_vec()))
4325    }
4326
4327    /// Device sigmoid router for Step-3.7 / DeepSeek-V3-class MoEs. `correction_bias` is added
4328    /// only to the top-k key; returned weights use the un-biased sigmoid score. `active` masks
4329    /// original expert ids before top-k. Exact key ties choose the smaller original id.
4330    #[allow(clippy::too_many_arguments)]
4331    pub fn moe_router_sigmoid_topk(
4332        &self,
4333        logits: &CudaSlice<f32>,
4334        t: usize,
4335        n_expert: usize,
4336        n_used: usize,
4337        active_count: usize,
4338        correction_bias: &CudaSlice<f32>,
4339        active: &CudaSlice<u8>,
4340        scaling_factor: f32,
4341        route_norm: bool,
4342    ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4343        crate::sigrouter_contract::validate_active_count(n_used, active_count)?;
4344        if n_expert == 0 || n_expert > 1024 || n_used == 0 || n_used > n_expert {
4345            return Err(format!(
4346                "sigmoid router shape unsupported: n_expert={n_expert}, n_used={n_used}",
4347            )
4348            .into());
4349        }
4350        if logits.len() < t * n_expert
4351            || correction_bias.len() != n_expert
4352            || active.len() != n_expert
4353        {
4354            return Err(format!(
4355                "sigmoid router buffer mismatch: logits={} bias={} active={} expected logits>={} row={}",
4356                logits.len(), correction_bias.len(), active.len(), t * n_expert, n_expert,
4357            ).into());
4358        }
4359        let f = self.func("moe_router_sigmoid_topk_f32");
4360        let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?;
4361        let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?;
4362        let threads = n_expert.div_ceil(32) * 32;
4363        let cfg = LaunchConfig {
4364            grid_dim: (t as u32, 1, 1),
4365            block_dim: (threads as u32, 1, 1),
4366            shared_mem_bytes: 0,
4367        };
4368        let (ne, nu, rn) = (n_expert as i32, n_used as i32, i32::from(route_norm));
4369        let __s_b = self.gpu.stream();
4370        let mut b = __s_b.launch_builder(&f);
4371        b.arg(logits)
4372            .arg(correction_bias)
4373            .arg(active)
4374            .arg(&mut sel_idx)
4375            .arg(&mut sel_w)
4376            .arg(&ne)
4377            .arg(&nu)
4378            .arg(&scaling_factor)
4379            .arg(&rn);
4380        unsafe {
4381            b.launch(cfg)?;
4382        }
4383        Ok((sel_idx, sel_w))
4384    }
4385
4386    /// Single-sync pinned readback twin of `moe_router_sigmoid_topk`. This preserves the existing
4387    /// grouped/staged dispatch contract while replacing the full-logit DtoH plus host sigmoid/sort.
4388    #[allow(clippy::too_many_arguments)]
4389    pub fn moe_router_sigmoid_topk_host(
4390        &self,
4391        logits: &CudaSlice<f32>,
4392        t: usize,
4393        n_expert: usize,
4394        n_used: usize,
4395        active_count: usize,
4396        correction_bias: &CudaSlice<f32>,
4397        active: &CudaSlice<u8>,
4398        scaling_factor: f32,
4399        route_norm: bool,
4400    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
4401        let (sel_idx, sel_w) = self.moe_router_sigmoid_topk(
4402            logits,
4403            t,
4404            n_expert,
4405            n_used,
4406            active_count,
4407            correction_bias,
4408            active,
4409            scaling_factor,
4410            route_norm,
4411        )?;
4412        let n = t * n_used;
4413        let bytes = n * 8;
4414        let mut guard = self.router_stage.lock().unwrap();
4415        if guard.as_ref().map(|p| p.cap < bytes).unwrap_or(true) {
4416            *guard = Some(PinnedStage::new(bytes.max(4096))?);
4417        }
4418        let stage = guard.as_mut().unwrap();
4419        let (si, sw) = unsafe {
4420            (
4421                std::slice::from_raw_parts_mut(stage.ptr as *mut i32, n),
4422                std::slice::from_raw_parts_mut(stage.ptr.add(n * 4) as *mut f32, n),
4423            )
4424        };
4425        self.gpu.stream().memcpy_dtoh(&sel_idx, si)?;
4426        self.gpu.stream().memcpy_dtoh(&sel_w, sw)?;
4427        self.gpu.stream().synchronize()?;
4428        Ok((si.iter().map(|&i| i as u32).collect(), sw.to_vec()))
4429    }
4430
4431    /// EDGE-1 §C.2: async H2D of `host_bytes` into `scratch[off..]` on the COPY stream, returning a
4432    /// recorded event the compute stream can `wait` on before the dependent GEMM. Used for in-token
4433    /// expert prefetch (pipeline by one). `host_bytes` should be pinned for a true DMA (§C.1).
4434    pub fn stage_expert_async(
4435        &self,
4436        host_bytes: &[u8],
4437        scratch: &mut CudaSlice<u8>,
4438        off: usize,
4439    ) -> Result<cudarc::driver::CudaEvent, Box<dyn std::error::Error>> {
4440        let mut dst = scratch.slice_mut(off..off + host_bytes.len());
4441        self.copy_stream.memcpy_htod(host_bytes, &mut dst)?;
4442        Ok(self.copy_stream.record_event(None)?)
4443    }
4444
4445    /// Make the compute stream wait for an async copy event (the consumer side of `stage_expert_async`).
4446    pub fn compute_wait(
4447        &self,
4448        ev: &cudarc::driver::CudaEvent,
4449    ) -> Result<(), Box<dyn std::error::Error>> {
4450        self.gpu.stream().wait(ev)?;
4451        Ok(())
4452    }
4453
4454    /// qmatvec over a byte sub-range of a (resident/scratch) CudaSlice<u8> holding ONE expert
4455    /// matrix. x is a CudaView<f32> (a sliced row of z, or a sliced activation). Reuses the
4456    /// validated qmatvec_f32 dequant path (NOT a fast path — the correctness gate). The
4457    /// CudaView base+offset pointer is honored by the launch arg.
4458    pub fn qmatvec_view(
4459        &self,
4460        w: &CudaSlice<u8>,
4461        range: std::ops::Range<usize>,
4462        x: &cudarc::driver::CudaView<f32>,
4463        m: usize,
4464        in_f: usize,
4465        out_f: usize,
4466        qtype: i32,
4467        row_bytes: usize,
4468    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4469        let f = self.func("qmatvec_f32");
4470        let wv = w.slice(range); // CudaView<u8>, offset honored
4471        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
4472        let cfg = LaunchConfig {
4473            grid_dim: (out_f as u32, m as u32, 1),
4474            block_dim: (256, 1, 1),
4475            shared_mem_bytes: 0,
4476        };
4477        let (inf, outf, mi, qt, rb) =
4478            (in_f as i32, out_f as i32, m as i32, qtype, row_bytes as i64);
4479        let __s_b = self.gpu.stream();
4480        let mut b = __s_b.launch_builder(&f);
4481        b.arg(&wv)
4482            .arg(x)
4483            .arg(&mut y)
4484            .arg(&inf)
4485            .arg(&outf)
4486            .arg(&mi)
4487            .arg(&qt)
4488            .arg(&rb);
4489        unsafe {
4490            b.launch(cfg)?;
4491        }
4492        Ok(y)
4493    }
4494
4495    /// STAGE-2 GROUPED DECODE (2026-07-04): one MoE layer's gate+up+SiLU for all `n_used` routed
4496    /// experts of ONE token in ONE launch (replaces 8x qmatvec(gate) + 8x qmatvec(up) + 8x
4497    /// silu_mul = 24 launches). `gp`/`up` are the 8 expert weight-block device pointers (SLRU
4498    /// cache slots — fixed-address, stable for the launch). Returns act [n_used, n_ff].
4499    /// BIT-IDENTICAL to the sequential chain: each dot reproduces qmatvec_f32's exact 256-thread
4500    /// reduction; the SiLU epilogue is silu_mul_f32's exact expression (see kernel header).
4501    #[allow(clippy::too_many_arguments)]
4502    /// dp4a q8 twins (MoE expert dp4a arc, 2026-07-06): same contract as the _f32 versions but
4503    /// consume a PRE-QUANTIZED q8_1 activation. FP-order differs from _f32 (int dot + warp tree)
4504    /// — the argmax/stream-identity battery arbitrates; MEMRA_MOE_Q8=0 restores f32.
4505    pub fn moe_gate_up_silu8_q8(
4506        &self,
4507        gp: WPtr8,
4508        up: WPtr8,
4509        aq: &CudaSlice<i8>,
4510        ad: &CudaSlice<f32>,
4511        in_f: usize,
4512        n_ff: usize,
4513        n_used: usize,
4514        qt_g: i32,
4515        qt_u: i32,
4516        rb_g: usize,
4517        rb_u: usize,
4518    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4519        let f = self.func("moe_gate_up_silu8_q8");
4520        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
4521        let cfg = LaunchConfig {
4522            grid_dim: (n_ff as u32, n_used as u32, 1),
4523            block_dim: (32, 1, 1),
4524            shared_mem_bytes: 0,
4525        };
4526        let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
4527        let __s_b = self.gpu.stream();
4528        let mut b = __s_b.launch_builder(&f);
4529        b.arg(&gp)
4530            .arg(&up)
4531            .arg(aq)
4532            .arg(ad)
4533            .arg(&mut act)
4534            .arg(&inf)
4535            .arg(&nff)
4536            .arg(&qt_g)
4537            .arg(&qt_u)
4538            .arg(&rbg)
4539            .arg(&rbu);
4540        unsafe {
4541            b.launch(cfg)?;
4542        }
4543        Ok(act)
4544    }
4545
4546    #[allow(clippy::too_many_arguments)]
4547    pub fn moe_down8_fma_q8(
4548        &self,
4549        dp: WPtr8,
4550        w: F32x8,
4551        aq2: &CudaSlice<i8>,
4552        ad2: &CudaSlice<f32>,
4553        dst: &mut cudarc::driver::CudaViewMut<f32>,
4554        in_f: usize,
4555        out_f: usize,
4556        n_used: usize,
4557        qt: i32,
4558        rb: usize,
4559    ) -> Result<(), Box<dyn std::error::Error>> {
4560        let f = self.func("moe_down8_fma_q8");
4561        let cfg = LaunchConfig {
4562            grid_dim: (out_f as u32, 1, 1),
4563            block_dim: (32, 1, 1),
4564            shared_mem_bytes: 0,
4565        };
4566        let (inf, outf, nu, rbi) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
4567        let __s_b = self.gpu.stream();
4568        let mut b = __s_b.launch_builder(&f);
4569        b.arg(&dp)
4570            .arg(&w)
4571            .arg(aq2)
4572            .arg(ad2)
4573            .arg(dst)
4574            .arg(&inf)
4575            .arg(&outf)
4576            .arg(&nu)
4577            .arg(&qt)
4578            .arg(&rbi);
4579        unsafe {
4580            b.launch(cfg)?;
4581        }
4582        Ok(())
4583    }
4584
4585    /// q8 sequential expert matvec (staged path twin of qmatvec_view for IQ3_S/IQ4_XS).
4586    pub fn qmatvec_expert_q8(
4587        &self,
4588        w: &CudaSlice<u8>,
4589        range: std::ops::Range<usize>,
4590        aq: &CudaSlice<i8>,
4591        ad: &CudaSlice<f32>,
4592        m: usize,
4593        in_f: usize,
4594        out_f: usize,
4595        qtype: i32,
4596        row_bytes: usize,
4597    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4598        let f = self.func("qmatvec_expert_q8");
4599        let wv = w.slice(range);
4600        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
4601        const ROWS: u32 = 4; // MEMRA_MMVQ_ROWS
4602        let cfg = LaunchConfig {
4603            grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, m as u32, 1),
4604            block_dim: (32, ROWS, 1),
4605            shared_mem_bytes: 0,
4606        };
4607        let (inf, outf, mi, rbi) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
4608        let __s_b = self.gpu.stream();
4609        let mut b = __s_b.launch_builder(&f);
4610        b.arg(&wv)
4611            .arg(aq)
4612            .arg(ad)
4613            .arg(&mut y)
4614            .arg(&inf)
4615            .arg(&outf)
4616            .arg(&mi)
4617            .arg(&qtype)
4618            .arg(&rbi);
4619        unsafe {
4620            b.launch(cfg)?;
4621        }
4622        Ok(y)
4623    }
4624
4625    pub fn moe_gate_up_silu8(
4626        &self,
4627        gp: WPtr8,
4628        up: WPtr8,
4629        x: &cudarc::driver::CudaView<f32>,
4630        in_f: usize,
4631        n_ff: usize,
4632        n_used: usize,
4633        qt_g: i32,
4634        qt_u: i32,
4635        rb_g: usize,
4636        rb_u: usize,
4637    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4638        let f = self.func("moe_gate_up_silu8_f32");
4639        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?; // fully overwritten
4640        let cfg = LaunchConfig {
4641            grid_dim: (n_ff as u32, n_used as u32, 1),
4642            block_dim: (256, 1, 1),
4643            shared_mem_bytes: 0,
4644        };
4645        let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
4646        let __s_b = self.gpu.stream();
4647        let mut b = __s_b.launch_builder(&f);
4648        b.arg(&gp)
4649            .arg(&up)
4650            .arg(x)
4651            .arg(&mut act)
4652            .arg(&inf)
4653            .arg(&nff)
4654            .arg(&qt_g)
4655            .arg(&qt_u)
4656            .arg(&rbg)
4657            .arg(&rbu);
4658        unsafe {
4659            b.launch(cfg)?;
4660        }
4661        Ok(act)
4662    }
4663
4664    /// STAGE-2 GROUPED DECODE: one MoE layer's down-proj + weighted accumulation for all `n_used`
4665    /// routed experts in ONE launch (replaces 8x qmatvec(down) + 8x axpy = 16 launches), writing
4666    /// the token's moe_out row DIRECTLY (`dst` is the zeroed row; the in-kernel slot-ordered
4667    /// __fmaf_rn chain starting at 0.0f reproduces the sequential axpy_f32 accumulation into the
4668    /// zeroed row bit-for-bit — the A2 byte-identity scheme at m=1).
4669    #[allow(clippy::too_many_arguments)]
4670    pub fn moe_down8_fma_into(
4671        &self,
4672        dp: WPtr8,
4673        w: F32x8,
4674        act: &CudaSlice<f32>,
4675        dst: &mut cudarc::driver::CudaViewMut<f32>,
4676        in_f: usize,
4677        out_f: usize,
4678        n_used: usize,
4679        qt: i32,
4680        rb: usize,
4681    ) -> Result<(), Box<dyn std::error::Error>> {
4682        let f = self.func("moe_down8_fma_f32");
4683        let cfg = LaunchConfig {
4684            grid_dim: (out_f as u32, 1, 1),
4685            block_dim: (256, 1, 1),
4686            shared_mem_bytes: 0,
4687        };
4688        let (inf, outf, nu, rbv) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
4689        let __s_b = self.gpu.stream();
4690        let mut b = __s_b.launch_builder(&f);
4691        b.arg(&dp)
4692            .arg(&w)
4693            .arg(act)
4694            .arg(dst)
4695            .arg(&inf)
4696            .arg(&outf)
4697            .arg(&nu)
4698            .arg(&qt)
4699            .arg(&rbv);
4700        unsafe {
4701            b.launch(cfg)?;
4702        }
4703        Ok(())
4704    }
4705
4706    /// LAUNCH-STRUCTURE STAGE 3: device-dispatch twin of `moe_gate_up_silu8` for FULLY-RESIDENT
4707    /// layers. The expert ids come from the router kernel's DEVICE `sel` output (no DtoH) and the
4708    /// weight pointers from the per-layer device table `[3, n_expert]` of slot base addresses.
4709    /// BIT-IDENTICAL math (same grid/block/reduction; only the pointer/id source differs).
4710    #[allow(clippy::too_many_arguments)]
4711    /// dp4a q8 twin of the _dev pair (resident-experts arc).
4712    ///
4713    /// GEOMETRY VARIANTS (multirow/occupancy arc 2026-07-05): all outputs are BIT-IDENTICAL to
4714    /// the base one-warp-per-(row,slot) kernel (same expert_dot_g g-order + warp tree per row;
4715    /// down's FMA chain stays slot-ordered serial). Seams:
4716    ///   MEMRA_MOE_DEVQ8_GU   = 0(base) | 1 | 2 | 4 -> _r{1,2,4} multirow twin (RPW rows/warp)
4717    ///                       | s2 (gate/up warp split) | s2z (s2 + WPB rows packed per block)
4718    ///                       | gs4 (gate/up x low/high-group 4-warp split, nsb==64 only)
4719    ///                       | u64 (nsb==64 unrolled ILP twin, geometry unchanged)
4720    ///   MEMRA_MOE_DEVQ8_WPB  = warps per block for _r twins / z-rows for s2z (default 4)
4721    ///   MEMRA_MOE_DEVQ8_DOWN = auto(default: w8h2 when in_f==512 & n_used<=8 — measured +3.8%
4722    ///                       decode on 35B/G7e) | 0 (base one-warp serial-slot) | 1 | 2 | 4 ->
4723    ///                       _w8r{1,2,4} slot-parallel twin | h2 (half-warp dual-row, nsb==16
4724    ///                       only) | w8h2 (h2 x slot-parallel)
4725    #[allow(clippy::too_many_arguments)]
4726    /// MoE PREFILL pair-batch matvec: one launch covers all (token,expert) pairs for one proj.
4727    #[allow(clippy::too_many_arguments)]
4728    pub fn moe_pairs_matvec_q8(
4729        &self,
4730        table: &CudaSlice<u64>,
4731        proj: i32,
4732        pair_tok: &CudaSlice<i32>,
4733        pair_ex: &CudaSlice<i32>,
4734        aq: &CudaSlice<i8>,
4735        ad: &CudaSlice<f32>,
4736        in_f: usize,
4737        out_f: usize,
4738        n_expert: usize,
4739        n_pairs: usize,
4740        qtype: i32,
4741        row_bytes: usize,
4742    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4743        let f = self.func("moe_pairs_matvec_q8");
4744        let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
4745        const ROWS: u32 = 4;
4746        let cfg = LaunchConfig {
4747            grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_pairs as u32, 1),
4748            block_dim: (32, ROWS, 1),
4749            shared_mem_bytes: 0,
4750        };
4751        let (inf, outf, ne, np, rbi) = (
4752            in_f as i32,
4753            out_f as i32,
4754            n_expert as i32,
4755            n_pairs as i32,
4756            row_bytes as i64,
4757        );
4758        let __s_b = self.gpu.stream();
4759        let mut b = __s_b.launch_builder(&f);
4760        b.arg(table)
4761            .arg(&proj)
4762            .arg(pair_tok)
4763            .arg(pair_ex)
4764            .arg(aq)
4765            .arg(ad)
4766            .arg(&mut y)
4767            .arg(&inf)
4768            .arg(&outf)
4769            .arg(&ne)
4770            .arg(&np)
4771            .arg(&qtype)
4772            .arg(&rbi);
4773        unsafe {
4774            b.launch(cfg)?;
4775        }
4776        Ok(y)
4777    }
4778
4779    /// Expert-major pair matvec (weight-reuse across each expert's token group).
4780    #[allow(clippy::too_many_arguments)]
4781    pub fn moe_pairs_matvec_q8_em(
4782        &self,
4783        table: &CudaSlice<u64>,
4784        proj: i32,
4785        ex_ids: &CudaSlice<i32>,
4786        ex_off: &CudaSlice<i32>,
4787        ex_pairs: &CudaSlice<i32>,
4788        pair_tok: &CudaSlice<i32>,
4789        aq: &CudaSlice<i8>,
4790        ad: &CudaSlice<f32>,
4791        in_f: usize,
4792        out_f: usize,
4793        n_expert: usize,
4794        n_active: usize,
4795        n_pairs: usize,
4796        qtype: i32,
4797        row_bytes: usize,
4798    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4799        let f = self.func("moe_pairs_matvec_q8_em");
4800        let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
4801        const ROWS: u32 = 4;
4802        let cfg = LaunchConfig {
4803            grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_active as u32, 1),
4804            block_dim: (32, ROWS, 1),
4805            shared_mem_bytes: 0,
4806        };
4807        let (inf, outf, ne, na, rbi) = (
4808            in_f as i32,
4809            out_f as i32,
4810            n_expert as i32,
4811            n_active as i32,
4812            row_bytes as i64,
4813        );
4814        let __s_b = self.gpu.stream();
4815        let mut b = __s_b.launch_builder(&f);
4816        b.arg(table)
4817            .arg(&proj)
4818            .arg(ex_ids)
4819            .arg(ex_off)
4820            .arg(ex_pairs)
4821            .arg(pair_tok)
4822            .arg(aq)
4823            .arg(ad)
4824            .arg(&mut y)
4825            .arg(&inf)
4826            .arg(&outf)
4827            .arg(&ne)
4828            .arg(&na)
4829            .arg(&qtype)
4830            .arg(&rbi);
4831        unsafe {
4832            b.launch(cfg)?;
4833        }
4834        Ok(y)
4835    }
4836
4837    // Decode-once expert-major MMQ (rung 3). Same CSR inputs/geometry as _em; kernel dequants each
4838    // weight group once per (row,group) then dp4a's across the expert's token group.
4839    #[allow(clippy::too_many_arguments)]
4840    pub fn moe_pairs_matvec_q8_dec(
4841        &self,
4842        table: &CudaSlice<u64>,
4843        proj: i32,
4844        ex_ids: &CudaSlice<i32>,
4845        ex_off: &CudaSlice<i32>,
4846        ex_pairs: &CudaSlice<i32>,
4847        pair_tok: &CudaSlice<i32>,
4848        aq: &CudaSlice<i8>,
4849        ad: &CudaSlice<f32>,
4850        in_f: usize,
4851        out_f: usize,
4852        n_expert: usize,
4853        n_active: usize,
4854        n_pairs: usize,
4855        qtype: i32,
4856        row_bytes: usize,
4857    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4858        let f = self.func("moe_pairs_matvec_q8_dec");
4859        let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
4860        const ROWS: u32 = 4;
4861        let cfg = LaunchConfig {
4862            grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_active as u32, 1),
4863            block_dim: (32, ROWS, 1),
4864            shared_mem_bytes: 0,
4865        };
4866        let (inf, outf, ne, na, rbi) = (
4867            in_f as i32,
4868            out_f as i32,
4869            n_expert as i32,
4870            n_active as i32,
4871            row_bytes as i64,
4872        );
4873        let __s_b = self.gpu.stream();
4874        let mut b = __s_b.launch_builder(&f);
4875        b.arg(table)
4876            .arg(&proj)
4877            .arg(ex_ids)
4878            .arg(ex_off)
4879            .arg(ex_pairs)
4880            .arg(pair_tok)
4881            .arg(aq)
4882            .arg(ad)
4883            .arg(&mut y)
4884            .arg(&inf)
4885            .arg(&outf)
4886            .arg(&ne)
4887            .arg(&na)
4888            .arg(&qtype)
4889            .arg(&rbi);
4890        unsafe {
4891            b.launch(cfg)?;
4892        }
4893        Ok(y)
4894    }
4895
4896    pub fn moe_pairs_gelu_mul(
4897        &self,
4898        gate: &CudaSlice<f32>,
4899        up: &CudaSlice<f32>,
4900        n: usize,
4901    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4902        let f = self.func("moe_pairs_gelu_mul");
4903        let mut act = self.alloc_uninit::<f32>(n)?;
4904        let cfg = LaunchConfig::for_num_elems(n as u32);
4905        let nl = n as i64;
4906        let __s_b = self.gpu.stream();
4907        let mut b = __s_b.launch_builder(&f);
4908        b.arg(gate).arg(up).arg(&mut act).arg(&nl);
4909        unsafe {
4910            b.launch(cfg)?;
4911        }
4912        Ok(act)
4913    }
4914
4915    pub fn moe_pairs_silu_mul(
4916        &self,
4917        gate: &CudaSlice<f32>,
4918        up: &CudaSlice<f32>,
4919        n: usize,
4920    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4921        let f = self.func("moe_pairs_silu_mul");
4922        let mut act = self.alloc_uninit::<f32>(n)?;
4923        let cfg = LaunchConfig::for_num_elems(n as u32);
4924        let nl = n as i64;
4925        let __s_b = self.gpu.stream();
4926        let mut b = __s_b.launch_builder(&f);
4927        b.arg(gate).arg(up).arg(&mut act).arg(&nl);
4928        unsafe {
4929            b.launch(cfg)?;
4930        }
4931        Ok(act)
4932    }
4933
4934    #[allow(clippy::too_many_arguments)]
4935    pub fn moe_pairs_scatter(
4936        &self,
4937        y_down: &CudaSlice<f32>,
4938        pair_w: &CudaSlice<f32>,
4939        tok_pair_off: &CudaSlice<i32>,
4940        tok_pair_ids: &CudaSlice<i32>,
4941        moe_out: &mut CudaSlice<f32>,
4942        t: usize,
4943        n_embd: usize,
4944    ) -> Result<(), Box<dyn std::error::Error>> {
4945        let f = self.func("moe_pairs_scatter");
4946        let cfg = LaunchConfig {
4947            grid_dim: (((n_embd + 255) / 256) as u32, t as u32, 1),
4948            block_dim: (256, 1, 1),
4949            shared_mem_bytes: 0,
4950        };
4951        let ne = n_embd as i32;
4952        let __s_b = self.gpu.stream();
4953        let mut b = __s_b.launch_builder(&f);
4954        b.arg(y_down)
4955            .arg(pair_w)
4956            .arg(tok_pair_off)
4957            .arg(tok_pair_ids)
4958            .arg(moe_out)
4959            .arg(&ne);
4960        unsafe {
4961            b.launch(cfg)?;
4962        }
4963        Ok(())
4964    }
4965
4966    /// gemma4 GELU twin of moe_gate_up_silu8_dev_q8 (base geometry — slot-packed j8/j8r2
4967    /// twins probed 2026-08-01 g26 decode dig: bit-identical rows, -2.5%/-2.9% whole-model
4968    /// decode x3 interleaved -> refuted and killed; research/g26-decode-20260801/receipts.md).
4969    #[allow(clippy::too_many_arguments)]
4970    pub fn moe_gate_up_gelu8_dev_q8(
4971        &self,
4972        table: &CudaSlice<u64>,
4973        sel: &cudarc::driver::CudaView<i32>,
4974        aq: &CudaSlice<i8>,
4975        ad: &CudaSlice<f32>,
4976        in_f: usize,
4977        n_ff: usize,
4978        n_used: usize,
4979        n_expert: usize,
4980        qt_g: i32,
4981        qt_u: i32,
4982        rb_g: usize,
4983        rb_u: usize,
4984    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4985        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
4986        let (inf, nff, ne, rbg, rbu) = (
4987            in_f as i32,
4988            n_ff as i32,
4989            n_expert as i32,
4990            rb_g as i64,
4991            rb_u as i64,
4992        );
4993        let f = self.func("moe_gate_up_gelu8_dev_q8");
4994        let cfg = LaunchConfig {
4995            grid_dim: (n_ff as u32, n_used as u32, 1),
4996            block_dim: (32, 1, 1),
4997            shared_mem_bytes: 0,
4998        };
4999        let __s_b = self.gpu.stream();
5000        let mut b = __s_b.launch_builder(&f);
5001        b.arg(table)
5002            .arg(sel)
5003            .arg(aq)
5004            .arg(ad)
5005            .arg(&mut act)
5006            .arg(&inf)
5007            .arg(&nff)
5008            .arg(&ne)
5009            .arg(&qt_g)
5010            .arg(&qt_u)
5011            .arg(&rbg)
5012            .arg(&rbu);
5013        unsafe {
5014            b.launch(cfg)?;
5015        }
5016        Ok(act)
5017    }
5018
5019    /// gemma4 GELU rows twin (verify): one launch over (n_ff, n_used, t).
5020    #[allow(clippy::too_many_arguments)]
5021    pub fn moe_gate_up_gelu8_dev_q8_rows(
5022        &self,
5023        table: &CudaSlice<u64>,
5024        sel: &CudaSlice<i32>,
5025        aq: &CudaSlice<i8>,
5026        ad: &CudaSlice<f32>,
5027        t: usize,
5028        in_f: usize,
5029        n_ff: usize,
5030        n_used: usize,
5031        n_expert: usize,
5032        qt_g: i32,
5033        qt_u: i32,
5034        rb_g: usize,
5035        rb_u: usize,
5036    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5037        let mut act = self.alloc_uninit::<f32>(t * n_used * n_ff)?;
5038        let (inf, nff, ne, rbg, rbu, nu) = (
5039            in_f as i32,
5040            n_ff as i32,
5041            n_expert as i32,
5042            rb_g as i64,
5043            rb_u as i64,
5044            n_used as i32,
5045        );
5046        let f = self.func("moe_gate_up_gelu8_dev_q8_rows");
5047        let cfg = LaunchConfig {
5048            grid_dim: (n_ff as u32, n_used as u32, t as u32),
5049            block_dim: (32, 1, 1),
5050            shared_mem_bytes: 0,
5051        };
5052        let __s_b = self.gpu.stream();
5053        let mut b = __s_b.launch_builder(&f);
5054        b.arg(table)
5055            .arg(sel)
5056            .arg(aq)
5057            .arg(ad)
5058            .arg(&mut act)
5059            .arg(&inf)
5060            .arg(&nff)
5061            .arg(&ne)
5062            .arg(&qt_g)
5063            .arg(&qt_u)
5064            .arg(&rbg)
5065            .arg(&rbu)
5066            .arg(&nu);
5067        unsafe {
5068            b.launch(cfg)?;
5069        }
5070        Ok(act)
5071    }
5072
5073    /// gemma4 GELU CSR twin (verify dedup: owner block serves every pair of its expert).
5074    #[allow(clippy::too_many_arguments)]
5075    pub fn moe_gate_up_gelu8_dev_q8_csr(
5076        &self,
5077        table: &CudaSlice<u64>,
5078        sel: &CudaSlice<i32>,
5079        aq: &CudaSlice<i8>,
5080        ad: &CudaSlice<f32>,
5081        n_pairs: usize,
5082        in_f: usize,
5083        n_ff: usize,
5084        n_used: usize,
5085        n_expert: usize,
5086        qt_g: i32,
5087        qt_u: i32,
5088        rb_g: usize,
5089        rb_u: usize,
5090    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5091        let mut act = self.alloc_uninit::<f32>(n_pairs * n_ff)?;
5092        let (inf, nff, ne, rbg, rbu, nu, npi) = (
5093            in_f as i32,
5094            n_ff as i32,
5095            n_expert as i32,
5096            rb_g as i64,
5097            rb_u as i64,
5098            n_used as i32,
5099            n_pairs as i32,
5100        );
5101        let f = self.func("moe_gate_up_gelu8_dev_q8_csr");
5102        let cfg = LaunchConfig {
5103            grid_dim: (n_ff as u32, n_pairs as u32, 1),
5104            block_dim: (32, 1, 1),
5105            shared_mem_bytes: 0,
5106        };
5107        let __s_b = self.gpu.stream();
5108        let mut b = __s_b.launch_builder(&f);
5109        b.arg(table)
5110            .arg(sel)
5111            .arg(aq)
5112            .arg(ad)
5113            .arg(&mut act)
5114            .arg(&inf)
5115            .arg(&nff)
5116            .arg(&ne)
5117            .arg(&qt_g)
5118            .arg(&qt_u)
5119            .arg(&rbg)
5120            .arg(&rbu)
5121            .arg(&nu)
5122            .arg(&npi);
5123        unsafe {
5124            b.launch(cfg)?;
5125        }
5126        Ok(act)
5127    }
5128
5129    /// gemma4 generic down rows twin (verify): one launch over (out_f, 1, t).
5130    #[allow(clippy::too_many_arguments)]
5131    pub fn moe_down8_fma_dev_q8_rows_g(
5132        &self,
5133        table: &CudaSlice<u64>,
5134        sel: &CudaSlice<i32>,
5135        w: &CudaSlice<f32>,
5136        aq2: &CudaSlice<i8>,
5137        ad2: &CudaSlice<f32>,
5138        dst: &mut CudaSlice<f32>,
5139        t: usize,
5140        in_f: usize,
5141        out_f: usize,
5142        n_used: usize,
5143        n_expert: usize,
5144        qt: i32,
5145        rb: usize,
5146    ) -> Result<(), Box<dyn std::error::Error>> {
5147        let (inf, outf, nu, ne, rbi) = (
5148            in_f as i32,
5149            out_f as i32,
5150            n_used as i32,
5151            n_expert as i32,
5152            rb as i64,
5153        );
5154        // Exact Step-3.7 B=1 shape: expose the eight independent slot dots as
5155        // eight warps, then replay the original slot-ordered FMA chain. Every
5156        // other shape retains the generic one-warp rows kernel.
5157        let step_b1_w8 = t == 1 && in_f == 1280 && out_f == 4096 && n_used == 8 && qt == QT_IQ4_XS;
5158        let f = self.func(if step_b1_w8 {
5159            "moe_down8_fma_dev_q8_rows_w8"
5160        } else {
5161            "moe_down8_fma_dev_q8_rows_g"
5162        });
5163        let cfg = LaunchConfig {
5164            grid_dim: (out_f as u32, 1, t as u32),
5165            block_dim: (32, if step_b1_w8 { 8 } else { 1 }, 1),
5166            shared_mem_bytes: 0,
5167        };
5168        let __s_b = self.gpu.stream();
5169        let mut b = __s_b.launch_builder(&f);
5170        b.arg(table)
5171            .arg(sel)
5172            .arg(w)
5173            .arg(aq2)
5174            .arg(ad2)
5175            .arg(dst)
5176            .arg(&inf)
5177            .arg(&outf)
5178            .arg(&nu)
5179            .arg(&ne)
5180            .arg(&qt)
5181            .arg(&rbi);
5182        unsafe {
5183            b.launch(cfg)?;
5184        }
5185        Ok(())
5186    }
5187
5188    /// rp_q4 microprobe (2026-07-10 verify-trunk lever): b4 GGUF-block layout vs the Q4_0
5189    /// split-plane twin on the wq-class shape. Returns (blk_us, rp_us) after asserting bitwise
5190    /// identity. Bench-only surface (rp_q4_probe bin); no production dispatch reads this.
5191    pub fn rp_probe_q4(&self, m: usize) -> Result<(f64, f64), Box<dyn std::error::Error>> {
5192        let (out_f, in_f) = (2048usize, 2816usize);
5193        let nblk = in_f / 32;
5194        let mut seed = 0x9E3779B97F4A7C15u64;
5195        let mut rng = move || {
5196            seed = seed
5197                .wrapping_mul(6364136223846793005)
5198                .wrapping_add(1442695040888963407);
5199            (seed >> 33) as u8
5200        };
5201        let mut w = vec![0u8; out_f * nblk * 18];
5202        for b in w.iter_mut() {
5203            *b = rng();
5204        }
5205        for r in 0..out_f {
5206            for g in 0..nblk {
5207                let off = (r * nblk + g) * 18;
5208                w[off] = 0x00;
5209                w[off + 1] = 0x2C; // sane half d
5210            }
5211        }
5212        let qplane = out_f * nblk * 16;
5213        let mut wrp = vec![0u8; w.len()];
5214        for r in 0..out_f {
5215            for g in 0..nblk {
5216                let src = &w[(r * nblk + g) * 18..(r * nblk + g) * 18 + 18];
5217                wrp[qplane + (r * nblk + g) * 2..qplane + (r * nblk + g) * 2 + 2]
5218                    .copy_from_slice(&src[0..2]);
5219                wrp[(r * nblk + g) * 16..(r * nblk + g) * 16 + 16].copy_from_slice(&src[2..18]);
5220            }
5221        }
5222        let w_d = self.htod_bytes(&w)?;
5223        let wrp_d = self.htod_bytes(&wrp)?;
5224        let mut aq = vec![0i8; m * in_f];
5225        for v in aq.iter_mut() {
5226            *v = rng() as i8;
5227        }
5228        let aq_d = self.htod_i8(&aq)?;
5229        let ad_d = self.htod(&vec![0.03125f32; m * nblk])?;
5230        let mut y0 = self.alloc_uninit::<f32>(m * out_f)?;
5231        let mut y1 = self.alloc_uninit::<f32>(m * out_f)?;
5232        const RPB: u32 = 4;
5233        let cfg = LaunchConfig {
5234            grid_dim: ((out_f as u32).div_ceil(RPB), 1, 1),
5235            block_dim: (32, RPB, 1),
5236            shared_mem_bytes: 0,
5237        };
5238        let (inf, outf, mi) = (in_f as i32, out_f as i32, m as i32);
5239        let (rb, qp) = ((nblk * 18) as i64, qplane as i64);
5240        let fb = self.func("qmatvec_q4_0_mmvq_b4");
5241        let fr = self.func("qmatvec_q4_0_mmvq_b4_rp");
5242        {
5243            let __s_b = self.gpu.stream();
5244            let mut b = __s_b.launch_builder(&fb);
5245            b.arg(&w_d)
5246                .arg(&aq_d)
5247                .arg(&ad_d)
5248                .arg(&mut y0)
5249                .arg(&inf)
5250                .arg(&outf)
5251                .arg(&mi)
5252                .arg(&rb);
5253            unsafe {
5254                b.launch(cfg)?;
5255            }
5256            let __s_b = self.gpu.stream();
5257            let mut b = __s_b.launch_builder(&fr);
5258            b.arg(&wrp_d)
5259                .arg(&aq_d)
5260                .arg(&ad_d)
5261                .arg(&mut y1)
5262                .arg(&inf)
5263                .arg(&outf)
5264                .arg(&mi)
5265                .arg(&qp);
5266            unsafe {
5267                b.launch(cfg)?;
5268            }
5269        }
5270        self.gpu.stream().synchronize()?;
5271        let (h0, h1) = (self.dtoh(&y0)?, self.dtoh(&y1)?);
5272        let nd = h0
5273            .iter()
5274            .zip(&h1)
5275            .filter(|(a, b)| a.to_bits() != b.to_bits())
5276            .count();
5277        if nd != 0 {
5278            return Err(format!("rp twin not bitwise: {nd}/{} diffs", h0.len()).into());
5279        }
5280        let mut time = |rp: bool| -> Result<f64, Box<dyn std::error::Error>> {
5281            self.gpu.stream().synchronize()?;
5282            let t0 = std::time::Instant::now();
5283            for _ in 0..500 {
5284                if rp {
5285                    let __s_b = self.gpu.stream();
5286                    let mut b = __s_b.launch_builder(&fr);
5287                    b.arg(&wrp_d)
5288                        .arg(&aq_d)
5289                        .arg(&ad_d)
5290                        .arg(&mut y1)
5291                        .arg(&inf)
5292                        .arg(&outf)
5293                        .arg(&mi)
5294                        .arg(&qp);
5295                    unsafe {
5296                        b.launch(cfg)?;
5297                    }
5298                } else {
5299                    let __s_b = self.gpu.stream();
5300                    let mut b = __s_b.launch_builder(&fb);
5301                    b.arg(&w_d)
5302                        .arg(&aq_d)
5303                        .arg(&ad_d)
5304                        .arg(&mut y0)
5305                        .arg(&inf)
5306                        .arg(&outf)
5307                        .arg(&mi)
5308                        .arg(&rb);
5309                    unsafe {
5310                        b.launch(cfg)?;
5311                    }
5312                }
5313            }
5314            self.gpu.stream().synchronize()?;
5315            Ok(t0.elapsed().as_secs_f64() * 1e6 / 500.0)
5316        };
5317        let _ = time(false)?;
5318        let _ = time(true)?; // warm
5319        Ok((time(false)?, time(true)?))
5320    }
5321
5322    /// Build the Q4_0 split-plane decode mirror for a 2D Quant tensor (device-side permutation,
5323    /// q4_0_split_rp_build). Raw bytes stay resident (prefill/gemm/Stage-A); the m<=8 decode
5324    /// dispatch prefers the mirror (_rp twins). No-op unless (Q4_0, 2D, mirror absent).
5325    /// VRAM cost == the tensor's weight size. MEMRA_Q4RP=0 disables at the call sites.
5326    pub fn build_q4_rp4(
5327        &self,
5328        t: &mut crate::model::GpuTensor,
5329    ) -> Result<(), Box<dyn std::error::Error>> {
5330        use crate::model::GpuTensor;
5331        let GpuTensor::Quant {
5332            bytes,
5333            qtype,
5334            row_bytes,
5335            ne,
5336            rp4,
5337            ..
5338        } = t
5339        else {
5340            return Ok(());
5341        };
5342        if *qtype != QT_Q4_0 || rp4.is_some() || ne.len() != 2 {
5343            return Ok(());
5344        }
5345        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5346        if in_f % 32 != 0 || *row_bytes != (in_f / 32) * 18 {
5347            return Ok(());
5348        }
5349        let nblk = in_f / 32;
5350        let mut dst = self.alloc_uninit::<u8>(out_f * nblk * 18)?;
5351        let f = self.func("q4_0_split_rp_build");
5352        let n = (out_f * nblk) as i32;
5353        let cfg = LaunchConfig {
5354            grid_dim: (((out_f * nblk) as u32).div_ceil(256), 1, 1),
5355            block_dim: (256, 1, 1),
5356            shared_mem_bytes: 0,
5357        };
5358        let (of, nb) = (out_f as i32, nblk as i32);
5359        let _ = n;
5360        let __s_b = self.gpu.stream();
5361        let mut b = __s_b.launch_builder(&f);
5362        b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
5363        unsafe {
5364            b.launch(cfg)?;
5365        }
5366        *rp4 = Some(dst);
5367        Ok(())
5368    }
5369
5370    /// Q8_0 twin of `build_q4_rp4` (H100 coalescing fix, 2026-07-26 ncu: GGUF 34B-stride
5371    /// weight loads hold Max Bandwidth at 41-46%; the split mirror makes them aligned 16B
5372    /// ldcs). Raw bytes stay resident (prefill GEMM/MMQ/fused m=1 launches read GGUF layout);
5373    /// the mmvq/batched decode arms prefer the mirror via `rp4`. Bit-identical outputs.
5374    pub fn build_q8_rp4(
5375        &self,
5376        t: &mut crate::model::GpuTensor,
5377    ) -> Result<(), Box<dyn std::error::Error>> {
5378        use crate::model::GpuTensor;
5379        let GpuTensor::Quant {
5380            bytes,
5381            qtype,
5382            row_bytes,
5383            ne,
5384            rp4,
5385            ..
5386        } = t
5387        else {
5388            return Ok(());
5389        };
5390        if *qtype != QT_Q8_0 || rp4.is_some() || ne.len() != 2 {
5391            return Ok(());
5392        }
5393        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5394        if in_f % 32 != 0 || *row_bytes != (in_f / 32) * 34 {
5395            return Ok(());
5396        }
5397        *rp4 = Some(self.build_q8_rp4_raw(bytes, in_f, out_f)?);
5398        Ok(())
5399    }
5400
5401    /// Raw rp-mirror build for gates/benches: split GGUF Q8_0 bytes into the qplane+dplane
5402    /// mirror without a GpuTensor (same kernel the loader path above uses).
5403    pub fn build_q8_rp4_raw(
5404        &self,
5405        bytes: &CudaSlice<u8>,
5406        in_f: usize,
5407        out_f: usize,
5408    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
5409        assert!(in_f % 32 == 0);
5410        let nblk = in_f / 32;
5411        let mut dst = self.alloc_uninit::<u8>(out_f * nblk * 34)?;
5412        let f = self.func("q8_0_split_rp_build");
5413        let cfg = LaunchConfig {
5414            grid_dim: (((out_f * nblk) as u32).div_ceil(256), 1, 1),
5415            block_dim: (256, 1, 1),
5416            shared_mem_bytes: 0,
5417        };
5418        let (of, nb) = (out_f as i32, nblk as i32);
5419        let __s_b = self.gpu.stream();
5420        let mut b = __s_b.launch_builder(&f);
5421        b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
5422        unsafe {
5423            b.launch(cfg)?;
5424        }
5425        Ok(dst)
5426    }
5427
5428    /// K-quant twins of `build_q8_rp4` (H100 K-quant coalescing fix, 2026-08-01 ncu on the
5429    /// q27 Q4_K_M decode: q4_K mmvq DRAM 41-54% with 65% excessive sectors, q6_K 40% with
5430    /// 78% — the 144B/210B superblock strides land every 4B weight load off-sector). The
5431    /// mirror re-packs each tensor into planes (q4_K: qs ++ 16B meta; q6_K: ql ++ qh ++
5432    /// scales ++ d — same total bytes) so every quant fetch is an aligned 16B ldcs. Raw
5433    /// bytes stay resident (prefill GEMM/dequant/Stage-A read GGUF layout); the mmvq/batched
5434    /// decode arms prefer the mirror via `rp4`. Bit-identical outputs.
5435    pub fn build_q4k_rp4(
5436        &self,
5437        t: &mut crate::model::GpuTensor,
5438    ) -> Result<(), Box<dyn std::error::Error>> {
5439        use crate::model::GpuTensor;
5440        let GpuTensor::Quant {
5441            bytes,
5442            qtype,
5443            row_bytes,
5444            ne,
5445            rp4,
5446            ..
5447        } = t
5448        else {
5449            return Ok(());
5450        };
5451        if *qtype != QT_Q4_K || rp4.is_some() || ne.len() != 2 {
5452            return Ok(());
5453        }
5454        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5455        if in_f % 256 != 0 || *row_bytes != (in_f / 256) * 144 {
5456            return Ok(());
5457        }
5458        *rp4 = Some(self.build_kq_rp4_raw(bytes, in_f, out_f, QT_Q4_K)?);
5459        Ok(())
5460    }
5461
5462    pub fn build_q6k_rp4(
5463        &self,
5464        t: &mut crate::model::GpuTensor,
5465    ) -> Result<(), Box<dyn std::error::Error>> {
5466        use crate::model::GpuTensor;
5467        let GpuTensor::Quant {
5468            bytes,
5469            qtype,
5470            row_bytes,
5471            ne,
5472            rp4,
5473            ..
5474        } = t
5475        else {
5476            return Ok(());
5477        };
5478        if *qtype != QT_Q6_K || rp4.is_some() || ne.len() != 2 {
5479            return Ok(());
5480        }
5481        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5482        if in_f % 256 != 0 || *row_bytes != (in_f / 256) * 210 {
5483            return Ok(());
5484        }
5485        *rp4 = Some(self.build_kq_rp4_raw(bytes, in_f, out_f, QT_Q6_K)?);
5486        Ok(())
5487    }
5488
5489    /// Raw K-quant rp-mirror build for gates/benches (same kernels the loader path uses).
5490    pub fn build_kq_rp4_raw(
5491        &self,
5492        bytes: &CudaSlice<u8>,
5493        in_f: usize,
5494        out_f: usize,
5495        qtype: i32,
5496    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
5497        assert!(in_f % 256 == 0);
5498        let nsbk = in_f / 256;
5499        let (sb_bytes, kname) = match qtype {
5500            QT_Q4_K => (144usize, "q4_K_split_rp_build"),
5501            QT_Q6_K => (210usize, "q6_K_split_rp_build"),
5502            _ => return Err(format!("build_kq_rp4_raw: qtype {qtype} has no rp mirror").into()),
5503        };
5504        let mut dst = self.alloc_uninit::<u8>(out_f * nsbk * sb_bytes)?;
5505        let f = self.func(kname);
5506        let cfg = LaunchConfig {
5507            grid_dim: (((out_f * nsbk) as u32).div_ceil(256), 1, 1),
5508            block_dim: (256, 1, 1),
5509            shared_mem_bytes: 0,
5510        };
5511        let (of, nb) = (out_f as i32, nsbk as i32);
5512        let __s_b = self.gpu.stream();
5513        let mut b = __s_b.launch_builder(&f);
5514        b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
5515        unsafe {
5516            b.launch(cfg)?;
5517        }
5518        Ok(dst)
5519    }
5520
5521    /// MEMRA_KQRP seam: the K-quant (q4_K/q6_K) split-plane decode mirrors at model load.
5522    /// Default follows the Q8RP convention — ON on the Hopper lane (80GB pays the mirror
5523    /// VRAM), OFF elsewhere (a 24GB card cannot hold model + mirror + KV for the big trunks).
5524    pub fn kqrp_enabled() -> bool {
5525        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5526        *ON.get_or_init(|| match std::env::var("MEMRA_KQRP").as_deref() {
5527            Ok("0") => false,
5528            Ok(_) => true,
5529            Err(_) => cfg!(memra_hopper_mma),
5530        })
5531    }
5532
5533    /// IN-PLACE split-plane swap (the 31B dense arc): build the split layout and REPLACE the
5534    /// GGUF bytes (zero extra steady-state VRAM — the transient peak is one tensor's size).
5535    /// The tensor's `rp` flag then routes every consumer (mmvq/batched `_rp` twins, the
5536    /// `qmatvec_gemm_q4_0_rp` prefill kernel). Callers gate on the fast path being active —
5537    /// the Stage-A f32 oracle (`MEMRA_FAST=0`) reads GGUF layout and must never see a swap.
5538    pub fn build_q4_rp_swap(
5539        &self,
5540        t: &mut crate::model::GpuTensor,
5541    ) -> Result<bool, Box<dyn std::error::Error>> {
5542        use crate::model::GpuTensor;
5543        // QTYPE GUARD IN THE SWAP ITSELF (gemma4 NVFP4mix prefill-NaN, 2026-08-17):
5544        // `rp4` is a SHARED field — the Q8RP walk parks Q8_0 MIRRORS there, and this
5545        // fn used to `take()` whatever it found. On a Q8_0-carrying gemma4-dense trunk
5546        // the swap hijacked those mirrors: `bytes` became split-plane in place, the
5547        // m<=16 `_rp` dispatch read them correctly (masking the corruption from every
5548        // decode pin), and every GGUF-layout prefill consumer (MMQ + GEMM) read the
5549        // fp16 d-plane as weights -> layer-0 NaN, <pad>-spam serving. Only a tensor
5550        // this fn's OWN builder serves may ever be swapped; everything else refuses
5551        // here, regardless of walk ordering.
5552        if !matches!(t, GpuTensor::Quant { qtype, .. } if *qtype == QT_Q4_0) {
5553            return Ok(false);
5554        }
5555        self.build_q4_rp4(t)?;
5556        self.gpu.stream().synchronize()?; // build kernel reads the GGUF bytes — drain BEFORE dropping them
5557        let GpuTensor::Quant { bytes, rp4, rp, .. } = t else {
5558            return Ok(false);
5559        };
5560        match rp4.take() {
5561            Some(split) => {
5562                *bytes = split; // the GGUF-layout buffer drops here
5563                *rp = true;
5564                Ok(true)
5565            }
5566            None => Ok(false),
5567        }
5568    }
5569
5570    /// MEMRA_Q4RP seam (default ON): the Q4_0 split-plane decode mirror at model load.
5571    pub fn q4rp_enabled() -> bool {
5572        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5573        *ON.get_or_init(|| {
5574            std::env::var("MEMRA_Q4RP")
5575                .map(|v| v != "0")
5576                .unwrap_or(true)
5577        })
5578    }
5579
5580    /// gemma4-E4B: dense [t][row_elems] gather of layer il's rows from the strided prologue
5581    /// buffer ([t][n_layer][n_epl]; off = il*n_epl, stride = n_layer*n_epl).
5582    pub fn copy_rows_strided(
5583        &self,
5584        src: &CudaSlice<f32>,
5585        dst: &mut CudaSlice<f32>,
5586        row_elems: usize,
5587        n_rows: usize,
5588        src_stride: usize,
5589        src_off: usize,
5590    ) -> Result<(), Box<dyn std::error::Error>> {
5591        let f = self.func("copy_rows_strided_f32");
5592        let cfg = LaunchConfig {
5593            grid_dim: (((row_elems as u32 + 255) / 256).max(1), n_rows as u32, 1),
5594            block_dim: (256, 1, 1),
5595            shared_mem_bytes: 0,
5596        };
5597        let (re, nr) = (row_elems as i32, n_rows as i32);
5598        let (st, off) = (src_stride as i64, src_off as i64);
5599        let __s_b = self.gpu.stream();
5600        let mut b = __s_b.launch_builder(&f);
5601        b.arg(src)
5602            .arg(&mut *dst)
5603            .arg(&re)
5604            .arg(&nr)
5605            .arg(&st)
5606            .arg(&off);
5607        unsafe {
5608            b.launch(cfg)?;
5609        }
5610        Ok(())
5611    }
5612
5613    /// Async device u32 store (value rides the kernel ARG — no host-memory transfer/sync).
5614    pub fn u32_set_k(
5615        &self,
5616        dst: &mut CudaSlice<u32>,
5617        v: u32,
5618        idx: usize,
5619    ) -> Result<(), Box<dyn std::error::Error>> {
5620        let f = self.func("u32_set_k");
5621        let cfg = LaunchConfig {
5622            grid_dim: (1, 1, 1),
5623            block_dim: (1, 1, 1),
5624            shared_mem_bytes: 0,
5625        };
5626        let ii = idx as i32;
5627        let __s_b = self.gpu.stream();
5628        let mut b = __s_b.launch_builder(&f);
5629        b.arg(dst).arg(&v).arg(&ii);
5630        unsafe {
5631            b.launch(cfg)?;
5632        }
5633        Ok(())
5634    }
5635
5636    /// counter += v (device-slot append advance; the +1 twin is `inc_seqlen`).
5637    pub fn i32_add_k(
5638        &self,
5639        d: &mut CudaSlice<i32>,
5640        v: i32,
5641    ) -> Result<(), Box<dyn std::error::Error>> {
5642        let f = self.func("i32_add_k");
5643        let cfg = LaunchConfig {
5644            grid_dim: (1, 1, 1),
5645            block_dim: (32, 1, 1),
5646            shared_mem_bytes: 0,
5647        };
5648        let __s_b = self.gpu.stream();
5649        let mut b = __s_b.launch_builder(&f);
5650        b.arg(d).arg(&v);
5651        unsafe {
5652            b.launch(cfg)?;
5653        }
5654        Ok(())
5655    }
5656
5657    /// pos rows from a device counter: dst[i] = ctr[0] + i (verify-stream rope positions).
5658    pub fn i32_iota_from(
5659        &self,
5660        ctr: &CudaSlice<i32>,
5661        dst: &mut CudaSlice<i32>,
5662        n: usize,
5663    ) -> Result<(), Box<dyn std::error::Error>> {
5664        let f = self.func("i32_iota_from");
5665        let cfg = LaunchConfig::for_num_elems(n as u32);
5666        let ni = n as i32;
5667        let __s_b = self.gpu.stream();
5668        let mut b = __s_b.launch_builder(&f);
5669        b.arg(ctr).arg(dst).arg(&ni);
5670        unsafe {
5671            b.launch(cfg)?;
5672        }
5673        Ok(())
5674    }
5675
5676    /// In-place trim-id translate: buf[idx] = map[buf[idx]] (FR-Spec d2t, async single-slot).
5677    pub fn u32_map_k(
5678        &self,
5679        buf: &mut CudaSlice<u32>,
5680        map: &CudaSlice<u32>,
5681        idx: usize,
5682    ) -> Result<(), Box<dyn std::error::Error>> {
5683        let f = self.func("u32_map_k");
5684        let cfg = LaunchConfig {
5685            grid_dim: (1, 1, 1),
5686            block_dim: (1, 1, 1),
5687            shared_mem_bytes: 0,
5688        };
5689        let ii = idx as i32;
5690        let __s_b = self.gpu.stream();
5691        let mut b = __s_b.launch_builder(&f);
5692        b.arg(buf).arg(map).arg(&ii);
5693        unsafe {
5694            b.launch(cfg)?;
5695        }
5696        Ok(())
5697    }
5698
5699    /// Pack a[off..off+n1] ++ b[0..n2] into one buffer (single dtoh follows).
5700    #[allow(clippy::too_many_arguments)]
5701    pub fn u32_pack2(
5702        &self,
5703        a: &CudaSlice<u32>,
5704        off_a: usize,
5705        n1: usize,
5706        b_in: &CudaSlice<u32>,
5707        n2: usize,
5708        out: &mut CudaSlice<u32>,
5709    ) -> Result<(), Box<dyn std::error::Error>> {
5710        let f = self.func("u32_pack2");
5711        let cfg = LaunchConfig::for_num_elems((n1 + n2) as u32);
5712        let (oa, i1, i2) = (off_a as i32, n1 as i32, n2 as i32);
5713        let __s_b = self.gpu.stream();
5714        let mut b = __s_b.launch_builder(&f);
5715        b.arg(a).arg(&oa).arg(&i1).arg(b_in).arg(&i2).arg(out);
5716        unsafe {
5717            b.launch(cfg)?;
5718        }
5719        Ok(())
5720    }
5721
5722    /// gemma4 R3 device fold: w[i] *= s[sel[i]] over the router's [n] (sel, w) pair.
5723    pub fn moe_w_exscale(
5724        &self,
5725        w: &mut CudaSlice<f32>,
5726        sel: &CudaSlice<i32>,
5727        s: &CudaSlice<f32>,
5728        n: usize,
5729    ) -> Result<(), Box<dyn std::error::Error>> {
5730        let f = self.func("moe_w_exscale");
5731        let cfg = LaunchConfig::for_num_elems(n as u32);
5732        let ni = n as i32;
5733        let __s_b = self.gpu.stream();
5734        let mut b = __s_b.launch_builder(&f);
5735        b.arg(w).arg(sel).arg(s).arg(&ni);
5736        unsafe {
5737            b.launch(cfg)?;
5738        }
5739        Ok(())
5740    }
5741
5742    /// Down-projection macro fold: w[i] *= macros[2*n_expert + sel[i]] on the device router
5743    /// weights (one launch per MoE layer, only for macro-carrying artifacts — see MoeWeights).
5744    pub fn moe_w_scale_by_expert(
5745        &self,
5746        w: &mut CudaSlice<f32>,
5747        sel: &CudaSlice<i32>,
5748        macros: &CudaSlice<f32>,
5749        n_expert: usize,
5750        n: usize,
5751    ) -> Result<(), Box<dyn std::error::Error>> {
5752        let f = self.func("moe_w_scale_by_expert");
5753        let cfg = LaunchConfig {
5754            grid_dim: (n.div_ceil(64) as u32, 1, 1),
5755            block_dim: (64, 1, 1),
5756            shared_mem_bytes: 0,
5757        };
5758        let (ne, nn) = (n_expert as i32, n as i32);
5759        let __s_b = self.gpu.stream();
5760        let mut b = __s_b.launch_builder(&f);
5761        b.arg(w).arg(sel).arg(macros).arg(&ne).arg(&nn);
5762        unsafe {
5763            b.launch(cfg)?;
5764        }
5765        Ok(())
5766    }
5767
5768    pub fn moe_gate_up_silu8_dev_q8(
5769        &self,
5770        table: &CudaSlice<u64>,
5771        sel: &cudarc::driver::CudaView<i32>,
5772        aq: &CudaSlice<i8>,
5773        ad: &CudaSlice<f32>,
5774        in_f: usize,
5775        n_ff: usize,
5776        n_used: usize,
5777        n_expert: usize,
5778        qt_g: i32,
5779        qt_u: i32,
5780        rb_g: usize,
5781        rb_u: usize,
5782        macros: &CudaSlice<f32>,
5783    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5784        static GU: std::sync::OnceLock<(String, u32)> = std::sync::OnceLock::new();
5785        let (mode, wpb) = GU.get_or_init(|| {
5786            let mode = std::env::var("MEMRA_MOE_DEVQ8_GU").unwrap_or_default();
5787            let wpb = std::env::var("MEMRA_MOE_DEVQ8_WPB")
5788                .ok()
5789                .and_then(|v| v.parse().ok())
5790                .unwrap_or(4u32)
5791                .clamp(1, 16);
5792            (mode, wpb)
5793        });
5794        let (mode, wpb) = (mode.as_str(), *wpb);
5795        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
5796        let (inf, nff, ne, rbg, rbu) = (
5797            in_f as i32,
5798            n_ff as i32,
5799            n_expert as i32,
5800            rb_g as i64,
5801            rb_u as i64,
5802        );
5803        let (f, cfg) = match mode {
5804            "1" | "2" | "4" => {
5805                let rpw: u32 = mode.parse().unwrap();
5806                let f = self.func(match rpw {
5807                    1 => "moe_gate_up_silu8_dev_q8_r1",
5808                    2 => "moe_gate_up_silu8_dev_q8_r2",
5809                    _ => "moe_gate_up_silu8_dev_q8_r4",
5810                });
5811                let rows_per_block = (rpw * wpb) as usize;
5812                let gx = n_ff.div_ceil(rows_per_block) as u32;
5813                (
5814                    f,
5815                    LaunchConfig {
5816                        grid_dim: (gx, n_used as u32, 1),
5817                        block_dim: (32, wpb, 1),
5818                        shared_mem_bytes: 0,
5819                    },
5820                )
5821            }
5822            "j8" if n_used <= 32 => (
5823                self.func("moe_gate_up_silu8_dev_q8_j8"),
5824                LaunchConfig {
5825                    grid_dim: (n_ff as u32, 1, 1),
5826                    block_dim: (32, n_used as u32, 1),
5827                    shared_mem_bytes: 0,
5828                },
5829            ),
5830            // SMEM-GRID twins (IQ3_S 2KB grid copied to shared, static smem — bit-identical dots)
5831            "vsm2" => {
5832                let f = self.func("moe_gate_up_silu8_dev_q8_vsm2");
5833                let sh = (rb_g + rb_u) as u32;
5834                use cudarc::driver::sys::CUfunction_attribute_enum as A;
5835                f.set_attribute(
5836                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
5837                    sh as i32,
5838                )?;
5839                (
5840                    f,
5841                    LaunchConfig {
5842                        grid_dim: (n_ff as u32, n_used as u32, 1),
5843                        block_dim: (32, 1, 1),
5844                        shared_mem_bytes: sh,
5845                    },
5846                )
5847            }
5848            "vsm" => {
5849                let f = self.func("moe_gate_up_silu8_dev_q8_vsm");
5850                let sh = (rb_g + rb_u) as u32;
5851                use cudarc::driver::sys::CUfunction_attribute_enum as A;
5852                f.set_attribute(
5853                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
5854                    sh as i32,
5855                )?;
5856                (
5857                    f,
5858                    LaunchConfig {
5859                        grid_dim: (n_ff as u32, n_used as u32, 1),
5860                        block_dim: (32, 1, 1),
5861                        shared_mem_bytes: sh,
5862                    },
5863                )
5864            }
5865            "sg" => (
5866                self.func("moe_gate_up_silu8_dev_q8_sg"),
5867                LaunchConfig {
5868                    grid_dim: (n_ff as u32, n_used as u32, 1),
5869                    block_dim: (32, 1, 1),
5870                    shared_mem_bytes: 0,
5871                },
5872            ),
5873            "j8sg" if n_used <= 32 => (
5874                self.func("moe_gate_up_silu8_dev_q8_j8sg"),
5875                LaunchConfig {
5876                    grid_dim: (n_ff as u32, 1, 1),
5877                    block_dim: (32, n_used as u32, 1),
5878                    shared_mem_bytes: 0,
5879                },
5880            ),
5881            "u64" if in_f == 2048 => (
5882                self.func("moe_gate_up_silu8_dev_q8_u64"),
5883                LaunchConfig {
5884                    grid_dim: (n_ff as u32, n_used as u32, 1),
5885                    block_dim: (32, 1, 1),
5886                    shared_mem_bytes: 0,
5887                },
5888            ),
5889            "gs4" if in_f == 2048 => (
5890                self.func("moe_gate_up_silu8_dev_q8_gs4"),
5891                LaunchConfig {
5892                    grid_dim: (n_ff as u32, n_used as u32, 1),
5893                    block_dim: (32, 4, 1),
5894                    shared_mem_bytes: 0,
5895                },
5896            ),
5897            // _v twin (down8 lane 2026-07-08): wide-load IQ4_XS dot, base geometry, bit-identical.
5898            "v" | "" => (
5899                self.func("moe_gate_up_silu8_dev_q8_v"),
5900                LaunchConfig {
5901                    grid_dim: (n_ff as u32, n_used as u32, 1),
5902                    block_dim: (32, 1, 1),
5903                    shared_mem_bytes: 0,
5904                },
5905            ),
5906            "s2" => (
5907                self.func("moe_gate_up_silu8_dev_q8_s2"),
5908                LaunchConfig {
5909                    grid_dim: (n_ff as u32, n_used as u32, 1),
5910                    block_dim: (32, 2, 1),
5911                    shared_mem_bytes: 0,
5912                },
5913            ),
5914            "s2z" => {
5915                let rz = wpb.min(16); // s2z smem tile is [16][2]
5916                (
5917                    self.func("moe_gate_up_silu8_dev_q8_s2z"),
5918                    LaunchConfig {
5919                        grid_dim: (n_ff.div_ceil(rz as usize) as u32, n_used as u32, 1),
5920                        block_dim: (32, 2, rz),
5921                        shared_mem_bytes: 0,
5922                    },
5923                )
5924            }
5925            _ => (
5926                self.func("moe_gate_up_silu8_dev_q8"),
5927                LaunchConfig {
5928                    grid_dim: (n_ff as u32, n_used as u32, 1),
5929                    block_dim: (32, 1, 1),
5930                    shared_mem_bytes: 0,
5931                },
5932            ),
5933        };
5934        let __s_b = self.gpu.stream();
5935        let mut b = __s_b.launch_builder(&f);
5936        b.arg(table)
5937            .arg(sel)
5938            .arg(aq)
5939            .arg(ad)
5940            .arg(&mut act)
5941            .arg(&inf)
5942            .arg(&nff)
5943            .arg(&ne)
5944            .arg(&qt_g)
5945            .arg(&qt_u)
5946            .arg(&rbg)
5947            .arg(&rbu)
5948            .arg(macros);
5949        unsafe {
5950            b.launch(cfg)?;
5951        }
5952        Ok(act)
5953    }
5954
5955    #[allow(clippy::too_many_arguments)]
5956    pub fn moe_down8_fma_dev_q8(
5957        &self,
5958        table: &CudaSlice<u64>,
5959        sel: &cudarc::driver::CudaView<i32>,
5960        w: &cudarc::driver::CudaView<f32>,
5961        aq2: &CudaSlice<i8>,
5962        ad2: &CudaSlice<f32>,
5963        dst: &mut cudarc::driver::CudaViewMut<f32>,
5964        in_f: usize,
5965        out_f: usize,
5966        n_used: usize,
5967        n_expert: usize,
5968        qt: i32,
5969        rb: usize,
5970    ) -> Result<(), Box<dyn std::error::Error>> {
5971        static DOWN: std::sync::OnceLock<String> = std::sync::OnceLock::new();
5972        let mode = DOWN.get_or_init(|| std::env::var("MEMRA_MOE_DEVQ8_DOWN").unwrap_or_default());
5973        let (inf, outf, nu, ne, rbi) = (
5974            in_f as i32,
5975            out_f as i32,
5976            n_used as i32,
5977            n_expert as i32,
5978            rb as i64,
5979        );
5980        // the w8 twins' smem tile is [RPW][8] — n_used must fit the 8-slot tile;
5981        // the h2 twins are nsb==16 (in_f==512) shape-gated.
5982        let (f, cfg) = match mode.as_str() {
5983            m @ ("1" | "2" | "4") if n_used <= 8 => {
5984                let rpw: usize = m.parse().unwrap();
5985                let f = self.func(match rpw {
5986                    1 => "moe_down8_fma_dev_q8_w8r1",
5987                    2 => "moe_down8_fma_dev_q8_w8r2",
5988                    _ => "moe_down8_fma_dev_q8_w8r4",
5989                });
5990                (
5991                    f,
5992                    LaunchConfig {
5993                        grid_dim: (out_f.div_ceil(rpw) as u32, 1, 1),
5994                        block_dim: (32, n_used as u32, 1),
5995                        shared_mem_bytes: 0,
5996                    },
5997                )
5998            }
5999            "h2" if in_f == 512 => (
6000                self.func("moe_down8_fma_dev_q8_h2"),
6001                LaunchConfig {
6002                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6003                    block_dim: (32, 1, 1),
6004                    shared_mem_bytes: 0,
6005                },
6006            ),
6007            // "" = AUTO gemma shape (in_f==704): w8r2 measured +1 tok/s vs base (sweep
6008            // 1/2/4 -> 133.6/134.2/133.6, 2026-07-10); slot-ordered chain preserved.
6009            "" if in_f == 704 && n_used <= 8 => (
6010                self.func("moe_down8_fma_dev_q8_w8r2"),
6011                LaunchConfig {
6012                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6013                    block_dim: (32, n_used as u32, 1),
6014                    shared_mem_bytes: 0,
6015                },
6016            ),
6017            // "" = AUTO: the measured winner for the 35B expert shape (arc 2026-07-05, +3.8%);
6018            // any shape the h2 kernels can't take (nsb!=16 / n_used>8) falls to base via `_`.
6019            // _v twins (down8 lane 2026-07-08): wide-load IQ4_XS dot, bit-identical outputs.
6020            "w8h2v" | "" if in_f == 512 && n_used <= 8 => (
6021                self.func("moe_down8_fma_dev_q8_w8h2v"),
6022                LaunchConfig {
6023                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6024                    block_dim: (32, n_used as u32, 1),
6025                    shared_mem_bytes: 0,
6026                },
6027            ),
6028            "w8h2r2v" if in_f == 512 && n_used <= 8 => (
6029                self.func("moe_down8_fma_dev_q8_w8h2r2v"),
6030                LaunchConfig {
6031                    grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
6032                    block_dim: (32, n_used as u32, 1),
6033                    shared_mem_bytes: 0,
6034                },
6035            ),
6036            "w8h2r2" if in_f == 512 && n_used <= 8 => (
6037                self.func("moe_down8_fma_dev_q8_w8h2r2"),
6038                LaunchConfig {
6039                    grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
6040                    block_dim: (32, n_used as u32, 1),
6041                    shared_mem_bytes: 0,
6042                },
6043            ),
6044            "w8h2" if in_f == 512 && n_used <= 8 => (
6045                self.func("moe_down8_fma_dev_q8_w8h2"),
6046                LaunchConfig {
6047                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6048                    block_dim: (32, n_used as u32, 1),
6049                    shared_mem_bytes: 0,
6050                },
6051            ),
6052            _ => (
6053                self.func("moe_down8_fma_dev_q8"),
6054                LaunchConfig {
6055                    grid_dim: (out_f as u32, 1, 1),
6056                    block_dim: (32, 1, 1),
6057                    shared_mem_bytes: 0,
6058                },
6059            ),
6060        };
6061        let __s_b = self.gpu.stream();
6062        let mut b = __s_b.launch_builder(&f);
6063        b.arg(table)
6064            .arg(sel)
6065            .arg(w)
6066            .arg(aq2)
6067            .arg(ad2)
6068            .arg(dst)
6069            .arg(&inf)
6070            .arg(&outf)
6071            .arg(&nu)
6072            .arg(&ne)
6073            .arg(&qt)
6074            .arg(&rbi);
6075        unsafe {
6076            b.launch(cfg)?;
6077        }
6078        Ok(())
6079    }
6080
6081    /// SMALL-M VERIFY rows twin (MEMRA_SPEC_M2, lane/spec-m2): ONE launch covers all `t` tokens
6082    /// of the spec verify's MoE dev gate/up (grid.z = token) — the _v geometry per token, with
6083    /// tok-offset sel/aq/ad/act pointers matching the serial loop's slices. BIT-IDENTICAL per
6084    /// token (see the kernel header). aq/ad are the BATCHED z-quantize ([t, in_f] rows —
6085    /// quantize_q8_1's per-32-block program is row-independent, so batched rows == the serial
6086    /// loop's per-token quantize_q8_1_view bytes). Returns act [t, n_used, n_ff].
6087    #[allow(clippy::too_many_arguments)]
6088    pub fn moe_gate_up_silu8_dev_q8_rows(
6089        &self,
6090        table: &CudaSlice<u64>,
6091        sel: &CudaSlice<i32>,
6092        aq: &CudaSlice<i8>,
6093        ad: &CudaSlice<f32>,
6094        t: usize,
6095        in_f: usize,
6096        n_ff: usize,
6097        n_used: usize,
6098        n_expert: usize,
6099        qt_g: i32,
6100        qt_u: i32,
6101        rb_g: usize,
6102        rb_u: usize,
6103        macros: &CudaSlice<f32>,
6104    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6105        let f = self.func("moe_gate_up_silu8_dev_q8_v_rows");
6106        let mut act = self.alloc_uninit::<f32>(t * n_used * n_ff)?;
6107        let cfg = LaunchConfig {
6108            grid_dim: (n_ff as u32, n_used as u32, t as u32),
6109            block_dim: (32, 1, 1),
6110            shared_mem_bytes: 0,
6111        };
6112        let (inf, nff, ne, nu, rbg, rbu) = (
6113            in_f as i32,
6114            n_ff as i32,
6115            n_expert as i32,
6116            n_used as i32,
6117            rb_g as i64,
6118            rb_u as i64,
6119        );
6120        let __s_b = self.gpu.stream();
6121        let mut b = __s_b.launch_builder(&f);
6122        b.arg(table)
6123            .arg(sel)
6124            .arg(aq)
6125            .arg(ad)
6126            .arg(&mut act)
6127            .arg(&inf)
6128            .arg(&nff)
6129            .arg(&ne)
6130            .arg(&qt_g)
6131            .arg(&qt_u)
6132            .arg(&rbg)
6133            .arg(&rbu)
6134            .arg(&nu)
6135            .arg(macros);
6136        unsafe {
6137            b.launch(cfg)?;
6138        }
6139        Ok(act)
6140    }
6141
6142    /// SMALL-M VERIFY rows twin of the down proj: w8h2v geometry per token on a grid.z token
6143    /// axis. Caller gates the w8h2v shape contract (in_f == 512, n_used <= 8) — same gate as
6144    /// the AUTO dispatch in `moe_down8_fma_dev_q8`. aq2/ad2 = batched act quantize
6145    /// ([t*n_used, in_f] rows). dst rows are FULLY overwritten per token.
6146    #[allow(clippy::too_many_arguments)]
6147    pub fn moe_down8_fma_dev_q8_rows(
6148        &self,
6149        table: &CudaSlice<u64>,
6150        sel: &CudaSlice<i32>,
6151        w: &CudaSlice<f32>,
6152        aq2: &CudaSlice<i8>,
6153        ad2: &CudaSlice<f32>,
6154        dst: &mut CudaSlice<f32>,
6155        t: usize,
6156        in_f: usize,
6157        out_f: usize,
6158        n_used: usize,
6159        n_expert: usize,
6160        qt: i32,
6161        rb: usize,
6162    ) -> Result<(), Box<dyn std::error::Error>> {
6163        assert!(
6164            in_f == 512 && n_used <= 8,
6165            "down rows twin is w8h2v shape-gated"
6166        );
6167        let f = self.func("moe_down8_fma_dev_q8_w8h2v_rows");
6168        let cfg = LaunchConfig {
6169            grid_dim: (out_f.div_ceil(2) as u32, 1, t as u32),
6170            block_dim: (32, n_used as u32, 1),
6171            shared_mem_bytes: 0,
6172        };
6173        let (inf, outf, nu, ne, rbi) = (
6174            in_f as i32,
6175            out_f as i32,
6176            n_used as i32,
6177            n_expert as i32,
6178            rb as i64,
6179        );
6180        let __s_b = self.gpu.stream();
6181        let mut b = __s_b.launch_builder(&f);
6182        b.arg(table)
6183            .arg(sel)
6184            .arg(w)
6185            .arg(aq2)
6186            .arg(ad2)
6187            .arg(dst)
6188            .arg(&inf)
6189            .arg(&outf)
6190            .arg(&nu)
6191            .arg(&ne)
6192            .arg(&qt)
6193            .arg(&rbi);
6194        unsafe {
6195            b.launch(cfg)?;
6196        }
6197        Ok(())
6198    }
6199
6200    /// CSR gate/up v3 (owner-scan dedup, no build kernel): qtypes {IQ4_XS, IQ3_S} (caller
6201    /// gates), grid.y = pair index; the first pair of each expert serves all its pairs.
6202    /// Bit-identical to moe_gate_up_silu8_dev_q8_v_rows (explicit-intrinsic accumulate).
6203    #[allow(clippy::too_many_arguments)]
6204    pub fn moe_gate_up_silu8_dev_q8_csr(
6205        &self,
6206        table: &CudaSlice<u64>,
6207        sel: &CudaSlice<i32>,
6208        aq: &CudaSlice<i8>,
6209        ad: &CudaSlice<f32>,
6210        n_pairs: usize,
6211        in_f: usize,
6212        n_ff: usize,
6213        n_used: usize,
6214        n_expert: usize,
6215        qt_g: i32,
6216        qt_u: i32,
6217        rb_g: usize,
6218        rb_u: usize,
6219    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6220        // NVFP4 experts take the NVFP4-specialized owner-scan twin (lane/moebatch-q35moe);
6221        // host gate guarantees qt_g == qt_u within a supported class.
6222        let f = if qt_g == crate::QT_NVFP4 {
6223            self.func("moe_gate_up_silu8_dev_q8_csr_nvfp4")
6224        } else {
6225            self.func("moe_gate_up_silu8_dev_q8_csr_iq4")
6226        };
6227        let mut act = self.alloc_uninit::<f32>(n_pairs * n_ff)?;
6228        let cfg = LaunchConfig {
6229            grid_dim: (n_ff as u32, n_pairs as u32, 1),
6230            block_dim: (32, 1, 1),
6231            shared_mem_bytes: 0,
6232        };
6233        let (inf, nff, ne, nu, npi, rbg, rbu) = (
6234            in_f as i32,
6235            n_ff as i32,
6236            n_expert as i32,
6237            n_used as i32,
6238            n_pairs as i32,
6239            rb_g as i64,
6240            rb_u as i64,
6241        );
6242        let __s_b = self.gpu.stream();
6243        let mut b = __s_b.launch_builder(&f);
6244        b.arg(table)
6245            .arg(sel)
6246            .arg(aq)
6247            .arg(ad)
6248            .arg(&mut act)
6249            .arg(&inf)
6250            .arg(&nff)
6251            .arg(&ne)
6252            .arg(&qt_g)
6253            .arg(&qt_u)
6254            .arg(&rbg)
6255            .arg(&rbu)
6256            .arg(&nu)
6257            .arg(&npi);
6258        unsafe {
6259            b.launch(cfg)?;
6260        }
6261        Ok(act)
6262    }
6263
6264    /// TEST SEAM (down8 lane 2026-07-08): launch a down dev_q8 variant BY NAME with its
6265    /// canonical geometry, bypassing the env-cached dispatch so moe-devq8-check can byte-
6266    /// compare variants in one process. Variants: "base", "w8h2", "w8h2r2", "w8h2v", "w8h2r2v".
6267    #[allow(clippy::too_many_arguments)]
6268    pub fn moe_down8_fma_dev_q8_variant(
6269        &self,
6270        variant: &str,
6271        table: &CudaSlice<u64>,
6272        sel: &cudarc::driver::CudaView<i32>,
6273        w: &cudarc::driver::CudaView<f32>,
6274        aq2: &CudaSlice<i8>,
6275        ad2: &CudaSlice<f32>,
6276        dst: &mut cudarc::driver::CudaViewMut<f32>,
6277        in_f: usize,
6278        out_f: usize,
6279        n_used: usize,
6280        n_expert: usize,
6281        qt: i32,
6282        rb: usize,
6283    ) -> Result<(), Box<dyn std::error::Error>> {
6284        let (inf, outf, nu, ne, rbi) = (
6285            in_f as i32,
6286            out_f as i32,
6287            n_used as i32,
6288            n_expert as i32,
6289            rb as i64,
6290        );
6291        let (f, cfg) = match variant {
6292            "w8h2" | "w8h2v" => (
6293                self.func(if variant == "w8h2" {
6294                    "moe_down8_fma_dev_q8_w8h2"
6295                } else {
6296                    "moe_down8_fma_dev_q8_w8h2v"
6297                }),
6298                LaunchConfig {
6299                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6300                    block_dim: (32, n_used as u32, 1),
6301                    shared_mem_bytes: 0,
6302                },
6303            ),
6304            "w8h2r2" | "w8h2r2v" => (
6305                self.func(if variant == "w8h2r2" {
6306                    "moe_down8_fma_dev_q8_w8h2r2"
6307                } else {
6308                    "moe_down8_fma_dev_q8_w8h2r2v"
6309                }),
6310                LaunchConfig {
6311                    grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
6312                    block_dim: (32, n_used as u32, 1),
6313                    shared_mem_bytes: 0,
6314                },
6315            ),
6316            _ => (
6317                self.func("moe_down8_fma_dev_q8"),
6318                LaunchConfig {
6319                    grid_dim: (out_f as u32, 1, 1),
6320                    block_dim: (32, 1, 1),
6321                    shared_mem_bytes: 0,
6322                },
6323            ),
6324        };
6325        let __s_b = self.gpu.stream();
6326        let mut b = __s_b.launch_builder(&f);
6327        b.arg(table)
6328            .arg(sel)
6329            .arg(w)
6330            .arg(aq2)
6331            .arg(ad2)
6332            .arg(dst)
6333            .arg(&inf)
6334            .arg(&outf)
6335            .arg(&nu)
6336            .arg(&ne)
6337            .arg(&qt)
6338            .arg(&rbi);
6339        unsafe {
6340            b.launch(cfg)?;
6341        }
6342        Ok(())
6343    }
6344
6345    /// TEST SEAM (down8 lane): gate_up twin of the above. Variants: "base", "v".
6346    #[allow(clippy::too_many_arguments)]
6347    pub fn moe_gate_up_silu8_dev_q8_variant(
6348        &self,
6349        variant: &str,
6350        table: &CudaSlice<u64>,
6351        sel: &cudarc::driver::CudaView<i32>,
6352        aq: &CudaSlice<i8>,
6353        ad: &CudaSlice<f32>,
6354        in_f: usize,
6355        n_ff: usize,
6356        n_used: usize,
6357        n_expert: usize,
6358        qt_g: i32,
6359        qt_u: i32,
6360        rb_g: usize,
6361        rb_u: usize,
6362    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6363        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
6364        let (inf, nff, ne, rbg, rbu) = (
6365            in_f as i32,
6366            n_ff as i32,
6367            n_expert as i32,
6368            rb_g as i64,
6369            rb_u as i64,
6370        );
6371        let f = self.func(if variant == "v" {
6372            "moe_gate_up_silu8_dev_q8_v"
6373        } else {
6374            "moe_gate_up_silu8_dev_q8"
6375        });
6376        let cfg = LaunchConfig {
6377            grid_dim: (n_ff as u32, n_used as u32, 1),
6378            block_dim: (32, 1, 1),
6379            shared_mem_bytes: 0,
6380        };
6381        let __s_b = self.gpu.stream();
6382        let mut b = __s_b.launch_builder(&f);
6383        b.arg(table)
6384            .arg(sel)
6385            .arg(aq)
6386            .arg(ad)
6387            .arg(&mut act)
6388            .arg(&inf)
6389            .arg(&nff)
6390            .arg(&ne)
6391            .arg(&qt_g)
6392            .arg(&qt_u)
6393            .arg(&rbg)
6394            .arg(&rbu);
6395        unsafe {
6396            b.launch(cfg)?;
6397        }
6398        Ok(act)
6399    }
6400
6401    pub fn moe_gate_up_silu8_dev(
6402        &self,
6403        table: &CudaSlice<u64>,
6404        sel: &cudarc::driver::CudaView<i32>,
6405        x: &cudarc::driver::CudaView<f32>,
6406        in_f: usize,
6407        n_ff: usize,
6408        n_used: usize,
6409        n_expert: usize,
6410        qt_g: i32,
6411        qt_u: i32,
6412        rb_g: usize,
6413        rb_u: usize,
6414        macros: &CudaSlice<f32>,
6415    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6416        let f = self.func("moe_gate_up_silu8_dev");
6417        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?; // fully overwritten
6418        let cfg = LaunchConfig {
6419            grid_dim: (n_ff as u32, n_used as u32, 1),
6420            block_dim: (256, 1, 1),
6421            shared_mem_bytes: 0,
6422        };
6423        let (inf, nff, ne, rbg, rbu) = (
6424            in_f as i32,
6425            n_ff as i32,
6426            n_expert as i32,
6427            rb_g as i64,
6428            rb_u as i64,
6429        );
6430        let __s_b = self.gpu.stream();
6431        let mut b = __s_b.launch_builder(&f);
6432        b.arg(table)
6433            .arg(sel)
6434            .arg(x)
6435            .arg(&mut act)
6436            .arg(&inf)
6437            .arg(&nff)
6438            .arg(&ne)
6439            .arg(&qt_g)
6440            .arg(&qt_u)
6441            .arg(&rbg)
6442            .arg(&rbu)
6443            .arg(macros);
6444        unsafe {
6445            b.launch(cfg)?;
6446        }
6447        Ok(act)
6448    }
6449
6450    /// LAUNCH-STRUCTURE STAGE 3: device-dispatch twin of `moe_down8_fma_into` — expert ids AND
6451    /// renormalized weights read from the router kernel's device output. BIT-IDENTICAL chain.
6452    #[allow(clippy::too_many_arguments)]
6453    pub fn moe_down8_fma_dev(
6454        &self,
6455        table: &CudaSlice<u64>,
6456        sel: &cudarc::driver::CudaView<i32>,
6457        w: &cudarc::driver::CudaView<f32>,
6458        act: &CudaSlice<f32>,
6459        dst: &mut cudarc::driver::CudaViewMut<f32>,
6460        in_f: usize,
6461        out_f: usize,
6462        n_used: usize,
6463        n_expert: usize,
6464        qt: i32,
6465        rb: usize,
6466    ) -> Result<(), Box<dyn std::error::Error>> {
6467        let f = self.func("moe_down8_fma_dev");
6468        let cfg = LaunchConfig {
6469            grid_dim: (out_f as u32, 1, 1),
6470            block_dim: (256, 1, 1),
6471            shared_mem_bytes: 0,
6472        };
6473        let (inf, outf, nu, ne, rbv) = (
6474            in_f as i32,
6475            out_f as i32,
6476            n_used as i32,
6477            n_expert as i32,
6478            rb as i64,
6479        );
6480        let __s_b = self.gpu.stream();
6481        let mut b = __s_b.launch_builder(&f);
6482        b.arg(table)
6483            .arg(sel)
6484            .arg(w)
6485            .arg(act)
6486            .arg(dst)
6487            .arg(&inf)
6488            .arg(&outf)
6489            .arg(&nu)
6490            .arg(&ne)
6491            .arg(&qt)
6492            .arg(&rbv);
6493        unsafe {
6494            b.launch(cfg)?;
6495        }
6496        Ok(())
6497    }
6498
6499    /// dst[i] += alpha * src[i], i in 0..n. dst is a CudaViewMut (a row of moe_out).
6500    pub fn axpy_into(
6501        &self,
6502        src: &CudaSlice<f32>,
6503        alpha: f32,
6504        dst: &mut cudarc::driver::CudaViewMut<f32>,
6505        n: usize,
6506    ) -> Result<(), Box<dyn std::error::Error>> {
6507        let f = self.func("axpy_f32");
6508        let cfg = LaunchConfig::for_num_elems(n as u32);
6509        let (a, ni) = (alpha, n as i32);
6510        let __s_b = self.gpu.stream();
6511        let mut b = __s_b.launch_builder(&f);
6512        b.arg(src).arg(dst).arg(&a).arg(&ni);
6513        unsafe {
6514            b.launch(cfg)?;
6515        }
6516        Ok(())
6517    }
6518
6519    /// dst[r*ncols + c] += src[r*ncols + c] * scale[r]. Per-row scalar accumulate (shared expert).
6520    pub fn add_scaled_rows(
6521        &self,
6522        src: &CudaSlice<f32>,
6523        scale: &CudaSlice<f32>,
6524        dst: &mut CudaSlice<f32>,
6525        ncols: usize,
6526        nrows: usize,
6527    ) -> Result<(), Box<dyn std::error::Error>> {
6528        let f = self.func("add_scaled_rows_f32");
6529        let cfg = LaunchConfig::for_num_elems((ncols * nrows) as u32);
6530        let (nc, nr) = (ncols as i32, nrows as i32);
6531        let __s_b = self.gpu.stream();
6532        let mut b = __s_b.launch_builder(&f);
6533        b.arg(src).arg(scale).arg(dst).arg(&nc).arg(&nr);
6534        unsafe {
6535            b.launch(cfg)?;
6536        }
6537        Ok(())
6538    }
6539
6540    // ======== A2 GROUPED MoE PREFILL KERNELS ========
6541
6542    /// Gather m_e rows from src[T, ncols] into dst[m_e, ncols] using index array idx[m_e].
6543    pub fn gather_rows(
6544        &self,
6545        src: &CudaSlice<f32>,
6546        idx: &CudaSlice<i32>,
6547        dst: &mut CudaSlice<f32>,
6548        ncols: usize,
6549        m_e: usize,
6550    ) -> Result<(), Box<dyn std::error::Error>> {
6551        let f = self.func("gather_rows_f32");
6552        let cfg = LaunchConfig::for_num_elems((m_e * ncols) as u32);
6553        let (nc, me) = (ncols as i32, m_e as i32);
6554        let __s_b = self.gpu.stream();
6555        let mut b = __s_b.launch_builder(&f);
6556        b.arg(src).arg(idx).arg(dst).arg(&nc).arg(&me);
6557        unsafe {
6558            b.launch(cfg)?;
6559        }
6560        Ok(())
6561    }
6562
6563    /// Scatter expert outputs into per-token slots: dst[tok_idx[r], slot_idx[r], :] = src[r, :] * weight[r].
6564    /// dst is [T, n_used, ncols], zero-initialized. Each (expert, token) pair maps to a unique slot.
6565    /// Scatter expert outputs into per-token slots (raw copy, no weight multiply).
6566    /// Weight stored into wbuf[tok*n_used + slot] for FMA in reduce step.
6567    pub fn scatter_slot(
6568        &self,
6569        src: &CudaSlice<f32>,
6570        tok_idx: &CudaSlice<i32>,
6571        slot_idx: &CudaSlice<i32>,
6572        weight: &CudaSlice<f32>,
6573        dst: &mut CudaSlice<f32>,
6574        wbuf: &mut CudaSlice<f32>,
6575        ncols: usize,
6576        n_used: usize,
6577        m_e: usize,
6578    ) -> Result<(), Box<dyn std::error::Error>> {
6579        let f = self.func("scatter_add_slot_f32");
6580        let cfg = LaunchConfig::for_num_elems((m_e * ncols) as u32);
6581        let (nc, nu, me) = (ncols as i32, n_used as i32, m_e as i32);
6582        let __s_b = self.gpu.stream();
6583        let mut b = __s_b.launch_builder(&f);
6584        b.arg(src)
6585            .arg(tok_idx)
6586            .arg(slot_idx)
6587            .arg(weight)
6588            .arg(dst)
6589            .arg(wbuf)
6590            .arg(&nc)
6591            .arg(&nu)
6592            .arg(&me);
6593        unsafe {
6594            b.launch(cfg)?;
6595        }
6596        Ok(())
6597    }
6598
6599    /// Reduce n_used slots per token: dst[t, col] = sum_s slots[t, s, col].
6600    /// Reduce n_used slots per token: dst[t, col] = sum_s FMA(wbuf[t,s], slots[t,s,col], acc).
6601    /// Uses FMA for bit-identity with the sequential axpy path.
6602    pub fn reduce_slots(
6603        &self,
6604        slots: &CudaSlice<f32>,
6605        wbuf: &CudaSlice<f32>,
6606        dst: &mut CudaSlice<f32>,
6607        ncols: usize,
6608        n_used: usize,
6609        t: usize,
6610    ) -> Result<(), Box<dyn std::error::Error>> {
6611        let f = self.func("reduce_slots_f32");
6612        let cfg = LaunchConfig::for_num_elems((t * ncols) as u32);
6613        let (nc, nu, ti) = (ncols as i32, n_used as i32, t as i32);
6614        let __s_b = self.gpu.stream();
6615        let mut b = __s_b.launch_builder(&f);
6616        b.arg(slots).arg(wbuf).arg(dst).arg(&nc).arg(&nu).arg(&ti);
6617        unsafe {
6618            b.launch(cfg)?;
6619        }
6620        Ok(())
6621    }
6622
6623    /// Stage-B: quantize activation [m,in] f32 -> q8_1 (int8 qs + per-block f32 scale).
6624    /// Quantize an activation [m, in_f] to q8_1 (int8 qs + per-32 f32 scale). Public so the
6625    /// forward can quantize a SHARED activation ONCE and feed it to several matmuls (gate+up
6626    /// share `z`; q/k/v and wqkv/gate/beta/alpha share `h`) — quantize_q8_1 was 13.5% of decode
6627    /// GPU time, ~half of it redundant re-quantization of the same row.
6628    /// quantize_q8_1 over a CudaView (a sliced z-row) — same kernel, offset-honoring arg.
6629    pub fn quantize_q8_1_view(
6630        &self,
6631        x: &cudarc::driver::CudaView<f32>,
6632        m: usize,
6633        in_f: usize,
6634    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6635        let f = self.func("quantize_q8_1");
6636        let nblk = in_f / 32;
6637        let mut q = self.alloc_uninit::<i8>(m * in_f)?;
6638        let mut d = self.alloc_uninit::<f32>(m * nblk)?;
6639        let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
6640        let (inf, mi) = (in_f as i32, m as i32);
6641        let __s_b = self.gpu.stream();
6642        let mut b = __s_b.launch_builder(&f);
6643        b.arg(x).arg(&mut q).arg(&mut d).arg(&inf).arg(&mi);
6644        unsafe {
6645            b.launch(cfg)?;
6646        }
6647        Ok((q, d))
6648    }
6649
6650    pub fn quantize_q8_1(
6651        &self,
6652        x: &CudaSlice<f32>,
6653        m: usize,
6654        in_f: usize,
6655    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6656        let nblk = in_f / 32;
6657        let mut q = self.alloc_uninit::<i8>(m * in_f)?; // full-overwrite output: skip memset
6658        let mut d = self.alloc_uninit::<f32>(m * nblk)?; // full-overwrite output: skip memset
6659        // WARP-PER-BLOCK kernel: one warp per 32-block -> m*in_f threads total.
6660        let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
6661        let (inf, mi) = (in_f as i32, m as i32);
6662        if Self::pdl_on() && Self::pdl_wb_on() {
6663            {
6664                use cudarc::driver::{DevicePtr, DevicePtrMut};
6665                let s = &self.gpu.stream();
6666                let (px, _g0) = x.device_ptr(s);
6667                let (pq, _g1) = q.device_ptr_mut(s);
6668                let (pd, _g2) = d.device_ptr_mut(s);
6669                let mut ps = [
6670                    &px as *const _ as *mut std::ffi::c_void,
6671                    &pq as *const _ as *mut _,
6672                    &pd as *const _ as *mut _,
6673                    &inf as *const _ as *mut _,
6674                    &mi as *const _ as *mut _,
6675                ];
6676                unsafe {
6677                    self.launch_pdl("quantize_q8_1", cfg.grid_dim, cfg.block_dim, &mut ps)?;
6678                }
6679            }
6680            return Ok((q, d));
6681        }
6682        let f = self.func("quantize_q8_1");
6683        let __s_b = self.gpu.stream();
6684        let mut b = __s_b.launch_builder(&f);
6685        b.arg(x).arg(&mut q).arg(&mut d).arg(&inf).arg(&mi);
6686        unsafe {
6687            b.launch(cfg)?;
6688        }
6689        Ok((q, d))
6690    }
6691
6692    /// Stage-C FP4: quantize activation [m,in] f32 -> e2m1 nibbles (aq4: u32 [m, in/8]) + per-16
6693    /// UE4M3 scale (ad4: u8 [m, in/16]), the layout the mxf4nvf4 block-scale GEMM B-operand wants.
6694    /// in_f must be a multiple of 64 (one NVFP4 K-block). One thread per (token, 16-block).
6695    pub fn quantize_fp4_act(
6696        &self,
6697        x: &CudaSlice<f32>,
6698        m: usize,
6699        in_f: usize,
6700    ) -> Result<(CudaSlice<u32>, CudaSlice<u8>), Box<dyn std::error::Error>> {
6701        let f = self.func("quantize_fp4_act");
6702        let nb16 = in_f / 16;
6703        let mut aq4 = self.alloc_uninit::<u32>(m * (in_f / 8))?; // full-overwrite output: skip memset
6704        let mut ad4 = self.alloc_uninit::<u8>(m * nb16)?; // full-overwrite output: skip memset
6705        let cfg = LaunchConfig::for_num_elems((m * nb16) as u32);
6706        let (inf, mi) = (in_f as i32, m as i32);
6707        let __s_b = self.gpu.stream();
6708        let mut b = __s_b.launch_builder(&f);
6709        b.arg(x).arg(&mut aq4).arg(&mut ad4).arg(&inf).arg(&mi);
6710        unsafe {
6711            b.launch(cfg)?;
6712        }
6713        Ok((aq4, ad4))
6714    }
6715
6716    /// Stage-C FP4 GEMM (NVFP4 weights): native mxf4nvf4 block-scale tensor-core matmul. Feeds raw
6717    /// e2m1 weight nibbles + raw UE4M3 micro-scales directly to mma.sync.m16n8k64 (762 TFLOP/s peak,
6718    /// 3.5x int8). Activation `x` is quantized to FP4 e2m1 here. NVFP4 per-tensor macro-scale applied
6719    /// post (scale==1.0 -> no-op). `bytes` = raw NVFP4 weight rows. Used by the MEMRA_FP4 prefill path.
6720    pub fn qmatvec_gemm_nvfp4_fp4(
6721        &self,
6722        bytes: &CudaSlice<u8>,
6723        x: &CudaSlice<f32>,
6724        m: usize,
6725        in_f: usize,
6726        out_f: usize,
6727        row_bytes: usize,
6728        scale: f32,
6729    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6730        assert!(
6731            in_f % 64 == 0,
6732            "FP4 GEMM requires in_f % 64 == 0, got {in_f}"
6733        );
6734        let (aq4, ad4) = self.quantize_fp4_act(x, m, in_f)?;
6735        let mut y = self.fp4_gemm_launch(bytes, &aq4, &ad4, m, in_f, out_f, row_bytes)?;
6736        if scale != 1.0 {
6737            self.scale_inplace(&mut y, scale, m * out_f)?;
6738        }
6739        Ok(y)
6740    }
6741
6742    /// Shared mxf4 GEMM launch (pre-quantized FP4 activation aq4/ad4). Same CTA tile as the int8 GEMM
6743    /// (BM=64 rows x BN=128 tokens, 4 warps). No macro-scale applied here.
6744    fn fp4_gemm_launch(
6745        &self,
6746        bytes: &CudaSlice<u8>,
6747        aq4: &CudaSlice<u32>,
6748        ad4: &CudaSlice<u8>,
6749        m: usize,
6750        in_f: usize,
6751        out_f: usize,
6752        row_bytes: usize,
6753    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6754        let f = self.func("qmatvec_gemm_nvfp4_fp4");
6755        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
6756        const BM: u32 = 64;
6757        const BN: u32 = 256;
6758        let cfg = LaunchConfig {
6759            grid_dim: ((out_f as u32 + BM - 1) / BM, (m as u32 + BN - 1) / BN, 1),
6760            block_dim: (32, 4, 1),
6761            shared_mem_bytes: 0,
6762        };
6763        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
6764        let __s_b = self.gpu.stream();
6765        let mut b = __s_b.launch_builder(&f);
6766        b.arg(bytes)
6767            .arg(aq4)
6768            .arg(ad4)
6769            .arg(&mut y)
6770            .arg(&inf)
6771            .arg(&outf)
6772            .arg(&mi)
6773            .arg(&rb);
6774        unsafe {
6775            b.launch(cfg)?;
6776        }
6777        Ok(y)
6778    }
6779
6780    /// Test entry (kernel_check): run the FP4 GEMM from raw bytes; NO macro-scale (caller compares bare).
6781    pub fn qmatvec_gemm_nvfp4_fp4_raw(
6782        &self,
6783        bytes: &CudaSlice<u8>,
6784        x: &CudaSlice<f32>,
6785        m: usize,
6786        in_f: usize,
6787        out_f: usize,
6788        row_bytes: usize,
6789    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6790        assert!(
6791            in_f % 64 == 0,
6792            "FP4 GEMM requires in_f % 64 == 0, got {in_f}"
6793        );
6794        let (aq4, ad4) = self.quantize_fp4_act(x, m, in_f)?;
6795        self.fp4_gemm_launch(bytes, &aq4, &ad4, m, in_f, out_f, row_bytes)
6796    }
6797
6798    /// Stage-B: Q8_0 weight x q8_1 activation int8 dp4a matmul. y[m,out]=x@W^T.
6799    pub fn qmatvec_q8_0_fast(
6800        &self,
6801        w: &CudaSlice<u8>,
6802        x: &CudaSlice<f32>,
6803        m: usize,
6804        in_f: usize,
6805        out_f: usize,
6806        row_bytes: usize,
6807    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6808        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
6809        let f = self.func("qmatvec_q8_0_dp4a");
6810        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
6811        let cfg = LaunchConfig {
6812            grid_dim: (out_f as u32, m as u32, 1),
6813            block_dim: (128, 1, 1),
6814            shared_mem_bytes: 0,
6815        };
6816        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
6817        let __s_b = self.gpu.stream();
6818        let mut b = __s_b.launch_builder(&f);
6819        b.arg(w)
6820            .arg(&aq)
6821            .arg(&ad)
6822            .arg(&mut y)
6823            .arg(&inf)
6824            .arg(&outf)
6825            .arg(&mi)
6826            .arg(&rb);
6827        unsafe {
6828            b.launch(cfg)?;
6829        }
6830        Ok(y)
6831    }
6832
6833    /// Stage-B: Q4_K weight x q8_1 activation int8 dp4a (decode). Min-offset via q8_1 sum term.
6834    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
6835    pub fn qmatvec_q4_K_fast(
6836        &self,
6837        w: &CudaSlice<u8>,
6838        x: &CudaSlice<f32>,
6839        m: usize,
6840        in_f: usize,
6841        out_f: usize,
6842        row_bytes: usize,
6843    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6844        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
6845        let f = self.func("qmatvec_q4_K_dp4a");
6846        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
6847        let cfg = LaunchConfig {
6848            grid_dim: (out_f as u32, m as u32, 1),
6849            block_dim: (128, 1, 1),
6850            shared_mem_bytes: 0,
6851        };
6852        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
6853        let __s_b = self.gpu.stream();
6854        let mut b = __s_b.launch_builder(&f);
6855        b.arg(w)
6856            .arg(&aq)
6857            .arg(&ad)
6858            .arg(&mut y)
6859            .arg(&inf)
6860            .arg(&outf)
6861            .arg(&mi)
6862            .arg(&rb);
6863        unsafe {
6864            b.launch(cfg)?;
6865        }
6866        Ok(y)
6867    }
6868
6869    /// Stage-B: Q6_K weight x q8_1 activation int8 dp4a (decode, symmetric).
6870    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
6871    pub fn qmatvec_q6_K_fast(
6872        &self,
6873        w: &CudaSlice<u8>,
6874        x: &CudaSlice<f32>,
6875        m: usize,
6876        in_f: usize,
6877        out_f: usize,
6878        row_bytes: usize,
6879    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6880        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
6881        let f = self.func("qmatvec_q6_K_dp4a");
6882        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
6883        let cfg = LaunchConfig {
6884            grid_dim: (out_f as u32, m as u32, 1),
6885            block_dim: (128, 1, 1),
6886            shared_mem_bytes: 0,
6887        };
6888        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
6889        let __s_b = self.gpu.stream();
6890        let mut b = __s_b.launch_builder(&f);
6891        b.arg(w)
6892            .arg(&aq)
6893            .arg(&ad)
6894            .arg(&mut y)
6895            .arg(&inf)
6896            .arg(&outf)
6897            .arg(&mi)
6898            .arg(&rb);
6899        unsafe {
6900            b.launch(cfg)?;
6901        }
6902        Ok(y)
6903    }
6904
6905    /// Stage-B: Q5_K weight x q8_1 activation int8 dp4a (decode). Min-offset via q8_1 sum term.
6906    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
6907    pub fn qmatvec_q5_K_fast(
6908        &self,
6909        w: &CudaSlice<u8>,
6910        x: &CudaSlice<f32>,
6911        m: usize,
6912        in_f: usize,
6913        out_f: usize,
6914        row_bytes: usize,
6915    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6916        self.qmatvec_dp4a_named("qmatvec_q5_K_dp4a", w, x, m, in_f, out_f, row_bytes)
6917    }
6918    /// Stage-B: Q3_K weight x q8_1 activation int8 dp4a (decode, symmetric).
6919    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
6920    pub fn qmatvec_q3_K_fast(
6921        &self,
6922        w: &CudaSlice<u8>,
6923        x: &CudaSlice<f32>,
6924        m: usize,
6925        in_f: usize,
6926        out_f: usize,
6927        row_bytes: usize,
6928    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6929        self.qmatvec_dp4a_named("qmatvec_q3_K_dp4a", w, x, m, in_f, out_f, row_bytes)
6930    }
6931    /// A6 split-plane twin of `qmatvec_nvfp4_fast` (weights repacked; used by the rp gates).
6932    pub fn qmatvec_nvfp4_fast_rp(
6933        &self,
6934        w: &CudaSlice<u8>,
6935        x: &CudaSlice<f32>,
6936        m: usize,
6937        in_f: usize,
6938        out_f: usize,
6939        row_bytes: usize,
6940    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6941        assert!(
6942            in_f % 64 == 0,
6943            "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
6944        );
6945        self.qmatvec_dp4a_named("qmatvec_nvfp4_dp4a_rp", w, x, m, in_f, out_f, row_bytes)
6946    }
6947    /// Stage-B: NVFP4 weight x q8_1 activation int8 dp4a (decode, symmetric, codebook lookup).
6948    pub fn qmatvec_nvfp4_fast(
6949        &self,
6950        w: &CudaSlice<u8>,
6951        x: &CudaSlice<f32>,
6952        m: usize,
6953        in_f: usize,
6954        out_f: usize,
6955        row_bytes: usize,
6956    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6957        // B1: the NVFP4 dp4a kernel maps two 32-elem q8_1 blocks onto one 64-elem block_nvfp4
6958        // (sblk = g >> 1). in_f must be a multiple of 64 or the last block reads a partial superblock.
6959        assert!(
6960            in_f % 64 == 0,
6961            "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
6962        );
6963        self.qmatvec_dp4a_named("qmatvec_nvfp4_dp4a", w, x, m, in_f, out_f, row_bytes)
6964    }
6965    /// Stage-B (optional perf): IQ4_XS codebook int8 dp4a.
6966    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
6967    pub fn qmatvec_iq4_XS_fast(
6968        &self,
6969        w: &CudaSlice<u8>,
6970        x: &CudaSlice<f32>,
6971        m: usize,
6972        in_f: usize,
6973        out_f: usize,
6974        row_bytes: usize,
6975    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6976        self.qmatvec_dp4a_named("qmatvec_iq4_XS_dp4a", w, x, m, in_f, out_f, row_bytes)
6977    }
6978
6979    /// Shared dp4a launcher: quantize_q8_1 then call the named kernel (grid (out,m), block 64).
6980    fn qmatvec_dp4a_named(
6981        &self,
6982        name: &str,
6983        w: &CudaSlice<u8>,
6984        x: &CudaSlice<f32>,
6985        m: usize,
6986        in_f: usize,
6987        out_f: usize,
6988        row_bytes: usize,
6989    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6990        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
6991        let f = self.func(name);
6992        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
6993        let cfg = LaunchConfig {
6994            grid_dim: (out_f as u32, m as u32, 1),
6995            block_dim: (128, 1, 1),
6996            shared_mem_bytes: 0,
6997        };
6998        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
6999        let __s_b = self.gpu.stream();
7000        let mut b = __s_b.launch_builder(&f);
7001        b.arg(w)
7002            .arg(&aq)
7003            .arg(&ad)
7004            .arg(&mut y)
7005            .arg(&inf)
7006            .arg(&outf)
7007            .arg(&mi)
7008            .arg(&rb);
7009        unsafe {
7010            b.launch(cfg)?;
7011        }
7012        Ok(y)
7013    }
7014
7015    pub fn htod(&self, v: &[f32]) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7016        Ok(self.gpu.stream().clone_htod(v)?)
7017    }
7018    pub fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
7019        Ok(self.gpu.stream().clone_htod(v)?)
7020    }
7021    /// i8 upload (moe-devq8-check: synthetic q8_1 activation bytes).
7022    pub fn htod_i8(&self, v: &[i8]) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
7023        Ok(self.gpu.stream().clone_htod(v)?)
7024    }
7025    pub fn htod_u64(&self, v: &[u64]) -> Result<CudaSlice<u64>, Box<dyn std::error::Error>> {
7026        Ok(self.gpu.stream().clone_htod(v)?)
7027    }
7028    /// View twin of `dtoh` (lean-logits component 3: D2H one row of a [B, n_vocab] stack).
7029    pub fn dtoh_view(
7030        &self,
7031        d: &cudarc::driver::CudaView<f32>,
7032    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
7033        let v = self.gpu.stream().clone_dtoh(d)?;
7034        self.gpu.stream().synchronize()?;
7035        Ok(v)
7036    }
7037    pub fn dtoh(&self, d: &CudaSlice<f32>) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
7038        let v = self.gpu.stream().clone_dtoh(d)?;
7039        self.gpu.stream().synchronize()?;
7040        Ok(v)
7041    }
7042    /// Queue two f32 device-to-host copies on the compute stream, then establish one host
7043    /// boundary for both. Hy3's CPU/GPU expert split needs the router logits and the MoE input;
7044    /// issuing them together avoids a second stream synchronization in every trunk layer.
7045    pub fn dtoh_pair(
7046        &self,
7047        a: &CudaSlice<f32>,
7048        b: &CudaSlice<f32>,
7049    ) -> Result<(Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
7050        let av = self.gpu.stream().clone_dtoh(a)?;
7051        let bv = self.gpu.stream().clone_dtoh(b)?;
7052        self.gpu.stream().synchronize()?;
7053        Ok((av, bv))
7054    }
7055    /// Device-to-host copy of an i32 buffer (fused-router sel_idx readback).
7056    pub fn dtoh_i32(&self, d: &CudaSlice<i32>) -> Result<Vec<i32>, Box<dyn std::error::Error>> {
7057        let v = self.gpu.stream().clone_dtoh(d)?;
7058        self.gpu.stream().synchronize()?;
7059        Ok(v)
7060    }
7061    /// Device-to-host copy of a u8 buffer (used to read back the quantized KV cache for validation).
7062    pub fn dtoh_u8(&self, d: &CudaSlice<u8>) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
7063        let v = self.gpu.stream().clone_dtoh(d)?;
7064        self.gpu.stream().synchronize()?;
7065        Ok(v)
7066    }
7067    pub fn dtoh_u8_view(
7068        &self,
7069        d: &cudarc::driver::CudaView<u8>,
7070    ) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
7071        let v = self.gpu.stream().clone_dtoh(d)?;
7072        self.gpu.stream().synchronize()?;
7073        Ok(v)
7074    }
7075    pub fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7076        let s = self.gpu.stream().alloc_zeros::<f32>(n)?;
7077        self.keep_if_capturing(&s);
7078        Ok(s)
7079    }
7080
7081    /// GPU-resident greedy argmax (CUDA-GRAPH-PLAN Phase 1): logits[n_vocab] -> token id in a
7082    /// resident device u32 [1]. PARALLEL 2-pass (RANK1 LEVER): the old single-CTA scan (one 256-thread
7083    /// block on one SM over 248K logits) was memory-starved at ~426us/token. Now pass 1 fans NB=256
7084    /// blocks across the SMs to saturate HBM, pass 2 reduces the NB partials. Bit-identical to host
7085    /// `argmax` (smallest index on tie). The whole point is NOT to dtoh logits — only a [1] u32 is read
7086    /// back (or kept resident for graph replay). Returns the device token buffer.
7087    /// Softmax probability of the (already-argmaxed) token `tok` under `logits` — the spec-decode
7088    /// p-min confidence signal. 2-pass like the parallel argmax; returns a device [1] f32.
7089    pub fn prob_of_token_device(
7090        &self,
7091        logits: &CudaSlice<f32>,
7092        tok: &CudaSlice<u32>,
7093        n_vocab: usize,
7094    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7095        let nb = ARGMAX_NB;
7096        let mut part = self.alloc_uninit::<f32>(nb)?;
7097        let mut p = self.alloc_uninit::<f32>(1)?;
7098        let f1 = self.func("prob_of_token_partial_f32");
7099        let cfg1 = LaunchConfig {
7100            grid_dim: (nb as u32, 1, 1),
7101            block_dim: (256, 1, 1),
7102            shared_mem_bytes: 0,
7103        };
7104        let nv = n_vocab as i32;
7105        let __s_b1 = self.gpu.stream();
7106        let mut b1 = __s_b1.launch_builder(&f1);
7107        b1.arg(logits).arg(tok).arg(&mut part).arg(&nv);
7108        unsafe {
7109            b1.launch(cfg1)?;
7110        }
7111        let f2 = self.func("prob_of_token_final_f32");
7112        let cfg2 = LaunchConfig {
7113            grid_dim: (1, 1, 1),
7114            block_dim: (256, 1, 1),
7115            shared_mem_bytes: 0,
7116        };
7117        let nbi = nb as i32;
7118        let __s_b2 = self.gpu.stream();
7119        let mut b2 = __s_b2.launch_builder(&f2);
7120        b2.arg(&part).arg(&mut p).arg(&nbi);
7121        unsafe {
7122            b2.launch(cfg2)?;
7123        }
7124        Ok(p)
7125    }
7126
7127    /// Like `prob_of_token_device` but writes into a PERSISTENT `p_out` buffer (stable pointer).
7128    /// Required for CUDA-graph capture of the draft chain: the captured prob kernels must write
7129    /// where the host reads the p-min confidence between replays. Same kernels, same math.
7130    /// Slot-addressed twin of `prob_of_token_device_into`: token read from `tok_all[tok_idx]`
7131    /// (a view at the slot), probability written to `p_out[p_idx]` — same two kernels, the
7132    /// pointers just land mid-buffer. Zero-sync (gemma confidence-adaptive draft depth).
7133    pub fn prob_of_token_device_col(
7134        &self,
7135        logits: &CudaSlice<f32>,
7136        tok_all: &CudaSlice<u32>,
7137        tok_idx: usize,
7138        p_out: &mut CudaSlice<f32>,
7139        p_idx: usize,
7140        n_vocab: usize,
7141    ) -> Result<(), Box<dyn std::error::Error>> {
7142        let tok_v = tok_all.slice(tok_idx..tok_idx + 1);
7143        let mut p_v = p_out.slice_mut(p_idx..p_idx + 1);
7144        let nb = ARGMAX_NB;
7145        let mut part = self.alloc_uninit::<f32>(nb)?;
7146        let f1 = self.func("prob_of_token_partial_f32");
7147        let cfg1 = LaunchConfig {
7148            grid_dim: (nb as u32, 1, 1),
7149            block_dim: (256, 1, 1),
7150            shared_mem_bytes: 0,
7151        };
7152        let nv = n_vocab as i32;
7153        let __s_b1 = self.gpu.stream();
7154        let mut b1 = __s_b1.launch_builder(&f1);
7155        b1.arg(logits).arg(&tok_v).arg(&mut part).arg(&nv);
7156        unsafe {
7157            b1.launch(cfg1)?;
7158        }
7159        let f2 = self.func("prob_of_token_final_f32");
7160        let cfg2 = LaunchConfig {
7161            grid_dim: (1, 1, 1),
7162            block_dim: (256, 1, 1),
7163            shared_mem_bytes: 0,
7164        };
7165        let nbi = nb as i32;
7166        let __s_b2 = self.gpu.stream();
7167        let mut b2 = __s_b2.launch_builder(&f2);
7168        b2.arg(&part).arg(&mut p_v).arg(&nbi);
7169        unsafe {
7170            b2.launch(cfg2)?;
7171        }
7172        Ok(())
7173    }
7174
7175    pub fn prob_of_token_device_into(
7176        &self,
7177        logits: &CudaSlice<f32>,
7178        tok: &CudaSlice<u32>,
7179        p_out: &mut CudaSlice<f32>,
7180        n_vocab: usize,
7181    ) -> Result<(), Box<dyn std::error::Error>> {
7182        let nb = ARGMAX_NB;
7183        let mut part = self.alloc_uninit::<f32>(nb)?;
7184        let f1 = self.func("prob_of_token_partial_f32");
7185        let cfg1 = LaunchConfig {
7186            grid_dim: (nb as u32, 1, 1),
7187            block_dim: (256, 1, 1),
7188            shared_mem_bytes: 0,
7189        };
7190        let nv = n_vocab as i32;
7191        let __s_b1 = self.gpu.stream();
7192        let mut b1 = __s_b1.launch_builder(&f1);
7193        b1.arg(logits).arg(tok).arg(&mut part).arg(&nv);
7194        unsafe {
7195            b1.launch(cfg1)?;
7196        }
7197        let f2 = self.func("prob_of_token_final_f32");
7198        let cfg2 = LaunchConfig {
7199            grid_dim: (1, 1, 1),
7200            block_dim: (256, 1, 1),
7201            shared_mem_bytes: 0,
7202        };
7203        let nbi = nb as i32;
7204        let __s_b2 = self.gpu.stream();
7205        let mut b2 = __s_b2.launch_builder(&f2);
7206        b2.arg(&part).arg(p_out).arg(&nbi);
7207        unsafe {
7208            b2.launch(cfg2)?;
7209        }
7210        Ok(())
7211    }
7212
7213    pub fn argmax_token_device(
7214        &self,
7215        logits: &CudaSlice<f32>,
7216        n_vocab: usize,
7217    ) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
7218        let mut tok = unsafe { self.gpu.stream().alloc::<u32>(1)? };
7219        self.argmax_token_device_into(logits, &mut tok, n_vocab)?;
7220        Ok(tok)
7221    }
7222    /// Like `argmax_token_device` but writes into a PERSISTENT `tok` buffer (stable pointer) instead
7223    /// of allocating a fresh one. Required for CUDA-graph capture: the captured argmax must write the
7224    /// next token into the SAME device buffer the next replay's embed_gather reads, so the buffer
7225    /// pointer is baked once and the token id never round-trips to host inside steady state. The
7226    /// pass-1 partials scratch (`argmax_partials`) is also a resident stable-pointer buffer so both
7227    /// captured passes bake fixed addresses.
7228    pub fn argmax_token_device_into(
7229        &self,
7230        logits: &CudaSlice<f32>,
7231        tok: &mut CudaSlice<u32>,
7232        n_vocab: usize,
7233    ) -> Result<(), Box<dyn std::error::Error>> {
7234        let nb = ARGMAX_NB;
7235        let f1 = self.func("argmax_partial_f32");
7236        let f2 = self.func("argmax_final_f32");
7237        let mut guard = self.argmax_partials.lock().unwrap();
7238        if guard.is_none() {
7239            // allocate ONCE; under generate_graph this runs in the tracking-off prime window so the
7240            // buffers carry no cudarc events (illegal inside capture).
7241            let pv = self.gpu.stream().alloc_zeros::<f32>(nb)?;
7242            let pi = self.gpu.stream().alloc_zeros::<i32>(nb)?;
7243            *guard = Some((pv, pi));
7244        }
7245        let (part_v, part_i) = guard.as_mut().unwrap();
7246        let nv = n_vocab as i32;
7247        let nbi = nb as i32;
7248        // pass 1: NB blocks x 256 threads grid-stride scan -> per-block (val, idx) partials.
7249        let cfg1 = LaunchConfig {
7250            grid_dim: (nb as u32, 1, 1),
7251            block_dim: (256, 1, 1),
7252            shared_mem_bytes: 0,
7253        };
7254        let __s_b1 = self.gpu.stream();
7255        let mut b1 = __s_b1.launch_builder(&f1);
7256        b1.arg(logits).arg(&mut *part_v).arg(&mut *part_i).arg(&nv);
7257        unsafe {
7258            b1.launch(cfg1)?;
7259        }
7260        // pass 2: one block reduces NB partials -> token_out[0].
7261        let cfg2 = LaunchConfig {
7262            grid_dim: (1, 1, 1),
7263            block_dim: (256, 1, 1),
7264            shared_mem_bytes: 0,
7265        };
7266        let __s_b2 = self.gpu.stream();
7267        let mut b2 = __s_b2.launch_builder(&f2);
7268        b2.arg(&*part_v).arg(&*part_i).arg(tok).arg(&nbi);
7269        unsafe {
7270            b2.launch(cfg2)?;
7271        }
7272        Ok(())
7273    }
7274    /// Column-`col` device argmax over a stacked verify-logits buffer [t, n_vocab] (spec accept
7275    /// walk): toks[out_idx] = argmax(logits[col*n_vocab .. (col+1)*n_vocab]). SAME 2-pass kernels
7276    /// and tie-break contract as `argmax_token_device_into` (bit-identical to host argmax,
7277    /// argmax_gate-validated) — only the input pointer (a column view) and the output slot differ.
7278    /// Lets the accept walk read ONE [t] u32 instead of dtoh'ing the full [t, n_vocab] logits.
7279    pub fn argmax_token_device_col(
7280        &self,
7281        logits: &CudaSlice<f32>,
7282        col: usize,
7283        n_vocab: usize,
7284        toks: &mut CudaSlice<u32>,
7285        out_idx: usize,
7286    ) -> Result<(), Box<dyn std::error::Error>> {
7287        let nb = ARGMAX_NB;
7288        let f1 = self.func("argmax_partial_f32");
7289        let f2 = self.func("argmax_final_f32");
7290        let mut guard = self.argmax_partials.lock().unwrap();
7291        if guard.is_none() {
7292            let pv = self.gpu.stream().alloc_zeros::<f32>(nb)?;
7293            let pi = self.gpu.stream().alloc_zeros::<i32>(nb)?;
7294            *guard = Some((pv, pi));
7295        }
7296        let (part_v, part_i) = guard.as_mut().unwrap();
7297        let col_view = logits.slice(col * n_vocab..(col + 1) * n_vocab);
7298        let nv = n_vocab as i32;
7299        let nbi = nb as i32;
7300        let cfg1 = LaunchConfig {
7301            grid_dim: (nb as u32, 1, 1),
7302            block_dim: (256, 1, 1),
7303            shared_mem_bytes: 0,
7304        };
7305        let __s_b1 = self.gpu.stream();
7306        let mut b1 = __s_b1.launch_builder(&f1);
7307        b1.arg(&col_view)
7308            .arg(&mut *part_v)
7309            .arg(&mut *part_i)
7310            .arg(&nv);
7311        unsafe {
7312            b1.launch(cfg1)?;
7313        }
7314        let mut tok_view = toks.slice_mut(out_idx..out_idx + 1);
7315        let cfg2 = LaunchConfig {
7316            grid_dim: (1, 1, 1),
7317            block_dim: (256, 1, 1),
7318            shared_mem_bytes: 0,
7319        };
7320        let __s_b2 = self.gpu.stream();
7321        let mut b2 = __s_b2.launch_builder(&f2);
7322        b2.arg(&*part_v).arg(&*part_i).arg(&mut tok_view).arg(&nbi);
7323        unsafe {
7324            b2.launch(cfg2)?;
7325        }
7326        Ok(())
7327    }
7328    /// Read back a device u32 buffer (the spec accept walk's [t] per-column argmax tokens).
7329    pub fn htod_u32_v(&self, v: &[u32]) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
7330        Ok(self.gpu.stream().clone_htod(v)?)
7331    }
7332    pub fn dtoh_u32(&self, d: &CudaSlice<u32>) -> Result<Vec<u32>, Box<dyn std::error::Error>> {
7333        let v = self.gpu.stream().clone_dtoh(d)?;
7334        self.gpu.stream().synchronize()?;
7335        Ok(v)
7336    }
7337    /// Allocate a zeroed device u32 buffer (persistent spec-loop prediction slots).
7338    /// H2D into an EXISTING u32 buffer (stable pointer — the per-step grammar-mask upload:
7339    /// contents change every step, the address must not, so a captured graph can read it).
7340    pub fn htod_u32_into(
7341        &self,
7342        dst: &mut CudaSlice<u32>,
7343        src: &[u32],
7344    ) -> Result<(), Box<dyn std::error::Error>> {
7345        let mut view = dst.slice_mut(0..src.len());
7346        self.gpu.stream().memcpy_htod(src, &mut view)?;
7347        Ok(())
7348    }
7349
7350    /// H2D into an existing i32 buffer. OPTIPIPE uses this to refresh a stage-local saved-len
7351    /// table without changing the device address its reconcile kernel consumes.
7352    pub fn htod_i32_into(
7353        &self,
7354        dst: &mut CudaSlice<i32>,
7355        src: &[i32],
7356    ) -> Result<(), Box<dyn std::error::Error>> {
7357        let mut view = dst.slice_mut(0..src.len());
7358        self.gpu.stream().memcpy_htod(src, &mut view)?;
7359        Ok(())
7360    }
7361
7362    pub fn alloc_u32_zeroed(&self, n: usize) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
7363        let s = self.gpu.stream().alloc_zeros::<u32>(n)?;
7364        self.keep_if_capturing(&s);
7365        Ok(s)
7366    }
7367    /// embed_gather into a PERSISTENT `x_out` buffer (stable pointer) for CUDA-graph capture (the
7368    /// embed output starts the per-step kernel chain and must be at a fixed address across replays).
7369    pub fn embed_gather_device_into(
7370        &self,
7371        embd: &CudaSlice<u8>,
7372        token_d: &CudaSlice<u32>,
7373        x_out: &mut CudaSlice<f32>,
7374        n_embd: usize,
7375        qtype: i32,
7376        row_bytes: usize,
7377    ) -> Result<(), Box<dyn std::error::Error>> {
7378        let f = self.func("embed_gather_u32");
7379        let cfg = LaunchConfig {
7380            grid_dim: (((n_embd as u32 + 255) / 256).max(1), 1, 1),
7381            block_dim: (256, 1, 1),
7382            shared_mem_bytes: 0,
7383        };
7384        let (ne, qt, rb) = (n_embd as i32, qtype, row_bytes as i64);
7385        let __s_b = self.gpu.stream();
7386        let mut b = __s_b.launch_builder(&f);
7387        b.arg(embd)
7388            .arg(token_d)
7389            .arg(x_out)
7390            .arg(&ne)
7391            .arg(&qt)
7392            .arg(&rb);
7393        unsafe {
7394            b.launch(cfg)?;
7395        }
7396        Ok(())
7397    }
7398    /// Read a [1] i32 device counter (pos / seqlen) back to host. Tiny D2H + sync.
7399    pub fn dtoh_i32_one(&self, d: &CudaSlice<i32>) -> Result<i32, Box<dyn std::error::Error>> {
7400        let v = self.gpu.stream().clone_dtoh(d)?;
7401        self.gpu.stream().synchronize()?;
7402        Ok(v[0])
7403    }
7404    /// Set a [1] i32 device counter IN PLACE (keeps the buffer pointer stable — required for the
7405    /// graph-resident pos/seqlen counters whose addresses are baked into captured graphs). Restores
7406    /// the counter value after the throwaway capture warmups corrupt it.
7407    /// ASYNC i32 single-slot store (value rides the kernel arg — no host-memory transfer/sync).
7408    /// The graph-arc device-len counters use this; set_i32_one below is the SYNCING pageable
7409    /// copy (fine at stream-idle boundaries, poison mid-round).
7410    pub fn i32_set_k(
7411        &self,
7412        dst: &mut CudaSlice<i32>,
7413        v: i32,
7414    ) -> Result<(), Box<dyn std::error::Error>> {
7415        let f = self.func("i32_set_k");
7416        let cfg = LaunchConfig {
7417            grid_dim: (1, 1, 1),
7418            block_dim: (1, 1, 1),
7419            shared_mem_bytes: 0,
7420        };
7421        let idx = 0i32;
7422        let __s_b = self.gpu.stream();
7423        let mut b = __s_b.launch_builder(&f);
7424        b.arg(dst).arg(&v).arg(&idx);
7425        unsafe {
7426            b.launch(cfg)?;
7427        }
7428        Ok(())
7429    }
7430
7431    pub fn set_i32_one(
7432        &self,
7433        d: &mut CudaSlice<i32>,
7434        v: i32,
7435    ) -> Result<(), Box<dyn std::error::Error>> {
7436        self.gpu.stream().memcpy_htod(&[v], d)?;
7437        Ok(())
7438    }
7439    /// Set a [1] u32 device buffer IN PLACE (stable pointer) — for the resident `token_d` counter
7440    /// during priming / capture-state restore.
7441    pub fn set_u32_one(
7442        &self,
7443        d: &mut CudaSlice<u32>,
7444        v: u32,
7445    ) -> Result<(), Box<dyn std::error::Error>> {
7446        self.gpu.stream().memcpy_htod(&[v], d)?;
7447        Ok(())
7448    }
7449    /// Read back a [1] u32 device buffer (the argmax token). One tiny D2H + sync.
7450    pub fn dtoh_u32_one(&self, d: &CudaSlice<u32>) -> Result<u32, Box<dyn std::error::Error>> {
7451        let v = self.gpu.stream().clone_dtoh(d)?;
7452        self.gpu.stream().synchronize()?;
7453        Ok(v[0])
7454    }
7455    /// Upload raw bytes to a resident device u8 buffer (e.g. the embed table for device gather).
7456    pub fn upload_u8(&self, bytes: &[u8]) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
7457        Ok(self.gpu.stream().clone_htod(bytes)?)
7458    }
7459    /// Embed-from-device (CUDA-GRAPH-PLAN Phase 1): gather+dequant the row for the token id in
7460    /// `token_d[0]` from the resident embed table `embd` -> x_out[n_embd]. Bit-identical to host
7461    /// EmbedHost::gather (same per-dtype `deq`). No host round-trip of the token id.
7462    pub fn embed_gather_device(
7463        &self,
7464        embd: &CudaSlice<u8>,
7465        token_d: &CudaSlice<u32>,
7466        n_embd: usize,
7467        qtype: i32,
7468        row_bytes: usize,
7469    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7470        let f = self.func("embed_gather_u32");
7471        let mut x = self.alloc_uninit::<f32>(n_embd)?;
7472        let cfg = LaunchConfig {
7473            grid_dim: (((n_embd as u32 + 255) / 256).max(1), 1, 1),
7474            block_dim: (256, 1, 1),
7475            shared_mem_bytes: 0,
7476        };
7477        let (ne, qt, rb) = (n_embd as i32, qtype, row_bytes as i64);
7478        let __s_b = self.gpu.stream();
7479        let mut b = __s_b.launch_builder(&f);
7480        b.arg(embd)
7481            .arg(token_d)
7482            .arg(&mut x)
7483            .arg(&ne)
7484            .arg(&qt)
7485            .arg(&rb);
7486        unsafe {
7487            b.launch(cfg)?;
7488        }
7489        Ok(x)
7490    }
7491
7492    /// T-token device embed gather (spec verify/replay): tokens uploaded as a tiny [T] u32 htod,
7493    /// rows dequanted on-device -> x[T, n_embd]. Replaces host per-row dequant + T*n_embd*4B htod
7494    /// (nsys: 84% of spec API time was HtoD). Bit-identical rows (same per-dtype deq).
7495    pub fn embed_gather_device_t(
7496        &self,
7497        embd: &CudaSlice<u8>,
7498        tokens: &[u32],
7499        n_embd: usize,
7500        qtype: i32,
7501        row_bytes: usize,
7502    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7503        let t = tokens.len();
7504        let tok_d = self.gpu.stream().clone_htod(tokens)?;
7505        let f = self.func("embed_gather_u32_t");
7506        let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
7507        let cfg = LaunchConfig {
7508            grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
7509            block_dim: (256, 1, 1),
7510            shared_mem_bytes: 0,
7511        };
7512        let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
7513        let __s_b = self.gpu.stream();
7514        let mut b = __s_b.launch_builder(&f);
7515        b.arg(embd)
7516            .arg(&tok_d)
7517            .arg(&mut x)
7518            .arg(&ne)
7519            .arg(&qt)
7520            .arg(&rb)
7521            .arg(&ti);
7522        unsafe {
7523            b.launch(cfg)?;
7524        }
7525        Ok(x)
7526    }
7527
7528    /// T-token embed gather from a DEVICE token buffer (round-stream stage c: the verify tokens
7529    /// are assembled on-device from the draft-chain pack slots; no host round trip). Same kernel
7530    /// as embed_gather_device_t — bit-identical rows.
7531    /// embed_gather over a token VIEW (spec round: tokens live in the round's batch buffer).
7532    pub fn embed_gather_device_tv(
7533        &self,
7534        embd: &CudaSlice<u8>,
7535        tok_v: &cudarc::driver::CudaView<u32>,
7536        t: usize,
7537        n_embd: usize,
7538        qtype: i32,
7539        row_bytes: usize,
7540    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7541        let f = self.func("embed_gather_u32_t");
7542        let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
7543        let cfg = LaunchConfig {
7544            grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
7545            block_dim: (256, 1, 1),
7546            shared_mem_bytes: 0,
7547        };
7548        let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
7549        let __s_b = self.gpu.stream();
7550        let mut b = __s_b.launch_builder(&f);
7551        b.arg(embd)
7552            .arg(tok_v)
7553            .arg(&mut x)
7554            .arg(&ne)
7555            .arg(&qt)
7556            .arg(&rb)
7557            .arg(&ti);
7558        unsafe {
7559            b.launch(cfg)?;
7560        }
7561        Ok(x)
7562    }
7563
7564    pub fn embed_gather_device_td(
7565        &self,
7566        embd: &CudaSlice<u8>,
7567        tok_d: &CudaSlice<u32>,
7568        t: usize,
7569        n_embd: usize,
7570        qtype: i32,
7571        row_bytes: usize,
7572    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7573        let f = self.func("embed_gather_u32_t");
7574        let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
7575        let cfg = LaunchConfig {
7576            grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
7577            block_dim: (256, 1, 1),
7578            shared_mem_bytes: 0,
7579        };
7580        let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
7581        let __s_b = self.gpu.stream();
7582        let mut b = __s_b.launch_builder(&f);
7583        b.arg(embd)
7584            .arg(tok_d)
7585            .arg(&mut x)
7586            .arg(&ne)
7587            .arg(&qt)
7588            .arg(&rb)
7589            .arg(&ti);
7590        unsafe {
7591            b.launch(cfg)?;
7592        }
7593        Ok(x)
7594    }
7595
7596    /// Uninitialized device buffer — SKIPS the memset that `alloc_zeros` always issues. Decode
7597    /// profile (nsys): ~1050 memsets/token = 6.5% of decode GPU time + ~half the launch count, the
7598    /// dominant contributor to the 19% inter-kernel idle gap and a blocker for clean CUDA-graph
7599    /// capture. Use ONLY for buffers a kernel FULLY overwrites (every element written, no `+=`).
7600    /// SAFETY: caller guarantees the producing kernel writes every element before any read.
7601    #[inline]
7602    /// Keep an allocation alive for the current capture (no-op when retain mode is off).
7603    fn keep_if_capturing<T: cudarc::driver::DeviceRepr + Send + 'static>(&self, s: &CudaSlice<T>) {
7604        if self
7605            .capture_keep_on
7606            .load(std::sync::atomic::Ordering::Relaxed)
7607        {
7608            self.capture_keep.lock().unwrap().push(Box::new(s.clone()));
7609        }
7610    }
7611
7612    fn alloc_uninit<T: cudarc::driver::DeviceRepr + Send + 'static>(
7613        &self,
7614        n: usize,
7615    ) -> Result<CudaSlice<T>, Box<dyn std::error::Error>> {
7616        let mut s = unsafe { self.gpu.stream().alloc::<T>(n)? };
7617        // MEMRA_DEBUG_ZERO_ALLOCS=1 (task #14 defect hunt): memset EVERY engine allocation —
7618        // the global uninit-read discriminator (the prime-fn-scoped zeroing experiment could
7619        // not cover engine-internal buffers). Debug-only: massive launch overhead.
7620        {
7621            static Z: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7622            if *Z.get_or_init(|| std::env::var("MEMRA_DEBUG_ZERO_ALLOCS").as_deref() == Ok("1")) {
7623                // raw D8 memset (T lacks ValidAsZeroBits in the generic bound)
7624                use cudarc::driver::DevicePtrMut;
7625                let n_bytes = s.len() * std::mem::size_of::<T>();
7626                let stream = self.gpu.stream();
7627                let (p_, _g) = s.device_ptr_mut(&stream);
7628                unsafe {
7629                    cudarc::driver::sys::cuMemsetD8Async(p_, 0, n_bytes, stream.cu_stream())
7630                        .result()?;
7631                }
7632            }
7633        }
7634        self.keep_if_capturing(&s);
7635        Ok(s)
7636    }
7637
7638    /// Public f32 uninitialized scratch (see `alloc_uninit`). For decode/forward scratch a kernel
7639    /// fully overwrites. SAFETY: producing kernel must write every element before any read.
7640    /// Uninitialized q8_1 activation pair (int8 + per-32 scales) — the fa combine q8-emit
7641    /// consumers alloc through this (m=1 decode arms).
7642    pub fn uninit_q8_pair(
7643        &self,
7644        n: usize,
7645    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7646        Ok((
7647            self.alloc_uninit::<i8>(n)?,
7648            self.alloc_uninit::<f32>(n / 32)?,
7649        ))
7650    }
7651
7652    pub fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7653        self.alloc_uninit::<f32>(n)
7654    }
7655
7656    /// i8 uninitialized scratch (same contract as `uninit`).
7657    pub fn alloc_i8_uninit(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
7658        self.alloc_uninit::<i8>(n)
7659    }
7660
7661    /// RMSNorm: x[ncols,nrows] row-major, weight[ncols] -> dst. One block/row, 256 threads.
7662    /// gemma4: 3 rms_norms of the SAME input in one launch (one reduction, three weights).
7663    /// Per-output bit-identical to three rms_norm calls (verbatim reduction/scale chain).
7664    #[allow(clippy::too_many_arguments)]
7665    pub fn rms_norm3(
7666        &self,
7667        x: &CudaSlice<f32>,
7668        w0: &CudaSlice<f32>,
7669        w1: &CudaSlice<f32>,
7670        w2: &CudaSlice<f32>,
7671        d0: &mut CudaSlice<f32>,
7672        d1: &mut CudaSlice<f32>,
7673        d2: &mut CudaSlice<f32>,
7674        ncols: usize,
7675        nrows: usize,
7676        eps: f32,
7677    ) -> Result<(), Box<dyn std::error::Error>> {
7678        let f = self.func("rms_norm3_f32");
7679        let cfg = LaunchConfig {
7680            grid_dim: (nrows as u32, 1, 1),
7681            block_dim: (rms_block(), 1, 1),
7682            shared_mem_bytes: 0,
7683        };
7684        let (nc, e) = (ncols as i32, eps);
7685        let __s_b = self.gpu.stream();
7686        let mut b = __s_b.launch_builder(&f);
7687        b.arg(x)
7688            .arg(w0)
7689            .arg(w1)
7690            .arg(w2)
7691            .arg(d0)
7692            .arg(d1)
7693            .arg(d2)
7694            .arg(&nc)
7695            .arg(&e);
7696        unsafe {
7697            b.launch(cfg)?;
7698        }
7699        Ok(())
7700    }
7701
7702    /// gemma4 fused q/k/v head norms (one launch, per-row rms_norm_f32-verbatim).
7703    #[allow(clippy::too_many_arguments)]
7704    /// True when the warp-per-row qkv norm would engage for (rows, ncols) — the emit lane
7705    /// piggybacks on the same conditions.
7706    pub fn qkvnorm_w_on_prefill(rows: usize, ncols: usize) -> bool {
7707        static WARP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7708        *WARP_ON.get_or_init(|| {
7709            std::env::var("MEMRA_QKVNORM_W")
7710                .map(|v| v != "0")
7711                .unwrap_or(true)
7712        }) && ncols % 4 == 0
7713            && rows >= 64
7714    }
7715
7716    /// w4 norm with bf16 V EMIT (31B glue lane): the v segment also writes its normed rows as
7717    /// bf16 (the FA V operand — bit-identical to a post-hoc f32_to_bf16). Prefill-depth only.
7718    #[allow(clippy::too_many_arguments)]
7719    pub fn rms_norm_qkv_w4b(
7720        &self,
7721        q: &CudaSlice<f32>,
7722        k: &CudaSlice<f32>,
7723        v: &CudaSlice<f32>,
7724        wq: &CudaSlice<f32>,
7725        wk: &CudaSlice<f32>,
7726        wv: &CudaSlice<f32>,
7727        dq: &mut CudaSlice<f32>,
7728        dk: &mut CudaSlice<f32>,
7729        dv: &mut CudaSlice<f32>,
7730        dvb: &mut CudaSlice<u8>,
7731        ncols: usize,
7732        rq: usize,
7733        rk: usize,
7734        eps: f32,
7735        vf16: bool,
7736    ) -> Result<(), Box<dyn std::error::Error>> {
7737        assert!(ncols % 4 == 0 && rq + 2 * rk >= 64);
7738        let f = self.func("rms_norm_qkv_w4b_f32");
7739        let rows = (rq + 2 * rk) as u32;
7740        let cfg = LaunchConfig {
7741            grid_dim: (rows.div_ceil(8), 1, 1),
7742            block_dim: (256, 1, 1),
7743            shared_mem_bytes: 0,
7744        };
7745        let (nc, rqi, rki, rvi, e) = (ncols as i32, rq as i32, rk as i32, rk as i32, eps);
7746        let vf = vf16 as i32;
7747        let __s_b = self.gpu.stream();
7748        let mut b = __s_b.launch_builder(&f);
7749        b.arg(q)
7750            .arg(k)
7751            .arg(v)
7752            .arg(wq)
7753            .arg(wk)
7754            .arg(wv)
7755            .arg(dq)
7756            .arg(dk)
7757            .arg(dv)
7758            .arg(&mut *dvb)
7759            .arg(&nc)
7760            .arg(&rqi)
7761            .arg(&rki)
7762            .arg(&rvi)
7763            .arg(&e)
7764            .arg(&vf);
7765        unsafe {
7766            b.launch(cfg)?;
7767        }
7768        Ok(())
7769    }
7770
7771    pub fn rms_norm_qkv(
7772        &self,
7773        q: &CudaSlice<f32>,
7774        k: &CudaSlice<f32>,
7775        v: &CudaSlice<f32>,
7776        wq: &CudaSlice<f32>,
7777        wk: &CudaSlice<f32>,
7778        wv: &CudaSlice<f32>,
7779        dq: &mut CudaSlice<f32>,
7780        dk: &mut CudaSlice<f32>,
7781        dv: &mut CudaSlice<f32>,
7782        ncols: usize,
7783        rq: usize,
7784        rk: usize,
7785        eps: f32,
7786    ) -> Result<(), Box<dyn std::error::Error>> {
7787        // Warp-per-row float4 twin (default; MEMRA_QKVNORM_W=0 reverts): the block-per-row form
7788        // spends 767us/launch on 17k+ 2KB rows at prefill depth (launch/reduce latency-bound,
7789        // ~92GB/s). Own numeric config (reduce order differs) — battery-gated.
7790        static WARP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7791        let warp_on = *WARP_ON.get_or_init(|| {
7792            std::env::var("MEMRA_QKVNORM_W")
7793                .map(|v| v != "0")
7794                .unwrap_or(true)
7795        });
7796        // rows >= 64 keeps decode (nh + 2*nkv rows) on the block-tree kernel — decode/verify/
7797        // replay numerics are untouched on every model; only prefill depth takes the new config.
7798        if warp_on && ncols % 4 == 0 && rq + 2 * rk >= 64 {
7799            let f = self.func("rms_norm_qkv_w4_f32");
7800            let rows = (rq + 2 * rk) as u32;
7801            let cfg = LaunchConfig {
7802                grid_dim: (rows.div_ceil(8), 1, 1),
7803                block_dim: (256, 1, 1),
7804                shared_mem_bytes: 0,
7805            };
7806            let (nc, rqi, rki, rvi, e) = (ncols as i32, rq as i32, rk as i32, rk as i32, eps);
7807            let __s_b = self.gpu.stream();
7808            let mut b = __s_b.launch_builder(&f);
7809            b.arg(q)
7810                .arg(k)
7811                .arg(v)
7812                .arg(wq)
7813                .arg(wk)
7814                .arg(wv)
7815                .arg(dq)
7816                .arg(dk)
7817                .arg(dv)
7818                .arg(&nc)
7819                .arg(&rqi)
7820                .arg(&rki)
7821                .arg(&rvi)
7822                .arg(&e);
7823            unsafe {
7824                b.launch(cfg)?;
7825            }
7826            return Ok(());
7827        }
7828        let f = self.func("rms_norm_qkv_f32");
7829        let grid = (rq + 2 * rk) as u32;
7830        let cfg = LaunchConfig {
7831            grid_dim: (grid, 1, 1),
7832            block_dim: (rms_block(), 1, 1),
7833            shared_mem_bytes: 0,
7834        };
7835        let (nc, rqi, rki, e) = (ncols as i32, rq as i32, rk as i32, eps);
7836        let __s_b = self.gpu.stream();
7837        let mut b = __s_b.launch_builder(&f);
7838        b.arg(q)
7839            .arg(k)
7840            .arg(v)
7841            .arg(wq)
7842            .arg(wk)
7843            .arg(wv)
7844            .arg(dq)
7845            .arg(dk)
7846            .arg(dv)
7847            .arg(&nc)
7848            .arg(&rqi)
7849            .arg(&rki)
7850            .arg(&e);
7851        unsafe {
7852            b.launch(cfg)?;
7853        }
7854        Ok(())
7855    }
7856
7857    /// gemma4 fused pair of rms_norms over two different inputs (same width).
7858    #[allow(clippy::too_many_arguments)]
7859    pub fn rms_norm2x(
7860        &self,
7861        a: &CudaSlice<f32>,
7862        bb: &CudaSlice<f32>,
7863        wa: &CudaSlice<f32>,
7864        wb: &CudaSlice<f32>,
7865        da: &mut CudaSlice<f32>,
7866        db: &mut CudaSlice<f32>,
7867        ncols: usize,
7868        nrows: usize,
7869        eps: f32,
7870    ) -> Result<(), Box<dyn std::error::Error>> {
7871        let f = self.func("rms_norm2x_f32");
7872        let cfg = LaunchConfig {
7873            grid_dim: (2 * nrows as u32, 1, 1),
7874            block_dim: (rms_block(), 1, 1),
7875            shared_mem_bytes: 0,
7876        };
7877        let (nc, nr, e) = (ncols as i32, nrows as i32, eps);
7878        let __s_b = self.gpu.stream();
7879        let mut b = __s_b.launch_builder(&f);
7880        b.arg(a)
7881            .arg(bb)
7882            .arg(wa)
7883            .arg(wb)
7884            .arg(da)
7885            .arg(db)
7886            .arg(&nc)
7887            .arg(&nr)
7888            .arg(&e);
7889        unsafe {
7890            b.launch(cfg)?;
7891        }
7892        Ok(())
7893    }
7894
7895    /// gemma4 R4: in-place final-logit softcap y = cap*tanh(y/cap).
7896    pub fn softcap(
7897        &self,
7898        y: &mut CudaSlice<f32>,
7899        cap: f32,
7900        n: usize,
7901    ) -> Result<(), Box<dyn std::error::Error>> {
7902        let f = self.func("softcap_f32");
7903        let cfg = LaunchConfig::for_num_elems(n as u32);
7904        let ni = n as i32;
7905        let __s_b = self.gpu.stream();
7906        let mut b = __s_b.launch_builder(&f);
7907        b.arg(y).arg(&cap).arg(&ni);
7908        unsafe {
7909            b.launch(cfg)?;
7910        }
7911        Ok(())
7912    }
7913
7914    /// gemma4 suppress-token mask: y[row][ids[j]] = -inf over t logits rows (fixed-arg launch —
7915    /// graph-capture safe; NOT monotonic like softcap, so it must run before any argmax).
7916    pub fn mask_ids_rows(
7917        &self,
7918        y: &mut CudaSlice<f32>,
7919        ids: &CudaSlice<i32>,
7920        n_ids: usize,
7921        n_vocab: usize,
7922        t: usize,
7923    ) -> Result<(), Box<dyn std::error::Error>> {
7924        let f = self.func("mask_ids_rows_f32");
7925        let cfg = LaunchConfig::for_num_elems((n_ids * t) as u32);
7926        let (ni, nv, ti) = (n_ids as i32, n_vocab as i32, t as i32);
7927        let __s_b = self.gpu.stream();
7928        let mut b = __s_b.launch_builder(&f);
7929        b.arg(y).arg(ids).arg(&ni).arg(&nv).arg(&ti);
7930        unsafe {
7931            b.launch(cfg)?;
7932        }
7933        Ok(())
7934    }
7935
7936    /// gemma4: res = (a+b)*c AND dst = rms_norm(res, w) in one launch.
7937    #[allow(clippy::too_many_arguments)]
7938    pub fn add_scale_rms_norm(
7939        &self,
7940        a: &CudaSlice<f32>,
7941        b_in: &CudaSlice<f32>,
7942        c: f32,
7943        w: &CudaSlice<f32>,
7944        res: &mut CudaSlice<f32>,
7945        dst: &mut CudaSlice<f32>,
7946        ncols: usize,
7947        nrows: usize,
7948        eps: f32,
7949    ) -> Result<(), Box<dyn std::error::Error>> {
7950        let f = self.func("add_scale_rms_norm_f32");
7951        let cfg = LaunchConfig {
7952            grid_dim: (nrows as u32, 1, 1),
7953            block_dim: (rms_block(), 1, 1),
7954            shared_mem_bytes: 0,
7955        };
7956        let (nc, e2) = (ncols as i32, eps);
7957        let __s_b = self.gpu.stream();
7958        let mut b = __s_b.launch_builder(&f);
7959        b.arg(a)
7960            .arg(b_in)
7961            .arg(&c)
7962            .arg(w)
7963            .arg(res)
7964            .arg(dst)
7965            .arg(&nc)
7966            .arg(&e2);
7967        unsafe {
7968            b.launch(cfg)?;
7969        }
7970        Ok(())
7971    }
7972
7973    /// gemma4: res = (a+b)*c AND the next layer's attn_norm EMITTED q8_1 in one launch.
7974    /// Quantize epilogue bit-identical to quantize_q8_1 (the rms_norm_q8_1 form).
7975    #[allow(clippy::too_many_arguments)]
7976    pub fn add_scale_rms_norm_q8_1(
7977        &self,
7978        a: &CudaSlice<f32>,
7979        b_in: &CudaSlice<f32>,
7980        c: f32,
7981        w: &CudaSlice<f32>,
7982        res: &mut CudaSlice<f32>,
7983        ncols: usize,
7984        nrows: usize,
7985        eps: f32,
7986    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7987        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
7988        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
7989        let (nc, e2) = (ncols as i32, eps);
7990        if Self::pdl_on() && Self::pdl_wb_on() {
7991            {
7992                use cudarc::driver::{DevicePtr, DevicePtrMut};
7993                let s = &self.gpu.stream();
7994                let (pa, _g0) = a.device_ptr(s);
7995                let (pb, _g1) = b_in.device_ptr(s);
7996                let (pw, _g2) = w.device_ptr(s);
7997                let (pr, _g3) = res.device_ptr_mut(s);
7998                let (pq, _g4) = out_q.device_ptr_mut(s);
7999                let (pd, _g5) = out_d.device_ptr_mut(s);
8000                let mut ps = [
8001                    &pa as *const _ as *mut std::ffi::c_void,
8002                    &pb as *const _ as *mut _,
8003                    &c as *const _ as *mut _,
8004                    &pw as *const _ as *mut _,
8005                    &pr as *const _ as *mut _,
8006                    &pq as *const _ as *mut _,
8007                    &pd as *const _ as *mut _,
8008                    &nc as *const _ as *mut _,
8009                    &e2 as *const _ as *mut _,
8010                ];
8011                unsafe {
8012                    self.launch_pdl(
8013                        "add_scale_rms_norm_q8_1",
8014                        (nrows as u32, 1, 1),
8015                        (rms_block(), 1, 1),
8016                        &mut ps,
8017                    )?;
8018                }
8019            }
8020            return Ok((out_q, out_d));
8021        }
8022        let f = self.func("add_scale_rms_norm_q8_1");
8023        let cfg = LaunchConfig {
8024            grid_dim: (nrows as u32, 1, 1),
8025            block_dim: (rms_block(), 1, 1),
8026            shared_mem_bytes: 0,
8027        };
8028        let __s_b = self.gpu.stream();
8029        let mut b = __s_b.launch_builder(&f);
8030        b.arg(a)
8031            .arg(b_in)
8032            .arg(&c)
8033            .arg(w)
8034            .arg(res)
8035            .arg(&mut out_q)
8036            .arg(&mut out_d)
8037            .arg(&nc)
8038            .arg(&e2);
8039        unsafe {
8040            b.launch(cfg)?;
8041        }
8042        Ok((out_q, out_d))
8043    }
8044
8045    /// Slot-fed add_scale_rms_norm_q8_1 twin (alloc-free capture lane).
8046    #[allow(clippy::too_many_arguments)]
8047    pub fn add_scale_rms_norm_q8_1_into(
8048        &self,
8049        a: &CudaSlice<f32>,
8050        b_in: &CudaSlice<f32>,
8051        c: f32,
8052        w: &CudaSlice<f32>,
8053        res: &mut CudaSlice<f32>,
8054        ncols: usize,
8055        nrows: usize,
8056        eps: f32,
8057        out_q: &mut CudaSlice<i8>,
8058        out_d: &mut CudaSlice<f32>,
8059    ) -> Result<(), Box<dyn std::error::Error>> {
8060        debug_assert!(out_q.len() >= nrows * ncols && out_d.len() >= nrows * (ncols / 32));
8061        let (nc, e2) = (ncols as i32, eps);
8062        if Self::pdl_on() && Self::pdl_wb_on() {
8063            use cudarc::driver::{DevicePtr, DevicePtrMut};
8064            let s = &self.gpu.stream();
8065            let (pa, _g0) = a.device_ptr(s);
8066            let (pb, _g1) = b_in.device_ptr(s);
8067            let (pw, _g2) = w.device_ptr(s);
8068            let (pr, _g3) = res.device_ptr_mut(s);
8069            let (pq, _g4) = out_q.device_ptr_mut(s);
8070            let (pd, _g5) = out_d.device_ptr_mut(s);
8071            let mut ps = [
8072                &pa as *const _ as *mut std::ffi::c_void,
8073                &pb as *const _ as *mut _,
8074                &c as *const _ as *mut _,
8075                &pw as *const _ as *mut _,
8076                &pr as *const _ as *mut _,
8077                &pq as *const _ as *mut _,
8078                &pd as *const _ as *mut _,
8079                &nc as *const _ as *mut _,
8080                &e2 as *const _ as *mut _,
8081            ];
8082            unsafe {
8083                self.launch_pdl(
8084                    "add_scale_rms_norm_q8_1",
8085                    (nrows as u32, 1, 1),
8086                    (rms_block(), 1, 1),
8087                    &mut ps,
8088                )?;
8089            }
8090            return Ok(());
8091        }
8092        let f = self.func("add_scale_rms_norm_q8_1");
8093        let cfg = LaunchConfig {
8094            grid_dim: (nrows as u32, 1, 1),
8095            block_dim: (rms_block(), 1, 1),
8096            shared_mem_bytes: 0,
8097        };
8098        let __s_b = self.gpu.stream();
8099        let mut b = __s_b.launch_builder(&f);
8100        b.arg(a)
8101            .arg(b_in)
8102            .arg(&c)
8103            .arg(w)
8104            .arg(res)
8105            .arg(&mut *out_q)
8106            .arg(&mut *out_d)
8107            .arg(&nc)
8108            .arg(&e2);
8109        unsafe {
8110            b.launch(cfg)?;
8111        }
8112        Ok(())
8113    }
8114
8115    /// E4B glue fusion: rms(a, wa) prologue + the add_scale_rms_norm_q8_1 program — one launch
8116    /// replaces the per-layer rms_norm_f32(y) + emit pair in the PLE tail.
8117    #[allow(clippy::too_many_arguments)]
8118    pub fn rms_pre_add_scale_rms_norm_q8_1(
8119        &self,
8120        a: &CudaSlice<f32>,
8121        wa: &CudaSlice<f32>,
8122        b_in: &CudaSlice<f32>,
8123        c: f32,
8124        w: &CudaSlice<f32>,
8125        res: &mut CudaSlice<f32>,
8126        ncols: usize,
8127        nrows: usize,
8128        eps: f32,
8129    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8130        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
8131        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
8132        let (nc, e2) = (ncols as i32, eps);
8133        if Self::pdl_on() {
8134            {
8135                use cudarc::driver::{DevicePtr, DevicePtrMut};
8136                let s = &self.gpu.stream();
8137                let (pa, _g0) = a.device_ptr(s);
8138                let (pwa, _g1) = wa.device_ptr(s);
8139                let (pb, _g2) = b_in.device_ptr(s);
8140                let (pw, _g3) = w.device_ptr(s);
8141                let (pr, _g4) = res.device_ptr_mut(s);
8142                let (pq, _g5) = out_q.device_ptr_mut(s);
8143                let (pd, _g6) = out_d.device_ptr_mut(s);
8144                let mut ps = [
8145                    &pa as *const _ as *mut std::ffi::c_void,
8146                    &pwa as *const _ as *mut _,
8147                    &pb as *const _ as *mut _,
8148                    &c as *const _ as *mut _,
8149                    &pw as *const _ as *mut _,
8150                    &pr as *const _ as *mut _,
8151                    &pq as *const _ as *mut _,
8152                    &pd as *const _ as *mut _,
8153                    &nc as *const _ as *mut _,
8154                    &e2 as *const _ as *mut _,
8155                ];
8156                unsafe {
8157                    self.launch_pdl(
8158                        "rms_pre_add_scale_rms_norm_q8_1",
8159                        (nrows as u32, 1, 1),
8160                        (rms_block(), 1, 1),
8161                        &mut ps,
8162                    )?;
8163                }
8164            }
8165            return Ok((out_q, out_d));
8166        }
8167        let f = self.func("rms_pre_add_scale_rms_norm_q8_1");
8168        let cfg = LaunchConfig {
8169            grid_dim: (nrows as u32, 1, 1),
8170            block_dim: (rms_block(), 1, 1),
8171            shared_mem_bytes: 0,
8172        };
8173        let __s_b = self.gpu.stream();
8174        let mut b = __s_b.launch_builder(&f);
8175        b.arg(a)
8176            .arg(wa)
8177            .arg(b_in)
8178            .arg(&c)
8179            .arg(w)
8180            .arg(res)
8181            .arg(&mut out_q)
8182            .arg(&mut out_d)
8183            .arg(&nc)
8184            .arg(&e2);
8185        unsafe {
8186            b.launch(cfg)?;
8187        }
8188        Ok((out_q, out_d))
8189    }
8190
8191    /// GELU(tanh)*up with the activation emitted q8_1 alongside f32 (glue-fusion lane): the
8192    /// consumer matmul rides matmul_pre, killing its standalone quantize_q8_1 launch.
8193    pub fn gelu_tanh_mul_q8_1(
8194        &self,
8195        gate: &CudaSlice<f32>,
8196        up: &cudarc::driver::CudaView<f32>,
8197        act: &mut CudaSlice<f32>,
8198        ncols: usize,
8199        nrows: usize,
8200    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8201        debug_assert!(ncols % 128 == 0);
8202        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
8203        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
8204        let nc = ncols as i32;
8205        if Self::pdl_on() {
8206            {
8207                use cudarc::driver::{DevicePtr, DevicePtrMut};
8208                let s = &self.gpu.stream();
8209                let (pg, _g0) = gate.device_ptr(s);
8210                let (pu, _g1) = up.device_ptr(s);
8211                let (pact, _g2) = act.device_ptr_mut(s);
8212                let (pq, _g3) = out_q.device_ptr_mut(s);
8213                let (pd, _g4) = out_d.device_ptr_mut(s);
8214                let mut ps = [
8215                    &pg as *const _ as *mut std::ffi::c_void,
8216                    &pu as *const _ as *mut _,
8217                    &pact as *const _ as *mut _,
8218                    &pq as *const _ as *mut _,
8219                    &pd as *const _ as *mut _,
8220                    &nc as *const _ as *mut _,
8221                ];
8222                unsafe {
8223                    self.launch_pdl(
8224                        "gelu_tanh_mul_q8_1",
8225                        (nrows as u32, 1, 1),
8226                        (rms_block(), 1, 1),
8227                        &mut ps,
8228                    )?;
8229                }
8230            }
8231            return Ok((out_q, out_d));
8232        }
8233        let f = self.func("gelu_tanh_mul_q8_1");
8234        let cfg = LaunchConfig {
8235            grid_dim: (nrows as u32, 1, 1),
8236            block_dim: (rms_block(), 1, 1),
8237            shared_mem_bytes: 0,
8238        };
8239        let __s_b = self.gpu.stream();
8240        let mut b = __s_b.launch_builder(&f);
8241        b.arg(gate)
8242            .arg(up)
8243            .arg(act)
8244            .arg(&mut out_q)
8245            .arg(&mut out_d)
8246            .arg(&nc);
8247        unsafe {
8248            b.launch(cfg)?;
8249        }
8250        Ok((out_q, out_d))
8251    }
8252
8253    /// Slot-fed gelu_tanh_mul_q8_1 twin (alloc-free capture lane; incl. the PDL arm).
8254    #[allow(clippy::too_many_arguments)]
8255    pub fn gelu_tanh_mul_q8_1_into(
8256        &self,
8257        gate: &CudaSlice<f32>,
8258        up: &cudarc::driver::CudaView<f32>,
8259        act: &mut CudaSlice<f32>,
8260        ncols: usize,
8261        nrows: usize,
8262        out_q: &mut CudaSlice<i8>,
8263        out_d: &mut CudaSlice<f32>,
8264    ) -> Result<(), Box<dyn std::error::Error>> {
8265        debug_assert!(ncols % 128 == 0);
8266        debug_assert!(out_q.len() >= nrows * ncols && out_d.len() >= nrows * (ncols / 32));
8267        let nc = ncols as i32;
8268        if Self::pdl_on() {
8269            use cudarc::driver::{DevicePtr, DevicePtrMut};
8270            let s = &self.gpu.stream();
8271            let (pg, _g0) = gate.device_ptr(s);
8272            let (pu, _g1) = up.device_ptr(s);
8273            let (pact, _g2) = act.device_ptr_mut(s);
8274            let (pq, _g3) = out_q.device_ptr_mut(s);
8275            let (pd, _g4) = out_d.device_ptr_mut(s);
8276            let mut ps = [
8277                &pg as *const _ as *mut std::ffi::c_void,
8278                &pu as *const _ as *mut _,
8279                &pact as *const _ as *mut _,
8280                &pq as *const _ as *mut _,
8281                &pd as *const _ as *mut _,
8282                &nc as *const _ as *mut _,
8283            ];
8284            unsafe {
8285                self.launch_pdl(
8286                    "gelu_tanh_mul_q8_1",
8287                    (nrows as u32, 1, 1),
8288                    (rms_block(), 1, 1),
8289                    &mut ps,
8290                )?;
8291            }
8292            return Ok(());
8293        }
8294        let f = self.func("gelu_tanh_mul_q8_1");
8295        let cfg = LaunchConfig {
8296            grid_dim: (nrows as u32, 1, 1),
8297            block_dim: (rms_block(), 1, 1),
8298            shared_mem_bytes: 0,
8299        };
8300        let __s_b = self.gpu.stream();
8301        let mut b = __s_b.launch_builder(&f);
8302        b.arg(gate)
8303            .arg(up)
8304            .arg(&mut *act)
8305            .arg(&mut *out_q)
8306            .arg(&mut *out_d)
8307            .arg(&nc);
8308        unsafe {
8309            b.launch(cfg)?;
8310        }
8311        Ok(())
8312    }
8313
8314    /// gemma4: add + rms_norm3 with outputs 0/2 emitted q8_1 (zsh + moe_in) and 1 f32 (router).
8315    #[allow(clippy::too_many_arguments)]
8316    pub fn add_rms_norm3_q8z(
8317        &self,
8318        a: &CudaSlice<f32>,
8319        b_in: &CudaSlice<f32>,
8320        w0: &CudaSlice<f32>,
8321        w1: &CudaSlice<f32>,
8322        w2: &CudaSlice<f32>,
8323        res: &mut CudaSlice<f32>,
8324        out1: &mut CudaSlice<f32>,
8325        ncols: usize,
8326        nrows: usize,
8327        eps: f32,
8328    ) -> Result<
8329        (
8330            (CudaSlice<i8>, CudaSlice<f32>),
8331            (CudaSlice<i8>, CudaSlice<f32>),
8332        ),
8333        Box<dyn std::error::Error>,
8334    > {
8335        let mut q0 = self.alloc_uninit::<i8>(nrows * ncols)?;
8336        let mut d0 = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
8337        let mut q2 = self.alloc_uninit::<i8>(nrows * ncols)?;
8338        let mut d2 = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
8339        let f = self.func("add_rms_norm3_q8z_f32");
8340        let cfg = LaunchConfig {
8341            grid_dim: (nrows as u32, 1, 1),
8342            block_dim: (rms_block(), 1, 1),
8343            shared_mem_bytes: 0,
8344        };
8345        let (nc, e2) = (ncols as i32, eps);
8346        let __s_b = self.gpu.stream();
8347        let mut b = __s_b.launch_builder(&f);
8348        b.arg(a)
8349            .arg(b_in)
8350            .arg(w0)
8351            .arg(w1)
8352            .arg(w2)
8353            .arg(res)
8354            .arg(&mut q0)
8355            .arg(&mut d0)
8356            .arg(out1)
8357            .arg(&mut q2)
8358            .arg(&mut d2)
8359            .arg(&nc)
8360            .arg(&e2);
8361        unsafe {
8362            b.launch(cfg)?;
8363        }
8364        Ok(((q0, d0), (q2, d2)))
8365    }
8366
8367    /// gemma4: res = a+b AND the three rms_norms of res in one launch.
8368    #[allow(clippy::too_many_arguments)]
8369    pub fn add_rms_norm3(
8370        &self,
8371        a: &CudaSlice<f32>,
8372        b_in: &CudaSlice<f32>,
8373        w0: &CudaSlice<f32>,
8374        w1: &CudaSlice<f32>,
8375        w2: &CudaSlice<f32>,
8376        res: &mut CudaSlice<f32>,
8377        d0: &mut CudaSlice<f32>,
8378        d1: &mut CudaSlice<f32>,
8379        d2: &mut CudaSlice<f32>,
8380        ncols: usize,
8381        nrows: usize,
8382        eps: f32,
8383    ) -> Result<(), Box<dyn std::error::Error>> {
8384        let f = self.func("add_rms_norm3_f32");
8385        let cfg = LaunchConfig {
8386            grid_dim: (nrows as u32, 1, 1),
8387            block_dim: (rms_block(), 1, 1),
8388            shared_mem_bytes: 0,
8389        };
8390        let (nc, e2) = (ncols as i32, eps);
8391        let __s_b = self.gpu.stream();
8392        let mut b = __s_b.launch_builder(&f);
8393        b.arg(a)
8394            .arg(b_in)
8395            .arg(w0)
8396            .arg(w1)
8397            .arg(w2)
8398            .arg(res)
8399            .arg(d0)
8400            .arg(d1)
8401            .arg(d2)
8402            .arg(&nc)
8403            .arg(&e2);
8404        unsafe {
8405            b.launch(cfg)?;
8406        }
8407        Ok(())
8408    }
8409
8410    /// dst = (a + b) * c (residual add + layer scale, one launch).
8411    pub fn add_scale(
8412        &self,
8413        a: &CudaSlice<f32>,
8414        b_in: &CudaSlice<f32>,
8415        c: f32,
8416        dst: &mut CudaSlice<f32>,
8417        n: usize,
8418    ) -> Result<(), Box<dyn std::error::Error>> {
8419        let f = self.func("add_scale_f32");
8420        let cfg = LaunchConfig::for_num_elems(n as u32);
8421        let ni = n as i32;
8422        let __s_b = self.gpu.stream();
8423        let mut b = __s_b.launch_builder(&f);
8424        b.arg(a).arg(b_in).arg(&c).arg(dst).arg(&ni);
8425        unsafe {
8426            b.launch(cfg)?;
8427        }
8428        Ok(())
8429    }
8430
8431    /// Vision-tower LayerNorm (with bias) over [nrows, ncols] — lane/vision.
8432    pub fn layer_norm_bias(
8433        &self,
8434        x: &CudaSlice<f32>,
8435        w: &CudaSlice<f32>,
8436        b: &CudaSlice<f32>,
8437        dst: &mut CudaSlice<f32>,
8438        ncols: usize,
8439        nrows: usize,
8440        eps: f32,
8441    ) -> Result<(), Box<dyn std::error::Error>> {
8442        let f = self.func("layer_norm_bias_f32");
8443        let (nc, e) = (ncols as i32, eps);
8444        let cfg = LaunchConfig {
8445            grid_dim: (nrows as u32, 1, 1),
8446            block_dim: (256, 1, 1),
8447            shared_mem_bytes: 0,
8448        };
8449        let __s_b = self.gpu.stream();
8450        let mut lb = __s_b.launch_builder(&f);
8451        lb.arg(x).arg(w).arg(b).arg(&mut *dst).arg(&nc).arg(&e);
8452        unsafe {
8453            lb.launch(cfg)?;
8454        }
8455        Ok(())
8456    }
8457
8458    /// gelu_pytorch_tanh elementwise (vision tower MLP activation).
8459    pub fn gelu_tanh(
8460        &self,
8461        x: &CudaSlice<f32>,
8462        dst: &mut CudaSlice<f32>,
8463        n: usize,
8464    ) -> Result<(), Box<dyn std::error::Error>> {
8465        let f = self.func("gelu_tanh_f32");
8466        let ni = n as i64;
8467        let cfg = LaunchConfig {
8468            grid_dim: (n.div_ceil(256) as u32, 1, 1),
8469            block_dim: (256, 1, 1),
8470            shared_mem_bytes: 0,
8471        };
8472        let __s_b = self.gpu.stream();
8473        let mut lb = __s_b.launch_builder(&f);
8474        lb.arg(x).arg(&mut *dst).arg(&ni);
8475        unsafe {
8476            lb.launch(cfg)?;
8477        }
8478        Ok(())
8479    }
8480
8481    /// In-place row softmax over [nrows, ncols] (bidirectional vision attention).
8482    pub fn row_softmax(
8483        &self,
8484        x: &mut CudaSlice<f32>,
8485        ncols: usize,
8486        nrows: usize,
8487    ) -> Result<(), Box<dyn std::error::Error>> {
8488        let f = self.func("row_softmax_f32");
8489        let nc = ncols as i32;
8490        let cfg = LaunchConfig {
8491            grid_dim: (nrows as u32, 1, 1),
8492            block_dim: (256, 1, 1),
8493            shared_mem_bytes: 0,
8494        };
8495        let __s_b = self.gpu.stream();
8496        let mut lb = __s_b.launch_builder(&f);
8497        lb.arg(&mut *x).arg(&nc);
8498        unsafe {
8499            lb.launch(cfg)?;
8500        }
8501        Ok(())
8502    }
8503
8504    pub fn rms_norm(
8505        &self,
8506        x: &CudaSlice<f32>,
8507        w: &CudaSlice<f32>,
8508        dst: &mut CudaSlice<f32>,
8509        ncols: usize,
8510        nrows: usize,
8511        eps: f32,
8512    ) -> Result<(), Box<dyn std::error::Error>> {
8513        let (nc, e) = (ncols as i32, eps);
8514        let kname = if Self::norm_ilp_on() {
8515            "rms_norm_f32_v2"
8516        } else {
8517            "rms_norm_f32"
8518        };
8519        if Self::pdl_on() && Self::pdl_wb_on() {
8520            use cudarc::driver::{DevicePtr, DevicePtrMut};
8521            let s = &self.gpu.stream();
8522            let (px, _g0) = x.device_ptr(s);
8523            let (pw, _g1) = w.device_ptr(s);
8524            let (pd, _g2) = dst.device_ptr_mut(s);
8525            let mut ps = [
8526                &px as *const _ as *mut std::ffi::c_void,
8527                &pw as *const _ as *mut _,
8528                &pd as *const _ as *mut _,
8529                &nc as *const _ as *mut _,
8530                &e as *const _ as *mut _,
8531            ];
8532            unsafe {
8533                self.launch_pdl(kname, (nrows as u32, 1, 1), (rms_block(), 1, 1), &mut ps)?;
8534            }
8535            return Ok(());
8536        }
8537        let f = self.func(kname);
8538        let cfg = LaunchConfig {
8539            grid_dim: (nrows as u32, 1, 1),
8540            block_dim: (rms_block(), 1, 1),
8541            shared_mem_bytes: 0,
8542        };
8543        let __s_b = self.gpu.stream();
8544        let mut b = __s_b.launch_builder(&f);
8545        b.arg(x).arg(w).arg(dst).arg(&nc).arg(&e);
8546        unsafe {
8547            b.launch(cfg)?;
8548        }
8549        Ok(())
8550    }
8551
8552    /// RMS-norm with blockDim=1024 — BIT-IDENTICAL to the fused `rms_norm_q8_1` and
8553    /// `add_rms_norm_q8_1` kernels' sum-of-squares reduction. The spec verify path MUST use this
8554    /// to match decode's FP accumulation order: the standard `rms_norm` at blockDim=256 has a
8555    /// different per-thread stride (ncols/256 partials vs ncols/1024 partials) and therefore a
8556    /// different shfl-tree reduction that can shift `scale = rsqrt(sum/n + eps)` by ULPs, causing
8557    /// divergence through the GDN scan and argmax flips on the 9B text prompt. The underlying
8558    /// `rms_norm_f32` kernel supports any blockDim (generic reduce with shared[32]).
8559    pub fn rms_norm_decode(
8560        &self,
8561        x: &CudaSlice<f32>,
8562        w: &CudaSlice<f32>,
8563        dst: &mut CudaSlice<f32>,
8564        ncols: usize,
8565        nrows: usize,
8566        eps: f32,
8567    ) -> Result<(), Box<dyn std::error::Error>> {
8568        let f = self.func(if Self::norm_ilp_on() {
8569            "rms_norm_f32_v2"
8570        } else {
8571            "rms_norm_f32"
8572        });
8573        let cfg = LaunchConfig {
8574            grid_dim: (nrows as u32, 1, 1),
8575            block_dim: (1024, 1, 1),
8576            shared_mem_bytes: 0,
8577        };
8578        let (nc, e) = (ncols as i32, eps);
8579        let __s_b = self.gpu.stream();
8580        let mut b = __s_b.launch_builder(&f);
8581        b.arg(x).arg(w).arg(dst).arg(&nc).arg(&e);
8582        unsafe {
8583            b.launch(cfg)?;
8584        }
8585        Ok(())
8586    }
8587
8588    /// DECODE GLUE-FUSION LEVER: `z = rms_norm(x)*w` emitted DIRECTLY as q8_1 (no f32 `z` materialized,
8589    /// no standalone quantize_q8_1 launch). Returns (out_q [nrows*ncols i8], out_d [nrows*nblk f32])
8590    /// ready to feed matmul_pre. BIT-IDENTICAL to rms_norm + quantize_q8_1. ncols % 32 == 0.
8591    pub fn rms_norm_q8_1(
8592        &self,
8593        x: &CudaSlice<f32>,
8594        w: &CudaSlice<f32>,
8595        ncols: usize,
8596        nrows: usize,
8597        eps: f32,
8598    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8599        let nblk = ncols / 32;
8600        let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
8601        let mut d = self.alloc_uninit::<f32>(nrows * nblk)?;
8602        let (nc, e) = (ncols as i32, eps);
8603        if Self::pdl_on() {
8604            {
8605                use cudarc::driver::{DevicePtr, DevicePtrMut};
8606                let s = &self.gpu.stream();
8607                let (px, _g0) = x.device_ptr(s);
8608                let (pw, _g1) = w.device_ptr(s);
8609                let (pq, _g2) = q.device_ptr_mut(s);
8610                let (pd, _g3) = d.device_ptr_mut(s);
8611                let mut ps = [
8612                    &px as *const _ as *mut std::ffi::c_void,
8613                    &pw as *const _ as *mut _,
8614                    &pq as *const _ as *mut _,
8615                    &pd as *const _ as *mut _,
8616                    &nc as *const _ as *mut _,
8617                    &e as *const _ as *mut _,
8618                ];
8619                unsafe {
8620                    self.launch_pdl("rms_norm_q8_1", (nrows as u32, 1, 1), (1024, 1, 1), &mut ps)?;
8621                }
8622            }
8623            return Ok((q, d));
8624        }
8625        let f = self.func("rms_norm_q8_1");
8626        // 1024 threads: decode is nrows=1 -> ONE CTA; 32 warps hide the pass1->pass2 latency
8627        // (s[32] reduce already sized for 32 warps). Same shape math at any blockDim.
8628        let cfg = LaunchConfig {
8629            grid_dim: (nrows as u32, 1, 1),
8630            block_dim: (1024, 1, 1),
8631            shared_mem_bytes: 0,
8632        };
8633        let __s_b = self.gpu.stream();
8634        let mut b = __s_b.launch_builder(&f);
8635        b.arg(x).arg(w).arg(&mut q).arg(&mut d).arg(&nc).arg(&e);
8636        unsafe {
8637            b.launch(cfg)?;
8638        }
8639        Ok((q, d))
8640    }
8641
8642    /// Slot-fed rms_norm_q8_1 twin (alloc-free capture lane): identical launch (incl. the
8643    /// PDL arm), caller-owned outputs.
8644    pub fn rms_norm_q8_1_into(
8645        &self,
8646        x: &CudaSlice<f32>,
8647        w: &CudaSlice<f32>,
8648        ncols: usize,
8649        nrows: usize,
8650        eps: f32,
8651        q: &mut CudaSlice<i8>,
8652        d: &mut CudaSlice<f32>,
8653    ) -> Result<(), Box<dyn std::error::Error>> {
8654        let nblk = ncols / 32;
8655        debug_assert!(q.len() >= nrows * ncols && d.len() >= nrows * nblk);
8656        let (nc, e) = (ncols as i32, eps);
8657        if Self::pdl_on() {
8658            use cudarc::driver::{DevicePtr, DevicePtrMut};
8659            let s = &self.gpu.stream();
8660            let (px, _g0) = x.device_ptr(s);
8661            let (pw, _g1) = w.device_ptr(s);
8662            let (pq, _g2) = q.device_ptr_mut(s);
8663            let (pd, _g3) = d.device_ptr_mut(s);
8664            let mut ps = [
8665                &px as *const _ as *mut std::ffi::c_void,
8666                &pw as *const _ as *mut _,
8667                &pq as *const _ as *mut _,
8668                &pd as *const _ as *mut _,
8669                &nc as *const _ as *mut _,
8670                &e as *const _ as *mut _,
8671            ];
8672            unsafe {
8673                self.launch_pdl("rms_norm_q8_1", (nrows as u32, 1, 1), (1024, 1, 1), &mut ps)?;
8674            }
8675            return Ok(());
8676        }
8677        let f = self.func("rms_norm_q8_1");
8678        let cfg = LaunchConfig {
8679            grid_dim: (nrows as u32, 1, 1),
8680            block_dim: (1024, 1, 1),
8681            shared_mem_bytes: 0,
8682        };
8683        let __s_b = self.gpu.stream();
8684        let mut b = __s_b.launch_builder(&f);
8685        b.arg(x).arg(w).arg(&mut *q).arg(&mut *d).arg(&nc).arg(&e);
8686        unsafe {
8687            b.launch(cfg)?;
8688        }
8689        Ok(())
8690    }
8691
8692    /// Slot-fed quantize_q8_1 twin (alloc-free capture lane).
8693    pub fn quantize_q8_1_into(
8694        &self,
8695        x: &CudaSlice<f32>,
8696        m: usize,
8697        in_f: usize,
8698        q: &mut CudaSlice<i8>,
8699        d: &mut CudaSlice<f32>,
8700    ) -> Result<(), Box<dyn std::error::Error>> {
8701        let nblk = in_f / 32;
8702        debug_assert!(q.len() >= m * in_f && d.len() >= m * nblk);
8703        let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
8704        let (inf, mi) = (in_f as i32, m as i32);
8705        if Self::pdl_on() && Self::pdl_wb_on() {
8706            use cudarc::driver::{DevicePtr, DevicePtrMut};
8707            let s = &self.gpu.stream();
8708            let (px, _g0) = x.device_ptr(s);
8709            let (pq, _g1) = q.device_ptr_mut(s);
8710            let (pd, _g2) = d.device_ptr_mut(s);
8711            let mut ps = [
8712                &px as *const _ as *mut std::ffi::c_void,
8713                &pq as *const _ as *mut _,
8714                &pd as *const _ as *mut _,
8715                &inf as *const _ as *mut _,
8716                &mi as *const _ as *mut _,
8717            ];
8718            unsafe {
8719                self.launch_pdl("quantize_q8_1", cfg.grid_dim, cfg.block_dim, &mut ps)?;
8720            }
8721            return Ok(());
8722        }
8723        let f = self.func("quantize_q8_1");
8724        let __s_b = self.gpu.stream();
8725        let mut b = __s_b.launch_builder(&f);
8726        b.arg(x).arg(&mut *q).arg(&mut *d).arg(&inf).arg(&mi);
8727        unsafe {
8728            b.launch(cfg)?;
8729        }
8730        Ok(())
8731    }
8732
8733    /// DECODE GLUE-FUSION LEVER: `res = a+b; z = rms_norm(res)*w` with z emitted as q8_1. `res` is
8734    /// still written (the post-ffn residual add reads it). Fuses add_rms_norm + quantize_q8_1.
8735    /// Returns (out_q, out_d) for matmul_pre. BIT-IDENTICAL. ncols % 32 == 0.
8736    pub fn add_rms_norm_q8_1(
8737        &self,
8738        a: &CudaSlice<f32>,
8739        b_in: &CudaSlice<f32>,
8740        w: &CudaSlice<f32>,
8741        res: &mut CudaSlice<f32>,
8742        ncols: usize,
8743        nrows: usize,
8744        eps: f32,
8745    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8746        let nblk = ncols / 32;
8747        let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
8748        let mut d = self.alloc_uninit::<f32>(nrows * nblk)?;
8749        let f = self.func("add_rms_norm_q8_1");
8750        // 1024 threads: same single-CTA-at-decode reasoning as rms_norm_q8_1.
8751        let cfg = LaunchConfig {
8752            grid_dim: (nrows as u32, 1, 1),
8753            block_dim: (1024, 1, 1),
8754            shared_mem_bytes: 0,
8755        };
8756        let (nc, e) = (ncols as i32, eps);
8757        let __s_bld = self.gpu.stream();
8758        let mut bld = __s_bld.launch_builder(&f);
8759        bld.arg(a)
8760            .arg(b_in)
8761            .arg(w)
8762            .arg(res)
8763            .arg(&mut q)
8764            .arg(&mut d)
8765            .arg(&nc)
8766            .arg(&e);
8767        unsafe {
8768            bld.launch(cfg)?;
8769        }
8770        Ok((q, d))
8771    }
8772
8773    /// RANK3 LEVER (add+rmsnorm fuse): `res = a + b; dst = rms_norm(res) * w` in ONE launch. Fuses
8774    /// e.add(a,b,res) + e.rms_norm(res,w,dst), removing one launch + one HBM read of the residual per
8775    /// residual+norm pair. BIT-IDENTICAL to the two-kernel sequence (same IEEE add, same reduction).
8776    pub fn add_rms_norm(
8777        &self,
8778        a: &CudaSlice<f32>,
8779        b: &CudaSlice<f32>,
8780        w: &CudaSlice<f32>,
8781        res: &mut CudaSlice<f32>,
8782        dst: &mut CudaSlice<f32>,
8783        ncols: usize,
8784        nrows: usize,
8785        eps: f32,
8786    ) -> Result<(), Box<dyn std::error::Error>> {
8787        let (nc, e) = (ncols as i32, eps);
8788        let kname = if Self::norm_ilp_on() {
8789            "add_rms_norm_f32_v2"
8790        } else {
8791            "add_rms_norm_f32"
8792        };
8793        if Self::pdl_on() && Self::pdl_wb_on() {
8794            use cudarc::driver::{DevicePtr, DevicePtrMut};
8795            let s = &self.gpu.stream();
8796            let (pa, _g0) = a.device_ptr(s);
8797            let (pb, _g1) = b.device_ptr(s);
8798            let (pw, _g2) = w.device_ptr(s);
8799            let (pr, _g3) = res.device_ptr_mut(s);
8800            let (pd, _g4) = dst.device_ptr_mut(s);
8801            let mut ps = [
8802                &pa as *const _ as *mut std::ffi::c_void,
8803                &pb as *const _ as *mut _,
8804                &pw as *const _ as *mut _,
8805                &pr as *const _ as *mut _,
8806                &pd as *const _ as *mut _,
8807                &nc as *const _ as *mut _,
8808                &e as *const _ as *mut _,
8809            ];
8810            unsafe {
8811                self.launch_pdl(kname, (nrows as u32, 1, 1), (rms_block(), 1, 1), &mut ps)?;
8812            }
8813            return Ok(());
8814        }
8815        let f = self.func(kname);
8816        let cfg = LaunchConfig {
8817            grid_dim: (nrows as u32, 1, 1),
8818            block_dim: (rms_block(), 1, 1),
8819            shared_mem_bytes: 0,
8820        };
8821        let __s_b2 = self.gpu.stream();
8822        let mut b2 = __s_b2.launch_builder(&f);
8823        b2.arg(a)
8824            .arg(b)
8825            .arg(w)
8826            .arg(&mut *res)
8827            .arg(&mut *dst)
8828            .arg(&nc)
8829            .arg(&e);
8830        unsafe {
8831            b2.launch(cfg)?;
8832        }
8833        Ok(())
8834    }
8835
8836    /// E4B glue fusion: rms(a, wa) prologue + add_rms_norm — folds the post-attn norm into
8837    /// the tail entry (res = rms(a)*wa + b; dst = rms(res)*w).
8838    #[allow(clippy::too_many_arguments)]
8839    pub fn rms_pre_add_rms_norm(
8840        &self,
8841        a: &CudaSlice<f32>,
8842        wa: &CudaSlice<f32>,
8843        b: &CudaSlice<f32>,
8844        w: &CudaSlice<f32>,
8845        res: &mut CudaSlice<f32>,
8846        dst: &mut CudaSlice<f32>,
8847        ncols: usize,
8848        nrows: usize,
8849        eps: f32,
8850    ) -> Result<(), Box<dyn std::error::Error>> {
8851        let f = self.func("rms_pre_add_rms_norm_f32");
8852        let cfg = LaunchConfig {
8853            grid_dim: (nrows as u32, 1, 1),
8854            block_dim: (rms_block(), 1, 1),
8855            shared_mem_bytes: 0,
8856        };
8857        let (nc, e) = (ncols as i32, eps);
8858        let __s_b2 = self.gpu.stream();
8859        let mut b2 = __s_b2.launch_builder(&f);
8860        b2.arg(a)
8861            .arg(wa)
8862            .arg(b)
8863            .arg(w)
8864            .arg(&mut *res)
8865            .arg(&mut *dst)
8866            .arg(&nc)
8867            .arg(&e);
8868        unsafe {
8869            b2.launch(cfg)?;
8870        }
8871        Ok(())
8872    }
8873
8874    /// wave-2 fold: rms(a,wa) + add + ffn-norm with zsh EMITTED q8_1 (fused2 consumes it).
8875    #[allow(clippy::too_many_arguments)]
8876    pub fn rms_pre_add_rms_norm_q8z(
8877        &self,
8878        a: &CudaSlice<f32>,
8879        wa: &CudaSlice<f32>,
8880        b: &CudaSlice<f32>,
8881        w: &CudaSlice<f32>,
8882        res: &mut CudaSlice<f32>,
8883        dst: &mut CudaSlice<f32>,
8884        ncols: usize,
8885        nrows: usize,
8886        eps: f32,
8887    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8888        debug_assert!(ncols % 128 == 0);
8889        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
8890        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
8891        let (nc, e) = (ncols as i32, eps);
8892        if Self::pdl_on() {
8893            {
8894                use cudarc::driver::{DevicePtr, DevicePtrMut};
8895                let s = &self.gpu.stream();
8896                let (pa, _g0) = a.device_ptr(s);
8897                let (pwa, _g1) = wa.device_ptr(s);
8898                let (pb, _g2) = b.device_ptr(s);
8899                let (pw, _g3) = w.device_ptr(s);
8900                let (pr, _g4) = res.device_ptr_mut(s);
8901                let (pdst, _g5) = dst.device_ptr_mut(s);
8902                let (pq, _g6) = out_q.device_ptr_mut(s);
8903                let (pd, _g7) = out_d.device_ptr_mut(s);
8904                let mut ps = [
8905                    &pa as *const _ as *mut std::ffi::c_void,
8906                    &pwa as *const _ as *mut _,
8907                    &pb as *const _ as *mut _,
8908                    &pw as *const _ as *mut _,
8909                    &pr as *const _ as *mut _,
8910                    &pdst as *const _ as *mut _,
8911                    &pq as *const _ as *mut _,
8912                    &pd as *const _ as *mut _,
8913                    &nc as *const _ as *mut _,
8914                    &e as *const _ as *mut _,
8915                ];
8916                unsafe {
8917                    self.launch_pdl(
8918                        "rms_pre_add_rms_norm_q8z_f32",
8919                        (nrows as u32, 1, 1),
8920                        (rms_block(), 1, 1),
8921                        &mut ps,
8922                    )?;
8923                }
8924            }
8925            return Ok((out_q, out_d));
8926        }
8927        let f = self.func("rms_pre_add_rms_norm_q8z_f32");
8928        let cfg = LaunchConfig {
8929            grid_dim: (nrows as u32, 1, 1),
8930            block_dim: (rms_block(), 1, 1),
8931            shared_mem_bytes: 0,
8932        };
8933        let __s_b2 = self.gpu.stream();
8934        let mut b2 = __s_b2.launch_builder(&f);
8935        b2.arg(a)
8936            .arg(wa)
8937            .arg(b)
8938            .arg(w)
8939            .arg(&mut *res)
8940            .arg(&mut *dst)
8941            .arg(&mut out_q)
8942            .arg(&mut out_d)
8943            .arg(&nc)
8944            .arg(&e);
8945        unsafe {
8946            b2.launch(cfg)?;
8947        }
8948        Ok((out_q, out_d))
8949    }
8950
8951    /// Slot-fed twin of `rms_pre_add_rms_norm_q8z` (gemma4 pn-fold, slotted/graph arm):
8952    /// identical kernel, caller-owned outputs, PLAIN launch only — the dc_slotted capture
8953    /// body must stay attribute-free (the fused2_into precedent).
8954    #[allow(clippy::too_many_arguments)]
8955    pub fn rms_pre_add_rms_norm_q8z_into(
8956        &self,
8957        a: &CudaSlice<f32>,
8958        wa: &CudaSlice<f32>,
8959        b: &CudaSlice<f32>,
8960        w: &CudaSlice<f32>,
8961        res: &mut CudaSlice<f32>,
8962        dst: &mut CudaSlice<f32>,
8963        ncols: usize,
8964        nrows: usize,
8965        eps: f32,
8966        out_q: &mut CudaSlice<i8>,
8967        out_d: &mut CudaSlice<f32>,
8968    ) -> Result<(), Box<dyn std::error::Error>> {
8969        debug_assert!(ncols % 128 == 0);
8970        let (nc, e) = (ncols as i32, eps);
8971        let f = self.func("rms_pre_add_rms_norm_q8z_f32");
8972        let cfg = LaunchConfig {
8973            grid_dim: (nrows as u32, 1, 1),
8974            block_dim: (rms_block(), 1, 1),
8975            shared_mem_bytes: 0,
8976        };
8977        let __s_b = self.gpu.stream();
8978        let mut b2 = __s_b.launch_builder(&f);
8979        b2.arg(a)
8980            .arg(wa)
8981            .arg(b)
8982            .arg(w)
8983            .arg(&mut *res)
8984            .arg(&mut *dst)
8985            .arg(&mut *out_q)
8986            .arg(&mut *out_d)
8987            .arg(&nc)
8988            .arg(&e);
8989        unsafe {
8990            b2.launch(cfg)?;
8991        }
8992        Ok(())
8993    }
8994
8995    /// Slot-fed twin of `rms_pre_add_scale_rms_norm_q8_1` (gemma4 pn-fold exit, slotted
8996    /// arm): identical kernel, caller-owned outputs, PLAIN launch (capture-safe).
8997    #[allow(clippy::too_many_arguments)]
8998    pub fn rms_pre_add_scale_rms_norm_q8_1_into(
8999        &self,
9000        a: &CudaSlice<f32>,
9001        wa: &CudaSlice<f32>,
9002        b_in: &CudaSlice<f32>,
9003        c: f32,
9004        w: &CudaSlice<f32>,
9005        res: &mut CudaSlice<f32>,
9006        ncols: usize,
9007        nrows: usize,
9008        eps: f32,
9009        out_q: &mut CudaSlice<i8>,
9010        out_d: &mut CudaSlice<f32>,
9011    ) -> Result<(), Box<dyn std::error::Error>> {
9012        debug_assert!(ncols % 128 == 0);
9013        let (nc, e2) = (ncols as i32, eps);
9014        let f = self.func("rms_pre_add_scale_rms_norm_q8_1");
9015        let cfg = LaunchConfig {
9016            grid_dim: (nrows as u32, 1, 1),
9017            block_dim: (rms_block(), 1, 1),
9018            shared_mem_bytes: 0,
9019        };
9020        let __s_b = self.gpu.stream();
9021        let mut b2 = __s_b.launch_builder(&f);
9022        b2.arg(a)
9023            .arg(wa)
9024            .arg(b_in)
9025            .arg(&c)
9026            .arg(w)
9027            .arg(&mut *res)
9028            .arg(&mut *out_q)
9029            .arg(&mut *out_d)
9030            .arg(&nc)
9031            .arg(&e2);
9032        unsafe {
9033            b2.launch(cfg)?;
9034        }
9035        Ok(())
9036    }
9037
9038    /// gemma4 pn-fold seam (GAP-DIAGNOSIS verdict 7, the E4B glue backport): the dense
9039    /// decode/verify/slotted trio folds post_attn_norm into the tail entry
9040    /// (rms_pre_add_rms_norm[_q8z]) and post_ffw_norm into the residual exit
9041    /// (rms_pre_add_scale_rms_norm_q8_1). BITS-CHANGING vs the two-launch chain (the
9042    /// single-phase reduction's expansion rounding — E4B receipts); every arm moves
9043    /// together so decode == verify == graph parity holds BY CONSTRUCTION within either
9044    /// seam value. MEMRA_G4_PNFOLD=0 restores the unfused chain everywhere.
9045    pub fn g4_pnfold_on() -> bool {
9046        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9047        *ON.get_or_init(|| {
9048            std::env::var("MEMRA_G4_PNFOLD")
9049                .map(|v| v != "0")
9050                .unwrap_or(true)
9051        })
9052    }
9053
9054    /// wave-4b: OUT-dim concat of three Q4_0 tensors (same in_features; rows are independent
9055    /// blocks, so the concat is a D2D byte concat of the GGUF-layout planes). Returns None
9056    /// off-class (non-Q4_0, mismatched widths, or any tensor already rp-swapped in place).
9057    pub fn build_q4_out_concat3(
9058        &self,
9059        w0: &crate::model::GpuTensor,
9060        w1: &crate::model::GpuTensor,
9061        w2: &crate::model::GpuTensor,
9062    ) -> Result<Option<crate::model::GpuTensor>, Box<dyn std::error::Error>> {
9063        use crate::model::GpuTensor;
9064        let part = |w: &GpuTensor| -> Option<(usize, usize)> {
9065            match w {
9066                GpuTensor::Quant {
9067                    qtype,
9068                    row_bytes,
9069                    rp,
9070                    ..
9071                } if *qtype == QT_Q4_0 && !*rp => Some((*row_bytes, w.out_features())),
9072                _ => None,
9073            }
9074        };
9075        let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (part(w0), part(w1), part(w2))
9076        else {
9077            return Ok(None);
9078        };
9079        if rb0 != rb1
9080            || rb0 != rb2
9081            || w0.in_features() != w1.in_features()
9082            || w0.in_features() != w2.in_features()
9083        {
9084            return Ok(None);
9085        }
9086        fn bytes_of(w: &crate::model::GpuTensor) -> &CudaSlice<u8> {
9087            match w {
9088                crate::model::GpuTensor::Quant { bytes, .. } => bytes,
9089                _ => unreachable!(),
9090            }
9091        }
9092        let (b0, b1, b2) = (bytes_of(w0), bytes_of(w1), bytes_of(w2));
9093        let total = rb0 * (o0 + o1 + o2);
9094        let mut cat = self.alloc_u8(total)?;
9095        self.copy_u8_into(&mut cat, 0, b0, rb0 * o0)?;
9096        self.copy_u8_into(&mut cat, rb0 * o0, b1, rb1 * o1)?;
9097        self.copy_u8_into(&mut cat, rb0 * (o0 + o1), b2, rb2 * o2)?;
9098        Ok(Some(GpuTensor::Quant {
9099            bytes: cat,
9100            qtype: QT_Q4_0,
9101            row_bytes: rb0,
9102            ne: vec![w0.in_features() as u64, (o0 + o1 + o2) as u64],
9103            scale: 1.0,
9104            rp: false,
9105            #[cfg(memra_cutlass)]
9106            cutlass: None,
9107            fp8: None,
9108            blk: None,
9109            rp4: None,
9110            f16: None,
9111        }))
9112    }
9113
9114    /// FULL-WIDTH-ROPE CONTRACT for the fused rms_norm+qkv+rope kernels
9115    /// (lane/graph-s-key-exactness-20260819, probe O-10 follow-up).
9116    ///
9117    /// `rms_norm_qkv_rope_f32` / `_cat_f32` (`cu/kernels.cu`) and
9118    /// `rms_norm_qkv_rope_append_body` (`cu/flash_attn.cu`) take NO `n_dims`/`n_rot` argument.
9119    /// They compute `int half = ncols / 2` and rotate the FULL head width by construction — the
9120    /// standalone `rope_neox*` kernels take `n_dims` and early-return above it, these do not.
9121    ///
9122    /// Every call site today is a gemma-4 arm, where `n_rot == head_dim`, so nothing is wrong
9123    /// now. But this is the n_rot bug ONE FUSION away: fuse a partial-rotary arch onto these
9124    /// kernels (qwen3.5 = 64 rotary dims of a 256-wide head; step35 full-attn = 64 of 128) and
9125    /// 192 dims that must pass through unrotated get rotated silently — no error, no NaN, just a
9126    /// wrong model. The n_rot lane already paid for that class once, in the config derivation.
9127    ///
9128    /// So the fusions now take the layer's DERIVED rope width and refuse anything but full
9129    /// width. A future partial-rotary caller fails at its first launch with the geometry named
9130    /// instead of serving quietly wrong logits.
9131    fn full_width_rope_only(
9132        kernel: &str,
9133        n_rot: usize,
9134        head_dim: usize,
9135    ) -> Result<(), Box<dyn std::error::Error>> {
9136        if n_rot == head_dim {
9137            return Ok(());
9138        }
9139        Err(format!(
9140            "{kernel}: PARTIAL ROTARY REFUSED — n_rot {n_rot} != head_dim {head_dim}. This fused \
9141             rms_norm+qkv+rope kernel carries no n_dims parameter and rotates the full head \
9142             width (half = ncols/2), so it would rotate dims {n_rot}..{head_dim} that must pass \
9143             through unrotated. Use the split path (rms_norm_qkv + rope_neox/rope_neox2 with \
9144             n_dims={n_rot}), or add an n_dims early-return to the kernel and widen this guard."
9145        )
9146        .into())
9147    }
9148
9149    /// wave-4b: the qkv-cat twin — one contiguous [rq+2*rk, hd] input from the concat matvec.
9150    /// `n_rot` is the layer's derived rotary width and MUST equal `head_dim`
9151    /// ([`Engine::full_width_rope_only`]).
9152    #[allow(clippy::too_many_arguments)]
9153    pub fn rms_norm_qkv_rope_cat(
9154        &self,
9155        qkv: &CudaSlice<f32>,
9156        wq: &CudaSlice<f32>,
9157        wk: &CudaSlice<f32>,
9158        wv: &CudaSlice<f32>,
9159        q: &mut CudaSlice<f32>,
9160        k: &mut CudaSlice<f32>,
9161        v: &mut CudaSlice<f32>,
9162        head_dim: usize,
9163        n_rot: usize,
9164        rq: usize,
9165        rk: usize,
9166        pos: &CudaSlice<i32>,
9167        nh_q: usize,
9168        nh_k: usize,
9169        base: f32,
9170        freq_scale: f32,
9171        ff: Option<&CudaSlice<f32>>,
9172        eps: f32,
9173    ) -> Result<(), Box<dyn std::error::Error>> {
9174        Self::full_width_rope_only("rms_norm_qkv_rope_cat", n_rot, head_dim)?;
9175        let rows = rq + rk + rk;
9176        let theta_scale = base.powf(-2.0 / head_dim as f32);
9177        let (nc, rqi, rki, nhq, nhk) = (
9178            head_dim as i32,
9179            rq as i32,
9180            rk as i32,
9181            nh_q as i32,
9182            nh_k as i32,
9183        );
9184        if Self::pdl_on() {
9185            use cudarc::driver::{DevicePtr, DevicePtrMut};
9186            let s = &self.gpu.stream();
9187            let (pqkv, _g0) = qkv.device_ptr(s);
9188            let (pwq, _g1) = wq.device_ptr(s);
9189            let (pwk, _g2) = wk.device_ptr(s);
9190            let (pwv, _g3) = wv.device_ptr(s);
9191            let (pq, _g4) = q.device_ptr_mut(s);
9192            let (pk, _g5) = k.device_ptr_mut(s);
9193            let (pv, _g6) = v.device_ptr_mut(s);
9194            let (ppos, _g7) = pos.device_ptr(s);
9195            let (pff, _g8) = match ff {
9196                Some(t) => {
9197                    let (p, g) = t.device_ptr(s);
9198                    (p, Some(g))
9199                }
9200                None => (0, None),
9201            };
9202            let mut ps = [
9203                &pqkv as *const _ as *mut std::ffi::c_void,
9204                &pwq as *const _ as *mut _,
9205                &pwk as *const _ as *mut _,
9206                &pwv as *const _ as *mut _,
9207                &pq as *const _ as *mut _,
9208                &pk as *const _ as *mut _,
9209                &pv as *const _ as *mut _,
9210                &nc as *const _ as *mut _,
9211                &rqi as *const _ as *mut _,
9212                &rki as *const _ as *mut _,
9213                &ppos as *const _ as *mut _,
9214                &nhq as *const _ as *mut _,
9215                &nhk as *const _ as *mut _,
9216                &theta_scale as *const _ as *mut _,
9217                &freq_scale as *const _ as *mut _,
9218                &pff as *const _ as *mut _,
9219                &eps as *const _ as *mut _,
9220            ];
9221            unsafe {
9222                self.launch_pdl(
9223                    "rms_norm_qkv_rope_cat_f32",
9224                    (rows as u32, 1, 1),
9225                    (rms_block(), 1, 1),
9226                    &mut ps,
9227                )?;
9228            }
9229            return Ok(());
9230        }
9231        let f = self.func("rms_norm_qkv_rope_cat_f32");
9232        let cfg = LaunchConfig {
9233            grid_dim: (rows as u32, 1, 1),
9234            block_dim: (rms_block(), 1, 1),
9235            shared_mem_bytes: 0,
9236        };
9237        let __s_b = self.gpu.stream();
9238        let mut b = __s_b.launch_builder(&f);
9239        match ff {
9240            Some(t) => {
9241                b.arg(qkv)
9242                    .arg(wq)
9243                    .arg(wk)
9244                    .arg(wv)
9245                    .arg(&mut *q)
9246                    .arg(&mut *k)
9247                    .arg(&mut *v)
9248                    .arg(&nc)
9249                    .arg(&rqi)
9250                    .arg(&rki)
9251                    .arg(pos)
9252                    .arg(&nhq)
9253                    .arg(&nhk)
9254                    .arg(&theta_scale)
9255                    .arg(&freq_scale)
9256                    .arg(t)
9257                    .arg(&eps);
9258                unsafe {
9259                    b.launch(cfg)?;
9260                }
9261            }
9262            None => {
9263                let null: u64 = 0;
9264                b.arg(qkv)
9265                    .arg(wq)
9266                    .arg(wk)
9267                    .arg(wv)
9268                    .arg(&mut *q)
9269                    .arg(&mut *k)
9270                    .arg(&mut *v)
9271                    .arg(&nc)
9272                    .arg(&rqi)
9273                    .arg(&rki)
9274                    .arg(pos)
9275                    .arg(&nhq)
9276                    .arg(&nhk)
9277                    .arg(&theta_scale)
9278                    .arg(&freq_scale)
9279                    .arg(&null)
9280                    .arg(&eps);
9281                unsafe {
9282                    b.launch(cfg)?;
9283                }
9284            }
9285        }
9286        Ok(())
9287    }
9288
9289    /// wave-3 fold: rms_norm_qkv + rope_neox2 in ONE launch (n_dims == head_dim; ff nullable).
9290    /// `n_rot` is the layer's derived rotary width and MUST equal `head_dim`
9291    /// ([`Engine::full_width_rope_only`]).
9292    #[allow(clippy::too_many_arguments)]
9293    pub fn rms_norm_qkv_rope(
9294        &self,
9295        q0: &CudaSlice<f32>,
9296        k0: &CudaSlice<f32>,
9297        v0: &CudaSlice<f32>,
9298        wq: &CudaSlice<f32>,
9299        wk: &CudaSlice<f32>,
9300        wv: &CudaSlice<f32>,
9301        q: &mut CudaSlice<f32>,
9302        k: &mut CudaSlice<f32>,
9303        v: &mut CudaSlice<f32>,
9304        head_dim: usize,
9305        n_rot: usize,
9306        rq: usize,
9307        rk: usize,
9308        pos: &CudaSlice<i32>,
9309        nh_q: usize,
9310        nh_k: usize,
9311        base: f32,
9312        freq_scale: f32,
9313        ff: Option<&CudaSlice<f32>>,
9314        eps: f32,
9315    ) -> Result<(), Box<dyn std::error::Error>> {
9316        Self::full_width_rope_only("rms_norm_qkv_rope", n_rot, head_dim)?;
9317        let f = self.func("rms_norm_qkv_rope_f32");
9318        let rows = rq + rk + rk; // q rows + k rows + v rows (rk == rv)
9319        let cfg = LaunchConfig {
9320            grid_dim: (rows as u32, 1, 1),
9321            block_dim: (rms_block(), 1, 1),
9322            shared_mem_bytes: 0,
9323        };
9324        let theta_scale = base.powf(-2.0 / head_dim as f32);
9325        let (nc, rqi, rki, nhq, nhk) = (
9326            head_dim as i32,
9327            rq as i32,
9328            rk as i32,
9329            nh_q as i32,
9330            nh_k as i32,
9331        );
9332        let __s_b = self.gpu.stream();
9333        let mut b = __s_b.launch_builder(&f);
9334        match ff {
9335            Some(t) => {
9336                b.arg(q0)
9337                    .arg(k0)
9338                    .arg(v0)
9339                    .arg(wq)
9340                    .arg(wk)
9341                    .arg(wv)
9342                    .arg(&mut *q)
9343                    .arg(&mut *k)
9344                    .arg(&mut *v)
9345                    .arg(&nc)
9346                    .arg(&rqi)
9347                    .arg(&rki)
9348                    .arg(pos)
9349                    .arg(&nhq)
9350                    .arg(&nhk)
9351                    .arg(&theta_scale)
9352                    .arg(&freq_scale)
9353                    .arg(t)
9354                    .arg(&eps);
9355                unsafe {
9356                    b.launch(cfg)?;
9357                }
9358            }
9359            None => {
9360                let null: u64 = 0;
9361                b.arg(q0)
9362                    .arg(k0)
9363                    .arg(v0)
9364                    .arg(wq)
9365                    .arg(wk)
9366                    .arg(wv)
9367                    .arg(&mut *q)
9368                    .arg(&mut *k)
9369                    .arg(&mut *v)
9370                    .arg(&nc)
9371                    .arg(&rqi)
9372                    .arg(&rki)
9373                    .arg(pos)
9374                    .arg(&nhq)
9375                    .arg(&nhk)
9376                    .arg(&theta_scale)
9377                    .arg(&freq_scale)
9378                    .arg(&null)
9379                    .arg(&eps);
9380                unsafe {
9381                    b.launch(cfg)?;
9382                }
9383            }
9384        }
9385        Ok(())
9386    }
9387
9388    /// FUSED norm+rope+APPEND (m=1 decode, 2026-07-23): one launch replaces the
9389    /// rms_norm_qkv_rope + append_kv_quantized_dc pair. Kernel lives in the flash fatbins
9390    /// (format-flavored quant tail) — `g` must mirror the append path's flavor exactly.
9391    /// `n_rot` is the layer's derived rotary width and MUST equal `head_dim`
9392    /// ([`Engine::full_width_rope_only`]).
9393    #[allow(clippy::too_many_arguments)]
9394    pub fn rms_norm_qkv_rope_append_dc(
9395        &self,
9396        q0: &CudaSlice<f32>,
9397        k0: &CudaSlice<f32>,
9398        v0: &CudaSlice<f32>,
9399        wq: &CudaSlice<f32>,
9400        wk: &CudaSlice<f32>,
9401        wv: &CudaSlice<f32>,
9402        q: &mut CudaSlice<f32>,
9403        k: &mut CudaSlice<f32>,
9404        v: &mut CudaSlice<f32>,
9405        head_dim: usize,
9406        n_rot: usize,
9407        rq: usize,
9408        rk: usize,
9409        pos: &CudaSlice<i32>,
9410        nh_q: usize,
9411        nh_k: usize,
9412        base: f32,
9413        freq_scale: f32,
9414        ff: Option<&CudaSlice<f32>>,
9415        eps: f32,
9416        kc: &mut CudaSlice<u8>,
9417        vc: &mut CudaSlice<u8>,
9418        t_dev: &CudaSlice<i32>,
9419        k_tok_bytes: usize,
9420        v_tok_bytes: usize,
9421        g: bool,
9422    ) -> Result<(), Box<dyn std::error::Error>> {
9423        Self::full_width_rope_only("rms_norm_qkv_rope_append_dc", n_rot, head_dim)?;
9424        let rows = rq + rk + rk;
9425        let theta_scale = base.powf(-2.0 / head_dim as f32);
9426        let (nc, rqi, rki, nhq, nhk) = (
9427            head_dim as i32,
9428            rq as i32,
9429            rk as i32,
9430            nh_q as i32,
9431            nh_k as i32,
9432        );
9433        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
9434        if Self::pdl_on() && Self::pdl_wb_on() {
9435            use cudarc::driver::{DevicePtr, DevicePtrMut};
9436            let s = &self.gpu.stream();
9437            let (p0, _a0) = q0.device_ptr(s);
9438            let (p1, _a1) = k0.device_ptr(s);
9439            let (p2, _a2) = v0.device_ptr(s);
9440            let (pwq, _a3) = wq.device_ptr(s);
9441            let (pwk, _a4) = wk.device_ptr(s);
9442            let (pwv, _a5) = wv.device_ptr(s);
9443            let (pq, _a6) = q.device_ptr_mut(s);
9444            let (pk, _a7) = k.device_ptr_mut(s);
9445            let (pv, _a8) = v.device_ptr_mut(s);
9446            let (pp, _a9) = pos.device_ptr(s);
9447            let pff: u64 = match ff {
9448                Some(t) => {
9449                    let (p, _gg) = t.device_ptr(s);
9450                    p as u64
9451                }
9452                None => 0,
9453            };
9454            let (pkc, _a10) = kc.device_ptr_mut(s);
9455            let (pvc, _a11) = vc.device_ptr_mut(s);
9456            let (pt, _a12) = t_dev.device_ptr(s);
9457            let mut ps = [
9458                &p0 as *const _ as *mut std::ffi::c_void,
9459                &p1 as *const _ as *mut _,
9460                &p2 as *const _ as *mut _,
9461                &pwq as *const _ as *mut _,
9462                &pwk as *const _ as *mut _,
9463                &pwv as *const _ as *mut _,
9464                &pq as *const _ as *mut _,
9465                &pk as *const _ as *mut _,
9466                &pv as *const _ as *mut _,
9467                &nc as *const _ as *mut _,
9468                &rqi as *const _ as *mut _,
9469                &rki as *const _ as *mut _,
9470                &pp as *const _ as *mut _,
9471                &nhq as *const _ as *mut _,
9472                &nhk as *const _ as *mut _,
9473                &theta_scale as *const _ as *mut _,
9474                &freq_scale as *const _ as *mut _,
9475                &pff as *const _ as *mut _,
9476                &eps as *const _ as *mut _,
9477                &pkc as *const _ as *mut _,
9478                &pvc as *const _ as *mut _,
9479                &pt as *const _ as *mut _,
9480                &ktb as *const _ as *mut _,
9481                &vtb as *const _ as *mut _,
9482            ];
9483            unsafe {
9484                self.launch_pdl_flash(
9485                    g,
9486                    "rms_norm_qkv_rope_append_dc_f32",
9487                    (rows as u32, 1, 1),
9488                    (rms_block(), 1, 1),
9489                    0,
9490                    &mut ps,
9491                )?;
9492            }
9493            return Ok(());
9494        }
9495        let f = if g {
9496            self.func_g("rms_norm_qkv_rope_append_dc_f32")
9497        } else {
9498            self.func("rms_norm_qkv_rope_append_dc_f32")
9499        };
9500        let cfg = LaunchConfig {
9501            grid_dim: (rows as u32, 1, 1),
9502            block_dim: (rms_block(), 1, 1),
9503            shared_mem_bytes: 0,
9504        };
9505        let __s_b = self.gpu.stream();
9506        let mut b = __s_b.launch_builder(&f);
9507        match ff {
9508            Some(t) => {
9509                b.arg(q0)
9510                    .arg(k0)
9511                    .arg(v0)
9512                    .arg(wq)
9513                    .arg(wk)
9514                    .arg(wv)
9515                    .arg(&mut *q)
9516                    .arg(&mut *k)
9517                    .arg(&mut *v)
9518                    .arg(&nc)
9519                    .arg(&rqi)
9520                    .arg(&rki)
9521                    .arg(pos)
9522                    .arg(&nhq)
9523                    .arg(&nhk)
9524                    .arg(&theta_scale)
9525                    .arg(&freq_scale)
9526                    .arg(t)
9527                    .arg(&eps)
9528                    .arg(&mut *kc)
9529                    .arg(&mut *vc)
9530                    .arg(t_dev)
9531                    .arg(&ktb)
9532                    .arg(&vtb);
9533                unsafe {
9534                    b.launch(cfg)?;
9535                }
9536            }
9537            None => {
9538                let null: u64 = 0;
9539                b.arg(q0)
9540                    .arg(k0)
9541                    .arg(v0)
9542                    .arg(wq)
9543                    .arg(wk)
9544                    .arg(wv)
9545                    .arg(&mut *q)
9546                    .arg(&mut *k)
9547                    .arg(&mut *v)
9548                    .arg(&nc)
9549                    .arg(&rqi)
9550                    .arg(&rki)
9551                    .arg(pos)
9552                    .arg(&nhq)
9553                    .arg(&nhk)
9554                    .arg(&theta_scale)
9555                    .arg(&freq_scale)
9556                    .arg(&null)
9557                    .arg(&eps)
9558                    .arg(&mut *kc)
9559                    .arg(&mut *vc)
9560                    .arg(t_dev)
9561                    .arg(&ktb)
9562                    .arg(&vtb);
9563                unsafe {
9564                    b.launch(cfg)?;
9565                }
9566            }
9567        }
9568        Ok(())
9569    }
9570
9571    /// wave-2 fold: a + b with the sum emitted q8_1 alongside f32.
9572    /// Host-len twin of `rms_norm_qkv_rope_append_dc` (zoo-fusion arc): the eager decode
9573    /// arm tracks kv length on host (`kvl.len`), so the append slot rides the launch arg
9574    /// instead of the device counter. Kernel body is SHARED with the _dc entry (one
9575    /// inlined body — bit-identical to the rms_norm_qkv_rope + append pair it replaces,
9576    /// same law as the dc fold). `n_rot` is the layer's derived rotary width and MUST equal
9577    /// `head_dim` ([`Engine::full_width_rope_only`]).
9578    #[allow(clippy::too_many_arguments)]
9579    pub fn rms_norm_qkv_rope_append(
9580        &self,
9581        q0: &CudaSlice<f32>,
9582        k0: &CudaSlice<f32>,
9583        v0: &CudaSlice<f32>,
9584        wq: &CudaSlice<f32>,
9585        wk: &CudaSlice<f32>,
9586        wv: &CudaSlice<f32>,
9587        q: &mut CudaSlice<f32>,
9588        k: &mut CudaSlice<f32>,
9589        v: &mut CudaSlice<f32>,
9590        head_dim: usize,
9591        n_rot: usize,
9592        rq: usize,
9593        rk: usize,
9594        pos: &CudaSlice<i32>,
9595        nh_q: usize,
9596        nh_k: usize,
9597        base: f32,
9598        freq_scale: f32,
9599        ff: Option<&CudaSlice<f32>>,
9600        eps: f32,
9601        kc: &mut CudaSlice<u8>,
9602        vc: &mut CudaSlice<u8>,
9603        t: usize,
9604        k_tok_bytes: usize,
9605        v_tok_bytes: usize,
9606        g: bool,
9607    ) -> Result<(), Box<dyn std::error::Error>> {
9608        Self::full_width_rope_only("rms_norm_qkv_rope_append", n_rot, head_dim)?;
9609        let rows = rq + rk + rk;
9610        let theta_scale = base.powf(-2.0 / head_dim as f32);
9611        let (nc, rqi, rki, nhq, nhk) = (
9612            head_dim as i32,
9613            rq as i32,
9614            rk as i32,
9615            nh_q as i32,
9616            nh_k as i32,
9617        );
9618        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
9619        let ti = t as i32;
9620        if Self::pdl_on() && Self::pdl_wb_on() {
9621            use cudarc::driver::{DevicePtr, DevicePtrMut};
9622            let s = &self.gpu.stream();
9623            let (p0, _a0) = q0.device_ptr(s);
9624            let (p1, _a1) = k0.device_ptr(s);
9625            let (p2, _a2) = v0.device_ptr(s);
9626            let (pwq, _a3) = wq.device_ptr(s);
9627            let (pwk, _a4) = wk.device_ptr(s);
9628            let (pwv, _a5) = wv.device_ptr(s);
9629            let (pq, _a6) = q.device_ptr_mut(s);
9630            let (pk, _a7) = k.device_ptr_mut(s);
9631            let (pv, _a8) = v.device_ptr_mut(s);
9632            let (pp, _a9) = pos.device_ptr(s);
9633            let pff: u64 = match ff {
9634                Some(t) => {
9635                    let (p, _gg) = t.device_ptr(s);
9636                    p as u64
9637                }
9638                None => 0,
9639            };
9640            let (pkc, _a10) = kc.device_ptr_mut(s);
9641            let (pvc, _a11) = vc.device_ptr_mut(s);
9642            let mut ps = [
9643                &p0 as *const _ as *mut std::ffi::c_void,
9644                &p1 as *const _ as *mut _,
9645                &p2 as *const _ as *mut _,
9646                &pwq as *const _ as *mut _,
9647                &pwk as *const _ as *mut _,
9648                &pwv as *const _ as *mut _,
9649                &pq as *const _ as *mut _,
9650                &pk as *const _ as *mut _,
9651                &pv as *const _ as *mut _,
9652                &nc as *const _ as *mut _,
9653                &rqi as *const _ as *mut _,
9654                &rki as *const _ as *mut _,
9655                &pp as *const _ as *mut _,
9656                &nhq as *const _ as *mut _,
9657                &nhk as *const _ as *mut _,
9658                &theta_scale as *const _ as *mut _,
9659                &freq_scale as *const _ as *mut _,
9660                &pff as *const _ as *mut _,
9661                &eps as *const _ as *mut _,
9662                &pkc as *const _ as *mut _,
9663                &pvc as *const _ as *mut _,
9664                &ti as *const _ as *mut _,
9665                &ktb as *const _ as *mut _,
9666                &vtb as *const _ as *mut _,
9667            ];
9668            unsafe {
9669                self.launch_pdl_flash(
9670                    g,
9671                    "rms_norm_qkv_rope_append_f32",
9672                    (rows as u32, 1, 1),
9673                    (rms_block(), 1, 1),
9674                    0,
9675                    &mut ps,
9676                )?;
9677            }
9678            return Ok(());
9679        }
9680        let f = if g {
9681            self.func_g("rms_norm_qkv_rope_append_f32")
9682        } else {
9683            self.func("rms_norm_qkv_rope_append_f32")
9684        };
9685        let cfg = LaunchConfig {
9686            grid_dim: (rows as u32, 1, 1),
9687            block_dim: (rms_block(), 1, 1),
9688            shared_mem_bytes: 0,
9689        };
9690        let __s_b = self.gpu.stream();
9691        let mut b = __s_b.launch_builder(&f);
9692        let null: u64 = 0;
9693        b.arg(q0)
9694            .arg(k0)
9695            .arg(v0)
9696            .arg(wq)
9697            .arg(wk)
9698            .arg(wv)
9699            .arg(&mut *q)
9700            .arg(&mut *k)
9701            .arg(&mut *v)
9702            .arg(&nc)
9703            .arg(&rqi)
9704            .arg(&rki)
9705            .arg(pos)
9706            .arg(&nhq)
9707            .arg(&nhk)
9708            .arg(&theta_scale)
9709            .arg(&freq_scale);
9710        match ff {
9711            Some(t) => {
9712                b.arg(t);
9713            }
9714            None => {
9715                b.arg(&null);
9716            }
9717        }
9718        b.arg(&eps)
9719            .arg(&mut *kc)
9720            .arg(&mut *vc)
9721            .arg(&ti)
9722            .arg(&ktb)
9723            .arg(&vtb);
9724        unsafe {
9725            b.launch(cfg)?;
9726        }
9727        Ok(())
9728    }
9729
9730    pub fn add_q8_1(
9731        &self,
9732        a: &CudaSlice<f32>,
9733        b: &CudaSlice<f32>,
9734        res: &mut CudaSlice<f32>,
9735        ncols: usize,
9736        nrows: usize,
9737    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9738        debug_assert!(ncols % 128 == 0);
9739        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
9740        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9741        let f = self.func("add_q8_1_f32");
9742        let cfg = LaunchConfig {
9743            grid_dim: (nrows as u32, 1, 1),
9744            block_dim: (rms_block(), 1, 1),
9745            shared_mem_bytes: 0,
9746        };
9747        let nc = ncols as i32;
9748        let __s_b2 = self.gpu.stream();
9749        let mut b2 = __s_b2.launch_builder(&f);
9750        b2.arg(a)
9751            .arg(b)
9752            .arg(&mut *res)
9753            .arg(&mut out_q)
9754            .arg(&mut out_d)
9755            .arg(&nc);
9756        unsafe {
9757            b2.launch(cfg)?;
9758        }
9759        Ok((out_q, out_d))
9760    }
9761
9762    /// E4B FFN-tail exit fusion (glue wave 5): resid = b + rms(a, wa) emitted f32 + q8_1 pair
9763    /// in ONE launch — replaces rms_norm(a,wa->sn) + add_q8_1(sn,b). Same rms_block() config
9764    /// as both parents (bit-identity: identical reduction + quad-walk quantize).
9765    pub fn rms_pre_add_q8_1(
9766        &self,
9767        a: &CudaSlice<f32>,
9768        wa: &CudaSlice<f32>,
9769        b: &CudaSlice<f32>,
9770        res: &mut CudaSlice<f32>,
9771        ncols: usize,
9772        nrows: usize,
9773        eps: f32,
9774    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9775        debug_assert!(ncols % 128 == 0);
9776        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
9777        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9778        let f = self.func("rms_pre_add_q8_1_f32");
9779        let cfg = LaunchConfig {
9780            grid_dim: (nrows as u32, 1, 1),
9781            block_dim: (rms_block(), 1, 1),
9782            shared_mem_bytes: 0,
9783        };
9784        let (nc, ep) = (ncols as i32, eps);
9785        let __s_b2 = self.gpu.stream();
9786        let mut b2 = __s_b2.launch_builder(&f);
9787        b2.arg(a)
9788            .arg(wa)
9789            .arg(b)
9790            .arg(&mut *res)
9791            .arg(&mut out_q)
9792            .arg(&mut out_d)
9793            .arg(&nc)
9794            .arg(&ep);
9795        unsafe {
9796            b2.launch(cfg)?;
9797        }
9798        Ok((out_q, out_d))
9799    }
9800
9801    /// L2 norm per row (head_dim), no weight.
9802    /// PREFILL l2 dispatch (round 27): the warp-per-row float4 v2 when the numeric-config
9803    /// seam allows (MEMRA_L2_V2, default ON, d_state==128 only); else the strided kernel.
9804    pub fn l2_v2_on(ncols: usize) -> bool {
9805        ncols == 128 && std::env::var("MEMRA_L2_V2").as_deref() != Ok("0")
9806    }
9807
9808    pub fn l2_norm_pp(
9809        &self,
9810        x: &CudaSlice<f32>,
9811        dst: &mut CudaSlice<f32>,
9812        dst16: Option<&mut CudaSlice<u8>>,
9813        ncols: usize,
9814        nrows: usize,
9815        eps: f32,
9816    ) -> Result<(), Box<dyn std::error::Error>> {
9817        if Self::l2_v2_on(ncols) {
9818            let f = self.func("l2_norm_pp_v2_f32");
9819            let rows_per_block = 8u32; // 256 threads = 8 warps = 8 rows
9820            let cfg = LaunchConfig {
9821                grid_dim: ((nrows as u32).div_ceil(rows_per_block), 1, 1),
9822                block_dim: (256, 1, 1),
9823                shared_mem_bytes: 0,
9824            };
9825            let (nc, nr, e) = (ncols as i32, nrows as i32, eps);
9826            // mirror-fold: bf16 twin address by value (0 = skip; matches the nullable param)
9827            let d16: u64 = match dst16 {
9828                Some(d) => self.addr_u8(d),
9829                None => 0,
9830            };
9831            let __s_b = self.gpu.stream();
9832            let mut b = __s_b.launch_builder(&f);
9833            b.arg(x).arg(dst).arg(&d16).arg(&nc).arg(&nr).arg(&e);
9834            unsafe {
9835                b.launch(cfg)?;
9836            }
9837            return Ok(());
9838        }
9839        self.l2_norm(x, dst, ncols, nrows, eps)
9840    }
9841
9842    pub fn l2_norm(
9843        &self,
9844        x: &CudaSlice<f32>,
9845        dst: &mut CudaSlice<f32>,
9846        ncols: usize,
9847        nrows: usize,
9848        eps: f32,
9849    ) -> Result<(), Box<dyn std::error::Error>> {
9850        let f = self.func("l2_norm_f32");
9851        let cfg = LaunchConfig {
9852            grid_dim: (nrows as u32, 1, 1),
9853            block_dim: (256, 1, 1),
9854            shared_mem_bytes: 0,
9855        };
9856        let (nc, e) = (ncols as i32, eps);
9857        let __s_b = self.gpu.stream();
9858        let mut b = __s_b.launch_builder(&f);
9859        b.arg(x).arg(dst).arg(&nc).arg(&e);
9860        unsafe {
9861            b.launch(cfg)?;
9862        }
9863        Ok(())
9864    }
9865
9866    /// L2-norm with blockDim=32 (warp-tree reduction) — BIT-IDENTICAL to gdn_prep_decode_f32's
9867    /// per-warp L2 norm. The verify path MUST use this to match decode's FP accumulation order:
9868    /// l2_norm at blockDim=256 produces a different shfl-tree reduction of the 128-element
9869    /// squared-sum (pairwise tree vs serial-4-then-warp-tree), causing ULP differences that
9870    /// propagate through gdn_scan and flip argmax on marginal logits.
9871    pub fn l2_norm_decode(
9872        &self,
9873        x: &CudaSlice<f32>,
9874        dst: &mut CudaSlice<f32>,
9875        ncols: usize,
9876        nrows: usize,
9877        eps: f32,
9878    ) -> Result<(), Box<dyn std::error::Error>> {
9879        let f = self.func("l2_norm_f32");
9880        let cfg = LaunchConfig {
9881            grid_dim: (nrows as u32, 1, 1),
9882            block_dim: (32, 1, 1),
9883            shared_mem_bytes: 0,
9884        };
9885        let (nc, e) = (ncols as i32, eps);
9886        let __s_b = self.gpu.stream();
9887        let mut b = __s_b.launch_builder(&f);
9888        b.arg(x).arg(dst).arg(&nc).arg(&e);
9889        unsafe {
9890            b.launch(cfg)?;
9891        }
9892        Ok(())
9893    }
9894
9895    /// RoPE NEOX in-place. x:[head_dim, n_heads, n_tokens], pos:[n_tokens].
9896    pub fn rope_neox(
9897        &self,
9898        x: &mut CudaSlice<f32>,
9899        pos: &CudaSlice<i32>,
9900        head_dim: usize,
9901        n_dims: usize,
9902        n_heads: usize,
9903        n_tokens: usize,
9904        freq_base: f32,
9905        freq_scale: f32,
9906    ) -> Result<(), Box<dyn std::error::Error>> {
9907        let f = self.func("rope_neox_f32");
9908        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
9909        let grid = (n_heads * n_tokens) as u32;
9910        let cfg = LaunchConfig {
9911            grid_dim: (grid, 1, 1),
9912            block_dim: ((head_dim / 2) as u32, 1, 1),
9913            shared_mem_bytes: 0,
9914        };
9915        let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
9916        let __s_b = self.gpu.stream();
9917        let mut b = __s_b.launch_builder(&f);
9918        b.arg(x)
9919            .arg(pos)
9920            .arg(&hd)
9921            .arg(&nd)
9922            .arg(&nh)
9923            .arg(&theta_scale)
9924            .arg(&freq_scale);
9925        unsafe {
9926            b.launch(cfg)?;
9927        }
9928        Ok(())
9929    }
9930
9931    /// RoPE NEOX with per-dim freq factors (gemma4 global layers, rope_freqs.weight [n_dims/2]).
9932    pub fn rope_neox_ff(
9933        &self,
9934        x: &mut CudaSlice<f32>,
9935        pos: &CudaSlice<i32>,
9936        head_dim: usize,
9937        n_dims: usize,
9938        n_heads: usize,
9939        n_tokens: usize,
9940        freq_base: f32,
9941        freq_scale: f32,
9942        ff: &CudaSlice<f32>,
9943    ) -> Result<(), Box<dyn std::error::Error>> {
9944        let f = self.func("rope_neox_ff_f32");
9945        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
9946        let grid = (n_heads * n_tokens) as u32;
9947        let cfg = LaunchConfig {
9948            grid_dim: (grid, 1, 1),
9949            block_dim: ((head_dim / 2) as u32, 1, 1),
9950            shared_mem_bytes: 0,
9951        };
9952        let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
9953        let __s_b = self.gpu.stream();
9954        let mut b = __s_b.launch_builder(&f);
9955        b.arg(x)
9956            .arg(pos)
9957            .arg(&hd)
9958            .arg(&nd)
9959            .arg(&nh)
9960            .arg(&theta_scale)
9961            .arg(&freq_scale)
9962            .arg(ff);
9963        unsafe {
9964            b.launch(cfg)?;
9965        }
9966        Ok(())
9967    }
9968
9969    /// gemma4: rope q and k in one launch (per-row chain = rope_neox / rope_neox_ff verbatim).
9970    #[allow(clippy::too_many_arguments)]
9971    pub fn rope_neox2(
9972        &self,
9973        q: &mut CudaSlice<f32>,
9974        k: &mut CudaSlice<f32>,
9975        pos: &CudaSlice<i32>,
9976        head_dim: usize,
9977        n_dims: usize,
9978        nh_q: usize,
9979        nh_k: usize,
9980        n_tokens: usize,
9981        freq_base: f32,
9982        freq_scale: f32,
9983        ff: Option<&CudaSlice<f32>>,
9984    ) -> Result<(), Box<dyn std::error::Error>> {
9985        let f = self.func("rope_neox2_f32");
9986        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
9987        let grid = ((nh_q + nh_k) * n_tokens) as u32;
9988        let cfg = LaunchConfig {
9989            grid_dim: (grid, 1, 1),
9990            block_dim: ((head_dim / 2) as u32, 1, 1),
9991            shared_mem_bytes: 0,
9992        };
9993        let (hd, nd, nq, nk, nt) = (
9994            head_dim as i32,
9995            n_dims as i32,
9996            nh_q as i32,
9997            nh_k as i32,
9998            n_tokens as i32,
9999        );
10000        let __s_b = self.gpu.stream();
10001        let mut b = __s_b.launch_builder(&f);
10002        b.arg(q)
10003            .arg(k)
10004            .arg(pos)
10005            .arg(&hd)
10006            .arg(&nd)
10007            .arg(&nq)
10008            .arg(&nk)
10009            .arg(&nt)
10010            .arg(&theta_scale)
10011            .arg(&freq_scale);
10012        match ff {
10013            Some(ffv) => {
10014                b.arg(ffv);
10015                unsafe {
10016                    b.launch(cfg)?;
10017                }
10018            }
10019            None => {
10020                let null: u64 = 0;
10021                b.arg(&null);
10022                unsafe {
10023                    b.launch(cfg)?;
10024                }
10025            }
10026        }
10027        Ok(())
10028    }
10029
10030    /// gemma4 R1: dst = GELU_tanh(gate) * up.
10031    pub fn gelu_tanh_mul(
10032        &self,
10033        gate: &CudaSlice<f32>,
10034        up: &CudaSlice<f32>,
10035        dst: &mut CudaSlice<f32>,
10036        n: usize,
10037    ) -> Result<(), Box<dyn std::error::Error>> {
10038        let f = self.func("gelu_tanh_mul_f32");
10039        let cfg = LaunchConfig::for_num_elems(n as u32);
10040        let ni = n as i32;
10041        let __s_b = self.gpu.stream();
10042        let mut b = __s_b.launch_builder(&f);
10043        b.arg(gate).arg(up).arg(dst).arg(&ni);
10044        unsafe {
10045            b.launch(cfg)?;
10046        }
10047        Ok(())
10048    }
10049
10050    pub fn silu_mul(
10051        &self,
10052        gate: &CudaSlice<f32>,
10053        up: &CudaSlice<f32>,
10054        dst: &mut CudaSlice<f32>,
10055        n: usize,
10056    ) -> Result<(), Box<dyn std::error::Error>> {
10057        let f = self.func("silu_mul_f32");
10058        // float4 kernel: one thread per 4 elements (tail handled in-kernel)
10059        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
10060        let ni = n as i32;
10061        let __s_b = self.gpu.stream();
10062        let mut b = __s_b.launch_builder(&f);
10063        b.arg(gate).arg(up).arg(dst).arg(&ni);
10064        unsafe {
10065            b.launch(cfg)?;
10066        }
10067        Ok(())
10068    }
10069
10070    /// f16out twin of `silu_mul` (task #17): the epilogue also emits the fp16 GEMM operand
10071    /// for the down projection — kills the standalone convert pass. Bit-identical class.
10072    pub fn silu_mul_f16out(
10073        &self,
10074        gate: &CudaSlice<f32>,
10075        up: &CudaSlice<f32>,
10076        dst: &mut CudaSlice<f32>,
10077        dst16: &mut CudaSlice<u8>,
10078        n: usize,
10079    ) -> Result<(), Box<dyn std::error::Error>> {
10080        let f = self.func("silu_mul_f16out_f32");
10081        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
10082        let ni = n as i32;
10083        let __s_b = self.gpu.stream();
10084        let mut b = __s_b.launch_builder(&f);
10085        b.arg(gate).arg(up).arg(dst).arg(dst16).arg(&ni);
10086        unsafe {
10087            b.launch(cfg)?;
10088        }
10089        Ok(())
10090    }
10091
10092    /// FFN SwiGLU epilogue fusion (RANK3 LEVER 2): `dst = silu(gate*gs) * (up*us)` in ONE launch,
10093    /// folding the per-tensor NVFP4 macro-scale (`gs`,`us`) that would otherwise be two separate
10094    /// `scale_inplace` launches on the gate/up matmul outputs. BIT-IDENTICAL to
10095    /// scale_inplace(gate,gs); scale_inplace(up,us); silu_mul(gate,up,dst) — identical float ops in
10096    /// identical order. For non-NVFP4 weights gs==us==1.0 -> identical to `silu_mul`. Net: -2
10097    /// launches per dense FFN layer (the gate+up post-matmul scales).
10098    pub fn silu_mul_scaled(
10099        &self,
10100        gate: &CudaSlice<f32>,
10101        up: &CudaSlice<f32>,
10102        gs: f32,
10103        us: f32,
10104        dst: &mut CudaSlice<f32>,
10105        n: usize,
10106    ) -> Result<(), Box<dyn std::error::Error>> {
10107        let f = self.func("silu_mul_scaled_f32");
10108        let cfg = LaunchConfig::for_num_elems(n as u32);
10109        let ni = n as i32;
10110        let (gsf, usf) = (gs, us);
10111        let __s_b = self.gpu.stream();
10112        let mut b = __s_b.launch_builder(&f);
10113        b.arg(gate).arg(up).arg(&gsf).arg(&usf).arg(dst).arg(&ni);
10114        unsafe {
10115            b.launch(cfg)?;
10116        }
10117        Ok(())
10118    }
10119
10120    /// swigluoai (MiniMax-M3 / GPT-OSS): clamped SwiGLU epilogue, math 1:1 vs llama.cpp
10121    /// ggml_cuda_op_swiglu_oai_single. `dst = swish_alpha(clamp(gate*gs)) * (1 + clamp(up*us))`.
10122    /// gs/us fold the NVFP4 macro-scales exactly like `silu_mul_scaled`.
10123    #[allow(clippy::too_many_arguments)]
10124    pub fn swigluoai_mul_scaled(
10125        &self,
10126        gate: &CudaSlice<f32>,
10127        up: &CudaSlice<f32>,
10128        gs: f32,
10129        us: f32,
10130        alpha: f32,
10131        limit: f32,
10132        dst: &mut CudaSlice<f32>,
10133        n: usize,
10134    ) -> Result<(), Box<dyn std::error::Error>> {
10135        let f = self.func("swigluoai_mul_scaled_f32");
10136        let cfg = LaunchConfig::for_num_elems(n as u32);
10137        let ni = n as i32;
10138        let __s_b = self.gpu.stream();
10139        let mut b = __s_b.launch_builder(&f);
10140        b.arg(gate)
10141            .arg(up)
10142            .arg(&gs)
10143            .arg(&us)
10144            .arg(&alpha)
10145            .arg(&limit)
10146            .arg(dst)
10147            .arg(&ni);
10148        unsafe {
10149            b.launch(cfg)?;
10150        }
10151        Ok(())
10152    }
10153
10154    /// RANK2 LEVER (q8_1 quant-fold): SwiGLU epilogue that EMITS the q8_1 quantization of `act`
10155    /// directly (aq int8 [n] + ad f32 [n/32]), so ffn_down's standalone `quantize_q8_1` launch is
10156    /// removed — the down-proj activation has one consumer, so the quant folds into the producer for
10157    /// free (no extra HBM read; no f32 `act` write). gs/us fold the gate/up NVFP4 macro-scales like
10158    /// `silu_mul_scaled`. BIT-IDENTICAL q8_1 to silu_mul_scaled(...) then quantize_q8_1(...). Only
10159    /// valid when ffn_down uses the q8_1 dp4a/mmvq path; the caller checks `uses_q8_1_fast(ffn_down)`.
10160    /// n must be a multiple of 32 (n_ff always is).
10161    pub fn silu_mul_scaled_q8_1(
10162        &self,
10163        gate: &CudaSlice<f32>,
10164        up: &CudaSlice<f32>,
10165        gs: f32,
10166        us: f32,
10167        n: usize,
10168    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10169        let f = self.func("silu_mul_scaled_q8_1");
10170        let nblk = n / 32;
10171        let mut aq = self.alloc_uninit::<i8>(n)?; // full-overwrite output
10172        let mut ad = self.alloc_uninit::<f32>(nblk)?; // full-overwrite output
10173        // WARP-PER-BLOCK kernel: one warp (32 lanes) per 32-block -> n threads total.
10174        let cfg = LaunchConfig::for_num_elems(n as u32);
10175        let (gsf, usf, ni) = (gs, us, n as i32);
10176        let __s_b = self.gpu.stream();
10177        let mut b = __s_b.launch_builder(&f);
10178        b.arg(gate)
10179            .arg(up)
10180            .arg(&gsf)
10181            .arg(&usf)
10182            .arg(&mut aq)
10183            .arg(&mut ad)
10184            .arg(&ni);
10185        unsafe {
10186            b.launch(cfg)?;
10187        }
10188        Ok((aq, ad))
10189    }
10190
10191    pub fn add(
10192        &self,
10193        a: &CudaSlice<f32>,
10194        b_in: &CudaSlice<f32>,
10195        dst: &mut CudaSlice<f32>,
10196        n: usize,
10197    ) -> Result<(), Box<dyn std::error::Error>> {
10198        let f = self.func("add_f32");
10199        // float4 kernel: one thread per 4 elements (tail handled in-kernel)
10200        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
10201        let ni = n as i32;
10202        let __s_bld = self.gpu.stream();
10203        let mut bld = __s_bld.launch_builder(&f);
10204        bld.arg(a).arg(b_in).arg(dst).arg(&ni);
10205        unsafe {
10206            bld.launch(cfg)?;
10207        }
10208        Ok(())
10209    }
10210
10211    pub fn mul(
10212        &self,
10213        a: &CudaSlice<f32>,
10214        b_in: &CudaSlice<f32>,
10215        dst: &mut CudaSlice<f32>,
10216        n: usize,
10217    ) -> Result<(), Box<dyn std::error::Error>> {
10218        let f = self.func("mul_f32");
10219        let cfg = LaunchConfig::for_num_elems(n as u32);
10220        let ni = n as i32;
10221        let __s_bld = self.gpu.stream();
10222        let mut bld = __s_bld.launch_builder(&f);
10223        bld.arg(a).arg(b_in).arg(dst).arg(&ni);
10224        unsafe {
10225            bld.launch(cfg)?;
10226        }
10227        Ok(())
10228    }
10229
10230    /// Unified weight-tensor matmul: dispatches quant tensors to qmatvec (weights packed) and
10231    /// float tensors to cuBLASLt. y[m,out] = x[m,in] @ W[out,in]^T.
10232    pub fn matmul(
10233        &self,
10234        w: &crate::model::GpuTensor,
10235        x: &CudaSlice<f32>,
10236        m: usize,
10237    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10238        use crate::model::GpuTensor;
10239        let in_f = w.in_features();
10240        let out_f = w.out_features();
10241        // PREFILL (T>1) ROOT FIX: batched tensor-core int8 GEMM. Decodes each weight tile to int8
10242        // in smem ONCE and reuses across all tokens via mma — vs the dp4a matvec's per-token weight
10243        // re-read. Only the 4 daily-hot dtypes; m=1 decode keeps dp4a (it's bandwidth-bound, mma
10244        // gives nothing). Quantize the activation once here then call the GEMM.
10245        // m cutoff FIXED at 16: the m=4 MMA-verify A/B (2026-07-06, was MEMRA_GEMM_M) measured
10246        // NEGATIVE — the MMA tile grid starves at m=4 (BN=256 -> grid.y=1) and its FP order
10247        // shifted verify argmax at tight margins. Do not lower without re-running that battery.
10248        #[allow(non_snake_case)]
10249        // VERIFY-EXACT scope pushes the GEMM crossover out of reach (usize::MAX) — the
10250        // t>=16 dflash verify must ride the decode-exact batched class (parity law).
10251        let GEMM_M_THRESHOLD = if self.verify_exact_on() {
10252            usize::MAX
10253        } else {
10254            16usize
10255        };
10256
10257        // PREFILL GEMM (m>=16). ACCURACY-FIRST dispatch (2026-06-28, prefill-gemm-beat-research wf
10258        // wllbyo6vc step 1): the int8 W4A8 GEMM (qmatvec_gemm, q8_1 activation, s32 accumulate) is
10259        // ACCURATE (prefill logit maxdiff 0.159, < dp4a 0.55) and the default. The FP4 W4A4 mxf4 path
10260        // (try_fp4_gemm) quantizes the ACTIVATION to e2m1 4-bit (8 magnitude levels) -> maxdiff 1.0
10261        // when combined — a real accuracy loss, NOT a math bug. So FP4-W4A4 is taken ONLY under the
10262        // explicit MEMRA_FP4 opt-in AND it must come SECOND (int8 W4A8 is the correct default for NVFP4).
10263        // The workflow plan rebuilds the FP4 path (kill per-K repack, widen K, deepen pipeline, TMA) to
10264        // be both fast AND accurate; until then NVFP4 prefill defaults to the accurate int8 GEMM.
10265        // TINY-OUT_F GUARD (2026-06-28, ncu trace): the tiling GEMM's grid is (ceil(out_f/BM=64),
10266        // ceil(m/BN=256)). For tiny out_f (ssm_beta/ssm_alpha out_f=num_v_heads~32), grid.x=1 -> only
10267        // ceil(m/256) CTAs (e.g. 2 for m=512) on 82 SMs = 0.39% SM throughput, 852us EACH (measured
10268        // worst offender). The dp4a path grids (out_f, m) = far more CTAs, filling the GPU. So route
10269        // out_f < 2*BM to dp4a (skip the tiling GEMM which structurally can't fill the SMs here).
10270        const GEMM_MIN_OUT_F: usize = 128; // 2*BM; below this the GEMM grid.x starves the 82 SMs
10271        // VENDORED llama MMQ prefill GEMMs. NVFP4 W4A8 is DEFAULT-ON (2026-07-05 flip: same int8
10272        // accuracy class as the int8 GEMM below at ~1.9x pp512, rp-loader coexists with the A6
10273        // repack; MEMRA_MMQ_W4A8=0 = escape hatch). W4A4 mxf4nvf4 + Q4_K/Q5_K stay behind MEMRA_MMQ=1.
10274        // The env policy lives in mmq_supports/qmatvec_mmq. Feeds raw f32 activation `x` (the
10275        // launcher quantizes internally). out_f>=MMQ_Y/2 keeps the tile grid from starving the SMs.
10276        // FP8-ACT PREFILL (MEMRA_PP_FP8=1, probe verdict 2026-07-08): F8-E4M3-origin projections
10277        // carry their raw e4m3 device bytes (the `fp8` operand stashed at load next to the Q8_0
10278        // re-encode) — cuBLASLt FP8 TN at 620-795 TF vs 47-72 TF for this class's int8 GEMM.
10279        // Weight side EXACT (checkpoint bytes); activation rides ONE per-batch e4m3 scale
10280        // (amax/448) folded with weight_scale in-GEMM. Prefill only; decode keeps Q8_0 untouched.
10281        if m >= GEMM_M_THRESHOLD {
10282            if let Some(y) = self.try_fp8_gemm(w, x, m)? {
10283                return Ok(y);
10284            }
10285            // PER-BLOCK FP8 MMQ (lane/fp8-mmq): the block-128 class try_fp8_gemm skips (cuBLASLt
10286            // takes no block grid on sm_120). Exact per block — the checkpoint's e4m3 bytes and its
10287            // f32 grid go into the tile unchanged. TWO SOURCES, TWO DEFAULTS: the load-time stash is
10288            // opt-in (MEMRA_FP8_MMQ=1), the native-resident QT_F8_E4M3_BLK grid is DEFAULT ON
10289            // (MEMRA_FP8_MMQ=0 reverts it to dequant-per-call) — see fp8_ffi.rs for why the same
10290            // tile defaults differently by operand source.
10291            if let Some(y) = self.try_fp8_blk_mmq(w, x, m)? {
10292                return Ok(y);
10293            }
10294            // FP16-mirror prefill (MEMRA_PP_F16=1, probe 2026-07-26: 3.2-3.7x the MMQ class).
10295            // Mirror presence IS the gate (only built under the env). Decode never reaches here.
10296            if let Some(y) = self.try_f16_gemm(w, x, m)? {
10297                return Ok(y);
10298            }
10299        }
10300        // F8-E4M3 BLOCK-128 (QT_F8_E4M3_BLK, lane/fp8-blk128-decode). TWO arms, split at the SAME
10301        // m threshold the rest of this method uses:
10302        //   * m >= threshold (prefill): dequant-per-call to the ARM B' Q8_0 slab and recurse, so
10303        //     prefill keeps the floor's kernels AND the floor's bits (try_e4m3_blk_prefill).
10304        //   * m <  threshold: the native per-block GEMV — m=1 decode and the m=2..15 verify tiers.
10305        //     grid.y=m runs the exact m=1 program per (token,row), so the decode-parity law holds
10306        //     across every tier by construction with no batched twin needed.
10307        //
10308        // NOT gated on `fast`: this dtype has no dp4a twin and no Stage-A f32-dequant oracle (the
10309        // generic `deq()` switch has no block-scale input), exactly as QT_F8_E4M3 has none, so
10310        // MEMRA_FAST=0 cannot route it anywhere else. Placed before every GEMM/MMQ arm below
10311        // because gemm_supports/mmq_supports/mmvq_supports all deliberately REFUSE this qtype —
10312        // reaching the generic tail would panic rather than produce wrong numbers, and this pair of
10313        // arms is what makes sure it never gets there.
10314        if let GpuTensor::Quant { qtype, .. } = w {
10315            if *qtype == QT_F8_E4M3_BLK {
10316                if m >= GEMM_M_THRESHOLD {
10317                    if let Some(y) = self.try_e4m3_blk_prefill(w, x, m)? {
10318                        return Ok(y);
10319                    }
10320                }
10321                let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10322                if let Some(y) = self.try_e4m3_blk_pre(w, &aq, &ad, m)? {
10323                    return Ok(y);
10324                }
10325            }
10326        }
10327        if m >= GEMM_M_THRESHOLD && out_f >= GEMM_MIN_OUT_F && self.mmq_supports(w) {
10328            return self.qmatvec_mmq(w, x, m);
10329        }
10330        if m >= GEMM_M_THRESHOLD && out_f >= GEMM_MIN_OUT_F && self.gemm_supports(w) {
10331            let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10332            return self.qmatvec_gemm(w, &aq, &ad, m);
10333        }
10334        // FP4 W4A4 only as an explicit speed/accuracy tradeoff opt-in, and only if the int8 GEMM
10335        // above didn't already handle this weight (e.g. NVFP4 with in_f%64!=0, or MEMRA_NO_GEMM set).
10336        if m >= GEMM_M_THRESHOLD {
10337            if let Some(y) = self.try_fp4_gemm(w, x, m, in_f, out_f)? {
10338                return Ok(y);
10339            }
10340        }
10341        // Stage-B fast int8 dp4a is the DEFAULT since 2026-07-08 (it has been the daily path
10342        // for weeks; the old opt-in flag was a silent-slow-path landmine). MEMRA_FAST=0 reverts
10343        // to Stage-A f32-dequant (the correctness oracle path).
10344        let fast = std::env::var("MEMRA_FAST").as_deref() != Ok("0");
10345        // PERF-3 decode-GEMV: m=1 warp-per-row MMVQ (MEMRA_MMVQ). The big decode matvecs reach
10346        // `matmul` directly (ffn_down, lm_head output, wo), so route them here too — not only the
10347        // matmul_pre siblings. qmatvec_mmvq_raw quantizes the activation internally (q8_1) like the
10348        // _fast paths; the NVFP4 macro-scale is applied by the `scale != 1.0` block below.
10349        if m == 1 && fast {
10350            if let GpuTensor::Quant {
10351                bytes,
10352                qtype,
10353                row_bytes,
10354                rp,
10355                rp4,
10356                scale,
10357                ..
10358            } = w
10359            {
10360                if self.mmvq_supports(*qtype) {
10361                    // NVFP4 macro-scale rides the kernel's fused epilogue arg (one launch total);
10362                    // non-NVFP4 has scale==1.0 so qmatvec_mmvq skips scale_inplace either way.
10363                    // Q4_0 split-plane mirror (rp4): the decode arm reads it via the _rp twins.
10364                    let (bytes, rp) = match rp4 {
10365                        Some(m4) => (m4, true),
10366                        None => (bytes, *rp),
10367                    };
10368                    let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10369                    return self.qmatvec_mmvq(
10370                        bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, *scale, rp,
10371                    );
10372                }
10373            }
10374        }
10375        // BATCHED weight-resident matvec for the m=2-4 band (the MTP/verify forward's ffn_down, wo, and
10376        // lm_head `output` reach `matmul` directly at m=T=2..4). Walks the weight ONCE, dp4a vs all m
10377        // activation columns -> 1 weight read for m tokens (vs grid.y=m re-reading m times below). Quant
10378        // the activation once here (q8_1) like the _fast paths; macro-scale applied via the scale!=1.0
10379        // block below. MEMRA_NO_BATCHED -> per-m path.
10380        //
10381        // DECODE-PARITY GATE (2026-07-07, the 9B synth K=3/4/6 spec FAIL root cause): the batched
10382        // kernels are bit-identical per (token,row) to MMVQ's 32-thread warp reduce, NOT to the
10383        // dp4a kernels' 128-thread two-level reduce. Without MEMRA_MMVQ the m=1 decode chain rides
10384        // dp4a, so a verify riding batched here has a DIFFERENT FP order than the decode it must
10385        // match bit-for-bit — greedy spec flips at tight-margin tokens (the old HANDOVER "ENV LAW:
10386        // FAST+MMVQ both required" footgun, closed here). Parity law: the m>1 kernel CLASS must be
10387        // a pure function of (dtype, env) equal to the m=1 class — batched iff MMVQ. Without MMVQ
10388        // the verify falls to the per-m grid.y=m dp4a path below (each column = the exact m=1
10389        // dp4a program). MEMRA_MMVQ=1 (the daily config) is dispatch-unchanged.
10390        if (2..=16).contains(&m)
10391            && fast
10392            && std::env::var("MEMRA_NO_BATCHED").is_err()
10393            && (m <= 4 || Self::b8_enabled())
10394        {
10395            // b16 tier (2026-07-11, spec K>7): Q4_0/Q6_K have base+_rp b16 kernels; Q8_0's
10396            // b16 exists only as the split-plane _rp twin, so it joins iff the q8rp mirror
10397            // is present (rp4) — the mirror pick below then routes to the _rp family.
10398            // QT_F8_E4M3 joins unconditionally (lane/rp-on-st): its b16 IS the base kernel,
10399            // because the native e4m3 row layout is already aligned and needs no mirror.
10400            // NVFP4/Q4_K/Q8_0 all join unconditionally now (lane/rp-on-st): each has base + _rp
10401            // b16 twins, so either residency layout has its aligned form at this width. Q8_0's
10402            // old `rp4.is_some()` precondition is GONE — the mirror is a bandwidth lever, not the
10403            // exact tier's admission ticket (it was refusing FP8-ST over 23.9 MiB of ssm_beta).
10404            let m_ok = m <= 8
10405                || matches!(w, GpuTensor::Quant { qtype, .. }
10406                if *qtype == QT_Q4_0 || *qtype == QT_Q6_K || *qtype == QT_F8_E4M3
10407                    || *qtype == QT_NVFP4 || *qtype == QT_Q4_K || *qtype == QT_Q5_K || *qtype == QT_Q8_0);
10408            if m_ok {
10409                if let GpuTensor::Quant {
10410                    bytes,
10411                    qtype,
10412                    row_bytes,
10413                    rp,
10414                    rp4,
10415                    ..
10416                } = w
10417                {
10418                    if self.batched_supports(*qtype) && self.mmvq_supports(*qtype) {
10419                        let (bytes, rp) = match rp4 {
10420                            Some(m4) => (m4, true),
10421                            None => (bytes, *rp),
10422                        };
10423                        let mcols = Self::batched_mcols(m);
10424                        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10425                        let mut y = self.qmatvec_mmvq_batched(
10426                            bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, mcols, 1.0, rp,
10427                        )?;
10428                        if let GpuTensor::Quant { scale, .. } = w {
10429                            if *scale != 1.0 {
10430                                self.scale_inplace(&mut y, *scale, m * out_f)?;
10431                            }
10432                        }
10433                        return Ok(y);
10434                    }
10435                }
10436            }
10437        }
10438        // F8-E4M3 (MEMRA_ST_E4M3) catch-all for the m<16 band the arms above didn't take (m=9..15,
10439        // the K=8 verify tier; or m=2..8 under MEMRA_NO_BATCHED/MEMRA_B8=0): grid.y=m e4m3 mmvq —
10440        // the SAME per-(token,row) program as the m=1 decode launch (bit-identical by construction),
10441        // weight re-read m times (rare tier; exactness over bandwidth here). There is no _dp4a twin
10442        // for this dtype, so the generic match below must never see it under `fast`.
10443        if fast {
10444            if let GpuTensor::Quant {
10445                bytes,
10446                qtype,
10447                row_bytes,
10448                scale,
10449                ..
10450            } = w
10451            {
10452                if *qtype == QT_F8_E4M3 {
10453                    let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10454                    return self.qmatvec_mmvq(
10455                        bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, *scale, false,
10456                    );
10457                }
10458            }
10459        }
10460        let mut y = match w {
10461            GpuTensor::Quant {
10462                bytes,
10463                qtype,
10464                row_bytes,
10465                ..
10466            } if fast && *qtype == QT_Q8_0 => {
10467                self.qmatvec_q8_0_fast(bytes, x, m, in_f, out_f, *row_bytes)?
10468            }
10469            GpuTensor::Quant {
10470                bytes,
10471                qtype,
10472                row_bytes,
10473                ..
10474            } if fast && *qtype == QT_Q4_K => {
10475                self.qmatvec_q4_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
10476            }
10477            GpuTensor::Quant {
10478                bytes,
10479                qtype,
10480                row_bytes,
10481                ..
10482            } if fast && *qtype == QT_Q6_K => {
10483                self.qmatvec_q6_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
10484            }
10485            GpuTensor::Quant {
10486                bytes,
10487                qtype,
10488                row_bytes,
10489                ..
10490            } if fast && *qtype == QT_Q5_K => {
10491                self.qmatvec_q5_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
10492            }
10493            GpuTensor::Quant {
10494                bytes,
10495                qtype,
10496                row_bytes,
10497                ..
10498            } if fast && *qtype == QT_Q3_K => {
10499                self.qmatvec_q3_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
10500            }
10501            GpuTensor::Quant {
10502                bytes,
10503                qtype,
10504                row_bytes,
10505                rp,
10506                ..
10507            } if fast && *qtype == QT_NVFP4 => self.qmatvec_dp4a_named(
10508                if *rp {
10509                    "qmatvec_nvfp4_dp4a_rp"
10510                } else {
10511                    "qmatvec_nvfp4_dp4a"
10512                },
10513                bytes,
10514                x,
10515                m,
10516                in_f,
10517                out_f,
10518                *row_bytes,
10519            )?,
10520            // IQ4_XS trunk fast path — DEFAULT ON since 2026-08-02 (MEMRA_IQ_FAST=0 reverts to
10521            // Stage-A; see iq_fast_enabled). The old opt-in default was the KAT-Coder decode
10522            // anomaly (research/kat-anomaly-20260802/).
10523            GpuTensor::Quant {
10524                bytes,
10525                qtype,
10526                row_bytes,
10527                ..
10528            } if fast && *qtype == QT_IQ4_XS && Self::iq_fast_enabled() => {
10529                self.qmatvec_iq4_XS_fast(bytes, x, m, in_f, out_f, *row_bytes)?
10530            }
10531            // B3: IQ3_S uses the Stage-A f32 dequant-in-kernel path. There is NO
10532            // qmatvec_iq3_s_dp4a kernel — do NOT add a `*qtype == QT_IQ3_S` fast guard here
10533            // without first writing the matching kernel, or func() will panic
10534            // "kernel ... not in any fatbin".
10535            GpuTensor::Quant {
10536                bytes,
10537                qtype,
10538                row_bytes,
10539                rp,
10540                ..
10541            } =>
10542            // Stage-A generic: repacked NVFP4 uses the device-side split-plane tag (the
10543            // deq(row,j) form cannot address the planes; same value/product order).
10544            {
10545                self.qmatvec(
10546                    bytes,
10547                    x,
10548                    m,
10549                    in_f,
10550                    out_f,
10551                    if *rp && *qtype == QT_NVFP4 {
10552                        QT_NVFP4_RP
10553                    } else {
10554                        *qtype
10555                    },
10556                    *row_bytes,
10557                )?
10558            }
10559            GpuTensor::Float { data, .. } => self.linear(x, data, m, in_f, out_f)?,
10560            // MEMRA_FULL_PREC bf16-resident weight: dequant-on-use to f32 scratch, then the same
10561            // cuBLASLt f32 GEMV as the Float arm.
10562            GpuTensor::FloatBf16 { data, .. } => {
10563                self.linear_bf16_chunked(x, data, m, in_f, out_f, false)?
10564            }
10565        };
10566        // NVFP4 per-tensor macro-scale (post-matmul). scale==1.0 for all other quants/float -> no-op.
10567        if let GpuTensor::Quant { scale, .. } = w {
10568            if *scale != 1.0 {
10569                self.scale_inplace(&mut y, *scale, m * out_f)?;
10570            }
10571        }
10572        Ok(y)
10573    }
10574
10575    /// True when `MEMRA_FAST=0`, i.e. the Stage-A f32 oracle is the requested arithmetic and every
10576    /// `matmul_pre` call will take the raw-f32 escape rather than the q8_1 pair.
10577    ///
10578    /// WHY THIS EXISTS AS ITS OWN PREDICATE: `uses_q8_1_fast` needs a weight, but the callers that
10579    /// have to DECIDE WHETHER TO MATERIALIZE an f32 activation sit one level above any weight — the
10580    /// gemma-4 decode trunk emits a q8_1 pair per layer and hands it down, so the f32 has to be
10581    /// produced (or not) by the loop that owns the residual. Cached in a OnceLock like every other
10582    /// `*_on()` flag here: this is read once per layer per token on the decode path, and the daily
10583    /// path must not pay an env lookup for a flag that is off.
10584    pub fn stage_a_raw_needed() -> bool {
10585        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
10586        *ON.get_or_init(|| std::env::var("MEMRA_FAST").as_deref() == Ok("0"))
10587    }
10588
10589    /// True if `w` would take the int8-dp4a fast path under MEMRA_FAST (so its activation can be
10590    /// pre-quantized once and shared across sibling matmuls via `matmul_pre`).
10591    pub fn uses_q8_1_fast(&self, w: &crate::model::GpuTensor) -> bool {
10592        use crate::model::GpuTensor;
10593        if std::env::var("MEMRA_FAST").as_deref() == Ok("0") {
10594            return false;
10595        }
10596        match w {
10597            // QT_F8_E4M3_BLK is admitted for the same reason QT_F8_E4M3 is: its ONLY kernel class
10598            // takes the shared q8_1 activation, so callers may pre-quantize once and share it
10599            // across siblings. It is NOT admitted to any of the fused/dual epilogue doors those
10600            // siblings can then open (`q8_fused_params`, `e4m3_fused_params` and
10601            // `matmul_pre_dual_noscale` all match on their own qtype and refuse this one) — the
10602            // block class has no fused twin yet, so each of its projections takes its own launch.
10603            GpuTensor::Quant { qtype, .. } => {
10604                matches!(
10605                    *qtype,
10606                    QT_Q8_0
10607                        | QT_Q4_K
10608                        | QT_Q6_K
10609                        | QT_Q5_K
10610                        | QT_Q3_K
10611                        | QT_NVFP4
10612                        | QT_F8_E4M3
10613                        | QT_F8_E4M3_BLK
10614                        | QT_Q4_0
10615                ) || (*qtype == QT_IQ4_XS && Self::iq_fast_enabled())
10616            }
10617            GpuTensor::Float { .. } | GpuTensor::FloatBf16 { .. } => false,
10618        }
10619    }
10620
10621    /// matmul with a PRE-QUANTIZED q8_1 activation (aq,ad from `quantize_q8_1`). Skips the
10622    /// per-matmul re-quantize so sibling matmuls that share an input (gate+up share `z`;
10623    /// q/k/v + wqkv/gate/beta/alpha share `h`) quantize ONCE. Caller MUST have checked
10624    /// `uses_q8_1_fast(w)`; falls back to plain `matmul` otherwise (Stage-A / Float / non-fast).
10625    pub fn matmul_pre(
10626        &self,
10627        w: &crate::model::GpuTensor,
10628        aq: &CudaSlice<i8>,
10629        ad: &CudaSlice<f32>,
10630        x_fallback: &CudaSlice<f32>,
10631        m: usize,
10632    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10633        use crate::model::GpuTensor;
10634        // Every raw-f32 arm below (fp8/f16/MMQ/fp4) reads m*in_f from x_fallback. Callers that
10635        // pre-quantized and dropped the f32 input pass an EMPTY x_fallback (E4B's fusion port:
10636        // h = zeros(0)) — the length guard keeps those on the aq/ad GEMM instead of feeding a
10637        // 0-byte buffer to a convert kernel (illegal address -> cublasLt status 13; the E4B
10638        // rc=30013 dig, 2026-07-31).
10639        let x_raw_ok = x_fallback.len() >= m * w.in_features();
10640        // FP8-ACT PREFILL (MEMRA_PP_FP8=1): same arm as `matmul` — the fp8 operand needs the RAW
10641        // f32 activation (per-batch e4m3 quant differs from q8_1), so x_fallback not aq/ad.
10642        if m >= 16 && x_raw_ok && !self.verify_exact_on() {
10643            if let Some(y) = self.try_fp8_gemm(w, x_fallback, m)? {
10644                return Ok(y);
10645            }
10646            // PER-BLOCK FP8 MMQ — same arm as `matmul` (stash opt-in, native-resident default ON);
10647            // its own quantizer wants the RAW f32 activation, so x_fallback not aq/ad.
10648            if let Some(y) = self.try_fp8_blk_mmq(w, x_fallback, m)? {
10649                return Ok(y);
10650            }
10651            // FP16-mirror prefill (same arm as `matmul` — fp16 wants the RAW f32 activation).
10652            if let Some(y) = self.try_f16_gemm(w, x_fallback, m)? {
10653                return Ok(y);
10654            }
10655        }
10656        // BLOCK-128 e4m3 (QT_F8_E4M3_BLK) — the same two arms as `matmul`, split at the same m, and
10657        // placed at the same point in the order (after the prefill GEMM hooks, before every arm
10658        // that refuses this qtype). The prefill arm needs the RAW f32 activation for the Q8_0
10659        // dispatch it recurses into, so it takes x_fallback and is skipped when that is empty
10660        // (a pre-quantized caller that dropped its f32 input never runs at prefill m anyway).
10661        if m >= 16 && x_raw_ok && !self.verify_exact_on() {
10662            if let Some(y) = self.try_e4m3_blk_prefill(w, x_fallback, m)? {
10663                return Ok(y);
10664            }
10665        }
10666        if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
10667            return Ok(y);
10668        }
10669        // VENDORED llama MMQ prefill GEMMs (NVFP4 W4A8 default-on; W4A4/k-quant behind MEMRA_MMQ=1
10670        // — policy in mmq_supports) — use the RAW f32 activation (their own internal quant:
10671        // q8_1 D4 for NVFP4 W4A8, FP8/UE4M3 for W4A4, q8_1 DS4 for Q4_K/Q5_K), so x_fallback not
10672        // aq/ad.
10673        if m >= 16
10674            && w.out_features() >= 128
10675            && self.mmq_supports(w)
10676            && !self.verify_exact_on()
10677            && x_raw_ok
10678        {
10679            return self.qmatvec_mmq(w, x_fallback, m);
10680        }
10681        // Stage-C FP4 prefill (MEMRA_FP4): native mxf4 GEMM needs the f32 activation (FP4-quant differs
10682        // from q8_1), so re-quantize from x_fallback rather than reuse aq/ad. NVFP4 only, m>=16.
10683        if m >= 16 && x_raw_ok && !self.verify_exact_on() {
10684            if let Some(y) =
10685                self.try_fp4_gemm(w, x_fallback, m, w.in_features(), w.out_features())?
10686            {
10687                return Ok(y);
10688            }
10689        }
10690        // Prefill GEMM root fix: if T>1 and the dtype has a GEMM kernel, batch via tensor cores
10691        // (reuses the already-quantized aq/ad — no extra quantize). m=1 falls through to dp4a.
10692        if m >= 16 && self.gemm_supports(w) && !self.verify_exact_on() {
10693            return self.qmatvec_gemm(w, aq, ad, m);
10694        }
10695        // THE SEVENTH ARM. Every raw-f32 arm above is guarded by `x_raw_ok`; this one — the
10696        // Stage-A / Float escape — was not, and it is the ONLY one `MEMRA_FAST=0` opens. So the
10697        // 2026-07-31 E4B rc=30013 fix (add the length guard, keep empty-fallback callers off the
10698        // raw-f32 arms) protected the six arms the FAST path can reach and missed the one the
10699        // ORACLE path reaches. Consequence before this guard: the gemma-4 decode arms, which pass
10700        // `e.zeros(0)` because their trunk emits only a q8_1 pair and never materializes an f32
10701        // attn-normed activation, fell through to Stage-A `qmatvec_f32` (cu/qmatvec.cu:5442),
10702        // which reads `m * in_f` floats out of a 0-byte allocation ->
10703        // CUDA_ERROR_ILLEGAL_ADDRESS at layer 0 of the first decode token. The fault is STICKY:
10704        // it poisons the context, so every LATER request in that process fails with an unrelated
10705        // message ("cache alloc failed: ...") and the true cause appears exactly once, in the
10706        // first failure. That is what made `MEMRA_FAST=0` — the reference named in
10707        // `tools/argmax-margin-gate.sh`'s own header — return an opaque HTTP 500 on the gemma
10708        // dense artifact and left the arm with no working truth instrument.
10709        //
10710        // Refuse loudly instead of reading out of bounds. A named error at the true call site is
10711        // strictly better than an illegal address surfacing later at an unrelated sync point, and
10712        // an oracle that cannot run must say so rather than corrupt the context it runs in.
10713        if !self.uses_q8_1_fast(w) {
10714            if !x_raw_ok {
10715                return Err(format!(
10716                    "matmul_pre: q8_1-fast is off for this weight but x_fallback holds {} f32 \
10717                     (need m*in_f = {}*{} = {}). This call site pre-quantized its activation and \
10718                     dropped the f32, so there is nothing to fall back to — pass the real f32 \
10719                     activation (see Engine::rms_norm_decode, which is bit-identical to \
10720                     rms_norm_q8_1's reduction) or keep the weight on the q8_1 path.",
10721                    x_fallback.len(),
10722                    m,
10723                    w.in_features(),
10724                    m * w.in_features()
10725                )
10726                .into());
10727            }
10728            return self.matmul(w, x_fallback, m);
10729        }
10730        let in_f = w.in_features();
10731        let out_f = w.out_features();
10732        let (bytes, qtype, row_bytes, scale, rp) = match w {
10733            GpuTensor::Quant {
10734                bytes,
10735                qtype,
10736                row_bytes,
10737                scale,
10738                rp,
10739                ..
10740            } => (bytes, *qtype, *row_bytes, *scale, *rp),
10741            _ => unreachable!("uses_q8_1_fast guaranteed Quant"),
10742        };
10743        // Q4_0 split-plane mirror: only the mmvq/batched decode arms read it (the _rp twins);
10744        // the dp4a/oracle tails below keep the raw GGUF bytes.
10745        let (mbytes, mrp) = match w {
10746            GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
10747            _ => (bytes, rp),
10748        };
10749        // PERF-3 decode-GEMV: warp-per-row MMVQ for the m=1 decode arm, gated behind MEMRA_MMVQ.
10750        // Only the 4 daily-hot dtypes have an _mmvq kernel (Q8_0/Q4_K/Q6_K/NVFP4); Q5_K/Q3_K/IQ4_XS
10751        // keep _dp4a (the oracle/fallback). Bit-equivalent to _dp4a up to f32 reduction order.
10752        if m == 1 && self.mmvq_supports(qtype) {
10753            return self.qmatvec_mmvq(mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, mrp);
10754        }
10755        // BATCHED weight-resident matvec for the m=2-4 band (the MTP/verify forward: full_attn_verify
10756        // and decode_step_t run their projections at m=T=k=2..4). The plain _dp4a path below launches
10757        // grid.y=m INDEPENDENT blocks per output row -> the weight row is re-read m times from HBM/L2.
10758        // The _b2/_b4 kernels walk the weight ONCE and dp4a vs all m activation columns, so m tokens
10759        // cost ~1 weight read instead of m (decode is weight-BW-bound). BIT-IDENTICAL per (token,row)
10760        // to the _mmvq path (32-thread warp reduce — NOT the dp4a 128-thread reduce below).
10761        // m=2 -> mcols=2; m∈{3,4} -> mcols=4; m∈{5..8} -> mcols=8 (kernel guards c>=m).
10762        // MEMRA_NO_BATCHED forces the per-m grid.y=m path (the A/B reference); MEMRA_B8=0 keeps
10763        // m=5..8 on the old per-m path (b8-tier-only seam).
10764        // DECODE-PARITY GATE (2026-07-07): batched iff mmvq_supports — see matmul's parity note.
10765        // Without MEMRA_MMVQ, m=1 decode rides dp4a (the arm below at m=1); the verify must ride
10766        // the SAME class per column (grid.y=m dp4a = the exact m=1 dp4a program per column).
10767        if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
10768            && std::env::var("MEMRA_NO_BATCHED").is_err()
10769            && (m <= 4 || Self::b8_enabled())
10770            // b16 tier: every class routed here now has base + _rp b16 kernels (Q4_0/Q6_K
10771            // pre-existing; NVFP4/Q4_K/Q8_0-base/F8_E4M3 added lane/rp-on-st 2026-08-06), so
10772            // there is no mirror precondition left — `mrp` still selects the LAYOUT below.
10773            && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || qtype == QT_NVFP4
10774                || qtype == QT_Q4_K || qtype == QT_Q5_K || qtype == QT_F8_E4M3 || qtype == QT_Q8_0)
10775        {
10776            let mcols = Self::batched_mcols(m);
10777            return self.qmatvec_mmvq_batched(
10778                mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, mrp,
10779            );
10780        }
10781        // F8-E4M3 catch-all (m=9..15 / batched-disabled seams): grid.y=m e4m3 mmvq — this dtype
10782        // has NO _dp4a twin, and per (token,row) the mmvq body is the exact m=1 decode program.
10783        // Q4_0 joins the catch-all (2026-07-11): adaptive-K cap 8 makes verify t=9 reachable
10784        // for the first time (past the b8 tier) and Q4_0 has no dp4a twin either. The mirror
10785        // (mbytes/mrp) keeps the rp layout consistent with the m=1 decode program.
10786        if qtype == QT_F8_E4M3 || qtype == QT_Q4_0 {
10787            let (b2, r2) = if qtype == QT_Q4_0 {
10788                (mbytes, mrp)
10789            } else {
10790                (bytes, rp)
10791            };
10792            return self.qmatvec_mmvq(b2, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, r2);
10793        }
10794        let name = match qtype {
10795            QT_Q8_0 => "qmatvec_q8_0_dp4a",
10796            QT_Q4_K => "qmatvec_q4_K_dp4a",
10797            QT_Q6_K => "qmatvec_q6_K_dp4a",
10798            QT_Q5_K => "qmatvec_q5_K_dp4a",
10799            QT_Q3_K => "qmatvec_q3_K_dp4a",
10800            QT_NVFP4 => {
10801                if rp {
10802                    "qmatvec_nvfp4_dp4a_rp"
10803                } else {
10804                    "qmatvec_nvfp4_dp4a"
10805                }
10806            }
10807            QT_IQ4_XS => "qmatvec_iq4_XS_dp4a",
10808            _ => unreachable!(),
10809        };
10810        let f = self.func(name);
10811        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output: skip memset
10812        let cfg = LaunchConfig {
10813            grid_dim: (out_f as u32, m as u32, 1),
10814            block_dim: (128, 1, 1),
10815            shared_mem_bytes: 0,
10816        };
10817        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
10818        let __s_b = self.gpu.stream();
10819        let mut b = __s_b.launch_builder(&f);
10820        b.arg(bytes)
10821            .arg(aq)
10822            .arg(ad)
10823            .arg(&mut y)
10824            .arg(&inf)
10825            .arg(&outf)
10826            .arg(&mi)
10827            .arg(&rb);
10828        unsafe {
10829            b.launch(cfg)?;
10830        }
10831        if scale != 1.0 {
10832            self.scale_inplace(&mut y, scale, m * out_f)?;
10833        }
10834        Ok(y)
10835    }
10836
10837    /// DECODE-EXACT matmul at any m: guarantees the SAME warp-per-row (MMVQ, 32-thread) FP
10838    /// accumulation order as the T=1 decode path for EVERY token row. The spec-decode verify MUST
10839    /// use this for linear-attn projections to be bit-identical to greedy decode. The dp4a kernel
10840    /// (128 threads, two-level reduction) used by `matmul`/`matmul_pre` at m>=5 has a different
10841    /// shfl-tree shape that produces ULP differences propagating through gdn_scan into argmax flips.
10842    /// The MMVQ kernel with grid.y=m already processes each row independently (same 32-thread warp
10843    /// reduce as m=1); this method just forces that path unconditionally.
10844    pub fn matmul_decode_exact(
10845        &self,
10846        w: &crate::model::GpuTensor,
10847        x: &CudaSlice<f32>,
10848        m: usize,
10849    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10850        use crate::model::GpuTensor;
10851        // FLOAT tensors (35B ssm_beta/ssm_alpha on every linear layer, F32 ne=[2048,32]): the
10852        // generic path is cuBLASLt, whose reduction splits are n-DEPENDENT — m=1 vs m=2 col-0
10853        // outputs differ in every bit (probe 2026-07-06: 32/32 bit-diff, maxdiff 3.5e-3), which
10854        // shifted 35B verify logits 0.26-0.56 vs eager and flipped greedy at tight margins (the
10855        // p3 spec FAIL). Decode-exact contract: per-COLUMN m=1 cuBLASLt calls — each column's
10856        // reduction is the exact kernel the T=1 decode path runs, so verify==decode bit-for-bit.
10857        // m<=10 here (K+2 verify tier), so the extra launches are a handful of 4us gemvs.
10858        if let GpuTensor::Float { data, .. } = w {
10859            return self.linear_decode_exact(x, data, m, w.in_features(), w.out_features());
10860        }
10861        // MEMRA_FULL_PREC bf16-resident weight: dequant-on-use, then the per-column decode-exact
10862        // float linear (same n-independent reduction contract as the Float arm above).
10863        if let GpuTensor::FloatBf16 { data, .. } = w {
10864            let (in_f, out_f) = (w.in_features(), w.out_features());
10865            return self.linear_bf16_chunked(x, data, m, in_f, out_f, true);
10866        }
10867        if !self.uses_q8_1_fast(w) {
10868            return self.matmul(w, x, m);
10869        }
10870        let in_f = w.in_features();
10871        let out_f = w.out_features();
10872        let (bytes, qtype, row_bytes, scale, rp) = match w {
10873            GpuTensor::Quant {
10874                bytes,
10875                qtype,
10876                row_bytes,
10877                scale,
10878                rp,
10879                ..
10880            } => (bytes, *qtype, *row_bytes, *scale, *rp),
10881            _ => return self.matmul(w, x, m),
10882        };
10883        // Q4_0 split-plane mirror for the mmvq/batched arms below (dp4a tail = matmul_pre,
10884        // which does its own mirror pick).
10885        let (bytes, rp) = match w {
10886            GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
10887            _ => (bytes, rp),
10888        };
10889        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10890        // BLOCK-128 e4m3 (QT_F8_E4M3_BLK): the same single kernel every other entry dispatches, so
10891        // the decode-exact contract needs nothing special — grid.y=m runs the m=1 program per
10892        // (token,row) by construction, which is exactly what this method exists to guarantee.
10893        if let Some(y) = self.try_e4m3_blk_pre(w, &aq, &ad, m)? {
10894            return Ok(y);
10895        }
10896        // Batched weight-resident matvec for m=2-8: BIT-IDENTICAL per (token,row) to MMVQ (exact
10897        // integer dp4a, same warp reduce — kernel-check gate rel=0.00e0), one weight read for m
10898        // tokens. The dispatch the divergence fix must avoid is dp4a's 128-thread two-level
10899        // reduce, NOT this. m=5..8 is the K=4..7 spec-verify tier (b8): pre-b8 T=5 fell to the
10900        // grid.y=m per-row MMVQ below = 5 full weight reads/launch — the measured 27B K=4 cliff.
10901        // DECODE-PARITY GATE (2026-07-07): batched (MMVQ-class order) only when the m=1 decode
10902        // chain rides MMVQ too — without MEMRA_MMVQ decode is dp4a, so the exact-contract here
10903        // must be per-column dp4a (matmul_pre fallthrough), not the MMVQ order.
10904        if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
10905            && std::env::var("MEMRA_NO_BATCHED").is_err()
10906            && (m <= 4 || Self::b8_enabled())
10907            // Every b16 class has base + _rp twins after lane/rp-on-st (see matmul_pre's note):
10908            // no mirror precondition, `rp` selects the layout only.
10909            && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || qtype == QT_F8_E4M3
10910                || qtype == QT_NVFP4 || qtype == QT_Q4_K || qtype == QT_Q5_K || qtype == QT_Q8_0)
10911        {
10912            let mcols = Self::batched_mcols(m);
10913            return self.qmatvec_mmvq_batched(
10914                bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, rp,
10915            );
10916        }
10917        if self.mmvq_supports(qtype) {
10918            // MMVQ at grid.y=m: each row is processed by its own warp independently — same 32-thread
10919            // accumulation + warp_reduce_sum as m=1 decode. Bit-identical per row.
10920            return self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, scale, rp);
10921        }
10922        // Fallback for non-MMVQ quant types (Q5_K, Q3_K): use dp4a (the only available kernel).
10923        // These types are not used in the 27B's linear-attn NVFP4+Q4_K layers.
10924        self.matmul_pre(w, &aq, &ad, x, m)
10925    }
10926
10927    /// DECODE-EXACT matmul from a PRE-QUANTIZED q8_1 activation (batched-verify epilogue
10928    /// re-fuse, lane/vt-fixes fix 2, 2026-08-03): the EXACT `matmul_decode_exact` dispatch for
10929    /// q8_1-fast Quant tensors, with the caller's (aq, ad) replacing the internal
10930    /// `quantize_q8_1`. quantize_q8_1 is deterministic (same input bytes -> same q8 bytes), so
10931    /// sharing one quantize across sibling matmuls of the same activation — or consuming the
10932    /// q8 emitted by a fused epilogue (rms_norm_q8_1 / add_rms_norm_q8_1 /
10933    /// silu_mul_scaled_q8_1 / gated_rmsnorm_q8_1, all kernel-check-pinned bit-identical to
10934    /// their unfused chains) — cannot change any dispatched kernel's input bytes.
10935    /// Caller MUST guarantee `uses_q8_1_fast(w)` (the fused epilogues only exist on that path).
10936    pub fn matmul_decode_exact_pre(
10937        &self,
10938        w: &crate::model::GpuTensor,
10939        aq: &CudaSlice<i8>,
10940        ad: &CudaSlice<f32>,
10941        m: usize,
10942    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10943        use crate::model::GpuTensor;
10944        debug_assert!(
10945            self.uses_q8_1_fast(w),
10946            "matmul_decode_exact_pre: caller must guarantee q8_1-fast"
10947        );
10948        // BLOCK-128 e4m3: same single kernel, all m — see matmul_decode_exact's note.
10949        if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
10950            return Ok(y);
10951        }
10952        let in_f = w.in_features();
10953        let out_f = w.out_features();
10954        let (bytes, qtype, row_bytes, scale, rp) = match w {
10955            GpuTensor::Quant {
10956                bytes,
10957                qtype,
10958                row_bytes,
10959                scale,
10960                rp,
10961                ..
10962            } => (bytes, *qtype, *row_bytes, *scale, *rp),
10963            _ => {
10964                return Err(
10965                    "matmul_decode_exact_pre: Quant tensor required (q8_1-fast contract)".into(),
10966                );
10967            }
10968        };
10969        // Q4_0 split-plane mirror — same pick as matmul_decode_exact.
10970        let (bytes, rp) = match w {
10971            GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
10972            _ => (bytes, rp),
10973        };
10974        // Dispatch mirror of matmul_decode_exact's q8_1-fast tail, condition for condition.
10975        if (2..=16).contains(&m)
10976            && self.batched_supports(qtype)
10977            && self.mmvq_supports(qtype)
10978            && std::env::var("MEMRA_NO_BATCHED").is_err()
10979            && (m <= 4 || Self::b8_enabled())
10980            && (m <= 8
10981                || qtype == QT_Q4_0
10982                || qtype == QT_Q6_K
10983                || qtype == QT_F8_E4M3
10984                || qtype == QT_NVFP4
10985                || qtype == QT_Q4_K
10986                || qtype == QT_Q5_K
10987                || qtype == QT_Q8_0)
10988        {
10989            let mcols = Self::batched_mcols(m);
10990            return self.qmatvec_mmvq_batched(
10991                bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, rp,
10992            );
10993        }
10994        if self.mmvq_supports(qtype) {
10995            return self.qmatvec_mmvq(bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, rp);
10996        }
10997        // Non-MMVQ quant types (Q5_K/Q3_K under MEMRA_MMVQ=0): dp4a via matmul_pre — the same
10998        // fallback matmul_decode_exact takes. m <= 16 on the verify tier never reads x_fallback.
10999        let x0 = self.zeros(0)?;
11000        self.matmul_pre(w, aq, ad, &x0, m)
11001    }
11002
11003    /// DUAL gate+up batched matvec from a PRE-QUANTIZED activation, macro-scales DEFERRED
11004    /// (lane/vt-fixes fix 2): same eligibility as `matmul_decode_exact_dual`, but the caller's
11005    /// (aq, ad) replaces the internal quantize and the NVFP4 per-tensor scales are RETURNED
11006    /// instead of applied via two `scale_inplace` launches — the fused SwiGLU epilogue
11007    /// (`silu_mul_scaled_q8_1`) folds them, exactly like the m=1 decode chain does. Deferring
11008    /// is value-exact: `y[i]*s` inline in the epilogue is the same IEEE multiply scale_inplace
11009    /// would store (f32 store/load round-trips are exact). None -> caller falls back to the
11010    /// per-tensor path.
11011    pub fn matmul_decode_exact_dual_pre(
11012        &self,
11013        w0: &crate::model::GpuTensor,
11014        w1: &crate::model::GpuTensor,
11015        aq: &CudaSlice<i8>,
11016        ad: &CudaSlice<f32>,
11017        m: usize,
11018    ) -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn std::error::Error>>
11019    {
11020        use crate::model::GpuTensor;
11021        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11022        let on = *ON.get_or_init(|| {
11023            std::env::var("MEMRA_SPEC_DUAL_T")
11024                .map(|v| v != "0")
11025                .unwrap_or(true)
11026        });
11027        if !on
11028            || !(2..=7).contains(&m)
11029            || std::env::var("MEMRA_NO_BATCHED").is_ok()
11030            || !self.uses_q8_1_fast(w0)
11031            || !self.uses_q8_1_fast(w1)
11032        {
11033            return Ok(None);
11034        }
11035        // DECODE-PARITY GATE (lane/nvfp4-strict, 2026-08-05): batched iff MMVQ — the dual
11036        // kernels are the MMVQ warp-reduce family, and without MEMRA_MMVQ the m=1 decode
11037        // chain this verify must match bit-for-bit rides dp4a (see matmul_decode_exact's
11038        // note). The singles enforce this via `mmvq_supports`; the dual door skipped it.
11039        if !self.mmvq_supports(QT_NVFP4) {
11040            return Ok(None);
11041        }
11042        let (in_f, out_f) = (w0.in_features(), w0.out_features());
11043        if w1.in_features() != in_f || w1.out_features() != out_f {
11044            return Ok(None);
11045        }
11046        let (b0, b1, row_bytes, s0, s1, rp) = match (w0, w1) {
11047            (
11048                GpuTensor::Quant {
11049                    bytes: b0,
11050                    qtype: q0,
11051                    row_bytes: rb0,
11052                    scale: s0,
11053                    rp: rp0,
11054                    rp4: None,
11055                    ..
11056                },
11057                GpuTensor::Quant {
11058                    bytes: b1,
11059                    qtype: q1,
11060                    row_bytes: rb1,
11061                    scale: s1,
11062                    rp: rp1,
11063                    rp4: None,
11064                    ..
11065                },
11066            ) if *q0 == QT_NVFP4 && *q1 == QT_NVFP4 && rb0 == rb1 && rp0 == rp1 => {
11067                (b0, b1, *rb0, *s0, *s1, *rp0)
11068            }
11069            _ => return Ok(None),
11070        };
11071        // m=5..7: only the exact-width rp duals exist (vt-fixes fix 1b); GGUF layout keeps
11072        // the singles. The b8 dual (MCOLS=8 at m=5..8) measured FLAT and stays dead.
11073        if m > 4 && !(rp && Self::b8_enabled() && std::env::var("MEMRA_B567").as_deref() != Ok("0"))
11074        {
11075            return Ok(None);
11076        }
11077        let (y0, y1) =
11078            self.qmatvec_batched_dual_raw(b0, b1, aq, ad, m, in_f, out_f, row_bytes, rp)?;
11079        Ok(Some(((y0, s0), (y1, s1))))
11080    }
11081
11082    /// GROUP-4 GDN-tuple BATCHED matvec (trunk-kernels slice C): the qwen35 linear-layer
11083    /// in-projection 4-tuple (wqkv / wqkv_gate / ssm_beta / ssm_alpha) from ONE pre-quantized
11084    /// activation in ONE launch. Blocks map to the concatenated row space; every out_f must be
11085    /// a multiple of 8 (rows_per_block) so each warp's row pair resolves to one tensor; per
11086    /// (tensor, token, row) the kernel body is `nvfp4_mmvq_batched_rp` VERBATIM with the
11087    /// tensor's macro-scale fused at the write (== the conditional scale_inplace pass,
11088    /// bit-identical) -> BIT-IDENTICAL to the four single launches. Split-plane rp NVFP4 only,
11089    /// m=2..8 (exact-width MCOLS at m=5..7 mirroring the B567 law; m=8 requires b8_enabled
11090    /// like the singles). None -> caller runs the four singles. MEMRA_TK_GDN_GROUP=0 rollback.
11091    pub fn matmul_decode_exact_group4_pre(
11092        &self,
11093        ws: [&crate::model::GpuTensor; 4],
11094        aq: &CudaSlice<i8>,
11095        ad: &CudaSlice<f32>,
11096        m: usize,
11097    ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
11098        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11099        let on = *ON.get_or_init(|| {
11100            std::env::var("MEMRA_TK_GDN_GROUP")
11101                .map(|v| v != "0")
11102                .unwrap_or(true)
11103        });
11104        self.matmul_decode_exact_group_pre(&ws, aq, ad, m, on, "GDN group4")
11105    }
11106
11107    /// GROUP-3 twin for the qwen35 full-attention q/k/v triple (trunk-kernels slice D):
11108    /// the SAME group4 kernels with n3=0 (blocks never reach the fourth range; W3/y3 are
11109    /// never dereferenced) — per (tensor, token, row) bit-identical to the three singles
11110    /// exactly as the group4 door is to its four. MEMRA_TK_FA_GROUP=0 rollback.
11111    pub fn matmul_decode_exact_group3_pre(
11112        &self,
11113        ws: [&crate::model::GpuTensor; 3],
11114        aq: &CudaSlice<i8>,
11115        ad: &CudaSlice<f32>,
11116        m: usize,
11117    ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
11118        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11119        let on = *ON.get_or_init(|| {
11120            std::env::var("MEMRA_TK_FA_GROUP")
11121                .map(|v| v != "0")
11122                .unwrap_or(true)
11123        });
11124        self.matmul_decode_exact_group_pre(&ws, aq, ad, m, on, "FA group3")
11125    }
11126
11127    /// Shared core of the group3/group4 doors: eligibility mirror of the singles' batched
11128    /// dispatch, then ONE `qmatvec_nvfp4_mmvq_group4_b*_rp` launch over the concatenated
11129    /// row space (3-tensor callers ride n3=0). Returns one output per input tensor.
11130    fn matmul_decode_exact_group_pre(
11131        &self,
11132        ws: &[&crate::model::GpuTensor],
11133        aq: &CudaSlice<i8>,
11134        ad: &CudaSlice<f32>,
11135        m: usize,
11136        on: bool,
11137        tag: &'static str,
11138    ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
11139        use crate::model::GpuTensor;
11140        if !on
11141            || !(2..=8).contains(&m)
11142            || std::env::var("MEMRA_NO_BATCHED").is_ok()
11143            || (m > 4 && !Self::b8_enabled())
11144            || !self.mmvq_supports(QT_NVFP4)
11145            || !self.batched_supports(QT_NVFP4)
11146        {
11147            return Ok(None);
11148        }
11149        let in_f = ws[0].in_features();
11150        let mut parts: Vec<(&CudaSlice<u8>, usize, f32)> = Vec::with_capacity(4);
11151        for w in ws {
11152            if !self.uses_q8_1_fast(w) || w.in_features() != in_f {
11153                return Ok(None);
11154            }
11155            match w {
11156                GpuTensor::Quant {
11157                    bytes,
11158                    qtype,
11159                    scale,
11160                    rp: true,
11161                    rp4: None,
11162                    ..
11163                } if *qtype == QT_NVFP4 && w.out_features() % 8 == 0 => {
11164                    parts.push((bytes, w.out_features(), *scale));
11165                }
11166                _ => return Ok(None),
11167            }
11168        }
11169        // MCOLS tier mirrors the singles: batched_mcols + the B567 exact-width law at m=5..7.
11170        static B567: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11171        let b567 = *B567.get_or_init(|| std::env::var("MEMRA_B567").as_deref() != Ok("0"));
11172        let mcols = if (5..=7).contains(&m) && b567 {
11173            m
11174        } else {
11175            Self::batched_mcols(m)
11176        };
11177        let kname: &'static str = match mcols {
11178            2 => "qmatvec_nvfp4_mmvq_group4_b2_rp",
11179            4 => "qmatvec_nvfp4_mmvq_group4_b4_rp",
11180            5 => "qmatvec_nvfp4_mmvq_group4_b5_rp",
11181            6 => "qmatvec_nvfp4_mmvq_group4_b6_rp",
11182            7 => "qmatvec_nvfp4_mmvq_group4_b7_rp",
11183            8 => "qmatvec_nvfp4_mmvq_group4_b8_rp",
11184            _ => return Ok(None),
11185        };
11186        // Engagement receipt PER DOOR (dead-arm lesson): one shared Once here suppressed
11187        // the second door's print on the slice-D battery — key the once-set by tag.
11188        if std::env::var("MEMRA_DEBUG").is_ok() {
11189            use std::sync::Mutex;
11190            static SEEN: Mutex<Vec<&'static str>> = Mutex::new(Vec::new());
11191            let mut seen = SEEN.lock().unwrap();
11192            if !seen.contains(&tag) {
11193                seen.push(tag);
11194                eprintln!("[memra] {tag} batched ENGAGED (m={m})");
11195            }
11196        }
11197        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
11198        let rows_per_block = ROWS_PER_BLOCK * 2; // WROWS=2 in the group kernel
11199        let total: usize = parts.iter().map(|p| p.1).sum();
11200        let three = parts.len() == 3;
11201        let mut y0 = self.alloc_uninit::<f32>(m * parts[0].1)?;
11202        let mut y1 = self.alloc_uninit::<f32>(m * parts[1].1)?;
11203        let mut y2 = self.alloc_uninit::<f32>(m * parts[2].1)?;
11204        // 3-tensor callers: n3=0 means no block ever resolves to the fourth range — W3/y3
11205        // are never dereferenced; a 1-element dummy keeps the launch ABI without aliasing y0.
11206        let mut y3 = self.alloc_uninit::<f32>(if three { 1 } else { m * parts[3].1 })?;
11207        let cfg = LaunchConfig {
11208            grid_dim: ((total as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
11209            block_dim: (32, ROWS_PER_BLOCK, 1),
11210            shared_mem_bytes: 0,
11211        };
11212        let (inf, mi) = (in_f as i32, m as i32);
11213        let (n0, n1, n2) = (parts[0].1 as i32, parts[1].1 as i32, parts[2].1 as i32);
11214        let n3 = if three { 0i32 } else { parts[3].1 as i32 };
11215        let (s0, s1, s2) = (parts[0].2, parts[1].2, parts[2].2);
11216        let s3 = if three { 1.0f32 } else { parts[3].2 };
11217        let w3 = if three { parts[0].0 } else { parts[3].0 };
11218        let f = self.func(kname);
11219        let __s_b = self.gpu.stream();
11220        let mut b = __s_b.launch_builder(&f);
11221        b.arg(parts[0].0)
11222            .arg(parts[1].0)
11223            .arg(parts[2].0)
11224            .arg(w3)
11225            .arg(aq)
11226            .arg(ad)
11227            .arg(&mut y0)
11228            .arg(&mut y1)
11229            .arg(&mut y2)
11230            .arg(&mut y3)
11231            .arg(&inf)
11232            .arg(&n0)
11233            .arg(&n1)
11234            .arg(&n2)
11235            .arg(&n3)
11236            .arg(&mi)
11237            .arg(&s0)
11238            .arg(&s1)
11239            .arg(&s2)
11240            .arg(&s3);
11241        unsafe {
11242            b.launch(cfg)?;
11243        }
11244        Ok(Some(if three {
11245            vec![y0, y1, y2]
11246        } else {
11247            vec![y0, y1, y2, y3]
11248        }))
11249    }
11250
11251    /// DUAL gate+up BATCHED matvec at verify t=2..8 (lane/verify-economics, 2026-08-02): ONE
11252    /// launch computes both FFN projections of a verify batch — same activation, same shape,
11253    /// blockIdx.y selects the tensor. Per (tensor, token, row) the kernel body is the single
11254    /// batched program on the SAME layout (split-plane rp: b2 rp / b4 rpr2 / b8 rpr2; GGUF:
11255    /// b2 base / b4 r2 / b8 r2) -> BIT-IDENTICAL to the two single `matmul_decode_exact`
11256    /// launches (kernel-check gates bitwise on both layouts; run-spec K=1..8 arbitrates e2e).
11257    /// The one activation quantize replaces two IDENTICAL quantizes of the same `x` (same
11258    /// kernel, same input -> same q8_1 bytes), and the two independent weight streams in one
11259    /// grid restore the memory-level parallelism the two-launch form loses to tail drain +
11260    /// launch gap (m=1 dual_mr2 precedent: DRAM 40% -> 47-50% on the 27B pair).
11261    /// `Some((y0, y1))` only when both tensors are NVFP4, the SAME layout (both rp or both
11262    /// GGUF, no rp4 mirror), identical (in_f, out_f, row_bytes), q8_1-fast, and m in 2..=4
11263    /// (the b2/b4 tiers = verify T for K=1..3, the profitable-K window — the b8 dual measured
11264    /// FLAT vs the rpsc singles x3 interleaved, research/verify-economics-20260802, and was
11265    /// killed per doctrine). None -> caller runs the two singles. MEMRA_SPEC_DUAL_T=0 rollback.
11266    pub fn matmul_decode_exact_dual(
11267        &self,
11268        w0: &crate::model::GpuTensor,
11269        w1: &crate::model::GpuTensor,
11270        x: &CudaSlice<f32>,
11271        m: usize,
11272    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
11273        use crate::model::GpuTensor;
11274        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11275        let on = *ON.get_or_init(|| {
11276            std::env::var("MEMRA_SPEC_DUAL_T")
11277                .map(|v| v != "0")
11278                .unwrap_or(true)
11279        });
11280        if !on
11281            || !(2..=4).contains(&m)
11282            || std::env::var("MEMRA_NO_BATCHED").is_ok()
11283            || !self.uses_q8_1_fast(w0)
11284            || !self.uses_q8_1_fast(w1)
11285        {
11286            return Ok(None);
11287        }
11288        // DECODE-PARITY GATE (lane/nvfp4-strict, 2026-08-05): batched iff MMVQ — same law as
11289        // the singles' `batched_supports && mmvq_supports` check in matmul_decode_exact,
11290        // which this dual door bypassed. Without MEMRA_MMVQ the m=1 decode is dp4a; the
11291        // verify must ride the per-column dp4a class, not the MMVQ-family dual.
11292        if !self.mmvq_supports(QT_NVFP4) {
11293            return Ok(None);
11294        }
11295        let (in_f, out_f) = (w0.in_features(), w0.out_features());
11296        if w1.in_features() != in_f || w1.out_features() != out_f {
11297            return Ok(None);
11298        }
11299        let (b0, b1, row_bytes, s0, s1, rp) = match (w0, w1) {
11300            (
11301                GpuTensor::Quant {
11302                    bytes: b0,
11303                    qtype: q0,
11304                    row_bytes: rb0,
11305                    scale: s0,
11306                    rp: rp0,
11307                    rp4: None,
11308                    ..
11309                },
11310                GpuTensor::Quant {
11311                    bytes: b1,
11312                    qtype: q1,
11313                    row_bytes: rb1,
11314                    scale: s1,
11315                    rp: rp1,
11316                    rp4: None,
11317                    ..
11318                },
11319            ) if *q0 == QT_NVFP4 && *q1 == QT_NVFP4 && rb0 == rb1 && rp0 == rp1 => {
11320                (b0, b1, *rb0, *s0, *s1, *rp0)
11321            }
11322            _ => return Ok(None),
11323        };
11324        // Engagement receipt (MEMRA_DEBUG=1): the first dead-arm A/B lesson — a `rp: false`
11325        // gate silently no-op'd the whole experiment; prove the arm is live in the log.
11326        if std::env::var("MEMRA_DEBUG").is_ok() {
11327            static ONCE: std::sync::Once = std::sync::Once::new();
11328            ONCE.call_once(|| eprintln!("[memra] dual gate+up batched ENGAGED (m={m} rp={rp})"));
11329        }
11330        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
11331        let (y0, y1) =
11332            self.qmatvec_batched_dual_raw(b0, b1, &aq, &ad, m, in_f, out_f, row_bytes, rp)?;
11333        let mut y0 = y0;
11334        let mut y1 = y1;
11335        if s0 != 1.0 {
11336            self.scale_inplace(&mut y0, s0, m * out_f)?;
11337        }
11338        if s1 != 1.0 {
11339            self.scale_inplace(&mut y1, s1, m * out_f)?;
11340        }
11341        Ok(Some((y0, y1)))
11342    }
11343
11344    /// Launch body of the dual batched twins from raw NVFP4 weight bytes + a pre-quantized q8_1
11345    /// activation (kernel-check's bit-equivalence entry; matmul_decode_exact_dual's core).
11346    /// mcols tier = batched_mcols(m); macro-scale NOT applied. `rp` selects the split-plane
11347    /// twins (both buffers must be the repacked layout).
11348    #[allow(clippy::too_many_arguments)]
11349    pub fn qmatvec_batched_dual_raw(
11350        &self,
11351        b0: &CudaSlice<u8>,
11352        b1: &CudaSlice<u8>,
11353        aq: &CudaSlice<i8>,
11354        ad: &CudaSlice<f32>,
11355        m: usize,
11356        in_f: usize,
11357        out_f: usize,
11358        row_bytes: usize,
11359        rp: bool,
11360    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11361        const ROWS_PER_BLOCK: u32 = 4;
11362        let mcols = Self::batched_mcols(m);
11363        // EXACT-WIDTH duals at m=5..7 (vt-fixes fix 1b): rp-only; bit-identical to the two
11364        // b5/b6/b7 singles (blockIdx.y selects the tensor, same template body).
11365        let tiny_rp1 = rp
11366            && mcols == 4
11367            && out_f <= 128
11368            && std::env::var("MEMRA_NVFP4_AUX_DUAL").as_deref() != Ok("0");
11369        let (name, rows_per_block) = if tiny_rp1 {
11370            ("qmatvec_nvfp4_mmvq_dual_b4_rp", ROWS_PER_BLOCK)
11371        } else {
11372            match (mcols, rp, m) {
11373                (2, false, _) => ("qmatvec_nvfp4_mmvq_dual_b2", ROWS_PER_BLOCK),
11374                (4, false, _) => ("qmatvec_nvfp4_mmvq_dual_b4_r2", ROWS_PER_BLOCK * 2),
11375                (2, true, _) => ("qmatvec_nvfp4_mmvq_dual_b2_rp", ROWS_PER_BLOCK),
11376                (4, true, _) => ("qmatvec_nvfp4_mmvq_dual_b4_rpr2", ROWS_PER_BLOCK * 2),
11377                (8, true, 5) => ("qmatvec_nvfp4_mmvq_dual_b5_rpr2", ROWS_PER_BLOCK * 2),
11378                (8, true, 6) => ("qmatvec_nvfp4_mmvq_dual_b6_rpr2", ROWS_PER_BLOCK * 2),
11379                (8, true, 7) => ("qmatvec_nvfp4_mmvq_dual_b7_rpr2", ROWS_PER_BLOCK * 2),
11380                _ => {
11381                    return Err(
11382                        format!("qmatvec_batched_dual_raw: no dual kernel for m {m}").into(),
11383                    );
11384                }
11385            }
11386        };
11387        let f = self.func(name);
11388        let mut y0 = self.alloc_uninit::<f32>(m * out_f)?;
11389        let mut y1 = self.alloc_uninit::<f32>(m * out_f)?;
11390        let cfg = LaunchConfig {
11391            grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 2, 1),
11392            block_dim: (32, ROWS_PER_BLOCK, 1),
11393            shared_mem_bytes: 0,
11394        };
11395        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
11396        let __s_b = self.gpu.stream();
11397        let mut b = __s_b.launch_builder(&f);
11398        b.arg(b0)
11399            .arg(b1)
11400            .arg(aq)
11401            .arg(ad)
11402            .arg(&mut y0)
11403            .arg(&mut y1)
11404            .arg(&inf)
11405            .arg(&outf)
11406            .arg(&mi)
11407            .arg(&rb);
11408        unsafe {
11409            b.launch(cfg)?;
11410        }
11411        Ok((y0, y1))
11412    }
11413
11414    /// Like `matmul_pre` but RETURNS THE RAW (un-macro-scaled) matmul output together with the
11415    /// per-tensor NVFP4 scale, instead of applying `scale_inplace` internally. Used by the fused
11416    /// SwiGLU epilogue (RANK3 LEVER 2) so the gate/up scales fold into one `silu_mul_scaled` launch.
11417    /// `Some((y_raw, scale))` only on the m==1 decode fast path (mmvq / dp4a) where the scale is a
11418    /// separate post-launch op we can defer; returns `None` for every other path (prefill GEMM, FP4
11419    /// GEMM, Stage-A, Float) so the caller falls back to the scaled `matmul_pre` + `silu_mul`.
11420    /// DUAL gate+up NVFP4 matvec (mm-fusion): ONE launch computes both projections (same
11421    /// activation, same shape) — grid.y selects the tensor. Bit-identical per element to two
11422    /// mr2 launches at m=1. Returns (gate_raw, up_raw) un-scaled (caller folds the two macro
11423    /// scales into the SwiGLU epilogue, same as the matmul_pre_noscale contract). None unless
11424    /// both tensors are NVFP4 q8_1-fast with identical (in_f, out_f, row_bytes) and m==1.
11425    pub fn matmul_pre_dual_noscale(
11426        &self,
11427        w0: &crate::model::GpuTensor,
11428        w1: &crate::model::GpuTensor,
11429        aq: &CudaSlice<i8>,
11430        ad: &CudaSlice<f32>,
11431        m: usize,
11432    ) -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn std::error::Error>>
11433    {
11434        use crate::model::GpuTensor;
11435        if m != 1 || !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) {
11436            return Ok(None);
11437        }
11438        // FP-ORDER LAW (lane/nvfp4-strict, 2026-08-05): every kernel this door can dispatch
11439        // (q8_0 fused2, nvfp4 dual_mr2) is the MMVQ family — 32-thread warp reduce. Without
11440        // MEMRA_MMVQ the m=1 singles ride dp4a (128-thread two-level reduce), so fusing here
11441        // would mix dispatch families across the pair — the exact class `q8_fused_params`
11442        // already refuses for Q8_0. The NVFP4 arm lacked this check, which is why
11443        // decode-batch-gate `--mode strict`'s equalizing env (MEMRA_MMVQ=0) never pinned
11444        // NVFP4 models: decode_step_h kept riding dual_mr2 while the batched body fell to
11445        // dp4a (gate1 maxdiff 1.639e-1 / gate2 step-8 divergence at the 2026-08-05 train
11446        // HEAD, research/nvfp4-strict-20260805/). Default env (MMVQ on) is dispatch-unchanged.
11447        if !self.mmvq_supports(QT_NVFP4) {
11448            return Ok(None);
11449        }
11450        let (in_f, out_f) = (w0.in_features(), w0.out_features());
11451        if w1.in_features() != in_f || w1.out_features() != out_f {
11452            return Ok(None);
11453        }
11454        // Q8_0 ARM (lane/q27-deepdive, 2026-08-05): the dense-FFN gate+up pair on a Q8_0 trunk fell
11455        // through this NVFP4-only gate to two `matmul_pre_noscale` launches — measured 128 of the
11456        // 1015 launches/token on q27-Q8_0 decode, the single largest un-fused class in the tick
11457        // (nsys `research/q27-deepdive-20260805/nsys/`). `q8_fused2_core` already serves the same
11458        // pair shape for the shared-expert gate/up, and its kernel body is `qmatvec_q8_0_mmvq`
11459        // VERBATIM per (tensor,row) -> BIT-IDENTICAL to the two separate launches. Q8_0 carries no
11460        // macro-scale (q8_fused_params requires scale==1.0), so the noscale contract is satisfied
11461        // by returning 1.0 for both: the SwiGLU epilogue's fold becomes the identity it already is
11462        // on this dtype today. Seam: MEMRA_Q8_FFN_FUSE2=0 rolls back to the two-launch pair.
11463        // rp4 guard: with MEMRA_Q8RP the singles route to the `_rp` split-plane twin over the
11464        // mirror buffer; the fused2 kernel has no `_rp` form, so fusing there would swap
11465        // dispatch families mid-model. Bail and let the two singles run (mirror lane unchanged).
11466        let no_mirror =
11467            |w: &crate::model::GpuTensor| !matches!(w, GpuTensor::Quant { rp4: Some(_), .. });
11468        if self.q8_ffn_fuse2_on()
11469            && no_mirror(w0)
11470            && no_mirror(w1)
11471            && let Some([p0, p1]) = self.q8_fused_params(&[w0, w1])
11472        {
11473            let (y0, y1) = self.q8_fused2_core(p0.0, p1.0, aq, ad, in_f, p0.1, p1.1, p0.2)?;
11474            return Ok(Some(((y0, 1.0), (y1, 1.0))));
11475        }
11476        // F8-E4M3 ARM (lane/fp8-decode-v1, 2026-08-05): with native e4m3 residency the FFN gate+up
11477        // pair (and the ssm beta+alpha dual, which routes through this same entry) fell through
11478        // both the NVFP4 gate below and the Q8_0 arm above to two `matmul_pre_noscale` launches —
11479        // native residency was UN-FUSING the trunk relative to the Q8_0 slab it replaces. The
11480        // fused2 kernel body is `qmatvec_e4m3_mmvq` VERBATIM per (tensor,row). Contract match:
11481        // `matmul_pre_noscale` on e4m3 launches with scale 1.0 and RETURNS the per-tensor
11482        // weight_scale for the caller to fold, so we pass ws=1.0 here and return (s0,s1) — same
11483        // bits, and the two macro-scale multiplies still fold into the SwiGLU epilogue.
11484        // MEMRA_E4M3_DUAL=0 rolls back to the two-launch pair.
11485        if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
11486            let (y0, y1) =
11487                self.e4m3_fused2_core(p0.0, p1.0, aq, ad, in_f, p0.1, p1.1, p0.2, 1.0, 1.0)?;
11488            return Ok(Some(((y0, p0.3), (y1, p1.3))));
11489        }
11490        let (b0, q0, rb0, s0, rp0) = match w0 {
11491            GpuTensor::Quant {
11492                bytes,
11493                qtype,
11494                row_bytes,
11495                scale,
11496                rp,
11497                ..
11498            } => (bytes, *qtype, *row_bytes, *scale, *rp),
11499            _ => return Ok(None),
11500        };
11501        let (b1, q1, rb1, s1, rp1) = match w1 {
11502            GpuTensor::Quant {
11503                bytes,
11504                qtype,
11505                row_bytes,
11506                scale,
11507                rp,
11508                ..
11509            } => (bytes, *qtype, *row_bytes, *scale, *rp),
11510            _ => return Ok(None),
11511        };
11512        if q0 != QT_NVFP4 || q1 != QT_NVFP4 || rb0 != rb1 || rp0 != rp1 {
11513            return Ok(None);
11514        }
11515        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
11516        const RPW: u32 = 2;
11517        let rows_per_block = ROWS_PER_BLOCK * RPW;
11518        let f = self.func(if rp0 {
11519            "qmatvec_nvfp4_mmvq_dual_mr2_rp"
11520        } else {
11521            "qmatvec_nvfp4_mmvq_dual_mr2"
11522        });
11523        let mut y0 = self.alloc_uninit::<f32>(out_f)?;
11524        let mut y1 = self.alloc_uninit::<f32>(out_f)?;
11525        let cfg = LaunchConfig {
11526            grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 2, 1),
11527            block_dim: (32, ROWS_PER_BLOCK, 1),
11528            shared_mem_bytes: 0,
11529        };
11530        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, 1i32, rb0 as i64);
11531        // noscale contract: the caller folds s0/s1 into the SwiGLU epilogue — the kernel's fused
11532        // yscale args stay 1.0 here (they exist for the single-tensor callers).
11533        let one = 1.0f32;
11534        let __s_b = self.gpu.stream();
11535        let mut b = __s_b.launch_builder(&f);
11536        b.arg(b0)
11537            .arg(b1)
11538            .arg(aq)
11539            .arg(ad)
11540            .arg(&mut y0)
11541            .arg(&mut y1)
11542            .arg(&inf)
11543            .arg(&outf)
11544            .arg(&mi)
11545            .arg(&rb)
11546            .arg(&one)
11547            .arg(&one);
11548        unsafe {
11549            b.launch(cfg)?;
11550        }
11551        Ok(Some(((y0, s0), (y1, s1))))
11552    }
11553
11554    /// FUSED NVFP4 matvec TRIPLE with unequal out_f (rig-native decode increment 1,
11555    /// lane/rig-native-nvfp4): wq+wk+wv in ONE launch via the q8_0 fused2 block-offset
11556    /// recipe. Per (tensor,row,t) the kernel body is nvfp4_mmvq_multirow_rp VERBATIM ->
11557    /// bit-identical to three separate `matmul_pre` launches; yscales fold in-kernel exactly
11558    /// as the singles do. grid.y = m (the t-parallel verify rows ride the same launch).
11559    /// None when ineligible (not all rp NVFP4 / in_f mismatch / mmvq off) — callers fall
11560    /// back to the three singles.
11561    #[allow(clippy::too_many_arguments)]
11562    pub fn matmul_nvfp4_fused3(
11563        &self,
11564        w0: &crate::model::GpuTensor,
11565        w1: &crate::model::GpuTensor,
11566        w2: &crate::model::GpuTensor,
11567        aq: &CudaSlice<i8>,
11568        ad: &CudaSlice<f32>,
11569        m: usize,
11570    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
11571    {
11572        use crate::model::GpuTensor;
11573        // m == 1 rides the original fused kernel; m = 2..=8 rides the BATCHED fused twin
11574        // (lane/samplat — the fused4_b8 pattern): the seg body is nvfp4_mmvq_batched_rp_sc
11575        // verbatim, weight rows read once for all m columns, bit-identical per
11576        // (tensor,row,column) to the three bN_rpsc singles. The old "at m>1 the fused
11577        // segments would re-read the weight per row" note described the grid.y=m lift,
11578        // which this twin deliberately is NOT.
11579        if !(1..=8).contains(&m)
11580            || !self.mmvq_supports(QT_NVFP4)
11581            || !self.uses_q8_1_fast(w0)
11582            || !self.uses_q8_1_fast(w1)
11583            || !self.uses_q8_1_fast(w2)
11584        {
11585            return Ok(None);
11586        }
11587        if m > 1 {
11588            let in_f = w0.in_features();
11589            if std::env::var("MEMRA_NVFP4_FUSED3B").as_deref() == Ok("0")
11590                || !self.batched_supports(QT_NVFP4)
11591                || std::env::var("MEMRA_NO_BATCHED").is_ok()
11592                || (m > 4 && !Self::b8_enabled())
11593                || in_f % 512 != 0
11594                || in_f / 64 > 272
11595            {
11596                return Ok(None);
11597            }
11598        }
11599        let unpack = |w: &crate::model::GpuTensor| match w {
11600            GpuTensor::Quant {
11601                bytes,
11602                qtype,
11603                scale,
11604                rp,
11605                ..
11606            } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
11607            _ => None,
11608        };
11609        let (Some(p0), Some(p1), Some(p2)) = (unpack(w0), unpack(w1), unpack(w2)) else {
11610            return Ok(None);
11611        };
11612        let in_f = w0.in_features();
11613        if w1.in_features() != in_f || w2.in_features() != in_f {
11614            return Ok(None);
11615        }
11616        let (o0, o1, o2) = (w0.out_features(), w1.out_features(), w2.out_features());
11617        const ROWS_PER_BLOCK: u32 = 4; // MEMRA_MMVQ_ROWS
11618        const RPW: u32 = 2;
11619        let rows_pb = ROWS_PER_BLOCK * RPW;
11620        let nb = |o: usize| (o as u32).div_ceil(rows_pb);
11621        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
11622        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
11623        let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
11624        let (inf, oi0, oi1, oi2, mi) = (in_f as i32, o0 as i32, o1 as i32, o2 as i32, m as i32);
11625        // SAFETY: the raw pointers come straight from the &GpuTensor borrows above and are
11626        // only dereferenced for the launch-arg build inside this call.
11627        let (b0, b1, b2) = unsafe { (&*p0.0, &*p1.0, &*p2.0) };
11628        if m > 1 {
11629            // batched twin has no in-kernel scale — refuse scale carriers (GGUF trunk = 1.0).
11630            if p0.1 != 1.0 || p1.1 != 1.0 || p2.1 != 1.0 {
11631                return Ok(None);
11632            }
11633            let f = self.func("qmatvec_nvfp4_mmvq_fused3_b8_rpsc");
11634            let cfg = LaunchConfig {
11635                grid_dim: (nb(o0) + nb(o1) + nb(o2), 1, 1),
11636                block_dim: (32, ROWS_PER_BLOCK, 1),
11637                shared_mem_bytes: 0,
11638            };
11639            let __s_b = self.gpu.stream();
11640            let mut b = __s_b.launch_builder(&f);
11641            b.arg(b0)
11642                .arg(b1)
11643                .arg(b2)
11644                .arg(aq)
11645                .arg(ad)
11646                .arg(&mut y0)
11647                .arg(&mut y1)
11648                .arg(&mut y2)
11649                .arg(&inf)
11650                .arg(&oi0)
11651                .arg(&oi1)
11652                .arg(&oi2)
11653                .arg(&mi);
11654            unsafe {
11655                b.launch(cfg)?;
11656            }
11657            return Ok(Some((y0, y1, y2)));
11658        }
11659        let f = self.func("qmatvec_nvfp4_mmvq_fused3_rp");
11660        let cfg = LaunchConfig {
11661            grid_dim: (nb(o0) + nb(o1) + nb(o2), m as u32, 1),
11662            block_dim: (32, ROWS_PER_BLOCK, 1),
11663            shared_mem_bytes: 0,
11664        };
11665        let __s_b = self.gpu.stream();
11666        let mut b = __s_b.launch_builder(&f);
11667        b.arg(b0)
11668            .arg(b1)
11669            .arg(b2)
11670            .arg(aq)
11671            .arg(ad)
11672            .arg(&mut y0)
11673            .arg(&mut y1)
11674            .arg(&mut y2)
11675            .arg(&inf)
11676            .arg(&oi0)
11677            .arg(&oi1)
11678            .arg(&oi2)
11679            .arg(&mi)
11680            .arg(&p0.1)
11681            .arg(&p1.1)
11682            .arg(&p2.1);
11683        unsafe {
11684            b.launch(cfg)?;
11685        }
11686        Ok(Some((y0, y1, y2)))
11687    }
11688
11689    /// fused2 twin of `matmul_nvfp4_fused3`, for MIXED-type weight groups: the gemma4
11690    /// dense NVFP4mix recipe keeps `attn_v` and `ffn_down` at Q8_0 (full-NVFP4 was
11691    /// measured garbage on this dense class), so its q/k/v trio and gate/up/down never
11692    /// satisfy an all-NVFP4 fused3 — the pairs that ARE uniformly NVFP4 (q,k and
11693    /// gate,up) fuse here instead. m==1 only, same law as fused3/fused4: per
11694    /// (tensor,row) the kernel seg body is VERBATIM, so the fusion is bit-identical to
11695    /// two separate launches. `MEMRA_NVFP4_FUSED2=0` is the rollback seam and the
11696    /// same-binary interleaved A/B arm.
11697    pub fn matmul_nvfp4_fused2(
11698        &self,
11699        w0: &crate::model::GpuTensor,
11700        w1: &crate::model::GpuTensor,
11701        aq: &CudaSlice<i8>,
11702        ad: &CudaSlice<f32>,
11703        m: usize,
11704    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
11705        use crate::model::GpuTensor;
11706        static FUSED2_OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11707        let off =
11708            *FUSED2_OFF.get_or_init(|| std::env::var("MEMRA_NVFP4_FUSED2").as_deref() == Ok("0"));
11709        // m==1 ONLY: at m>1 the singles ride the _b16 weight-once column program (one weight
11710        // read serves all m rows); the fused segments would re-read the weight per row.
11711        if off
11712            || m != 1
11713            || !self.mmvq_supports(QT_NVFP4)
11714            || !self.uses_q8_1_fast(w0)
11715            || !self.uses_q8_1_fast(w1)
11716        {
11717            return Ok(None);
11718        }
11719        let unpack = |w: &crate::model::GpuTensor| match w {
11720            GpuTensor::Quant {
11721                bytes,
11722                qtype,
11723                scale,
11724                rp,
11725                ..
11726            } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
11727            _ => None,
11728        };
11729        let (Some(p0), Some(p1)) = (unpack(w0), unpack(w1)) else {
11730            return Ok(None);
11731        };
11732        let in_f = w0.in_features();
11733        if w1.in_features() != in_f {
11734            return Ok(None);
11735        }
11736        let (o0, o1) = (w0.out_features(), w1.out_features());
11737        const ROWS_PER_BLOCK: u32 = 4; // MEMRA_MMVQ_ROWS
11738        const RPW: u32 = 2;
11739        let rows_pb = ROWS_PER_BLOCK * RPW;
11740        let nb = |o: usize| (o as u32).div_ceil(rows_pb);
11741        let f = self.func("qmatvec_nvfp4_mmvq_fused2_rp");
11742        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
11743        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
11744        let cfg = LaunchConfig {
11745            grid_dim: (nb(o0) + nb(o1), m as u32, 1),
11746            block_dim: (32, ROWS_PER_BLOCK, 1),
11747            shared_mem_bytes: 0,
11748        };
11749        let (inf, oi0, oi1, mi) = (in_f as i32, o0 as i32, o1 as i32, m as i32);
11750        // SAFETY: the raw pointers come straight from the &GpuTensor borrows above and are
11751        // only dereferenced for the launch-arg build inside this call.
11752        let (b0, b1) = unsafe { (&*p0.0, &*p1.0) };
11753        // PDL wave-B: the kernel carries MEMRA_PDL_ENTRY — launch overlaps the producer's
11754        // drain (quantize/norm emitting aq/ad). Same math, same order, bit-identical.
11755        if Self::pdl_on() && Self::pdl_mmvq_on() && Self::pdl_nvfp4q8_on() {
11756            {
11757                use cudarc::driver::{DevicePtr, DevicePtrMut};
11758                let s = &self.gpu.stream();
11759                let (pw0, _g0) = b0.device_ptr(s);
11760                let (pw1, _g1) = b1.device_ptr(s);
11761                let (paq, _g2) = aq.device_ptr(s);
11762                let (pad, _g3) = ad.device_ptr(s);
11763                let (py0, _g4) = y0.device_ptr_mut(s);
11764                let (py1, _g5) = y1.device_ptr_mut(s);
11765                let (s0, s1) = (p0.1, p1.1);
11766                let mut ps = [
11767                    &pw0 as *const _ as *mut std::ffi::c_void,
11768                    &pw1 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                    &inf as *const _ as *mut _,
11774                    &oi0 as *const _ as *mut _,
11775                    &oi1 as *const _ as *mut _,
11776                    &mi as *const _ as *mut _,
11777                    &s0 as *const _ as *mut _,
11778                    &s1 as *const _ as *mut _,
11779                ];
11780                unsafe {
11781                    self.launch_pdl(
11782                        "qmatvec_nvfp4_mmvq_fused2_rp",
11783                        cfg.grid_dim,
11784                        cfg.block_dim,
11785                        &mut ps,
11786                    )?;
11787                }
11788            }
11789            return Ok(Some((y0, y1)));
11790        }
11791        let __s_b = self.gpu.stream();
11792        let mut b = __s_b.launch_builder(&f);
11793        b.arg(b0)
11794            .arg(b1)
11795            .arg(aq)
11796            .arg(ad)
11797            .arg(&mut y0)
11798            .arg(&mut y1)
11799            .arg(&inf)
11800            .arg(&oi0)
11801            .arg(&oi1)
11802            .arg(&mi)
11803            .arg(&p0.1)
11804            .arg(&p1.1);
11805        unsafe {
11806            b.launch(cfg)?;
11807        }
11808        Ok(Some((y0, y1)))
11809    }
11810
11811    /// Slot-fed fused2 twin (alloc-free capture lane): identical launch to
11812    /// `matmul_nvfp4_fused2`, caller-owned outputs — the gemma4 dc_slotted graph body
11813    /// needs zero mem nodes, so the allocating wrapper can't serve it. Returns false
11814    /// (decline) on any non-NVFP4/rp pair; the caller chains or refuses.
11815    pub fn matmul_nvfp4_fused2_into(
11816        &self,
11817        w0: &crate::model::GpuTensor,
11818        w1: &crate::model::GpuTensor,
11819        aq: &CudaSlice<i8>,
11820        ad: &CudaSlice<f32>,
11821        y0: &mut CudaSlice<f32>,
11822        y1: &mut CudaSlice<f32>,
11823    ) -> Result<bool, Box<dyn std::error::Error>> {
11824        use crate::model::GpuTensor;
11825        static FUSED2_OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11826        let off =
11827            *FUSED2_OFF.get_or_init(|| std::env::var("MEMRA_NVFP4_FUSED2").as_deref() == Ok("0"));
11828        if off
11829            || !self.mmvq_supports(QT_NVFP4)
11830            || !self.uses_q8_1_fast(w0)
11831            || !self.uses_q8_1_fast(w1)
11832        {
11833            return Ok(false);
11834        }
11835        let unpack = |w: &crate::model::GpuTensor| match w {
11836            GpuTensor::Quant {
11837                bytes,
11838                qtype,
11839                scale,
11840                rp,
11841                ..
11842            } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
11843            _ => None,
11844        };
11845        let (Some(p0), Some(p1)) = (unpack(w0), unpack(w1)) else {
11846            return Ok(false);
11847        };
11848        let in_f = w0.in_features();
11849        if w1.in_features() != in_f {
11850            return Ok(false);
11851        }
11852        let (o0, o1) = (w0.out_features(), w1.out_features());
11853        if y0.len() < o0 || y1.len() < o1 {
11854            return Ok(false);
11855        }
11856        const ROWS_PER_BLOCK: u32 = 4; // MEMRA_MMVQ_ROWS
11857        const RPW: u32 = 2;
11858        let rows_pb = ROWS_PER_BLOCK * RPW;
11859        let nb = |o: usize| (o as u32).div_ceil(rows_pb);
11860        let f = self.func("qmatvec_nvfp4_mmvq_fused2_rp");
11861        let cfg = LaunchConfig {
11862            grid_dim: (nb(o0) + nb(o1), 1, 1),
11863            block_dim: (32, ROWS_PER_BLOCK, 1),
11864            shared_mem_bytes: 0,
11865        };
11866        let (inf, oi0, oi1, mi) = (in_f as i32, o0 as i32, o1 as i32, 1i32);
11867        // SAFETY: the raw pointers come straight from the &GpuTensor borrows above and are
11868        // only dereferenced for the launch-arg build inside this call.
11869        let (b0, b1) = unsafe { (&*p0.0, &*p1.0) };
11870        let __s_b = self.gpu.stream();
11871        let mut b = __s_b.launch_builder(&f);
11872        b.arg(b0)
11873            .arg(b1)
11874            .arg(aq)
11875            .arg(ad)
11876            .arg(&mut *y0)
11877            .arg(&mut *y1)
11878            .arg(&inf)
11879            .arg(&oi0)
11880            .arg(&oi1)
11881            .arg(&mi)
11882            .arg(&p0.1)
11883            .arg(&p1.1);
11884        unsafe {
11885            b.launch(cfg)?;
11886        }
11887        Ok(true)
11888    }
11889
11890    /// fused4 twin of `matmul_nvfp4_fused3`: the Linear-mixer projection quartet
11891    /// (wqkv + wqkv_gate + ssm_beta + ssm_alpha) in one launch, m==1 only. Per
11892    /// (tensor,row) the kernel body is nvfp4_mmvq_multirow_rp VERBATIM — bit-identical
11893    /// to four separate launches (rig-native decode increment 2, RIG-NATIVE-DECODE.md).
11894    #[allow(clippy::type_complexity)]
11895    pub fn matmul_nvfp4_fused4(
11896        &self,
11897        w0: &crate::model::GpuTensor,
11898        w1: &crate::model::GpuTensor,
11899        w2: &crate::model::GpuTensor,
11900        w3: &crate::model::GpuTensor,
11901        aq: &CudaSlice<i8>,
11902        ad: &CudaSlice<f32>,
11903        m: usize,
11904    ) -> Result<
11905        Option<(
11906            CudaSlice<f32>,
11907            CudaSlice<f32>,
11908            CudaSlice<f32>,
11909            CudaSlice<f32>,
11910        )>,
11911        Box<dyn std::error::Error>,
11912    > {
11913        use crate::model::GpuTensor;
11914        // MEMRA_NVFP4_FUSED4=0: rollback seam + the same-binary interleaved A/B arm.
11915        // m == 1 rides the original fused kernel; m = 2..=8 rides the BATCHED fused twin
11916        // (lane/samplat, 2026-08-21): same quartet-in-one-launch shape, seg body =
11917        // nvfp4_mmvq_batched_rp_sc verbatim (weight rows read once for all m columns) —
11918        // bit-identical per (tensor,row,column) to the four bN_rpsc singles it replaces.
11919        // Admission mirrors the singles' batched gates below.
11920        if !(1..=8).contains(&m)
11921            || std::env::var("MEMRA_NVFP4_FUSED4").as_deref() == Ok("0")
11922            || !self.mmvq_supports(QT_NVFP4)
11923            || !self.uses_q8_1_fast(w0)
11924            || !self.uses_q8_1_fast(w1)
11925            || !self.uses_q8_1_fast(w2)
11926            || !self.uses_q8_1_fast(w3)
11927        {
11928            return Ok(None);
11929        }
11930        if m > 1 {
11931            // the batched-twin gates: the bN_rpsc program this must stay byte-identical to
11932            // (matmul_pre's batched arm), plus the rp-sc dispatch shape requirements.
11933            let in_f = w0.in_features();
11934            if !self.batched_supports(QT_NVFP4)
11935                || std::env::var("MEMRA_NO_BATCHED").is_ok()
11936                || (m > 4 && !Self::b8_enabled())
11937                || in_f % 512 != 0
11938                || in_f / 64 > 272
11939            {
11940                return Ok(None);
11941            }
11942        }
11943        let unpack = |w: &crate::model::GpuTensor| match w {
11944            GpuTensor::Quant {
11945                bytes,
11946                qtype,
11947                scale,
11948                rp,
11949                ..
11950            } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
11951            _ => None,
11952        };
11953        let (Some(p0), Some(p1), Some(p2), Some(p3)) =
11954            (unpack(w0), unpack(w1), unpack(w2), unpack(w3))
11955        else {
11956            return Ok(None);
11957        };
11958        let in_f = w0.in_features();
11959        if w1.in_features() != in_f || w2.in_features() != in_f || w3.in_features() != in_f {
11960            return Ok(None);
11961        }
11962        let (o0, o1, o2, o3) = (
11963            w0.out_features(),
11964            w1.out_features(),
11965            w2.out_features(),
11966            w3.out_features(),
11967        );
11968        const ROWS_PER_BLOCK: u32 = 4; // MEMRA_MMVQ_ROWS
11969        const RPW: u32 = 2;
11970        let rows_pb = ROWS_PER_BLOCK * RPW;
11971        let nb = |o: usize| (o as u32).div_ceil(rows_pb);
11972        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
11973        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
11974        let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
11975        let mut y3 = self.alloc_uninit::<f32>(m * o3)?;
11976        let (inf, oi0, oi1, oi2, oi3, mi) = (
11977            in_f as i32,
11978            o0 as i32,
11979            o1 as i32,
11980            o2 as i32,
11981            o3 as i32,
11982            m as i32,
11983        );
11984        // SAFETY: the raw pointers come straight from the &GpuTensor borrows above and are
11985        // only dereferenced for the launch-arg build inside this call.
11986        let (b0, b1, b2, b3) = unsafe { (&*p0.0, &*p1.0, &*p2.0, &*p3.0) };
11987        if m > 1 {
11988            // Batched fused twin: no in-kernel scale (the bN_rpsc program has none) — refuse
11989            // scale-carrying tensors so the singles path keeps them (GGUF trunk scales are 1.0).
11990            if p0.1 != 1.0 || p1.1 != 1.0 || p2.1 != 1.0 || p3.1 != 1.0 {
11991                return Ok(None);
11992            }
11993            let f = self.func("qmatvec_nvfp4_mmvq_fused4_b8_rpsc");
11994            let cfg = LaunchConfig {
11995                grid_dim: (nb(o0) + nb(o1) + nb(o2) + nb(o3), 1, 1),
11996                block_dim: (32, ROWS_PER_BLOCK, 1),
11997                shared_mem_bytes: 0,
11998            };
11999            let __s_b = self.gpu.stream();
12000            let mut b = __s_b.launch_builder(&f);
12001            b.arg(b0)
12002                .arg(b1)
12003                .arg(b2)
12004                .arg(b3)
12005                .arg(aq)
12006                .arg(ad)
12007                .arg(&mut y0)
12008                .arg(&mut y1)
12009                .arg(&mut y2)
12010                .arg(&mut y3)
12011                .arg(&inf)
12012                .arg(&oi0)
12013                .arg(&oi1)
12014                .arg(&oi2)
12015                .arg(&oi3)
12016                .arg(&mi);
12017            unsafe {
12018                b.launch(cfg)?;
12019            }
12020            return Ok(Some((y0, y1, y2, y3)));
12021        }
12022        let f = self.func("qmatvec_nvfp4_mmvq_fused4_rp");
12023        let cfg = LaunchConfig {
12024            grid_dim: (nb(o0) + nb(o1) + nb(o2) + nb(o3), m as u32, 1),
12025            block_dim: (32, ROWS_PER_BLOCK, 1),
12026            shared_mem_bytes: 0,
12027        };
12028        let __s_b = self.gpu.stream();
12029        let mut b = __s_b.launch_builder(&f);
12030        b.arg(b0)
12031            .arg(b1)
12032            .arg(b2)
12033            .arg(b3)
12034            .arg(aq)
12035            .arg(ad)
12036            .arg(&mut y0)
12037            .arg(&mut y1)
12038            .arg(&mut y2)
12039            .arg(&mut y3)
12040            .arg(&inf)
12041            .arg(&oi0)
12042            .arg(&oi1)
12043            .arg(&oi2)
12044            .arg(&oi3)
12045            .arg(&mi)
12046            .arg(&p0.1)
12047            .arg(&p1.1)
12048            .arg(&p2.1)
12049            .arg(&p3.1);
12050        unsafe {
12051            b.launch(cfg)?;
12052        }
12053        Ok(Some((y0, y1, y2, y3)))
12054    }
12055
12056    /// FUSED Q8_0 m=1 matvec PAIR with UNEQUAL out_f (trunk launch-fusion, 2026-07-05). Folds two
12057    /// same-input q8_0 projections (35B trunk: wqkv+wqkv_gate 8192/4096, gate_shexp+up_shexp
12058    /// 512/512) into ONE launch via a block-offset split (blocks [0,nb0) -> w0, rest -> w1) — the
12059    /// dual-mr2 recipe with the same-out_f restriction lifted. Per (tensor,row) the kernel body is
12060    /// qmatvec_q8_0_mmvq VERBATIM -> BIT-IDENTICAL to two separate m=1 launches. Returns None when
12061    /// ineligible (not both Q8_0 / in_f mismatch / MEMRA_MMVQ off / MEMRA_Q8_DUAL=0) — caller falls
12062    /// back to the per-tensor path.
12063    pub fn matmul_q8_fused2(
12064        &self,
12065        w0: &crate::model::GpuTensor,
12066        w1: &crate::model::GpuTensor,
12067        aq: &CudaSlice<i8>,
12068        ad: &CudaSlice<f32>,
12069    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
12070        // e4m3 twin (lane/fp8-decode-v1): this entry is the trunk's generic m=1 pair door
12071        // (wqkv+wqkv_gate, ssm_beta+alpha, gate_shexp+up_shexp), so admitting QT_F8_E4M3 here
12072        // fuses the NATIVE-RESIDENCY FP8 trunk at every existing call site with no call-site
12073        // change. Scale is folded in-kernel per range -> the returned buffers are already scaled,
12074        // exactly like the per-tensor `matmul_pre` e4m3 dispatch this replaces.
12075        if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
12076            return Ok(Some(self.e4m3_fused2_core(
12077                p0.0,
12078                p1.0,
12079                aq,
12080                ad,
12081                w0.in_features(),
12082                p0.1,
12083                p1.1,
12084                p0.2,
12085                p0.3,
12086                p1.3,
12087            )?));
12088        }
12089        let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
12090            return Ok(None);
12091        };
12092        Ok(Some(self.q8_fused2_core(
12093            p0.0,
12094            p1.0,
12095            aq,
12096            ad,
12097            w0.in_features(),
12098            p0.1,
12099            p1.1,
12100            p0.2,
12101        )?))
12102    }
12103
12104    #[allow(clippy::too_many_arguments)]
12105    fn q8_fused2_core(
12106        &self,
12107        b0: &CudaSlice<u8>,
12108        b1: &CudaSlice<u8>,
12109        aq: &CudaSlice<i8>,
12110        ad: &CudaSlice<f32>,
12111        in_f: usize,
12112        out0: usize,
12113        out1: usize,
12114        row_bytes: usize,
12115    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12116        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
12117        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
12118        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
12119        let f = self.func("qmatvec_q8_0_mmvq_fused2");
12120        let mut y0 = self.alloc_uninit::<f32>(out0)?;
12121        let mut y1 = self.alloc_uninit::<f32>(out1)?;
12122        let cfg = LaunchConfig {
12123            grid_dim: (nb0 + nb1, 1, 1),
12124            block_dim: (32, ROWS_PER_BLOCK, 1),
12125            shared_mem_bytes: 0,
12126        };
12127        let (inf, o0, o1, rbl) = (in_f as i32, out0 as i32, out1 as i32, row_bytes as i64);
12128        let __s_b = self.gpu.stream();
12129        let mut b = __s_b.launch_builder(&f);
12130        b.arg(b0)
12131            .arg(b1)
12132            .arg(aq)
12133            .arg(ad)
12134            .arg(&mut y0)
12135            .arg(&mut y1)
12136            .arg(&inf)
12137            .arg(&o0)
12138            .arg(&o1)
12139            .arg(&rbl);
12140        unsafe {
12141            b.launch(cfg)?;
12142        }
12143        Ok((y0, y1))
12144    }
12145
12146    /// f32-activation entry for the fused2 pair: quantizes x to q8_1 ONCE then runs the fused
12147    /// launch — replaces two `matmul(w, x, 1)` calls that would each re-quantize the same x
12148    /// (35B shared-expert gate+up per MoE layer per token). Same bits: quantize_q8_1 is
12149    /// deterministic, the fused body is the MMVQ kernel verbatim. None when ineligible (the
12150    /// callers' m==1-under-MEMRA_FAST dispatch would take MMVQ; anything else falls back).
12151    pub fn matmul_q8_fused2_x(
12152        &self,
12153        w0: &crate::model::GpuTensor,
12154        w1: &crate::model::GpuTensor,
12155        x: &CudaSlice<f32>,
12156    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
12157        if !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) {
12158            return Ok(None);
12159        }
12160        if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
12161            let (aq, ad) = self.quantize_q8_1(x, 1, w0.in_features())?;
12162            return Ok(Some(self.e4m3_fused2_core(
12163                p0.0,
12164                p1.0,
12165                &aq,
12166                &ad,
12167                w0.in_features(),
12168                p0.1,
12169                p1.1,
12170                p0.2,
12171                p0.3,
12172                p1.3,
12173            )?));
12174        }
12175        let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
12176            return Ok(None);
12177        };
12178        let (aq, ad) = self.quantize_q8_1(x, 1, w0.in_features())?;
12179        Ok(Some(self.q8_fused2_core(
12180            p0.0,
12181            p1.0,
12182            &aq,
12183            &ad,
12184            w0.in_features(),
12185            p0.1,
12186            p1.1,
12187            p0.2,
12188        )?))
12189    }
12190
12191    /// Test entry for the kernel_check gate: launch the fused2 kernel from raw weight bytes,
12192    /// quantizing the f32 activation internally (mirrors qmatvec_mmvq_raw; no env gating).
12193    #[allow(clippy::too_many_arguments)]
12194    pub fn qmatvec_q8_fused2_raw(
12195        &self,
12196        b0: &CudaSlice<u8>,
12197        b1: &CudaSlice<u8>,
12198        x: &CudaSlice<f32>,
12199        in_f: usize,
12200        out0: usize,
12201        out1: usize,
12202        row_bytes: usize,
12203    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12204        let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
12205        self.q8_fused2_core(b0, b1, &aq, &ad, in_f, out0, out1, row_bytes)
12206    }
12207
12208    /// FUSED Q8_0 m=1 matvec TRIPLE (wq+wk+wv on the 35B full-attn layers: out_f 8192/512/512).
12209    /// Same block-offset recipe as `matmul_q8_fused2` with three ranges. BIT-IDENTICAL per
12210    /// (tensor,row) to three separate m=1 MMVQ launches.
12211    /// FUSED Q4_0 m=1 TRIPLE (gemma q/k/v — same quantized input; per (tensor,row) chain
12212    /// identical to the mr2 kernel). Returns None unless all three are Q4_0 with equal in_f.
12213    pub fn matmul_q4_fused3(
12214        &self,
12215        w0: &crate::model::GpuTensor,
12216        w1: &crate::model::GpuTensor,
12217        w2: &crate::model::GpuTensor,
12218        aq: &CudaSlice<i8>,
12219        ad: &CudaSlice<f32>,
12220    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
12221    {
12222        use crate::model::GpuTensor;
12223        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
12224            match w {
12225                GpuTensor::Quant {
12226                    qtype, row_bytes, ..
12227                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
12228                _ => None,
12229            }
12230        };
12231        let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (q4(w0), q4(w1), q4(w2)) else {
12232            return Ok(None);
12233        };
12234        if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
12235            return Ok(None);
12236        }
12237        // Effective (bytes, rp) per tensor: mirror (rp4) OR the in-place swap (rp flag,
12238        // bytes already split). Mixed layouts cannot share one fused launch -> fall back to
12239        // the separate matvecs (each routes its own rp).
12240        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
12241            match w {
12242                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
12243                    Some(m) => (m, true),
12244                    None => (bytes, *rp),
12245                },
12246                _ => unreachable!(),
12247            }
12248        }
12249        let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
12250        if rp0 != rp1 || rp1 != rp2 {
12251            return Ok(None);
12252        }
12253        let rp = rp0;
12254        let rpb: u32 = 4;
12255        // mr1 (one row/warp, 2026-07-14): follows the singles' MEMRA_Q40_MR default — the
12256        // fused t=1 kernels were left on mr2 when the singles flipped (DRAM-duty map:
12257        // fused3 57% / fused2 86%; small qkv segments starve under mr2's half grid).
12258        let mr1 = rp && Self::q40_mr1_on();
12259        let nb = |o: usize| {
12260            if mr1 {
12261                (o as u32).div_ceil(rpb)
12262            } else {
12263                (o as u32).div_ceil(2).div_ceil(rpb)
12264            }
12265        };
12266        let grid = nb(o0) + nb(o1) + nb(o2);
12267        let mut y0 = self.alloc_uninit::<f32>(o0)?;
12268        let mut y1 = self.alloc_uninit::<f32>(o1)?;
12269        let mut y2 = self.alloc_uninit::<f32>(o2)?;
12270        let f = self.func(if mr1 {
12271            "qmatvec_q4_0_mmvq_fused3_mr1_rp"
12272        } else if rp {
12273            "qmatvec_q4_0_mmvq_fused3_rp"
12274        } else {
12275            "qmatvec_q4_0_mmvq_fused3"
12276        });
12277        let cfg = LaunchConfig {
12278            grid_dim: (grid, 1, 1),
12279            block_dim: (32, rpb, 1),
12280            shared_mem_bytes: 0,
12281        };
12282        let inf = w0.in_features() as i32;
12283        let (oo0, oo1, oo2) = (o0 as i32, o1 as i32, o2 as i32);
12284        let (r0, r1, r2) = (rb0 as i64, rb1 as i64, rb2 as i64);
12285        // PDL wave-A (2026-07-23): the mr1 kernel carries MEMRA_PDL_ENTRY; only that
12286        // variant may take the programmatic-serialization launch.
12287        if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
12288            {
12289                use cudarc::driver::{DevicePtr, DevicePtrMut};
12290                let s = &self.gpu.stream();
12291                let (p0, _g0) = b0.device_ptr(s);
12292                let (p1, _g1) = b1.device_ptr(s);
12293                let (p2, _g2) = b2.device_ptr(s);
12294                let (paq, _g3) = aq.device_ptr(s);
12295                let (pad, _g4) = ad.device_ptr(s);
12296                let (py0, _g5) = y0.device_ptr_mut(s);
12297                let (py1, _g6) = y1.device_ptr_mut(s);
12298                let (py2, _g7) = y2.device_ptr_mut(s);
12299                let mut ps = [
12300                    &p0 as *const _ as *mut std::ffi::c_void,
12301                    &p1 as *const _ as *mut _,
12302                    &p2 as *const _ as *mut _,
12303                    &paq as *const _ as *mut _,
12304                    &pad as *const _ as *mut _,
12305                    &py0 as *const _ as *mut _,
12306                    &py1 as *const _ as *mut _,
12307                    &py2 as *const _ as *mut _,
12308                    &inf as *const _ as *mut _,
12309                    &oo0 as *const _ as *mut _,
12310                    &oo1 as *const _ as *mut _,
12311                    &oo2 as *const _ as *mut _,
12312                    &r0 as *const _ as *mut _,
12313                    &r1 as *const _ as *mut _,
12314                    &r2 as *const _ as *mut _,
12315                ];
12316                unsafe {
12317                    self.launch_pdl(
12318                        "qmatvec_q4_0_mmvq_fused3_mr1_rp",
12319                        (grid, 1, 1),
12320                        (32, rpb, 1),
12321                        &mut ps,
12322                    )?;
12323                }
12324            }
12325            return Ok(Some((y0, y1, y2)));
12326        }
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(&r0)
12342            .arg(&r1)
12343            .arg(&r2);
12344        unsafe {
12345            b.launch(cfg)?;
12346        }
12347        Ok(Some((y0, y1, y2)))
12348    }
12349
12350    /// Slot-fed fused3 twin (alloc-free capture lane): identical launch, caller-owned outputs.
12351    /// Returns Ok(false) when the fused path is unavailable (caller falls back).
12352    #[allow(clippy::too_many_arguments)]
12353    pub fn matmul_q4_fused3_into(
12354        &self,
12355        w0: &crate::model::GpuTensor,
12356        w1: &crate::model::GpuTensor,
12357        w2: &crate::model::GpuTensor,
12358        aq: &CudaSlice<i8>,
12359        ad: &CudaSlice<f32>,
12360        y0: &mut CudaSlice<f32>,
12361        y1: &mut CudaSlice<f32>,
12362        y2: &mut CudaSlice<f32>,
12363    ) -> Result<bool, Box<dyn std::error::Error>> {
12364        use crate::model::GpuTensor;
12365        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
12366            match w {
12367                GpuTensor::Quant {
12368                    qtype, row_bytes, ..
12369                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
12370                _ => None,
12371            }
12372        };
12373        let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (q4(w0), q4(w1), q4(w2)) else {
12374            return Ok(false);
12375        };
12376        if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
12377            return Ok(false);
12378        }
12379        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
12380            match w {
12381                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
12382                    Some(m) => (m, true),
12383                    None => (bytes, *rp),
12384                },
12385                _ => unreachable!(),
12386            }
12387        }
12388        let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
12389        if rp0 != rp1 || rp1 != rp2 {
12390            return Ok(false);
12391        }
12392        let rp = rp0;
12393        let rpb: u32 = 4;
12394        let mr1 = rp && Self::q40_mr1_on();
12395        let nb = |o: usize| {
12396            if mr1 {
12397                (o as u32).div_ceil(rpb)
12398            } else {
12399                (o as u32).div_ceil(2).div_ceil(rpb)
12400            }
12401        };
12402        let grid = nb(o0) + nb(o1) + nb(o2);
12403        debug_assert!(y0.len() >= o0 && y1.len() >= o1 && y2.len() >= o2);
12404        let f = self.func(if mr1 {
12405            "qmatvec_q4_0_mmvq_fused3_mr1_rp"
12406        } else if rp {
12407            "qmatvec_q4_0_mmvq_fused3_rp"
12408        } else {
12409            "qmatvec_q4_0_mmvq_fused3"
12410        });
12411        let cfg = LaunchConfig {
12412            grid_dim: (grid, 1, 1),
12413            block_dim: (32, rpb, 1),
12414            shared_mem_bytes: 0,
12415        };
12416        let inf = w0.in_features() as i32;
12417        let (oo0, oo1, oo2) = (o0 as i32, o1 as i32, o2 as i32);
12418        let (r0, r1, r2) = (rb0 as i64, rb1 as i64, rb2 as i64);
12419        // PDL wave-A: identical to the owned twin (capture-lane parity).
12420        if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
12421            use cudarc::driver::{DevicePtr, DevicePtrMut};
12422            let s = &self.gpu.stream();
12423            let (p0, _g0) = b0.device_ptr(s);
12424            let (p1, _g1) = b1.device_ptr(s);
12425            let (p2, _g2) = b2.device_ptr(s);
12426            let (paq, _g3) = aq.device_ptr(s);
12427            let (pad, _g4) = ad.device_ptr(s);
12428            let (py0, _g5) = y0.device_ptr_mut(s);
12429            let (py1, _g6) = y1.device_ptr_mut(s);
12430            let (py2, _g7) = y2.device_ptr_mut(s);
12431            let mut ps = [
12432                &p0 as *const _ as *mut std::ffi::c_void,
12433                &p1 as *const _ as *mut _,
12434                &p2 as *const _ as *mut _,
12435                &paq as *const _ as *mut _,
12436                &pad as *const _ as *mut _,
12437                &py0 as *const _ as *mut _,
12438                &py1 as *const _ as *mut _,
12439                &py2 as *const _ as *mut _,
12440                &inf as *const _ as *mut _,
12441                &oo0 as *const _ as *mut _,
12442                &oo1 as *const _ as *mut _,
12443                &oo2 as *const _ as *mut _,
12444                &r0 as *const _ as *mut _,
12445                &r1 as *const _ as *mut _,
12446                &r2 as *const _ as *mut _,
12447            ];
12448            unsafe {
12449                self.launch_pdl(
12450                    "qmatvec_q4_0_mmvq_fused3_mr1_rp",
12451                    (grid, 1, 1),
12452                    (32, rpb, 1),
12453                    &mut ps,
12454                )?;
12455            }
12456            return Ok(true);
12457        }
12458        let __s_b = self.gpu.stream();
12459        let mut b = __s_b.launch_builder(&f);
12460        b.arg(b0)
12461            .arg(b1)
12462            .arg(b2)
12463            .arg(aq)
12464            .arg(ad)
12465            .arg(&mut *y0)
12466            .arg(&mut *y1)
12467            .arg(&mut *y2)
12468            .arg(&inf)
12469            .arg(&oo0)
12470            .arg(&oo1)
12471            .arg(&oo2)
12472            .arg(&r0)
12473            .arg(&r1)
12474            .arg(&r2);
12475        unsafe {
12476            b.launch(cfg)?;
12477        }
12478        Ok(true)
12479    }
12480
12481    /// FUSED Q4_0 m=1 PAIR (gemma shared gate+up).
12482    pub fn matmul_q4_fused2(
12483        &self,
12484        w0: &crate::model::GpuTensor,
12485        w1: &crate::model::GpuTensor,
12486        aq: &CudaSlice<i8>,
12487        ad: &CudaSlice<f32>,
12488    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
12489        use crate::model::GpuTensor;
12490        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
12491            match w {
12492                GpuTensor::Quant {
12493                    qtype, row_bytes, ..
12494                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
12495                _ => None,
12496            }
12497        };
12498        let (Some((rb0, o0)), Some((rb1, o1))) = (q4(w0), q4(w1)) else {
12499            return Ok(None);
12500        };
12501        if w0.in_features() != w1.in_features() {
12502            return Ok(None);
12503        }
12504        // Effective (bytes, rp) per tensor (mirror or in-place swap); mixed -> separate matvecs.
12505        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
12506            match w {
12507                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
12508                    Some(m) => (m, true),
12509                    None => (bytes, *rp),
12510                },
12511                _ => unreachable!(),
12512            }
12513        }
12514        let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
12515        if rp0 != rp1 {
12516            return Ok(None);
12517        }
12518        let rp = rp0;
12519        let rpb: u32 = 4;
12520        // mr1 twin — see matmul_q4_fused3.
12521        let mr1 = rp && Self::q40_mr1_on();
12522        let nb = |o: usize| {
12523            if mr1 {
12524                (o as u32).div_ceil(rpb)
12525            } else {
12526                (o as u32).div_ceil(2).div_ceil(rpb)
12527            }
12528        };
12529        let grid = nb(o0) + nb(o1);
12530        let mut y0 = self.alloc_uninit::<f32>(o0)?;
12531        let mut y1 = self.alloc_uninit::<f32>(o1)?;
12532        let f = self.func(if mr1 {
12533            "qmatvec_q4_0_mmvq_fused2_mr1_rp"
12534        } else if rp {
12535            "qmatvec_q4_0_mmvq_fused2_rp"
12536        } else {
12537            "qmatvec_q4_0_mmvq_fused2"
12538        });
12539        let cfg = LaunchConfig {
12540            grid_dim: (grid, 1, 1),
12541            block_dim: (32, rpb, 1),
12542            shared_mem_bytes: 0,
12543        };
12544        let inf = w0.in_features() as i32;
12545        let (oo0, oo1) = (o0 as i32, o1 as i32);
12546        let (r0, r1) = (rb0 as i64, rb1 as i64);
12547        // PDL wave-A: mr1 kernel carries MEMRA_PDL_ENTRY.
12548        if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
12549            {
12550                use cudarc::driver::{DevicePtr, DevicePtrMut};
12551                let s = &self.gpu.stream();
12552                let (p0, _g0) = b0.device_ptr(s);
12553                let (p1, _g1) = b1.device_ptr(s);
12554                let (paq, _g2) = aq.device_ptr(s);
12555                let (pad, _g3) = ad.device_ptr(s);
12556                let (py0, _g4) = y0.device_ptr_mut(s);
12557                let (py1, _g5) = y1.device_ptr_mut(s);
12558                let mut ps = [
12559                    &p0 as *const _ as *mut std::ffi::c_void,
12560                    &p1 as *const _ as *mut _,
12561                    &paq as *const _ as *mut _,
12562                    &pad as *const _ as *mut _,
12563                    &py0 as *const _ as *mut _,
12564                    &py1 as *const _ as *mut _,
12565                    &inf as *const _ as *mut _,
12566                    &oo0 as *const _ as *mut _,
12567                    &oo1 as *const _ as *mut _,
12568                    &r0 as *const _ as *mut _,
12569                    &r1 as *const _ as *mut _,
12570                ];
12571                unsafe {
12572                    self.launch_pdl(
12573                        "qmatvec_q4_0_mmvq_fused2_mr1_rp",
12574                        (grid, 1, 1),
12575                        (32, rpb, 1),
12576                        &mut ps,
12577                    )?;
12578                }
12579            }
12580            return Ok(Some((y0, y1)));
12581        }
12582        let __s_b = self.gpu.stream();
12583        let mut b = __s_b.launch_builder(&f);
12584        b.arg(b0)
12585            .arg(b1)
12586            .arg(aq)
12587            .arg(ad)
12588            .arg(&mut y0)
12589            .arg(&mut y1)
12590            .arg(&inf)
12591            .arg(&oo0)
12592            .arg(&oo1)
12593            .arg(&r0)
12594            .arg(&r1);
12595        unsafe {
12596            b.launch(cfg)?;
12597        }
12598        Ok(Some((y0, y1)))
12599    }
12600
12601    /// Slot-fed fused2 twin (alloc-free capture lane): identical launch, caller-owned outputs.
12602    pub fn matmul_q4_fused2_into(
12603        &self,
12604        w0: &crate::model::GpuTensor,
12605        w1: &crate::model::GpuTensor,
12606        aq: &CudaSlice<i8>,
12607        ad: &CudaSlice<f32>,
12608        y0: &mut CudaSlice<f32>,
12609        y1: &mut CudaSlice<f32>,
12610    ) -> Result<bool, Box<dyn std::error::Error>> {
12611        use crate::model::GpuTensor;
12612        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
12613            match w {
12614                GpuTensor::Quant {
12615                    qtype, row_bytes, ..
12616                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
12617                _ => None,
12618            }
12619        };
12620        let (Some((rb0, o0)), Some((rb1, o1))) = (q4(w0), q4(w1)) else {
12621            return Ok(false);
12622        };
12623        if w0.in_features() != w1.in_features() {
12624            return Ok(false);
12625        }
12626        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
12627            match w {
12628                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
12629                    Some(m) => (m, true),
12630                    None => (bytes, *rp),
12631                },
12632                _ => unreachable!(),
12633            }
12634        }
12635        let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
12636        if rp0 != rp1 {
12637            return Ok(false);
12638        }
12639        let rp = rp0;
12640        let rpb: u32 = 4;
12641        let mr1 = rp && Self::q40_mr1_on();
12642        let nb = |o: usize| {
12643            if mr1 {
12644                (o as u32).div_ceil(rpb)
12645            } else {
12646                (o as u32).div_ceil(2).div_ceil(rpb)
12647            }
12648        };
12649        let grid = nb(o0) + nb(o1);
12650        debug_assert!(y0.len() >= o0 && y1.len() >= o1);
12651        let f = self.func(if mr1 {
12652            "qmatvec_q4_0_mmvq_fused2_mr1_rp"
12653        } else if rp {
12654            "qmatvec_q4_0_mmvq_fused2_rp"
12655        } else {
12656            "qmatvec_q4_0_mmvq_fused2"
12657        });
12658        let cfg = LaunchConfig {
12659            grid_dim: (grid, 1, 1),
12660            block_dim: (32, rpb, 1),
12661            shared_mem_bytes: 0,
12662        };
12663        let inf = w0.in_features() as i32;
12664        let (oo0, oo1) = (o0 as i32, o1 as i32);
12665        let (r0, r1) = (rb0 as i64, rb1 as i64);
12666        // PDL wave-A: identical to the owned twin (capture-lane parity).
12667        if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
12668            use cudarc::driver::{DevicePtr, DevicePtrMut};
12669            let s = &self.gpu.stream();
12670            let (p0, _g0) = b0.device_ptr(s);
12671            let (p1, _g1) = b1.device_ptr(s);
12672            let (paq, _g2) = aq.device_ptr(s);
12673            let (pad, _g3) = ad.device_ptr(s);
12674            let (py0, _g4) = y0.device_ptr_mut(s);
12675            let (py1, _g5) = y1.device_ptr_mut(s);
12676            let mut ps = [
12677                &p0 as *const _ as *mut std::ffi::c_void,
12678                &p1 as *const _ as *mut _,
12679                &paq as *const _ as *mut _,
12680                &pad as *const _ as *mut _,
12681                &py0 as *const _ as *mut _,
12682                &py1 as *const _ as *mut _,
12683                &inf as *const _ as *mut _,
12684                &oo0 as *const _ as *mut _,
12685                &oo1 as *const _ as *mut _,
12686                &r0 as *const _ as *mut _,
12687                &r1 as *const _ as *mut _,
12688            ];
12689            unsafe {
12690                self.launch_pdl(
12691                    "qmatvec_q4_0_mmvq_fused2_mr1_rp",
12692                    (grid, 1, 1),
12693                    (32, rpb, 1),
12694                    &mut ps,
12695                )?;
12696            }
12697            return Ok(true);
12698        }
12699        let __s_b = self.gpu.stream();
12700        let mut b = __s_b.launch_builder(&f);
12701        b.arg(b0)
12702            .arg(b1)
12703            .arg(aq)
12704            .arg(ad)
12705            .arg(&mut *y0)
12706            .arg(&mut *y1)
12707            .arg(&inf)
12708            .arg(&oo0)
12709            .arg(&oo1)
12710            .arg(&r0)
12711            .arg(&r1);
12712        unsafe {
12713            b.launch(cfg)?;
12714        }
12715        Ok(true)
12716    }
12717
12718    /// BATCHED fused2 (2026-07-13, megakernel-microcosm probe): gate+up b-tier matvecs in
12719    /// ONE segmented-grid launch — the up segment fills SMs as the gate segment drains
12720    /// (the per-launch tail waves behind the 6x-falsified b-tier plateau). Bit-identical
12721    /// per row to two mr2_rp launches. rp layout required; m in 2..=8 (b16 has no twin).
12722    pub fn matmul_q4_fused2_batched(
12723        &self,
12724        w0: &crate::model::GpuTensor,
12725        w1: &crate::model::GpuTensor,
12726        aq: &CudaSlice<i8>,
12727        ad: &CudaSlice<f32>,
12728        m: usize,
12729    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
12730        use crate::model::GpuTensor;
12731        if m < 2 || m > 8 {
12732            return Ok(None);
12733        }
12734        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
12735            match w {
12736                GpuTensor::Quant {
12737                    qtype, row_bytes, ..
12738                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
12739                _ => None,
12740            }
12741        };
12742        let (Some((rb0, o0)), Some((_rb1, o1))) = (q4(w0), q4(w1)) else {
12743            return Ok(None);
12744        };
12745        if w0.in_features() != w1.in_features() {
12746            return Ok(None);
12747        }
12748        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
12749            match w {
12750                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
12751                    Some(mr) => (mr, true),
12752                    None => (bytes, *rp),
12753                },
12754                _ => unreachable!(),
12755            }
12756        }
12757        let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
12758        if !rp0 || !rp1 {
12759            return Ok(None);
12760        }
12761        let mcols = Self::batched_mcols(m);
12762        let rpb: u32 = 4;
12763        let nb = |o: usize| (o as u32).div_ceil(2 * rpb);
12764        let grid = nb(o0) + nb(o1);
12765        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
12766        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
12767        let f = self.func(match mcols {
12768            2 => "qmatvec_q4_0_mmvq_b2_f2_rp",
12769            4 => "qmatvec_q4_0_mmvq_b4_f2_rp",
12770            _ => "qmatvec_q4_0_mmvq_b8_f2_rp",
12771        });
12772        let cfg = LaunchConfig {
12773            grid_dim: (grid, 1, 1),
12774            block_dim: (32, rpb, 1),
12775            shared_mem_bytes: 0,
12776        };
12777        let inf = w0.in_features() as i32;
12778        let (oo0, oo1, mi) = (o0 as i32, o1 as i32, m as i32);
12779        let rb = rb0 as i64;
12780        let __s_b = self.gpu.stream();
12781        let mut b = __s_b.launch_builder(&f);
12782        b.arg(b0)
12783            .arg(b1)
12784            .arg(aq)
12785            .arg(ad)
12786            .arg(&mut y0)
12787            .arg(&mut y1)
12788            .arg(&inf)
12789            .arg(&oo0)
12790            .arg(&oo1)
12791            .arg(&mi)
12792            .arg(&rb);
12793        unsafe {
12794            b.launch(cfg)?;
12795        }
12796        Ok(Some((y0, y1)))
12797    }
12798
12799    /// BATCHED fused3 (see matmul_q4_fused2_batched): three-segment single launch for the
12800    /// verify qkv triple. Same-in_f q4_0 rp tensors, m in 2..=8. Bit-identical per row.
12801    #[allow(clippy::too_many_arguments)]
12802    pub fn matmul_q4_fused3_batched(
12803        &self,
12804        w0: &crate::model::GpuTensor,
12805        w1: &crate::model::GpuTensor,
12806        w2: &crate::model::GpuTensor,
12807        aq: &CudaSlice<i8>,
12808        ad: &CudaSlice<f32>,
12809        m: usize,
12810    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
12811    {
12812        use crate::model::GpuTensor;
12813        if m < 2 || m > 8 {
12814            return Ok(None);
12815        }
12816        let q4 = |w: &GpuTensor| -> Option<usize> {
12817            match w {
12818                GpuTensor::Quant { qtype, .. } if *qtype == QT_Q4_0 => Some(w.out_features()),
12819                _ => None,
12820            }
12821        };
12822        let (Some(o0), Some(o1), Some(o2)) = (q4(w0), q4(w1), q4(w2)) else {
12823            return Ok(None);
12824        };
12825        if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
12826            return Ok(None);
12827        }
12828        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
12829            match w {
12830                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
12831                    Some(mr) => (mr, true),
12832                    None => (bytes, *rp),
12833                },
12834                _ => unreachable!(),
12835            }
12836        }
12837        let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
12838        if !rp0 || !rp1 || !rp2 {
12839            return Ok(None);
12840        }
12841        let mcols = Self::batched_mcols(m);
12842        let rpb: u32 = 4;
12843        let nb = |o: usize| (o as u32).div_ceil(2 * rpb);
12844        let grid = nb(o0) + nb(o1) + nb(o2);
12845        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
12846        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
12847        let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
12848        let f = self.func(match mcols {
12849            2 => "qmatvec_q4_0_mmvq_b2_f3_rp",
12850            4 => "qmatvec_q4_0_mmvq_b4_f3_rp",
12851            _ => "qmatvec_q4_0_mmvq_b8_f3_rp",
12852        });
12853        let cfg = LaunchConfig {
12854            grid_dim: (grid, 1, 1),
12855            block_dim: (32, rpb, 1),
12856            shared_mem_bytes: 0,
12857        };
12858        let inf = w0.in_features() as i32;
12859        let (oo0, oo1, oo2, mi) = (o0 as i32, o1 as i32, o2 as i32, m as i32);
12860        let rb = 0i64;
12861        let __s_b = self.gpu.stream();
12862        let mut b = __s_b.launch_builder(&f);
12863        b.arg(b0)
12864            .arg(b1)
12865            .arg(b2)
12866            .arg(aq)
12867            .arg(ad)
12868            .arg(&mut y0)
12869            .arg(&mut y1)
12870            .arg(&mut y2)
12871            .arg(&inf)
12872            .arg(&oo0)
12873            .arg(&oo1)
12874            .arg(&oo2)
12875            .arg(&mi)
12876            .arg(&rb);
12877        unsafe {
12878            b.launch(cfg)?;
12879        }
12880        Ok(Some((y0, y1, y2)))
12881    }
12882
12883    pub fn matmul_q8_fused3(
12884        &self,
12885        w0: &crate::model::GpuTensor,
12886        w1: &crate::model::GpuTensor,
12887        w2: &crate::model::GpuTensor,
12888        aq: &CudaSlice<i8>,
12889        ad: &CudaSlice<f32>,
12890    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
12891    {
12892        // e4m3 twin (lane/fp8-decode-v1): the full-attn wq/wk/wv triple — on the NV-27B those three
12893        // are per-tensor FP8, so native residency without this arm meant three separate launches.
12894        if let Some([p0, p1, p2]) = self.e4m3_fused_params(&[w0, w1, w2]) {
12895            return Ok(Some(self.e4m3_fused3_core(
12896                p0.0,
12897                p1.0,
12898                p2.0,
12899                aq,
12900                ad,
12901                w0.in_features(),
12902                p0.1,
12903                p1.1,
12904                p2.1,
12905                p0.2,
12906                p0.3,
12907                p1.3,
12908                p2.3,
12909            )?));
12910        }
12911        let Some([p0, p1, p2]) = self.q8_fused_params(&[w0, w1, w2]) else {
12912            return Ok(None);
12913        };
12914        Ok(Some(self.q8_fused3_core(
12915            p0.0,
12916            p1.0,
12917            p2.0,
12918            aq,
12919            ad,
12920            w0.in_features(),
12921            p0.1,
12922            p1.1,
12923            p2.1,
12924            p0.2,
12925        )?))
12926    }
12927
12928    #[allow(clippy::too_many_arguments)]
12929    fn q8_fused3_core(
12930        &self,
12931        b0: &CudaSlice<u8>,
12932        b1: &CudaSlice<u8>,
12933        b2: &CudaSlice<u8>,
12934        aq: &CudaSlice<i8>,
12935        ad: &CudaSlice<f32>,
12936        in_f: usize,
12937        out0: usize,
12938        out1: usize,
12939        out2: usize,
12940        row_bytes: usize,
12941    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12942        const ROWS_PER_BLOCK: u32 = 4;
12943        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
12944        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
12945        let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
12946        let f = self.func("qmatvec_q8_0_mmvq_fused3");
12947        let mut y0 = self.alloc_uninit::<f32>(out0)?;
12948        let mut y1 = self.alloc_uninit::<f32>(out1)?;
12949        let mut y2 = self.alloc_uninit::<f32>(out2)?;
12950        let cfg = LaunchConfig {
12951            grid_dim: (nb0 + nb1 + nb2, 1, 1),
12952            block_dim: (32, ROWS_PER_BLOCK, 1),
12953            shared_mem_bytes: 0,
12954        };
12955        let (inf, o0, o1, o2, rbl) = (
12956            in_f as i32,
12957            out0 as i32,
12958            out1 as i32,
12959            out2 as i32,
12960            row_bytes as i64,
12961        );
12962        let __s_b = self.gpu.stream();
12963        let mut b = __s_b.launch_builder(&f);
12964        b.arg(b0)
12965            .arg(b1)
12966            .arg(b2)
12967            .arg(aq)
12968            .arg(ad)
12969            .arg(&mut y0)
12970            .arg(&mut y1)
12971            .arg(&mut y2)
12972            .arg(&inf)
12973            .arg(&o0)
12974            .arg(&o1)
12975            .arg(&o2)
12976            .arg(&rbl);
12977        unsafe {
12978            b.launch(cfg)?;
12979        }
12980        Ok((y0, y1, y2))
12981    }
12982
12983    /// Test entry for the kernel_check gate: fused3 from raw weight bytes (internal q8_1 quant).
12984    #[allow(clippy::too_many_arguments)]
12985    pub fn qmatvec_q8_fused3_raw(
12986        &self,
12987        b0: &CudaSlice<u8>,
12988        b1: &CudaSlice<u8>,
12989        b2: &CudaSlice<u8>,
12990        x: &CudaSlice<f32>,
12991        in_f: usize,
12992        out0: usize,
12993        out1: usize,
12994        out2: usize,
12995        row_bytes: usize,
12996    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12997        let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
12998        self.q8_fused3_core(b0, b1, b2, &aq, &ad, in_f, out0, out1, out2, row_bytes)
12999    }
13000
13001    /// BATCHED twin of `matmul_q8_fused2` for the verify t=2-4 tier (MEMRA_SPEC_FUSED_T call
13002    /// sites, lane/close35b): ONE launch computes both same-input Q8_0 projections for m tokens.
13003    /// Per (tensor,token,row) the kernel body is q8_0_mmvq_batched VERBATIM with the identical
13004    /// row mapping (Q8_0's batched_variant is always "base") -> BIT-IDENTICAL to the two
13005    /// per-tensor _b2/_b4 launches `matmul_decode_exact` dispatches at m=2-4, with the caller's
13006    /// single shared q8_1 activation replacing two per-call re-quantizes (quantize_q8_1 is
13007    /// deterministic -> same bytes). None when ineligible (m outside 2..=4 / not both Q8_0 /
13008    /// in_f mismatch / MEMRA_MMVQ=0 / MEMRA_Q8_DUAL=0 / MEMRA_NO_BATCHED set — the last keeps
13009    /// dispatch parity: without batched kernels decode-exact runs grid.y=m MMVQ, and the fused
13010    /// twin must not introduce a batched program the reference path would not run).
13011    pub fn matmul_q8_fused2_t(
13012        &self,
13013        w0: &crate::model::GpuTensor,
13014        w1: &crate::model::GpuTensor,
13015        aq: &CudaSlice<i8>,
13016        ad: &CudaSlice<f32>,
13017        m: usize,
13018    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
13019        // m<=8 (lane/q27-deepdive, 2026-08-05): was 2..=4 (the verify tier's mcols 2/4). The
13020        // serving tick's mcols-8 tier now has its fused2_b8 wrapper, so c=5..8 batched decode
13021        // fuses too — same template body, still bit-identical to the two _b8 launches.
13022        if !(2..=8).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok() {
13023            return Ok(None);
13024        }
13025        // e4m3 twin: MEMRA_B8 parity — without it m=5..8 e4m3 decode runs the per-m grid.y=m path,
13026        // so the fused b8 launch would introduce a batched program the reference path would not run.
13027        if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
13028            if m > 4 && !Self::b8_enabled() {
13029                return Ok(None);
13030            }
13031            return Ok(Some(self.e4m3_fused2_t_core(
13032                p0.0,
13033                p1.0,
13034                aq,
13035                ad,
13036                m,
13037                w0.in_features(),
13038                p0.1,
13039                p1.1,
13040                p0.2,
13041                p0.3,
13042                p1.3,
13043            )?));
13044        }
13045        let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
13046            return Ok(None);
13047        };
13048        Ok(Some(self.q8_fused2_t_core(
13049            p0.0,
13050            p1.0,
13051            aq,
13052            ad,
13053            m,
13054            w0.in_features(),
13055            p0.1,
13056            p1.1,
13057            p0.2,
13058        )?))
13059    }
13060
13061    #[allow(clippy::too_many_arguments)]
13062    fn q8_fused2_t_core(
13063        &self,
13064        b0: &CudaSlice<u8>,
13065        b1: &CudaSlice<u8>,
13066        aq: &CudaSlice<i8>,
13067        ad: &CudaSlice<f32>,
13068        m: usize,
13069        in_f: usize,
13070        out0: usize,
13071        out1: usize,
13072        row_bytes: usize,
13073    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13074        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
13075        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
13076        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
13077        let f = self.func(match Self::batched_mcols(m) {
13078            2 => "qmatvec_q8_0_mmvq_fused2_b2",
13079            4 => "qmatvec_q8_0_mmvq_fused2_b4",
13080            // b8 = the SERVING tier (lane/q27-deepdive): c=5..8 batched decode.
13081            _ => "qmatvec_q8_0_mmvq_fused2_b8",
13082        });
13083        let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
13084        let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
13085        let cfg = LaunchConfig {
13086            grid_dim: (nb0 + nb1, 1, 1),
13087            block_dim: (32, ROWS_PER_BLOCK, 1),
13088            shared_mem_bytes: 0,
13089        };
13090        let (inf, o0, o1, mi, rbl) = (
13091            in_f as i32,
13092            out0 as i32,
13093            out1 as i32,
13094            m as i32,
13095            row_bytes as i64,
13096        );
13097        let __s_b = self.gpu.stream();
13098        let mut b = __s_b.launch_builder(&f);
13099        b.arg(b0)
13100            .arg(b1)
13101            .arg(aq)
13102            .arg(ad)
13103            .arg(&mut y0)
13104            .arg(&mut y1)
13105            .arg(&inf)
13106            .arg(&o0)
13107            .arg(&o1)
13108            .arg(&mi)
13109            .arg(&rbl);
13110        unsafe {
13111            b.launch(cfg)?;
13112        }
13113        Ok((y0, y1))
13114    }
13115
13116    /// Test entry for the kernel_check gate: fused2 batched from raw weight bytes (internal
13117    /// q8_1 quant of the [m, in_f] activation), no env gating.
13118    #[allow(clippy::too_many_arguments)]
13119    pub fn qmatvec_q8_fused2_t_raw(
13120        &self,
13121        b0: &CudaSlice<u8>,
13122        b1: &CudaSlice<u8>,
13123        x: &CudaSlice<f32>,
13124        m: usize,
13125        in_f: usize,
13126        out0: usize,
13127        out1: usize,
13128        row_bytes: usize,
13129    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13130        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
13131        self.q8_fused2_t_core(b0, b1, &aq, &ad, m, in_f, out0, out1, row_bytes)
13132    }
13133
13134    /// BATCHED twin of `matmul_q8_fused3` (wq+wk+wv at verify t=2-4). Same contract as
13135    /// `matmul_q8_fused2_t` with three ranges.
13136    #[allow(clippy::too_many_arguments)]
13137    pub fn matmul_q8_fused3_t(
13138        &self,
13139        w0: &crate::model::GpuTensor,
13140        w1: &crate::model::GpuTensor,
13141        w2: &crate::model::GpuTensor,
13142        aq: &CudaSlice<i8>,
13143        ad: &CudaSlice<f32>,
13144        m: usize,
13145    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
13146    {
13147        if !(2..=4).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok() {
13148            return Ok(None);
13149        }
13150        if let Some([p0, p1, p2]) = self.e4m3_fused_params(&[w0, w1, w2]) {
13151            return Ok(Some(self.e4m3_fused3_t_core(
13152                p0.0,
13153                p1.0,
13154                p2.0,
13155                aq,
13156                ad,
13157                m,
13158                w0.in_features(),
13159                p0.1,
13160                p1.1,
13161                p2.1,
13162                p0.2,
13163                p0.3,
13164                p1.3,
13165                p2.3,
13166            )?));
13167        }
13168        let Some([p0, p1, p2]) = self.q8_fused_params(&[w0, w1, w2]) else {
13169            return Ok(None);
13170        };
13171        Ok(Some(self.q8_fused3_t_core(
13172            p0.0,
13173            p1.0,
13174            p2.0,
13175            aq,
13176            ad,
13177            m,
13178            w0.in_features(),
13179            p0.1,
13180            p1.1,
13181            p2.1,
13182            p0.2,
13183        )?))
13184    }
13185
13186    #[allow(clippy::too_many_arguments)]
13187    fn q8_fused3_t_core(
13188        &self,
13189        b0: &CudaSlice<u8>,
13190        b1: &CudaSlice<u8>,
13191        b2: &CudaSlice<u8>,
13192        aq: &CudaSlice<i8>,
13193        ad: &CudaSlice<f32>,
13194        m: usize,
13195        in_f: usize,
13196        out0: usize,
13197        out1: usize,
13198        out2: usize,
13199        row_bytes: usize,
13200    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13201        const ROWS_PER_BLOCK: u32 = 4;
13202        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
13203        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
13204        let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
13205        let f = self.func(if Self::batched_mcols(m) == 2 {
13206            "qmatvec_q8_0_mmvq_fused3_b2"
13207        } else {
13208            "qmatvec_q8_0_mmvq_fused3_b4"
13209        });
13210        let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
13211        let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
13212        let mut y2 = self.alloc_uninit::<f32>(m * out2)?;
13213        let cfg = LaunchConfig {
13214            grid_dim: (nb0 + nb1 + nb2, 1, 1),
13215            block_dim: (32, ROWS_PER_BLOCK, 1),
13216            shared_mem_bytes: 0,
13217        };
13218        let (inf, o0, o1, o2, mi, rbl) = (
13219            in_f as i32,
13220            out0 as i32,
13221            out1 as i32,
13222            out2 as i32,
13223            m as i32,
13224            row_bytes as i64,
13225        );
13226        let __s_b = self.gpu.stream();
13227        let mut b = __s_b.launch_builder(&f);
13228        b.arg(b0)
13229            .arg(b1)
13230            .arg(b2)
13231            .arg(aq)
13232            .arg(ad)
13233            .arg(&mut y0)
13234            .arg(&mut y1)
13235            .arg(&mut y2)
13236            .arg(&inf)
13237            .arg(&o0)
13238            .arg(&o1)
13239            .arg(&o2)
13240            .arg(&mi)
13241            .arg(&rbl);
13242        unsafe {
13243            b.launch(cfg)?;
13244        }
13245        Ok((y0, y1, y2))
13246    }
13247
13248    /// Test entry for the kernel_check gate: fused3 batched from raw weight bytes.
13249    #[allow(clippy::too_many_arguments)]
13250    pub fn qmatvec_q8_fused3_t_raw(
13251        &self,
13252        b0: &CudaSlice<u8>,
13253        b1: &CudaSlice<u8>,
13254        b2: &CudaSlice<u8>,
13255        x: &CudaSlice<f32>,
13256        m: usize,
13257        in_f: usize,
13258        out0: usize,
13259        out1: usize,
13260        out2: usize,
13261        row_bytes: usize,
13262    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13263        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
13264        self.q8_fused3_t_core(b0, b1, b2, &aq, &ad, m, in_f, out0, out1, out2, row_bytes)
13265    }
13266
13267    /// Rollback seam for the Q8_0 dense-FFN gate+up fusion arm in `matmul_pre_dual_noscale`
13268    /// (lane/q27-deepdive, 2026-08-05). Default ON; `MEMRA_Q8_FFN_FUSE2=0` restores the
13269    /// two-`matmul_pre_noscale` pair. Read once — the dispatch must not vary within a run.
13270    pub fn q8_ffn_fuse2_on(&self) -> bool {
13271        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13272        *ON.get_or_init(|| std::env::var("MEMRA_Q8_FFN_FUSE2").as_deref() != Ok("0"))
13273    }
13274
13275    /// Eligibility + param extraction for the fused q8_0 launches: every tensor must be Quant Q8_0
13276    /// with macro-scale 1.0 (always true for GGUF q8_0; only NVFP4 carries scale) and share w[0]'s
13277    /// in_f (q8_0 row_bytes is a pure function of in_f, so equal in_f => equal row_bytes). MEMRA_MMVQ
13278    /// must be on: the fused body is the MMVQ kernel; without it decode m=1 runs dp4a and fusing
13279    /// would mix dispatch families (FP-order law). MEMRA_Q8_DUAL=0 = rollback seam.
13280    #[allow(clippy::type_complexity)]
13281    fn q8_fused_params<'w, const N: usize>(
13282        &self,
13283        ws: &[&'w crate::model::GpuTensor; N],
13284    ) -> Option<[(&'w CudaSlice<u8>, usize, usize); N]> {
13285        use crate::model::GpuTensor;
13286        if std::env::var("MEMRA_MMVQ").as_deref() == Ok("0") {
13287            return None;
13288        }
13289        if std::env::var("MEMRA_Q8_DUAL").is_ok_and(|v| v == "0") {
13290            return None;
13291        }
13292        let in_f = ws[0].in_features();
13293        let mut out: [Option<(&CudaSlice<u8>, usize, usize)>; N] = [None; N];
13294        for (i, w) in ws.iter().enumerate() {
13295            match w {
13296                GpuTensor::Quant {
13297                    bytes,
13298                    qtype,
13299                    row_bytes,
13300                    scale,
13301                    ..
13302                } if *qtype == QT_Q8_0 && *scale == 1.0 && w.in_features() == in_f => {
13303                    out[i] = Some((bytes, w.out_features(), *row_bytes))
13304                }
13305                _ => return None,
13306            }
13307        }
13308        Some(out.map(|o| o.unwrap()))
13309    }
13310
13311    /// Rollback seam for the F8-E4M3 launch-fusion arm (lane/fp8-decode-v1, 2026-08-05).
13312    /// Default ON; `MEMRA_E4M3_DUAL=0` restores the per-tensor m=1/batched launches.
13313    pub fn e4m3_dual_on(&self) -> bool {
13314        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13315        *ON.get_or_init(|| std::env::var("MEMRA_E4M3_DUAL").as_deref() != Ok("0"))
13316    }
13317
13318    /// Eligibility + param extraction for the FUSED e4m3 launches — the QT_F8_E4M3 twin of
13319    /// `q8_fused_params`. Differences that are inherent to the dtype, not policy:
13320    ///   * each tensor carries its OWN per-tensor `weight_scale` (returned as the 4th field);
13321    ///     Q8_0 hard-requires scale==1.0 because it has no macro-scale at all.
13322    ///   * no MEMRA_MMVQ gate: `mmvq_supports` exempts QT_F8_E4M3 (the e4m3 mmvq family is that
13323    ///     dtype's ONLY int8-act kernel class), so the per-tensor fallback these fused kernels
13324    ///     replace is ALWAYS the same mmvq body under every env — the FP-order law holds.
13325    ///   * `row_bytes == in_f` is asserted rather than derived: the native-residency load arm keeps
13326    ///     the checkpoint's raw [out_f, in_f] rows, and a re-encoded slab must never reach here.
13327    /// Rejects any split-plane mirror (`rp`/`rp4`): there is no `_rp` e4m3 fused form, so fusing
13328    /// there would swap dispatch families mid-model. MEMRA_E4M3_DUAL=0 = rollback seam.
13329    #[allow(clippy::type_complexity)]
13330    fn e4m3_fused_params<'w, const N: usize>(
13331        &self,
13332        ws: &[&'w crate::model::GpuTensor; N],
13333    ) -> Option<[(&'w CudaSlice<u8>, usize, usize, f32); N]> {
13334        use crate::model::GpuTensor;
13335        if !self.e4m3_dual_on() {
13336            return None;
13337        }
13338        let in_f = ws[0].in_features();
13339        let mut out: [Option<(&CudaSlice<u8>, usize, usize, f32)>; N] = [None; N];
13340        for (i, w) in ws.iter().enumerate() {
13341            match w {
13342                GpuTensor::Quant {
13343                    bytes,
13344                    qtype,
13345                    row_bytes,
13346                    scale,
13347                    rp,
13348                    rp4,
13349                    ..
13350                } if *qtype == QT_F8_E4M3
13351                    && w.in_features() == in_f
13352                    && *row_bytes == in_f
13353                    && !*rp
13354                    && rp4.is_none() =>
13355                {
13356                    out[i] = Some((bytes, w.out_features(), *row_bytes, *scale))
13357                }
13358                _ => return None,
13359            }
13360        }
13361        Some(out.map(|o| o.unwrap()))
13362    }
13363
13364    /// FUSED e4m3 m=1 PAIR. Block-offset split (`qmatvec_e4m3_mmvq_fused2`), per-tensor
13365    /// weight_scale folded at the write like the single-tensor `qmatvec_e4m3_mmvq` — so per
13366    /// (tensor,row) this is BIT-IDENTICAL to two separate m=1 launches, scale included.
13367    #[allow(clippy::too_many_arguments)]
13368    fn e4m3_fused2_core(
13369        &self,
13370        b0: &CudaSlice<u8>,
13371        b1: &CudaSlice<u8>,
13372        aq: &CudaSlice<i8>,
13373        ad: &CudaSlice<f32>,
13374        in_f: usize,
13375        out0: usize,
13376        out1: usize,
13377        row_bytes: usize,
13378        ws0: f32,
13379        ws1: f32,
13380    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13381        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
13382        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
13383        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
13384        let f = self.func("qmatvec_e4m3_mmvq_fused2");
13385        let mut y0 = self.alloc_uninit::<f32>(out0)?;
13386        let mut y1 = self.alloc_uninit::<f32>(out1)?;
13387        let cfg = LaunchConfig {
13388            grid_dim: (nb0 + nb1, 1, 1),
13389            block_dim: (32, ROWS_PER_BLOCK, 1),
13390            shared_mem_bytes: 0,
13391        };
13392        let (inf, o0, o1, rbl) = (in_f as i32, out0 as i32, out1 as i32, row_bytes as i64);
13393        let __s_b = self.gpu.stream();
13394        let mut b = __s_b.launch_builder(&f);
13395        b.arg(b0)
13396            .arg(b1)
13397            .arg(aq)
13398            .arg(ad)
13399            .arg(&mut y0)
13400            .arg(&mut y1)
13401            .arg(&inf)
13402            .arg(&o0)
13403            .arg(&o1)
13404            .arg(&rbl)
13405            .arg(&ws0)
13406            .arg(&ws1);
13407        unsafe {
13408            b.launch(cfg)?;
13409        }
13410        Ok((y0, y1))
13411    }
13412
13413    /// FUSED e4m3 m=1 TRIPLE (`qmatvec_e4m3_mmvq_fused3`). Same contract as the pair.
13414    #[allow(clippy::too_many_arguments)]
13415    fn e4m3_fused3_core(
13416        &self,
13417        b0: &CudaSlice<u8>,
13418        b1: &CudaSlice<u8>,
13419        b2: &CudaSlice<u8>,
13420        aq: &CudaSlice<i8>,
13421        ad: &CudaSlice<f32>,
13422        in_f: usize,
13423        out0: usize,
13424        out1: usize,
13425        out2: usize,
13426        row_bytes: usize,
13427        ws0: f32,
13428        ws1: f32,
13429        ws2: f32,
13430    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13431        const ROWS_PER_BLOCK: u32 = 4;
13432        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
13433        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
13434        let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
13435        let f = self.func("qmatvec_e4m3_mmvq_fused3");
13436        let mut y0 = self.alloc_uninit::<f32>(out0)?;
13437        let mut y1 = self.alloc_uninit::<f32>(out1)?;
13438        let mut y2 = self.alloc_uninit::<f32>(out2)?;
13439        let cfg = LaunchConfig {
13440            grid_dim: (nb0 + nb1 + nb2, 1, 1),
13441            block_dim: (32, ROWS_PER_BLOCK, 1),
13442            shared_mem_bytes: 0,
13443        };
13444        let (inf, o0, o1, o2, rbl) = (
13445            in_f as i32,
13446            out0 as i32,
13447            out1 as i32,
13448            out2 as i32,
13449            row_bytes as i64,
13450        );
13451        let __s_b = self.gpu.stream();
13452        let mut b = __s_b.launch_builder(&f);
13453        b.arg(b0)
13454            .arg(b1)
13455            .arg(b2)
13456            .arg(aq)
13457            .arg(ad)
13458            .arg(&mut y0)
13459            .arg(&mut y1)
13460            .arg(&mut y2)
13461            .arg(&inf)
13462            .arg(&o0)
13463            .arg(&o1)
13464            .arg(&o2)
13465            .arg(&rbl)
13466            .arg(&ws0)
13467            .arg(&ws1)
13468            .arg(&ws2);
13469        unsafe {
13470            b.launch(cfg)?;
13471        }
13472        Ok((y0, y1, y2))
13473    }
13474
13475    /// BATCHED FUSED e4m3 pair (m=2..8). The batched kernels carry no `ws` arg (every batched
13476    /// kernel in the tree is scale-free), so each output takes its own `scale_inplace` — the
13477    /// SAME post-op the per-tensor batched dispatch applies, hence still bit-identical.
13478    #[allow(clippy::too_many_arguments)]
13479    fn e4m3_fused2_t_core(
13480        &self,
13481        b0: &CudaSlice<u8>,
13482        b1: &CudaSlice<u8>,
13483        aq: &CudaSlice<i8>,
13484        ad: &CudaSlice<f32>,
13485        m: usize,
13486        in_f: usize,
13487        out0: usize,
13488        out1: usize,
13489        row_bytes: usize,
13490        ws0: f32,
13491        ws1: f32,
13492    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13493        const ROWS_PER_BLOCK: u32 = 4;
13494        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
13495        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
13496        let f = self.func(match Self::batched_mcols(m) {
13497            2 => "qmatvec_e4m3_mmvq_fused2_b2",
13498            4 => "qmatvec_e4m3_mmvq_fused2_b4",
13499            _ => "qmatvec_e4m3_mmvq_fused2_b8",
13500        });
13501        let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
13502        let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
13503        let cfg = LaunchConfig {
13504            grid_dim: (nb0 + nb1, 1, 1),
13505            block_dim: (32, ROWS_PER_BLOCK, 1),
13506            shared_mem_bytes: 0,
13507        };
13508        let (inf, o0, o1, mi, rbl) = (
13509            in_f as i32,
13510            out0 as i32,
13511            out1 as i32,
13512            m as i32,
13513            row_bytes as i64,
13514        );
13515        let __s_b = self.gpu.stream();
13516        let mut b = __s_b.launch_builder(&f);
13517        b.arg(b0)
13518            .arg(b1)
13519            .arg(aq)
13520            .arg(ad)
13521            .arg(&mut y0)
13522            .arg(&mut y1)
13523            .arg(&inf)
13524            .arg(&o0)
13525            .arg(&o1)
13526            .arg(&mi)
13527            .arg(&rbl);
13528        unsafe {
13529            b.launch(cfg)?;
13530        }
13531        if ws0 != 1.0 {
13532            self.scale_inplace(&mut y0, ws0, m * out0)?;
13533        }
13534        if ws1 != 1.0 {
13535            self.scale_inplace(&mut y1, ws1, m * out1)?;
13536        }
13537        Ok((y0, y1))
13538    }
13539
13540    /// BATCHED FUSED e4m3 triple (m=2..4). Same contract as the batched pair.
13541    #[allow(clippy::too_many_arguments)]
13542    fn e4m3_fused3_t_core(
13543        &self,
13544        b0: &CudaSlice<u8>,
13545        b1: &CudaSlice<u8>,
13546        b2: &CudaSlice<u8>,
13547        aq: &CudaSlice<i8>,
13548        ad: &CudaSlice<f32>,
13549        m: usize,
13550        in_f: usize,
13551        out0: usize,
13552        out1: usize,
13553        out2: usize,
13554        row_bytes: usize,
13555        ws0: f32,
13556        ws1: f32,
13557        ws2: f32,
13558    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13559        const ROWS_PER_BLOCK: u32 = 4;
13560        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
13561        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
13562        let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
13563        let f = self.func(if Self::batched_mcols(m) == 2 {
13564            "qmatvec_e4m3_mmvq_fused3_b2"
13565        } else {
13566            "qmatvec_e4m3_mmvq_fused3_b4"
13567        });
13568        let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
13569        let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
13570        let mut y2 = self.alloc_uninit::<f32>(m * out2)?;
13571        let cfg = LaunchConfig {
13572            grid_dim: (nb0 + nb1 + nb2, 1, 1),
13573            block_dim: (32, ROWS_PER_BLOCK, 1),
13574            shared_mem_bytes: 0,
13575        };
13576        let (inf, o0, o1, o2, mi, rbl) = (
13577            in_f as i32,
13578            out0 as i32,
13579            out1 as i32,
13580            out2 as i32,
13581            m as i32,
13582            row_bytes as i64,
13583        );
13584        let __s_b = self.gpu.stream();
13585        let mut b = __s_b.launch_builder(&f);
13586        b.arg(b0)
13587            .arg(b1)
13588            .arg(b2)
13589            .arg(aq)
13590            .arg(ad)
13591            .arg(&mut y0)
13592            .arg(&mut y1)
13593            .arg(&mut y2)
13594            .arg(&inf)
13595            .arg(&o0)
13596            .arg(&o1)
13597            .arg(&o2)
13598            .arg(&mi)
13599            .arg(&rbl);
13600        unsafe {
13601            b.launch(cfg)?;
13602        }
13603        if ws0 != 1.0 {
13604            self.scale_inplace(&mut y0, ws0, m * out0)?;
13605        }
13606        if ws1 != 1.0 {
13607            self.scale_inplace(&mut y1, ws1, m * out1)?;
13608        }
13609        if ws2 != 1.0 {
13610            self.scale_inplace(&mut y2, ws2, m * out2)?;
13611        }
13612        Ok((y0, y1, y2))
13613    }
13614
13615    /// BLOCK-128 e4m3 MMVQ launcher (`qmatvec_e4m3_blk_mmvq`, lane/fp8-blk128-decode 2026-08-05).
13616    /// The per-block-dequant twin of `qmatvec_mmvq`'s QT_F8_E4M3 arm: same grid/block decomposition
13617    /// (warp per output row, ROWS_PER_BLOCK warps per block, grid.y = m), same q8_1 activation, but
13618    /// the weight scale is a resident [rows, cols] f32 grid read per k128 block inside the kernel
13619    /// instead of one scalar folded at the write. It cannot share `qmatvec_mmvq`'s body because
13620    /// that launcher's arg list is fixed at (bytes, aq, ad, y, in_f, out_f, m, row_bytes [, scale]).
13621    ///
13622    /// `mr` and `rp` have no analogue here (no split-plane e4m3 layout exists), so there is exactly
13623    /// one kernel and no name table — a shape this cannot serve must be refused at LOAD, not here.
13624    pub fn qmatvec_e4m3_blk_mmvq(
13625        &self,
13626        bytes: &CudaSlice<u8>,
13627        aq: &CudaSlice<i8>,
13628        ad: &CudaSlice<f32>,
13629        scales: &CudaSlice<f32>,
13630        m: usize,
13631        in_f: usize,
13632        out_f: usize,
13633        row_bytes: usize,
13634        scale_cols: usize,
13635    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13636        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
13637        self.qmatvec_e4m3_blk_mmvq_into(
13638            bytes, aq, ad, scales, m, in_f, out_f, row_bytes, scale_cols, &mut y,
13639        )?;
13640        Ok(y)
13641    }
13642
13643    /// Slot-fed twin of `qmatvec_e4m3_blk_mmvq` (caller-owned output; the alloc-free capture lane).
13644    #[allow(clippy::too_many_arguments)]
13645    pub fn qmatvec_e4m3_blk_mmvq_into(
13646        &self,
13647        bytes: &CudaSlice<u8>,
13648        aq: &CudaSlice<i8>,
13649        ad: &CudaSlice<f32>,
13650        scales: &CudaSlice<f32>,
13651        m: usize,
13652        in_f: usize,
13653        out_f: usize,
13654        row_bytes: usize,
13655        scale_cols: usize,
13656        y: &mut CudaSlice<f32>,
13657    ) -> Result<(), Box<dyn std::error::Error>> {
13658        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
13659        let f = self.func("qmatvec_e4m3_blk_mmvq");
13660        let cfg = LaunchConfig {
13661            grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), m as u32, 1),
13662            block_dim: (32, ROWS_PER_BLOCK, 1), // warp-per-row
13663            shared_mem_bytes: 0,                // warp-only reduce
13664        };
13665        let (inf, outf, mi, rb, sc) = (
13666            in_f as i32,
13667            out_f as i32,
13668            m as i32,
13669            row_bytes as i64,
13670            scale_cols as i32,
13671        );
13672        let __s_b = self.gpu.stream();
13673        let mut b = __s_b.launch_builder(&f);
13674        b.arg(bytes)
13675            .arg(aq)
13676            .arg(ad)
13677            .arg(scales)
13678            .arg(&mut *y)
13679            .arg(&inf)
13680            .arg(&outf)
13681            .arg(&mi)
13682            .arg(&rb)
13683            .arg(&sc);
13684        unsafe {
13685            b.launch(cfg)?;
13686        }
13687        Ok(())
13688    }
13689
13690    /// BLOCK-128 e4m3 BATCHED matvec (lane/rp-on-st, 2026-08-06): the weight-read-once twin of
13691    /// `qmatvec_e4m3_blk_mmvq` for m=2..16. Per (token,row) BIT-IDENTICAL to the grid.y=m launch
13692    /// (same fmaf chain, same per-k32 `s * ad` fold, same warp reduce), so it inherits the
13693    /// decode-exactness contract while reading the weight ONCE for up to `mcols` columns instead
13694    /// of `m` times. `mcols` must be one of {2,4,8,16} and satisfy `mcols >= m`.
13695    #[allow(clippy::too_many_arguments)]
13696    pub fn qmatvec_e4m3_blk_mmvq_batched(
13697        &self,
13698        bytes: &CudaSlice<u8>,
13699        aq: &CudaSlice<i8>,
13700        ad: &CudaSlice<f32>,
13701        scales: &CudaSlice<f32>,
13702        m: usize,
13703        in_f: usize,
13704        out_f: usize,
13705        row_bytes: usize,
13706        scale_cols: usize,
13707        mcols: usize,
13708    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13709        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
13710        debug_assert!(mcols >= m, "blk batched: mcols {mcols} < m {m}");
13711        let name = match mcols {
13712            2 => "qmatvec_e4m3_blk_mmvq_b2",
13713            4 => "qmatvec_e4m3_blk_mmvq_b4",
13714            8 => "qmatvec_e4m3_blk_mmvq_b8",
13715            16 => "qmatvec_e4m3_blk_mmvq_b16",
13716            _ => {
13717                return Err(
13718                    format!("qmatvec_e4m3_blk_mmvq_batched: no kernel for mcols {mcols}").into(),
13719                );
13720            }
13721        };
13722        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
13723        let f = self.func(name);
13724        let cfg = LaunchConfig {
13725            grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
13726            block_dim: (32, ROWS_PER_BLOCK, 1),
13727            shared_mem_bytes: 0,
13728        };
13729        let (inf, outf, mi, rb, sc) = (
13730            in_f as i32,
13731            out_f as i32,
13732            m as i32,
13733            row_bytes as i64,
13734            scale_cols as i32,
13735        );
13736        let __s_b = self.gpu.stream();
13737        let mut b = __s_b.launch_builder(&f);
13738        b.arg(bytes)
13739            .arg(aq)
13740            .arg(ad)
13741            .arg(scales)
13742            .arg(&mut y)
13743            .arg(&inf)
13744            .arg(&outf)
13745            .arg(&mi)
13746            .arg(&rb)
13747            .arg(&sc);
13748        unsafe {
13749            b.launch(cfg)?;
13750        }
13751        Ok(y)
13752    }
13753
13754    /// Test entry for the kernel_check exactness gate: the block-128 e4m3 batched MMVQ from raw
13755    /// bytes with an internal q8_1 quantize (mirrors `qmatvec_batched_raw`).
13756    #[allow(clippy::too_many_arguments)]
13757    pub fn qmatvec_e4m3_blk_batched_raw(
13758        &self,
13759        bytes: &CudaSlice<u8>,
13760        x: &CudaSlice<f32>,
13761        scales: &CudaSlice<f32>,
13762        m: usize,
13763        in_f: usize,
13764        out_f: usize,
13765        row_bytes: usize,
13766        scale_cols: usize,
13767        mcols: usize,
13768    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13769        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
13770        self.qmatvec_e4m3_blk_mmvq_batched(
13771            bytes, &aq, &ad, scales, m, in_f, out_f, row_bytes, scale_cols, mcols,
13772        )
13773    }
13774
13775    /// Test entry for the kernel_check exactness gate: the block-128 e4m3 MMVQ from raw bytes with
13776    /// an internal q8_1 quantize (mirrors `qmatvec_mmvq_raw`).
13777    #[allow(clippy::too_many_arguments)]
13778    pub fn qmatvec_e4m3_blk_mmvq_raw(
13779        &self,
13780        bytes: &CudaSlice<u8>,
13781        x: &CudaSlice<f32>,
13782        scales: &CudaSlice<f32>,
13783        m: usize,
13784        in_f: usize,
13785        out_f: usize,
13786        row_bytes: usize,
13787        scale_cols: usize,
13788    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13789        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
13790        self.qmatvec_e4m3_blk_mmvq(
13791            bytes, &aq, &ad, scales, m, in_f, out_f, row_bytes, scale_cols,
13792        )
13793    }
13794
13795    /// Test entries for the kernel_check bit-parity gate: fused e4m3 launches from raw weight
13796    /// bytes with internal q8_1 quantize, no env gating (mirrors `qmatvec_q8_fused*_raw`).
13797    #[allow(clippy::too_many_arguments)]
13798    pub fn qmatvec_e4m3_fused2_raw(
13799        &self,
13800        b0: &CudaSlice<u8>,
13801        b1: &CudaSlice<u8>,
13802        x: &CudaSlice<f32>,
13803        in_f: usize,
13804        out0: usize,
13805        out1: usize,
13806        row_bytes: usize,
13807        ws0: f32,
13808        ws1: f32,
13809    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13810        let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
13811        self.e4m3_fused2_core(b0, b1, &aq, &ad, in_f, out0, out1, row_bytes, ws0, ws1)
13812    }
13813
13814    #[allow(clippy::too_many_arguments)]
13815    pub fn qmatvec_e4m3_fused3_raw(
13816        &self,
13817        b0: &CudaSlice<u8>,
13818        b1: &CudaSlice<u8>,
13819        b2: &CudaSlice<u8>,
13820        x: &CudaSlice<f32>,
13821        in_f: usize,
13822        out0: usize,
13823        out1: usize,
13824        out2: usize,
13825        row_bytes: usize,
13826        ws0: f32,
13827        ws1: f32,
13828        ws2: f32,
13829    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13830        let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
13831        self.e4m3_fused3_core(
13832            b0, b1, b2, &aq, &ad, in_f, out0, out1, out2, row_bytes, ws0, ws1, ws2,
13833        )
13834    }
13835
13836    #[allow(clippy::too_many_arguments)]
13837    pub fn qmatvec_e4m3_fused2_t_raw(
13838        &self,
13839        b0: &CudaSlice<u8>,
13840        b1: &CudaSlice<u8>,
13841        x: &CudaSlice<f32>,
13842        m: usize,
13843        in_f: usize,
13844        out0: usize,
13845        out1: usize,
13846        row_bytes: usize,
13847        ws0: f32,
13848        ws1: f32,
13849    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13850        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
13851        self.e4m3_fused2_t_core(b0, b1, &aq, &ad, m, in_f, out0, out1, row_bytes, ws0, ws1)
13852    }
13853
13854    #[allow(clippy::too_many_arguments)]
13855    pub fn qmatvec_e4m3_fused3_t_raw(
13856        &self,
13857        b0: &CudaSlice<u8>,
13858        b1: &CudaSlice<u8>,
13859        b2: &CudaSlice<u8>,
13860        x: &CudaSlice<f32>,
13861        m: usize,
13862        in_f: usize,
13863        out0: usize,
13864        out1: usize,
13865        out2: usize,
13866        row_bytes: usize,
13867        ws0: f32,
13868        ws1: f32,
13869        ws2: f32,
13870    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13871        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
13872        self.e4m3_fused3_t_core(
13873            b0, b1, b2, &aq, &ad, m, in_f, out0, out1, out2, row_bytes, ws0, ws1, ws2,
13874        )
13875    }
13876
13877    /// THE single dispatch point for `QT_F8_E4M3_BLK` from a PRE-QUANTIZED q8_1 activation
13878    /// (lane/fp8-blk128-decode). Every `matmul_pre`-family entry calls this first, so the block-128
13879    /// class has exactly ONE code path across `matmul`, `matmul_pre`, `matmul_pre_noscale`,
13880    /// `matmul_decode_exact` and `matmul_decode_exact_pre` — the same kernel at the same grid for
13881    /// every m, which is what makes verify == decode bit-for-bit at every tier for free.
13882    ///
13883    /// Returns None for any other qtype (the caller continues its normal dispatch). The `blk: Some`
13884    /// pattern is part of the match, not an unwrap: qtype and grid presence are set together in the
13885    /// one residency arm that builds this tensor, and a qtype-without-grid would be a construction
13886    /// bug — better to fall through and hit a loud refusal than to unwrap a None here.
13887    fn try_e4m3_blk_pre(
13888        &self,
13889        w: &crate::model::GpuTensor,
13890        aq: &CudaSlice<i8>,
13891        ad: &CudaSlice<f32>,
13892        m: usize,
13893    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
13894        use crate::model::GpuTensor;
13895        if let GpuTensor::Quant {
13896            bytes,
13897            qtype,
13898            row_bytes,
13899            blk: Some(g),
13900            ..
13901        } = w
13902        {
13903            if *qtype == QT_F8_E4M3_BLK {
13904                // BATCHED tier m=2..16 (lane/rp-on-st): weight read ONCE for up to mcols columns
13905                // instead of m grid.y re-reads. Bit-identical per (token,row) to the grid.y=m form
13906                // below, so the decode-exactness contract is preserved at every width. Gated by
13907                // the same seams the other batched families honor (MEMRA_NO_BATCHED, MEMRA_B8) so
13908                // one rollback door covers every dtype's batched tier.
13909                if (2..=16).contains(&m)
13910                    && std::env::var("MEMRA_NO_BATCHED").is_err()
13911                    && (m <= 4 || Self::b8_enabled())
13912                {
13913                    let mcols = Self::batched_mcols(m);
13914                    return Ok(Some(self.qmatvec_e4m3_blk_mmvq_batched(
13915                        bytes,
13916                        aq,
13917                        ad,
13918                        &g.scales,
13919                        m,
13920                        w.in_features(),
13921                        w.out_features(),
13922                        *row_bytes,
13923                        g.cols,
13924                        mcols,
13925                    )?));
13926                }
13927                return Ok(Some(self.qmatvec_e4m3_blk_mmvq(
13928                    bytes,
13929                    aq,
13930                    ad,
13931                    &g.scales,
13932                    m,
13933                    w.in_features(),
13934                    w.out_features(),
13935                    *row_bytes,
13936                    g.cols,
13937                )?));
13938            }
13939        }
13940        Ok(None)
13941    }
13942
13943    /// PREFILL (m >= GEMM_M_THRESHOLD) for `QT_F8_E4M3_BLK` — DEQUANT-PER-CALL to the Q8_0 slab
13944    /// this class's residency replaced, then the ordinary Q8_0 prefill dispatch on the transient.
13945    ///
13946    /// WHY THIS EXISTS AT ALL, i.e. the regression it prevents: the decode kernel is a warp-per-row
13947    /// GEMV. At grid.y=m it re-reads the whole weight once PER TOKEN, so letting a 512-token prefill
13948    /// chunk reach it would be a ~500x weight-traffic blowup on the single most bandwidth-bound part
13949    /// of the forward. Native residency is a DECODE win and must not be paid for in prefill, so
13950    /// prefill keeps the floor's arithmetic and the floor's kernels.
13951    ///
13952    /// WHY DEQUANT-PER-CALL rather than a second resident slab: a resident slab is dual residency —
13953    /// it gives back the entire 1.0-vs-1.0625 B/weight win this lane exists to capture (and then
13954    /// some, since the e4m3 copy stays too). The transient costs one linear device pass per
13955    /// (projection, prefill call) and frees immediately.
13956    ///
13957    /// NUMERICALLY IT IS THE FLOOR, EXACTLY: `fp8_blk_dequant_q8_0` is the merged ARM B' kernel,
13958    /// gate-proven BYTE-IDENTICAL to the host dequant+re-encode (kernel-check `fp8-blk-gpu`). So the
13959    /// slab these bytes form is bit-for-bit the slab the `MEMRA_ST_E4M3_BLK=0` arm makes resident,
13960    /// and every prefill kernel downstream sees identical input — prefill logits under this lane are
13961    /// bit-identical to prefill logits under the floor, which is what makes the decode A/B a clean
13962    /// single-variable comparison instead of a two-variable one.
13963    ///
13964    /// WHAT IT COSTS, MEASURED, AND WHY THAT COST IS MOSTLY STRUCTURAL (27B block-128 ckpt, pp512,
13965    /// this rig = RTX 5090 Laptop, ~896 GB/s GDDR7). This arm makes prefill move the weight THREE
13966    /// times instead of once: read 6.88 GB of e4m3, write 7.31 GB of Q8_0, then the MMQ reads that
13967    /// 7.31 GB back. The two extra passes are 14.19 GB = 15.8 ms at this card's roofline against a
13968    /// ~332 ms pp512, i.e. **~-4.5% pp is a floor no kernel tuning can remove** — only deleting the
13969    /// dequant can. Measured: the dequant kernel costs 27.9 ms/pass (nsys, 208 projections) after
13970    /// the 2026-08-05 vector rewrite (was 66.5 ms at one byte per thread), and e2e pp512 is
13971    /// 1451.4 vs the slab arm's 1541.6 tok/s = -5.8% (N=3 interleaved pairs). So ~1.3pp of the
13972    /// -5.8% is residual kernel inefficiency and ~4.5pp is the extra traffic itself.
13973    ///
13974    /// SO THE DEQUANT IS NO LONGER THE DEFAULT ROUTE — it is the FALLBACK. The per-block FP8 MMQ
13975    /// tile (`try_fp8_blk_mmq`) consumes the resident e4m3 bytes + grid DIRECTLY, deleting both extra
13976    /// passes, and since 2026-08-05 it runs FIRST and by default for the native-resident source
13977    /// (`fp8_blk_mmq_native_enabled`; `MEMRA_FP8_MMQ=0` is the seam back to this dequant). On paper
13978    /// the trade was unassumable — lane/fp8-mmq-v2 measured that tile at 0.85-1.09x the Q8_0 MMQ
13979    /// floor GEMM-only, so it swapped a -4.5% traffic cost for a 0-to-15% GEMM cost of unknown sign.
13980    /// Measured on the 27B (3 arms interleaved, N=3, research/fp8blk-20260805/VERDICT.md): slab
13981    /// 1540.5 / this dequant 1449.1 / the tile 1553.3 tok/s, min(tile) > max(slab). The tile wins
13982    /// because v2's denominator had its slab already resident while this class's floor must build it
13983    /// every call; same tile, opposite sign, because the question changed.
13984    ///
13985    /// THIS ARM STILL RUNS, and is not dead code: every `try_fp8_blk_mmq` precondition (in_f % 16,
13986    /// grid dims vs shape, per-tensor scale == 1.0, the e4m3-NaN scan) refuses by falling through to
13987    /// here, so a checkpoint the tile cannot take keeps exact prefill on the floor's own bits rather
13988    /// than losing the class. It is also what `MEMRA_FP8_MMQ=0` reverts to.
13989    fn try_e4m3_blk_prefill(
13990        &self,
13991        w: &crate::model::GpuTensor,
13992        x: &CudaSlice<f32>,
13993        m: usize,
13994    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
13995        use crate::model::GpuTensor;
13996        let GpuTensor::Quant {
13997            bytes,
13998            qtype,
13999            blk: Some(g),
14000            ..
14001        } = w
14002        else {
14003            return Ok(None);
14004        };
14005        if *qtype != QT_F8_E4M3_BLK {
14006            return Ok(None);
14007        }
14008        // NO-DEQUANT ROUTE, THE DEFAULT (MEMRA_FP8_MMQ=0 reverts): the per-block MMQ tile eats the
14009        // resident e4m3 bytes and grid as-is, so neither extra weight pass happens. Its own
14010        // preconditions (in_f % 16, grid dims, scale == 1.0, no e4m3 NaN code) can refuse — fall
14011        // through to the dequant below when they do, never silently produce nothing.
14012        if let Some(y) = self.try_fp8_blk_mmq(w, x, m)? {
14013            return Ok(Some(y));
14014        }
14015        let (in_f, out_f) = (w.in_features(), w.out_features());
14016        let slab = self.fp8_blk_dequant_q8_0_dev(bytes, &g.scales, out_f, in_f)?;
14017        let tmp = GpuTensor::Quant {
14018            bytes: slab,
14019            qtype: QT_Q8_0,
14020            row_bytes: in_f / 32 * 34,
14021            ne: vec![in_f as u64, out_f as u64],
14022            scale: 1.0,
14023            rp: false,
14024            #[cfg(memra_cutlass)]
14025            cutlass: None,
14026            fp8: None,
14027            blk: None,
14028            f16: None,
14029            rp4: None,
14030        };
14031        // Recursion terminates: `tmp` is QT_Q8_0 with `blk: None`, so it cannot re-enter this arm.
14032        Ok(Some(self.matmul(&tmp, x, m)?))
14033    }
14034
14035    pub fn matmul_pre_noscale(
14036        &self,
14037        w: &crate::model::GpuTensor,
14038        aq: &CudaSlice<i8>,
14039        ad: &CudaSlice<f32>,
14040        m: usize,
14041    ) -> Result<Option<(CudaSlice<f32>, f32)>, Box<dyn std::error::Error>> {
14042        use crate::model::GpuTensor;
14043        // BLOCK-128 e4m3: every scale factor is folded inside the kernel per k128, so the
14044        // "separable post-op scale" this entry exists to defer is 1.0 — return it explicitly
14045        // rather than let the tail below refuse and cost the caller a re-dispatch.
14046        if m == 1 {
14047            if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
14048                return Ok(Some((y, 1.0)));
14049            }
14050        }
14051        // Only the m==1 fast path applies the scale as a separable post-op; bail everywhere else.
14052        if m != 1 || !self.uses_q8_1_fast(w) {
14053            return Ok(None);
14054        }
14055        let in_f = w.in_features();
14056        let out_f = w.out_features();
14057        let (bytes, qtype, row_bytes, scale, rp) = match w {
14058            GpuTensor::Quant {
14059                bytes,
14060                qtype,
14061                row_bytes,
14062                scale,
14063                rp,
14064                ..
14065            } => (bytes, *qtype, *row_bytes, *scale, *rp),
14066            _ => return Ok(None),
14067        };
14068        // MMVQ warp-per-row (scale==1.0 passed -> kernel skips its internal scale; we return scale).
14069        if self.mmvq_supports(qtype) {
14070            // Q4_0 split-plane mirror (dp4a fallback below keeps the raw GGUF bytes).
14071            let (mbytes, mrp) = match w {
14072                GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
14073                _ => (bytes, rp),
14074            };
14075            let y = self.qmatvec_mmvq(
14076                mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, /*scale*/ 1.0, mrp,
14077            )?;
14078            return Ok(Some((y, scale)));
14079        }
14080        // dp4a fallback: same launch as matmul_pre but WITHOUT the post scale_inplace.
14081        let name = match qtype {
14082            QT_Q8_0 => "qmatvec_q8_0_dp4a",
14083            QT_Q4_K => "qmatvec_q4_K_dp4a",
14084            QT_Q6_K => "qmatvec_q6_K_dp4a",
14085            QT_Q5_K => "qmatvec_q5_K_dp4a",
14086            QT_Q3_K => "qmatvec_q3_K_dp4a",
14087            QT_NVFP4 => {
14088                if rp {
14089                    "qmatvec_nvfp4_dp4a_rp"
14090                } else {
14091                    "qmatvec_nvfp4_dp4a"
14092                }
14093            }
14094            QT_IQ4_XS => "qmatvec_iq4_XS_dp4a",
14095            _ => return Ok(None),
14096        };
14097        let f = self.func(name);
14098        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
14099        let cfg = LaunchConfig {
14100            grid_dim: (out_f as u32, m as u32, 1),
14101            block_dim: (128, 1, 1),
14102            shared_mem_bytes: 0,
14103        };
14104        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
14105        let __s_b = self.gpu.stream();
14106        let mut b = __s_b.launch_builder(&f);
14107        b.arg(bytes)
14108            .arg(aq)
14109            .arg(ad)
14110            .arg(&mut y)
14111            .arg(&inf)
14112            .arg(&outf)
14113            .arg(&mi)
14114            .arg(&rb);
14115        unsafe {
14116            b.launch(cfg)?;
14117        }
14118        Ok(Some((y, scale)))
14119    }
14120
14121    /// True if `qtype` has a warp-per-row MMVQ decode kernel AND MEMRA_MMVQ is set. Only the 4
14122    /// daily-hot dtypes (Q8_0, Q4_K, Q6_K, NVFP4) — others keep the _dp4a matvec (oracle/fallback).
14123    pub fn mmvq_supports(&self, qtype: i32) -> bool {
14124        // DEFAULT ON since 2026-07-08 (MEMRA_MMVQ=0 reverts to the _dp4a matvec class).
14125        // QT_F8_E4M3 is exempt from the MEMRA_MMVQ=0 escape: the e4m3 mmvq family is that dtype's
14126        // ONLY int8-act kernel class (there is no _dp4a twin), so its m=1/verify/batched dispatch
14127        // is a pure function of the dtype — the decode-parity law holds under every env.
14128        if qtype == QT_F8_E4M3 {
14129            return true;
14130        }
14131        if std::env::var("MEMRA_MMVQ").as_deref() == Ok("0") {
14132            return false;
14133        }
14134        matches!(
14135            qtype,
14136            QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q6_K | QT_NVFP4 | QT_Q4_0
14137        )
14138    }
14139
14140    /// PERF-3 warp-per-row MMVQ launcher (decode m=1 hot path). block=(32,ROWS_PER_BLOCK,1):
14141    /// one warp owns one output row, warp-only __shfl reduction (no smem barrier). Bit-equivalent
14142    /// to qmatvec_*_dp4a up to f32 reduction order. Pre-quantized q8_1 activation (aq,ad). NVFP4
14143    /// per-tensor macro-scale applied post (scale==1.0 for other dtypes -> no-op).
14144    pub fn qmatvec_mmvq(
14145        &self,
14146        bytes: &CudaSlice<u8>,
14147        aq: &CudaSlice<i8>,
14148        ad: &CudaSlice<f32>,
14149        m: usize,
14150        in_f: usize,
14151        out_f: usize,
14152        qtype: i32,
14153        row_bytes: usize,
14154        scale: f32,
14155        rp: bool,
14156    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14157        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output: skip memset
14158        self.qmatvec_mmvq_into(
14159            bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, rp, &mut y,
14160        )?;
14161        Ok(y)
14162    }
14163
14164    /// Slot-fed MMVQ twin (alloc-free capture lane): full policy body, caller-owned output.
14165    #[allow(clippy::too_many_arguments)]
14166    pub fn qmatvec_mmvq_into(
14167        &self,
14168        bytes: &CudaSlice<u8>,
14169        aq: &CudaSlice<i8>,
14170        ad: &CudaSlice<f32>,
14171        m: usize,
14172        in_f: usize,
14173        out_f: usize,
14174        qtype: i32,
14175        row_bytes: usize,
14176        scale: f32,
14177        rp: bool,
14178        y: &mut CudaSlice<f32>,
14179    ) -> Result<(), Box<dyn std::error::Error>> {
14180        debug_assert!(y.len() >= m * out_f);
14181        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
14182        // SMALL-SHAPE GRID FILL (H100 lane, 2026-07-26 microbench: attn qkv out_f=2048 =
14183        // 0.97 waves at the 4-warp block -> 66% of peak). The g2 twin (2 warps/block)
14184        // doubles the grid when the 4-warp launch would be sub-wave; per-row program
14185        // identical -> bit-identical. MEMRA_Q80_G2=0 reverts.
14186        if qtype == QT_Q8_0
14187            && rp
14188            && m == 1
14189            && out_f >= 64
14190            && (out_f as u32).div_ceil(ROWS_PER_BLOCK) < 4 * self.sm_count() as u32
14191            && {
14192                static G2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
14193                *G2.get_or_init(|| std::env::var("MEMRA_Q80_G2").as_deref() != Ok("0"))
14194            }
14195        {
14196            let f = self.func("qmatvec_q8_0_mmvq_rp_g2");
14197            let cfg = LaunchConfig {
14198                grid_dim: ((out_f as u32).div_ceil(2), 1, 1),
14199                block_dim: (32, 2, 1),
14200                shared_mem_bytes: 0,
14201            };
14202            let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, 1i32, row_bytes as i64);
14203            let __s_b = self.gpu.stream();
14204            let mut b = __s_b.launch_builder(&f);
14205            b.arg(bytes)
14206                .arg(aq)
14207                .arg(ad)
14208                .arg(&mut *y)
14209                .arg(&inf)
14210                .arg(&outf)
14211                .arg(&mi)
14212                .arg(&rb);
14213            unsafe {
14214                b.launch(cfg)?;
14215            }
14216            if scale != 1.0 {
14217                self.scale_inplace(y, scale, out_f)?;
14218            }
14219            return Ok(());
14220        }
14221        // Multi-row-per-warp (mr2) policy, fixed since the 2026-07 sweeps (the MEMRA_MMVQ_MR
14222        // override + mr4 kernel were retired 2026-07-08 — mr4 regressed on register pressure and
14223        // crashed under rp; q4_K/q6_K mr2 measured flat, "no gain = no change"):
14224        //   NVFP4 m=1 -> mr2 (clean +1-2% on 9B: RPW acc chains hide the weight-load latency
14225        //     that pins the single-row kernel at 30-46% DRAM). Bit-identical per row.
14226        //   Q5_K m=1 -> mr2 (2026-07-05: the FR-Spec trimmed draft head is Q5_K 32768 rows = 8%
14227        //     of the 27B p3 spec wall; latency-bound like the other k-quants pre-fix).
14228        //   Q4_K/Q6_K m=1 -> single-row (mr2 measured +0.7% / flat — weight-bandwidth-bound).
14229        let mut mr: u32 = if m == 1 && (qtype == QT_NVFP4 || qtype == QT_Q5_K) {
14230            2
14231        } else {
14232            1
14233        };
14234        // Q4_0 mr (gemma trunk): DEFAULT 1 since 2026-07-13 (MEMRA_Q40_MR=2 reverts) — the
14235        // mr1 rp twin doubles the block count and wins the tail-quantization/latency battle
14236        // on every gemma model (E4B +3.75%: 198.9 vs 191.7; 26B +0.7%; 31B +0.9%; N=2-3
14237        // valid-window interleaved, bit-identical per row — same dot program).
14238        if m == 1 && qtype == QT_Q4_0 {
14239            static Q40MR: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
14240            // shape policy PROBED NEGATIVE (2026-07-13): tall-only mr1 197.2 vs
14241            // mr1-everywhere 198.7 — mr1 wins wide-output shapes too; arm removed.
14242            mr = *Q40MR.get_or_init(|| {
14243                std::env::var("MEMRA_Q40_MR")
14244                    .ok()
14245                    .and_then(|v| v.parse().ok())
14246                    .unwrap_or(1)
14247            });
14248        }
14249        // q5issue lane (2026-07-08): MEMRA_Q5K_ISSUE swaps the q5_K m=1 mmvq kernels for the
14250        // issue-reduced `_il` bodies (uint4 header/qh/qs loads + branchless scale decode —
14251        // cuts ~34 LDG.U16 + ~5 LDG.U8 + a warp-divergent scale branch per 32-elem group-row
14252        // to 5 LDG.128). Bit-identical per (token,row) to the reference kernels.
14253        // `1` = shape-aware policy (N=3 clock-locked micro-bench, mem P0, synthetic real shapes):
14254        //   out_f <= 65536 (trunk/frspec regime): il at the default mr — mr2_il -9.5%/-10.5%
14255        //     on 4096x4096/4096x8192, -3.1% on the 32768 frspec head vs the mr2-ref default;
14256        //   out_f > 65536 (the 248320-row 27B lm_head, already ~97% of the mem wall): mr2_il
14257        //     REGRESSES +22% there but mr1_il wins -2.1% vs the mr2-ref default -> force mr=1.
14258        // `2` = force il at the current mr for EVERY shape (A/B probe seam). Default OFF.
14259        let q5_mode = std::env::var("MEMRA_Q5K_ISSUE").ok();
14260        let q5_force = q5_mode.as_deref() == Some("2");
14261        // DEFAULT ON since 2026-07-08 (MEMRA_Q5K_ISSUE=0 reverts): +1.8% 9B plain e2e N=3
14262        // (128.2 -> 130.4), 27B flat (its big head is already at the mem wall), all gates green.
14263        let q5_il = qtype == QT_Q5_K
14264            && m == 1
14265            && (q5_force || q5_mode.as_deref().map(|v| v != "0").unwrap_or(true));
14266        if q5_il && !q5_force && out_f > 65536 {
14267            mr = 1;
14268        }
14269        // Q4_0 split-plane rp: mr2 default; MEMRA_Q40_MR=1 reaches the mr1 rp twin
14270        // (2026-07-13 — the tall-input/short-output tail-quantization probe).
14271        if qtype == QT_Q4_0 && rp && mr != 1 {
14272            mr = 2;
14273        }
14274        // Q8_0 rp (H100 lane): mr1 default — the q4_0 mr2 recipe MEASURED NEGATIVE on H100
14275        // (2026-07-26 N=3: mr1 186.2 vs mr2 171.5 tok/s; halving the grid on 132 SMs costs
14276        // more than 2-row ILP buys). mr2 kernel stays behind MEMRA_Q80_MR=2 for the corpus.
14277        if qtype == QT_Q8_0 && rp {
14278            static Q80MR: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
14279            mr = *Q80MR.get_or_init(|| {
14280                std::env::var("MEMRA_Q80_MR")
14281                    .ok()
14282                    .and_then(|v| v.parse().ok())
14283                    .unwrap_or(1)
14284            });
14285        }
14286        let name = match (qtype, mr, rp) {
14287            (QT_NVFP4, 2, false) => "qmatvec_nvfp4_mmvq_mr2",
14288            (QT_NVFP4, 2, true) => "qmatvec_nvfp4_mmvq_mr2_rp",
14289            (QT_NVFP4, _, true) => "qmatvec_nvfp4_mmvq_rp",
14290            (QT_Q4_0, 1, true) => "qmatvec_q4_0_mmvq_rp",
14291            (QT_Q4_0, _, true) => "qmatvec_q4_0_mmvq_mr2_rp",
14292            (QT_Q5_K, 2, _) => {
14293                if q5_il {
14294                    "qmatvec_q5_K_mmvq_mr2_il"
14295                } else {
14296                    "qmatvec_q5_K_mmvq_mr2"
14297                }
14298            }
14299            (QT_Q8_0, 2, true) => "qmatvec_q8_0_mmvq_mr2_rp",
14300            // rpca (cp.async-staged weight ring): MEASURED NEGATIVE on H100 for Q8_0
14301            // (2026-07-26 N=3: 181.8 vs plain rp 185.5 — the smem round-trip exceeds the
14302            // latency it hides for 8-bit direct-dp4a; the NVFP4 win case overlaps table
14303            // decode with half the bytes). OPT-IN via MEMRA_Q80_CA=1 for the corpus.
14304            (QT_Q8_0, _, true)
14305                if in_f % 1024 == 0 && {
14306                    static CA: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
14307                    *CA.get_or_init(|| std::env::var("MEMRA_Q80_CA").as_deref() == Ok("1"))
14308                } =>
14309            {
14310                "qmatvec_q8_0_mmvq_rpca"
14311            }
14312            (QT_Q8_0, _, true) => "qmatvec_q8_0_mmvq_rp",
14313            (QT_Q8_0, _, _) => "qmatvec_q8_0_mmvq",
14314            // K-quant split-plane twins (H100 K-quant coalescing fix, 2026-08-01): the rp4
14315            // mirror routes here; GGUF layout keeps the plain kernels. rp bytes MUST never
14316            // reach a GGUF-layout kernel or vice versa.
14317            (QT_Q4_K, _, true) => "qmatvec_q4_K_mmvq_rp",
14318            (QT_Q6_K, _, true) => "qmatvec_q6_K_mmvq_rp",
14319            (QT_Q4_K, _, _) => "qmatvec_q4_K_mmvq",
14320            (QT_Q4_0, 2, false) => "qmatvec_q4_0_mmvq_mr2",
14321            (QT_Q4_0, _, false) => "qmatvec_q4_0_mmvq",
14322            (QT_Q5_K, _, _) => {
14323                if q5_il {
14324                    "qmatvec_q5_K_mmvq_il"
14325                } else {
14326                    "qmatvec_q5_K_mmvq"
14327                }
14328            }
14329            (QT_Q6_K, _, _) => "qmatvec_q6_K_mmvq",
14330            (QT_NVFP4, _, false) => "qmatvec_nvfp4_mmvq",
14331            (QT_F8_E4M3, _, _) => "qmatvec_e4m3_mmvq",
14332            _ => panic!("qmatvec_mmvq: qtype {qtype} has no MMVQ kernel"),
14333        };
14334        let f = self.func(name);
14335        // each block still has ROWS_PER_BLOCK warps; with mr rows/warp it covers ROWS_PER_BLOCK*mr rows.
14336        let rows_per_block = ROWS_PER_BLOCK * mr;
14337        let cfg = LaunchConfig {
14338            grid_dim: (
14339                (out_f as u32 + rows_per_block - 1) / rows_per_block,
14340                m as u32,
14341                1,
14342            ),
14343            block_dim: (32, ROWS_PER_BLOCK, 1), // warp-per-row (x mr rows each)
14344            shared_mem_bytes: 0,                // warp-only reduce at m=1
14345        };
14346        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
14347        let __s_b = self.gpu.stream();
14348        let mut b = __s_b.launch_builder(&f);
14349        // NVFP4 + e4m3 mmvq kernels take the macro-scale as a fused epilogue arg (applied at the
14350        // write — bit-identical to the old separate scale_inplace pass, minus one launch per matvec:
14351        // 53 scale launches/token on the 9B; for e4m3 the scale is the checkpoint's per-tensor f32
14352        // weight_scale). Other mmvq kernels keep the 8-arg signature.
14353        if qtype == QT_NVFP4 || qtype == QT_F8_E4M3 {
14354            // PDL wave-B: the nvfp4 mr2_rp single (gemma wo / generic rp singles) joins
14355            // the wave-A launch class — 9-arg flavor (fused macro-scale epilogue).
14356            if Self::pdl_on()
14357                && Self::pdl_mmvq_on()
14358                && Self::pdl_nvfp4q8_on()
14359                && name == "qmatvec_nvfp4_mmvq_mr2_rp"
14360            {
14361                use cudarc::driver::{DevicePtr, DevicePtrMut};
14362                let s = &self.gpu.stream();
14363                let (pw, _g0) = bytes.device_ptr(s);
14364                let (paq, _g1) = aq.device_ptr(s);
14365                let (pad, _g2) = ad.device_ptr(s);
14366                let (py, _g3) = y.device_ptr_mut(s);
14367                let mut ps = [
14368                    &pw as *const _ as *mut std::ffi::c_void,
14369                    &paq as *const _ as *mut _,
14370                    &pad as *const _ as *mut _,
14371                    &py as *const _ as *mut _,
14372                    &inf as *const _ as *mut _,
14373                    &outf as *const _ as *mut _,
14374                    &mi as *const _ as *mut _,
14375                    &rb as *const _ as *mut _,
14376                    &scale as *const _ as *mut _,
14377                ];
14378                unsafe {
14379                    self.launch_pdl(name, cfg.grid_dim, cfg.block_dim, &mut ps)?;
14380                }
14381                return Ok(());
14382            }
14383            b.arg(bytes)
14384                .arg(aq)
14385                .arg(ad)
14386                .arg(&mut *y)
14387                .arg(&inf)
14388                .arg(&outf)
14389                .arg(&mi)
14390                .arg(&rb)
14391                .arg(&scale);
14392            unsafe {
14393                b.launch(cfg)?;
14394            }
14395        } else if Self::pdl_on()
14396            && Self::pdl_mmvq_on()
14397            && (matches!(
14398                name,
14399                "qmatvec_q4_0_mmvq_rp" | "qmatvec_q6_K_mmvq" | "qmatvec_q6_K_mmvq_rp"
14400            ) || (Self::pdl_nvfp4q8_on()
14401                && matches!(name, "qmatvec_q8_0_mmvq_rp" | "qmatvec_q8_0_mmvq_mr2_rp")))
14402        {
14403            // PDL wave-A (2026-07-23): the two decode-hot single-matvec kernels carry
14404            // MEMRA_PDL_ENTRY — grid launches while the producer drains. ONLY the marked
14405            // names may take this launch (unmarked kernels would read unordered).
14406            {
14407                use cudarc::driver::{DevicePtr, DevicePtrMut};
14408                let s = &self.gpu.stream();
14409                let (pw, _g0) = bytes.device_ptr(s);
14410                let (paq, _g1) = aq.device_ptr(s);
14411                let (pad, _g2) = ad.device_ptr(s);
14412                let (py, _g3) = y.device_ptr_mut(s);
14413                let mut ps = [
14414                    &pw as *const _ as *mut std::ffi::c_void,
14415                    &paq as *const _ as *mut _,
14416                    &pad as *const _ as *mut _,
14417                    &py as *const _ as *mut _,
14418                    &inf as *const _ as *mut _,
14419                    &outf as *const _ as *mut _,
14420                    &mi as *const _ as *mut _,
14421                    &rb as *const _ as *mut _,
14422                ];
14423                unsafe {
14424                    self.launch_pdl(name, cfg.grid_dim, cfg.block_dim, &mut ps)?;
14425                }
14426            }
14427            if scale != 1.0 {
14428                self.scale_inplace(y, scale, m * out_f)?;
14429            }
14430        } else {
14431            b.arg(bytes)
14432                .arg(aq)
14433                .arg(ad)
14434                .arg(&mut *y)
14435                .arg(&inf)
14436                .arg(&outf)
14437                .arg(&mi)
14438                .arg(&rb);
14439            unsafe {
14440                b.launch(cfg)?;
14441            }
14442            if scale != 1.0 {
14443                self.scale_inplace(y, scale, m * out_f)?;
14444            }
14445        }
14446        Ok(())
14447    }
14448
14449    /// Test entry for the kernel_check bit-equivalence gate: run the warp-per-row MMVQ directly
14450    /// from raw weight bytes (quantize the f32 activation `x` to q8_1 internally). NVFP4 per-tensor
14451    /// macro-scale is NOT applied (caller compares bare, like qmatvec_*_fast). Mirrors qmatvec_gemm_raw.
14452    pub fn qmatvec_mmvq_raw(
14453        &self,
14454        bytes: &CudaSlice<u8>,
14455        x: &CudaSlice<f32>,
14456        m: usize,
14457        in_f: usize,
14458        out_f: usize,
14459        qtype: i32,
14460        row_bytes: usize,
14461        rp: bool,
14462    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14463        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
14464        self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, 1.0, rp)
14465    }
14466
14467    /// True if `qtype` has a batched weight-resident (`_b2`/`_b4`) matvec kernel. These mirror the
14468    /// `_mmvq` kernels but iterate the m token columns INSIDE one warp/row, so the weight bytes leave
14469    /// HBM/L2 once for m tokens (vs grid.y=m re-reading m times). The 5 daily-hot dtypes have them.
14470    pub fn batched_supports(&self, qtype: i32) -> bool {
14471        matches!(
14472            qtype,
14473            QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q6_K | QT_NVFP4 | QT_F8_E4M3 | QT_Q4_0
14474        )
14475    }
14476
14477    /// IQ4_XS trunk fast seam: MEMRA_IQ_FAST=0 reverts non-expert IQ4_XS matmuls to the Stage-A
14478    /// f32 oracle path. Default ON since 2026-08-02 (research/kat-anomaly-20260802/): the old
14479    /// opt-in default left every IQ4_XS-trunk artifact (KAT-Coder IQ4_XS: attn_qkv/attn_gate/
14480    /// ssm_out/shexp, ~0.52GB re-read per decode tick) on the oracle kernel — decode 106.7 ->
14481    /// 193.4 tok/s (x5 interleaved), pp512 228 -> 697, same bytes, via qmatvec_iq4_XS_dp4a. The
14482    /// supported artifacts carry IQ4_XS only in EXPERT banks (their own dispatch, not this seam),
14483    /// so this admission is dispatch-unchanged for every non-IQ4_XS-trunk model.
14484    pub fn iq_fast_enabled() -> bool {
14485        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
14486        *ON.get_or_init(|| {
14487            std::env::var("MEMRA_IQ_FAST")
14488                .map(|v| v != "0")
14489                .unwrap_or(true)
14490        })
14491    }
14492
14493    /// b8 tier seam: MEMRA_B8=0 keeps m=5..8 on the per-m grid.y=m path (m=2..4 batched dispatch
14494    /// unaffected). Default ON — the K=4..7 spec-verify weight-read-once fix.
14495    pub fn b8_enabled() -> bool {
14496        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
14497        *ON.get_or_init(|| std::env::var("MEMRA_B8").map(|v| v != "0").unwrap_or(true))
14498    }
14499
14500    /// Compile-time column batch for a runtime m: 2 -> b2, 3..4 -> b4, 5..8 -> b8.
14501    pub fn batched_mcols(m: usize) -> usize {
14502        if m == 2 {
14503            2
14504        } else if m <= 4 {
14505            4
14506        } else if m <= 8 {
14507            8
14508        } else {
14509            16
14510        }
14511    }
14512
14513    /// Kernel name for the batched matvec of `(qtype, mcols)`. mcols ∈ {2,4,8}. The b8 tier is the
14514    /// K=4..7 spec-verify fix (T=5..8): pre-b8 those T fell to grid.y=m per-row MMVQ = m full
14515    /// weight reads/launch — the measured 27B K=4 cliff (101 -> 73 tok/s at p3 despite acceptance
14516    /// holding 54%). One b8 launch reads the weight ONCE for up to 8 columns (c >= m masked).
14517    fn batched_kernel_name(qtype: i32, mcols: usize) -> Option<&'static str> {
14518        Some(match (qtype, mcols) {
14519            (QT_Q8_0, 2) => "qmatvec_q8_0_mmvq_b2",
14520            (QT_Q8_0, 4) => "qmatvec_q8_0_mmvq_b4",
14521            (QT_Q8_0, 8) => "qmatvec_q8_0_mmvq_b8",
14522            // b16 now has BOTH forms (lane/rp-on-st, 2026-08-06). It used to be rp-ONLY, which
14523            // made the q8rp mirror the exact-16 tier's admission ticket for any model carrying a
14524            // single Q8_0 matmul — measured as the FP8-ST refusal (`L0.ssm_beta qtype=0
14525            // rp4=false`, 96 t / 23.9 MiB = 0.143% of resident weight). The mirror stays a
14526            // BANDWIDTH lever on Q8_0-dominant GGUFs; it is no longer a correctness prerequisite.
14527            (QT_Q8_0, 16) => "qmatvec_q8_0_mmvq_b16",
14528            (QT_Q4_K, 2) => "qmatvec_q4_K_mmvq_b2",
14529            (QT_Q4_K, 4) => "qmatvec_q4_K_mmvq_b4",
14530            (QT_Q4_K, 8) => "qmatvec_q4_K_mmvq_b8",
14531            // b16 base + _rp (lane/rp-on-st): the 9B NVFP4 GGUF's blocker — real NVFP4 GGUFs keep
14532            // Q4_K attention next to NVFP4 MLP, and the tier's predicate is an ALL.
14533            (QT_Q4_K, 16) => "qmatvec_q4_K_mmvq_b16",
14534            (QT_Q5_K, 2) => "qmatvec_q5_K_mmvq_b2",
14535            (QT_Q5_K, 4) => "qmatvec_q5_K_mmvq_b4",
14536            (QT_Q5_K, 8) => "qmatvec_q5_K_mmvq_b8",
14537            // b16 base only (lane/rp-on-st): Q5_K has no rp twins at any width, so there is
14538            // nothing to mirror. Named by the diagnostic as `L0.wqkv_gate qtype=3` on the 9B.
14539            (QT_Q5_K, 16) => "qmatvec_q5_K_mmvq_b16",
14540            (QT_Q6_K, 2) => "qmatvec_q6_K_mmvq_b2",
14541            (QT_Q6_K, 4) => "qmatvec_q6_K_mmvq_b4",
14542            (QT_Q6_K, 8) => "qmatvec_q6_K_mmvq_b8",
14543            (QT_Q6_K, 16) => "qmatvec_q6_K_mmvq_b16",
14544            (QT_NVFP4, 2) => "qmatvec_nvfp4_mmvq_b2",
14545            (QT_NVFP4, 4) => "qmatvec_nvfp4_mmvq_b4",
14546            (QT_NVFP4, 8) => "qmatvec_nvfp4_mmvq_b8",
14547            // b16 (lane/rp-on-st): no mirror needed — NVFP4's 36 B/k32 block is already the
14548            // aligned form its own kernel walks. Unlocks the exact-16 tier for every NVFP4 model
14549            // AND for the mixed FP8-ST artifact, whose 193 NVFP4 tensors were refusing it.
14550            (QT_NVFP4, 16) => "qmatvec_nvfp4_mmvq_b16",
14551            (QT_F8_E4M3, 2) => "qmatvec_e4m3_mmvq_b2",
14552            (QT_F8_E4M3, 4) => "qmatvec_e4m3_mmvq_b4",
14553            (QT_F8_E4M3, 8) => "qmatvec_e4m3_mmvq_b8",
14554            // b16 tier (lane/rp-on-st): e4m3 needs NO split-plane mirror to reach it — its native
14555            // row-major layout is already 32B-aligned per k32 block, so the base kernel IS the
14556            // aligned form. Contrast Q8_0, whose b16 exists only as the `_rp` twin (hence q8rp).
14557            (QT_F8_E4M3, 16) => "qmatvec_e4m3_mmvq_b16",
14558            (QT_Q4_0, 2) => "qmatvec_q4_0_mmvq_b2",
14559            (QT_Q4_0, 4) => "qmatvec_q4_0_mmvq_b4",
14560            (QT_Q4_0, 8) => "qmatvec_q4_0_mmvq_b8",
14561            (QT_Q4_0, 16) => "qmatvec_q4_0_mmvq_b16",
14562            _ => return None,
14563        })
14564    }
14565
14566    /// BATCHED weight-tile-resident matvec from a PRE-QUANTIZED q8_1 activation (the m=2-8 verify/MTP
14567    /// win). One warp walks the weight row ONCE, dp4a vs all m activation columns -> weight HBM/L2
14568    /// traffic 1x for m tokens (vs grid.y=m re-reading it m times). `mcols` ∈ {2,4,8} is the
14569    /// compile-time batch; m must be <= mcols (the c >= m columns are masked in-kernel). y is
14570    /// [m, out_f] token-major. NVFP4 per-tensor macro-scale applied post
14571    /// (scale==1.0 for other dtypes -> no-op). BIT-IDENTICAL per (token,row) to qmatvec_*_mmvq.
14572    ///
14573    /// NVFP4 VARIANT DISPATCH: the batched NVFP4 kernel measured memory-LATENCY bound on the real
14574    /// 27B verify (ncu --set full, 12 steady launches: long_scoreboard 18-30 stalls/issue vs <=1.7
14575    /// for every other reason, DRAM only 41-51% active, lg_throttle 0.7, L1 hit 94% — ONE 6-LDG
14576    /// weight wavefront in flight per warp is the binding constraint, NOT bandwidth and NOT the
14577    /// column-unroll break). Two exactness-free fixes, chosen PER SHAPE from the DRAM-cold 8-copy
14578    /// msweep on all six 27B shapes (2026-07-03):
14579    ///   `pf` = next-g weight-prefetch double-buffer (48 regs, occupancy intact) — wins everywhere
14580    ///          it applies for b4 (-3..-14%), never loses;
14581    ///   `r2` = two rows/warp (67 regs -> 7 resident blocks/SM) — the bigger win (-8.5..-30%) but
14582    ///          wave-quantization-sensitive: with the grid halved to ceil(out_f/8) blocks, a
14583    ///          fractional straggler wave (waves in ~1.05-1.5) costs a full extra latency round on
14584    ///          a latency-bound kernel (27B ffn_down 640 blocks / 574 resident = 1.11 waves: +17%),
14585    ///          while <=1 wave (9B ffn_down 0.89: -30%) or >=2 waves (tail amortized; qkv 2.2:
14586    ///          -8.5%, ffn_gate 3.8: -12.5%) win. For b2, r2 wins on DEEP k-loops (in_f>=6144:
14587    ///          -8..-19%) where the 2-col body starves weight MLP hardest; pf measured negative.
14588    /// b4: r2 when waves(out_f) <= 1 (and grid fills >=half the SMs) or >= 2, else pf.
14589    /// b2: in_f>=6144 -> r2, else base.
14590    /// MEMRA_MMVQ_BV=base|pf|r2|pfr2 forces one variant everywhere (A/B + rollback seam).
14591    /// All variants BIT-IDENTICAL per (token,row): same dp4a order, scales, adg factor, reduce —
14592    /// only load issue time and the row->warp mapping change (kernel-check gates all of them).
14593    /// `rp` = the weight buffer is the A6 SPLIT-PLANE repacked layout (NVFP4 only): the same
14594    /// wave-aware auto rule applies, mapped onto the `_rp` twins (rp/rpr2/rpr2w8 mirror
14595    /// pf/r2/r2w8 — regs 44/67/64 land in the same residency classes).
14596    /// The variant the batched dispatch will pick for this (shape, m, mcols, layout) — exposed so
14597    /// gates can distinguish bit-identical variants (bit-bad==0 required) from the k-split family
14598    /// (deterministic but k-reduce-order-shifted: rel<1e-3 + run-to-run bit-identity required).
14599    /// Device SM count (cached) — grid-fill policy input.
14600    pub fn sm_count(&self) -> i32 {
14601        static SMS: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
14602        *SMS.get_or_init(|| {
14603            use cudarc::driver::sys::CUdevice_attribute_enum as A;
14604            self.gpu
14605                .ctx
14606                .attribute(A::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)
14607                .unwrap_or(82)
14608        })
14609    }
14610
14611    pub fn batched_variant(
14612        &self,
14613        _m: usize,
14614        in_f: usize,
14615        out_f: usize,
14616        qtype: i32,
14617        row_bytes: usize,
14618        mcols: usize,
14619        rp: bool,
14620    ) -> &'static str {
14621        // Q8_0 never joined the auto variant machinery (on sm_120 its only batched shapes
14622        // were tiny aux tensors). On Q8_0-trunk models the layout is the whole game: the
14623        // split-plane mirror (rp) routes to the _rp twins (H100 coalescing fix, 2026-07-26);
14624        // GGUF layout stays "base". rp bytes MUST never reach the base kernel or vice versa.
14625        if qtype == QT_Q8_0 {
14626            return if rp { "rp" } else { "base" };
14627        }
14628        static BV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
14629        let bv = *BV.get_or_init(|| match std::env::var("MEMRA_MMVQ_BV").as_deref() {
14630            Ok("base") => "base",
14631            Ok("pf") => "pf",
14632            Ok("r2") => "r2",
14633            Ok("r2w8") => "r2w8",
14634            Ok("pfr2") => "pfr2",
14635            Ok("ca") => "ca",
14636            Ok("car2") => "car2",
14637            // rp* = SPLIT-PLANE REPACKED layout kernels (A6 prototype): W must already be the
14638            // repacked buffer (msweep MSWEEP_RP harness) — never valid on GGUF-layout weights.
14639            Ok("rp") => "rp",
14640            Ok("rpr2") => "rpr2",
14641            Ok("rpr2w8") => "rpr2w8",
14642            // rpca* = cp.async software-pipelined split-plane (2026-07-05): hides the _rp
14643            // long_scoreboard load stall. rp-layout only; b4/b2 (no b8 twin).
14644            Ok("rpca") => "rpca",
14645            Ok("rpcar2") => "rpcar2",
14646            // 2026-07-06 m-small latency arc: rpsc = rpr2 + per-warp smem scale prestage (kills
14647            // the scale-plane global dependency, zero reg growth); rpms/rpmsc = m-split x2
14648            // across warp pairs (2x blocks of rpr2, column halves per warp, BIT-identical to
14649            // _rp); rpks/rpksc = k-split x2 (fastest microbench cells but k-reduce-order-shifted:
14650            // run-spec self-consistency FAILED on the 27B daily driver — verify logits must be
14651            // bit-identical to the decode path — measurement corpus ONLY, never auto).
14652            Ok("rpsc") => "rpsc",
14653            Ok("rpms") => "rpms",
14654            Ok("rpmsc") => "rpmsc",
14655            Ok("rpks") => "rpks",
14656            Ok("rpksc") => "rpksc",
14657            _ => "auto",
14658        });
14659        // cp.async ring variants need 16B-aligned rows (in_f%256==0 -> (in_f/64)*36 % 16 == 0)
14660        // and whole 32-group warp iterations (nsb%32==0 <=> in_f%1024==0). All 27B/9B trunk
14661        // shapes qualify; anything else falls back to the register variants.
14662        let ca_ok = qtype == QT_NVFP4 && (row_bytes % 16 == 0) && (in_f % 1024 == 0);
14663        // rpsc: smem scale plane fits (nsb64 <= 272) + int4-aligned staging (nsb64 % 4 == 0).
14664        // rpks/rpksc: half-plane staging alignment needs nsb64 % 8 == 0 (in_f % 512 == 0).
14665        // MEMRA_KS=0 removes the 2026-07-06 rpsc/rpks/rpksc entries from AUTO (rollback seam;
14666        // forced MEMRA_MMVQ_BV values still work).
14667        static KS_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
14668        let ks_on = *KS_ON.get_or_init(|| std::env::var("MEMRA_KS").as_deref() != Ok("0"));
14669        let sc_ok = ks_on && qtype == QT_NVFP4 && (in_f % 256 == 0) && (in_f / 64 <= 272);
14670        let ks_ok = ks_on && qtype == QT_NVFP4 && (in_f % 512 == 0) && (in_f / 64 <= 272);
14671        static SMS: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
14672        let sms = *SMS.get_or_init(|| {
14673            use cudarc::driver::sys::CUdevice_attribute_enum as A;
14674            self.gpu
14675                .ctx
14676                .attribute(A::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)
14677                .unwrap_or(82)
14678        });
14679        // k-quant r2 port (2026-07-04): q4_K/q5_K/q6_K have _r2/_r2w8 twins. ncu on the DRAM-cold
14680        // 9B msweep showed q4_K/q5_K b4 memory-latency bound like NVFP4 pre-fix (long_scoreboard
14681        // 19.6/16.4 per issue, DRAM 47.7/38.2%, L2 weight hit ~13%); q6_K lm_head is the exception
14682        // at DRAM 90-91% = wall-bound (yet r2 still wins -8%: deeper MLP raises achieved DRAM).
14683        // No _pf port (a k-quant group stages 10+ words vs NVFP4's 5 — register cost outweighs;
14684        // r2 covers the same MLP) and no rp (GGUF layout only). Q8_0 stays base: its only real
14685        // batched shapes are the tiny out_f=32 ssm_alpha/beta (8-block grids never fill one SM).
14686        // AUTO RULE = the measured winners table (differs from NVFP4's!):
14687        //   r2w8 NEVER in auto — the reg squeeze (72 -> 64 regs = stack spill) loses to unbounded
14688        //     r2 on every measured k-quant cell, incl. the wave-crossing lm_heads (q6_K 1316 vs
14689        //     r2 1258us) — kernels kept behind the force seam for the corpus;
14690        //   q4_K: r2 whenever the halved grid fills the SMs (blocks >= 4*SMs), INCLUDING the
14691        //     1.05-2.0 straggler window where NVFP4's r2 lost (qkv 1.78 waves: r2 -15% here; the
14692        //     k-quant base kernel leaves more latency on the table than a straggler wave costs);
14693        //   q5_K/q6_K: r2 only at waves >= 2 (the 248320-row lm_heads, 48+ waves: q6_K -8%, q5_K
14694        //     -2%); mid shapes measured base-or-flat (q5_K qkv 49.1 base vs 49.7 r2, attn_gate
14695        //     flat, attn_k base) — the 5/6-bit two-stream unpack makes r2's staging pricier.
14696        //   b2 same table with 8-row blocks: q4_K r2 when filled (-3..-22% all measured shapes),
14697        //     q5_K/q6_K r2 at waves >= 2 (27B lm_head -2.9%; 9B q6_K flat, harmless).
14698        let kq_r2 = matches!(qtype, QT_Q4_K | QT_Q5_K | QT_Q6_K);
14699        // MEMRA_KQ_BV=base|r2|r2w8 forces the k-quant variant WITHOUT touching the NVFP4 dispatch
14700        // (MEMRA_MMVQ_BV is global — an interleaved k-quant-only e2e A/B needs this narrower seam).
14701        static KQBV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
14702        let kq_bv = *KQBV.get_or_init(|| match std::env::var("MEMRA_KQ_BV").as_deref() {
14703            Ok("base") => "base",
14704            Ok("r2") => "r2",
14705            Ok("r2w8") => "r2w8",
14706            _ => "auto",
14707        });
14708        let variant: &'static str = if qtype == QT_Q4_0 {
14709            // Q4_0 r2 (gemma verify trunk, 2026-07-10): shared activation loads + the
14710            // row-independent ones-sum computed once per (col,group) for 2 rows. Same
14711            // fill rule as q4_K: r2 when the halved grid still fills the SMs.
14712            static Q40BV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
14713            let q40 = *Q40BV.get_or_init(|| match std::env::var("MEMRA_Q40_BV").as_deref() {
14714                // ms/sm/la = force-only measurement seams (ALL FLAT/NEGATIVE 2026-07-13,
14715                // never auto): m-split flat (nvcc keeps 72 regs); smem-slab −11% (staging
14716                // + syncs cost more than the stalls, bank-pad made no difference);
14717                // register load-ahead flat (nvcc already reorders). The b-tier limiter
14718                // is still unidentified — see the jsonl row.
14719                Ok("base") => "base",
14720                Ok("r2") => "r2",
14721                Ok("ms") => "ms",
14722                Ok("sm") => "sm",
14723                Ok("la") => "la",
14724                _ => "auto",
14725            });
14726            let v = if q40 != "auto" {
14727                q40
14728            } else if (out_f as u32).div_ceil(8) >= 4 * sms as u32 {
14729                "r2"
14730            } else {
14731                "base"
14732            };
14733            // split-plane mirror twins (2026-07-10): same fill rule, _rp names.
14734            // (m-split r2 pair twin PROBED FLAT 2026-07-13 — nvcc kept 72 regs either way
14735            // and the limiter is the per-column activation load chain (long_scoreboard
14736            // 42.5%), not occupancy; arm killed per doctrine, jsonl row is the record.)
14737            if rp {
14738                match v {
14739                    "ms" => "r2ms_rp",
14740                    "sm" => "r2sm_rp",
14741                    "la" => "r2la_rp",
14742                    "r2" => "r2_rp",
14743                    _ => "rp",
14744                }
14745            } else if matches!(v, "ms" | "sm" | "la") {
14746                "r2"
14747            } else {
14748                v
14749            }
14750        } else if qtype != QT_NVFP4 && !kq_r2 {
14751            "base"
14752        } else if kq_r2 && rp {
14753            // K-quant split-plane mirror (2026-08-01): only the plain _rp batched twins are
14754            // compiled for q4_K/q6_K — rp is a LAYOUT, it must survive every heuristic
14755            // (split-plane bytes through a GGUF-layout kernel = NaN). q5_K never mirrors.
14756            "rp"
14757        } else if kq_r2 {
14758            // k-quant r2w8 only exists at b4 (b2_r2 already 8-resident; b8 has no w8 twin) ->
14759            // mcols != 4 forced r2w8 falls to unbounded r2.
14760            if kq_bv != "auto" {
14761                if kq_bv == "r2w8" && mcols != 4 {
14762                    "r2"
14763                } else {
14764                    kq_bv
14765                }
14766            } else if bv != "auto" {
14767                match bv {
14768                    "r2" | "pfr2" | "rpr2" | "car2" => "r2",
14769                    "r2w8" | "rpr2w8" => {
14770                        if mcols != 4 {
14771                            "r2"
14772                        } else {
14773                            "r2w8"
14774                        }
14775                    }
14776                    _ => "base", // base/pf/ca/rp forced -> base (no such k-quant kernels)
14777                }
14778            } else {
14779                let blocks = (out_f + 7) / 8;
14780                let waves = blocks as f64 / (7 * sms as usize) as f64;
14781                let filled = blocks >= 4 * sms as usize;
14782                let use_r2 = if qtype == QT_Q4_K {
14783                    filled
14784                } else {
14785                    waves >= 2.0
14786                };
14787                if use_r2 { "r2" } else { "base" }
14788            }
14789        } else if bv != "auto" {
14790            // r2w8 only exists for b4/b8 (the b2_r2 kernel is already 8-blocks-resident at 60 regs).
14791            // ca/car2 need the alignment gate AND have no b8 twins; pfr2 has no b8 twin either —
14792            // unsupported (shape, mcols) combos fall back to pf/r2.
14793            // On rp buffers, forced legacy names map to their rp twins (layout law).
14794            let v = if bv == "r2w8" && mcols == 2 {
14795                "r2"
14796            } else if bv == "ca" && (!ca_ok || mcols == 8) {
14797                "pf"
14798            } else if bv == "car2" && (!ca_ok || mcols == 8) {
14799                "r2"
14800            } else if bv == "pfr2" && mcols == 8 {
14801                "r2"
14802            } else if (bv == "rpr2w8" || bv == "rpr2") && mcols == 2 {
14803                "rpr2"
14804            }
14805            // rpca* has no b8 twin (falls to rpr2w8/rpr2); needs the ca alignment gate.
14806            else if (bv == "rpca" || bv == "rpcar2") && (!ca_ok || mcols == 8) {
14807                if mcols == 8 { "rpr2w8" } else { "rpr2" }
14808            } else if bv == "rpcar2" && mcols == 2 {
14809                "rpca"
14810            }
14811            // rpsc/rpmsc/rpks* gate on smem-fit + alignment; fall to rpr2 outside it
14812            // (rpms has no smem and no alignment need — always valid on rp buffers).
14813            else if (bv == "rpsc" || bv == "rpmsc") && !sc_ok {
14814                "rpr2"
14815            } else if (bv == "rpks" || bv == "rpksc") && !ks_ok {
14816                "rpr2"
14817            } else {
14818                bv
14819            };
14820            if rp {
14821                match v {
14822                    "base" | "pf" | "ca" | "rp" => "rp",
14823                    "r2" | "pfr2" | "car2" | "rpr2" => "rpr2",
14824                    "r2w8" | "rpr2w8" => {
14825                        if mcols == 2 {
14826                            "rpr2"
14827                        } else {
14828                            "rpr2w8"
14829                        }
14830                    }
14831                    other => other, // rpca/rpcar2/rpsc/rpks/rpksc pass through (already rp-layout)
14832                }
14833            } else {
14834                v
14835            }
14836        } else if mcols == 8 {
14837            // b8 AUTO (2026-07-06 m-small latency arc, g7e DRAM-cold rp msweep m=5/6/8 all five
14838            // 27B shapes): rpsc — the rpr2w8 schedule with the warp's scale rows prestaged to
14839            // smem, leaving ONE global dependency (the quant stream) in the k-loop at zero reg
14840            // growth. BIT-identical to rpr2w8 and wins or ties EVERY b8 cell: ffn_gate m5
14841            // 50.7->46.9 m8 64.1->57.1 (-11%), qkv m8 34.6->33.0, ssm_out m8 29.7->28.8,
14842            // attn_gate m8 26.9->26.1, ffn_down m5 58.2->56.9. The faster split-grid twins are
14843            // OUT: rpksc (k-split, ffn_down m5 -21%) broke run-spec self-consistency (k-reduce
14844            // order shifts verify argmax at tie margins — verify must stay bit-identical to the
14845            // m=1 decode chain); rpmsc (m-split, bit-identical) measured NEGATIVE everywhere
14846            // (twin warp's duplicated weight stream: ffn_down m5 85.7 vs 56.9).
14847            if rp {
14848                if sc_ok { "rpsc" } else { "rpr2w8" }
14849            } else {
14850                "r2w8"
14851            }
14852        } else if mcols >= 4 {
14853            // r2 runs 7 resident blocks/SM (67 regs); its __launch_bounds__(128,8) twin `r2w8`
14854            // (64 regs) runs 8. grid = ceil(out_f/8) for both. rp twins land in the same
14855            // residency classes (rp 44 regs ~ pf-class occupancy, rpr2 67, rpr2w8 64).
14856            let blocks = (out_f + 7) / 8;
14857            let r7 = 7 * sms as usize;
14858            let r8 = 8 * sms as usize;
14859            let waves = blocks as f64 / r7 as f64;
14860            let filled = blocks >= 4 * sms as usize;
14861            // 2026-07-06 m-small latency arc: b4 keeps the wave rule (rpms/rpmsc measured
14862            // flat-to-negative at m=3/4 on every shape — the m-split twin duplicates the weight
14863            // stream; rpsc b4 also negative on r2-class picks, ffn_down m4 51.1 vs 46.5).
14864            if filled && blocks.div_ceil(r8) < blocks.div_ceil(r7) {
14865                // the extra residency drops the INTEGER wave count -> the straggler wave a
14866                // latency-bound kernel pays in full disappears (ffn_down 1.11 -> 0.98 waves:
14867                // 112.5 -> 81.6us, beats pf 90.1; qkv 2.23 -> 1.95: 58.1 -> 51.1).
14868                if rp { "rpr2w8" } else { "r2w8" }
14869            } else if waves >= 2.0 || (waves <= 1.0 && filled) {
14870                // tail amortized (>=2 waves) or single wave: unbounded r2 (no reg-squeeze tax —
14871                // gate/up 81.1 vs 83.9 bounded, attn_q 61.0 vs 63.4).
14872                if rp { "rpr2" } else { "r2" }
14873            } else {
14874                // fractional straggler-wave window with no crossing, or grid too small to fill
14875                // the SMs (tiny out_f<=1024 shapes want max row-parallelism): prefetch variant
14876                // (rp = the r1 split-plane twin — measured the attn_gate winner, 35.4 vs pf 36.4).
14877                if rp { "rp" } else { "pf" }
14878            }
14879        } else if in_f >= 6144 {
14880            // b2 deep-k (2026-07-06): every new twin measured flat-to-negative here (rpms 44.1
14881            // vs rpr2 40.8 ffn_down; rpsc 43.6; the winning rpks is banned on k-order) — rpr2
14882            // stays.
14883            if rp { "rpr2" } else { "r2" }
14884        } else if rp {
14885            // b2 shallow-k: qkv (out_f=10240, 0.97 waves at 7-resident) is the one measured cell
14886            // where the r2-schedule scale-prestage twin beats the r1 rp pick (24.7 vs 28.9us
14887            // -15%); the wider (ffn_gate 1.65 waves) and smaller (attn_gate 0.58) shapes LOSE
14888            // (41.8 vs 38.2 / 16.6 vs 14.6) — gate on the single-wave window.
14889            let waves = ((out_f + 7) / 8) as f64 / (7 * sms as usize) as f64;
14890            if sc_ok && waves >= 0.9 && waves <= 1.1 {
14891                "rpsc"
14892            } else {
14893                "rp"
14894            }
14895        } else {
14896            "base"
14897        };
14898        variant
14899    }
14900
14901    pub fn qmatvec_mmvq_batched(
14902        &self,
14903        bytes: &CudaSlice<u8>,
14904        aq: &CudaSlice<i8>,
14905        ad: &CudaSlice<f32>,
14906        m: usize,
14907        in_f: usize,
14908        out_f: usize,
14909        qtype: i32,
14910        row_bytes: usize,
14911        mcols: usize,
14912        scale: f32,
14913        rp: bool,
14914    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14915        const ROWS_PER_BLOCK: u32 = 4;
14916        // TUNE SEAM (H100 lane): MEMRA_BVAR forces the batched-variant pick for the whole
14917        // process — the auto heuristics were tuned on sm_120 (82 SMs / 858 GB/s) and the
14918        // sm_90a re-tune sweeps this seam empirically. Layout variants stay safe: an rp
14919        // weight keeps its rp-layout kernel family regardless of the override.
14920        let forced: Option<&'static str> = {
14921            static V: std::sync::OnceLock<Option<String>> = std::sync::OnceLock::new();
14922            V.get_or_init(|| std::env::var("MEMRA_BVAR").ok())
14923                .as_deref()
14924                .map(|s| Box::leak(s.to_string().into_boxed_str()) as &'static str)
14925        };
14926        let variant = match forced {
14927            Some(v) if !rp || v.contains("rp") => v,
14928            _ => self.batched_variant(m, in_f, out_f, qtype, row_bytes, mcols, rp),
14929        };
14930        let base_name = Self::batched_kernel_name(qtype, mcols).ok_or_else(|| {
14931            format!("qmatvec_mmvq_batched: no kernel for qtype {qtype} mcols {mcols}")
14932        })?;
14933        // b16 tier (t=9..16 verify): only base/_rp b16 kernels are compiled — the b2..b8
14934        // per-shape perf variants (r2/pf/...) do not apply at this width. rp is a LAYOUT,
14935        // not a perf variant: it must survive (base kernel on split-plane bytes = NaN).
14936        let variant = if mcols == 16 {
14937            if rp { "rp" } else { "base" }
14938        } else {
14939            variant
14940        };
14941        // EXACT-WIDTH b5/b6/b7 twins (lane/vt-fixes fix 1, 2026-08-03): the b8 kernels
14942        // allocate acc[WROWS][8] at ANY m, so T=5..7 verify paid the full 8-wide register
14943        // tax — the measured T=4->5 cliff. The same template at MCOLS=m runs the identical
14944        // per-(token,row) chain (columns c >= m never execute in either form) ->
14945        // BIT-IDENTICAL to the b8 launch. NVFP4 split-plane only (the sm_120 default trunk);
14946        // covers both b8 auto schedules (rpsc, rpr2w8). MEMRA_B567=0 rollback.
14947        static B567: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
14948        let b567 = *B567.get_or_init(|| std::env::var("MEMRA_B567").as_deref() != Ok("0"));
14949        if b567
14950            && qtype == QT_NVFP4
14951            && rp
14952            && mcols == 8
14953            && (5..=7).contains(&m)
14954            && matches!(variant, "rpsc" | "rpr2w8")
14955        {
14956            let f = self.func(&format!("qmatvec_nvfp4_mmvq_b{m}_{variant}"));
14957            let rows_per_block = ROWS_PER_BLOCK * 2; // r2-class schedules: 2 rows/warp
14958            let mut y = self.alloc_uninit::<f32>(m * out_f)?;
14959            let cfg = LaunchConfig {
14960                grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
14961                block_dim: (32, ROWS_PER_BLOCK, 1),
14962                shared_mem_bytes: 0,
14963            };
14964            let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
14965            let __s_b = self.gpu.stream();
14966            let mut b = __s_b.launch_builder(&f);
14967            b.arg(bytes)
14968                .arg(aq)
14969                .arg(ad)
14970                .arg(&mut y)
14971                .arg(&inf)
14972                .arg(&outf)
14973                .arg(&mi)
14974                .arg(&rb);
14975            unsafe {
14976                b.launch(cfg)?;
14977            }
14978            if scale != 1.0 {
14979                self.scale_inplace(&mut y, scale, m * out_f)?;
14980            }
14981            return Ok(y);
14982        }
14983        let (name, rows_per_block): (std::borrow::Cow<'static, str>, u32) = match variant {
14984            "base" => (base_name.into(), ROWS_PER_BLOCK),
14985            "pf" => (format!("{base_name}_pf").into(), ROWS_PER_BLOCK),
14986            "ca" => (format!("{base_name}_ca").into(), ROWS_PER_BLOCK),
14987            "rp" => (format!("{base_name}_rp").into(), ROWS_PER_BLOCK),
14988            "rpca" => (format!("{base_name}_rpca").into(), ROWS_PER_BLOCK), // 1 row/warp cp.async
14989            // split families: 2 warp-pairs x 2 rows = 4 rows/block (the k-range or column set
14990            // splits across the pair's two warps; grid.x doubles vs rpr2 at the same regs).
14991            "rpks" => (format!("{base_name}_rpks").into(), ROWS_PER_BLOCK),
14992            "rpksc" => (format!("{base_name}_rpksc").into(), ROWS_PER_BLOCK),
14993            "rpms" => (format!("{base_name}_rpms").into(), ROWS_PER_BLOCK),
14994            "rpmsc" => (format!("{base_name}_rpmsc").into(), ROWS_PER_BLOCK),
14995            "r2ms_rp" => (format!("{base_name}_r2ms_rp").into(), ROWS_PER_BLOCK),
14996            "r2sm_rp" => (format!("{base_name}_r2sm_rp").into(), ROWS_PER_BLOCK * 2),
14997            "r2la_rp" => (format!("{base_name}_r2la_rp").into(), ROWS_PER_BLOCK * 2),
14998            v => (format!("{base_name}_{v}").into(), ROWS_PER_BLOCK * 2), // r2-class: 2 rows/warp
14999        };
15000        debug_assert!(
15001            !rp || name.contains("_rp"),
15002            "rp weight dispatched to a GGUF-layout kernel"
15003        );
15004        let f = self.func(&name);
15005        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
15006        // r2sm_rp: [MCOLS][32 blk][8 int] activation slab + [MCOLS][32] f32 scales.
15007        let smem = if name.contains("_r2sm_rp") {
15008            (mcols * 32 * 9 * 4 + mcols * 32 * 4) as u32
15009        } else {
15010            0
15011        };
15012        let cfg = LaunchConfig {
15013            grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
15014            block_dim: (32, ROWS_PER_BLOCK, 1),
15015            shared_mem_bytes: smem,
15016        };
15017        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
15018        let __s_b = self.gpu.stream();
15019        let mut b = __s_b.launch_builder(&f);
15020        b.arg(bytes)
15021            .arg(aq)
15022            .arg(ad)
15023            .arg(&mut y)
15024            .arg(&inf)
15025            .arg(&outf)
15026            .arg(&mi)
15027            .arg(&rb);
15028        unsafe {
15029            b.launch(cfg)?;
15030        }
15031        if scale != 1.0 {
15032            self.scale_inplace(&mut y, scale, m * out_f)?;
15033        }
15034        Ok(y)
15035    }
15036
15037    /// BATCHED weight-tile-resident matvec from raw weight bytes (quantizes the f32 activation `x` to
15038    /// q8_1 internally; macro-scale NOT applied — caller compares bare, like qmatvec_*_fast). For the
15039    /// kernel_check bit-equivalence gate. `mcols` ∈ {2,4,8}. Works for Q8_0/Q4_K/Q5_K/Q6_K/NVFP4.
15040    pub fn qmatvec_batched_raw(
15041        &self,
15042        bytes: &CudaSlice<u8>,
15043        x: &CudaSlice<f32>,
15044        m: usize,
15045        in_f: usize,
15046        out_f: usize,
15047        qtype: i32,
15048        row_bytes: usize,
15049        mcols: usize,
15050        rp: bool,
15051    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15052        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15053        self.qmatvec_mmvq_batched(
15054            bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, mcols, 1.0, rp,
15055        )
15056    }
15057
15058    /// Back-compat NVFP4-only batched raw launcher (used by older gates). Delegates to the generic one.
15059    pub fn qmatvec_nvfp4_batched_raw(
15060        &self,
15061        bytes: &CudaSlice<u8>,
15062        x: &CudaSlice<f32>,
15063        m: usize,
15064        in_f: usize,
15065        out_f: usize,
15066        row_bytes: usize,
15067        mcols: usize,
15068        rp: bool,
15069    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15070        self.qmatvec_batched_raw(bytes, x, m, in_f, out_f, QT_NVFP4, row_bytes, mcols, rp)
15071    }
15072
15073    /// Stage-C FP4 gate (MEMRA_FP4): if `w` is an NVFP4 weight with in_f%64==0, run the native mxf4
15074    /// block-scale GEMM and apply the per-tensor macro-scale, returning Some(y). Else None (caller
15075    /// falls through to the int8 GEMM / dp4a). Strict opt-in over the proven int8 path; m>=16 only.
15076    fn try_fp4_gemm(
15077        &self,
15078        w: &crate::model::GpuTensor,
15079        x: &CudaSlice<f32>,
15080        m: usize,
15081        in_f: usize,
15082        out_f: usize,
15083    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
15084        use crate::model::GpuTensor;
15085        if cfg!(memra_portable_cuda) {
15086            return Ok(None);
15087        }
15088        if std::env::var("MEMRA_FP4").is_err() {
15089            return Ok(None);
15090        }
15091        // CUTLASS prefill branch (m>=128 + MEMRA_FP4_CUTLASS + a repacked CutlassWeight present): route
15092        // to the CUTLASS sm120 NVFP4 GEMM, folding the per-tensor macro-scale into the epilogue alpha
15093        // (1/scale) — no post-matmul scale_inplace. Decode (m<128) and the m∈[16,128) middle band keep
15094        // the hand-roll below: CUTLASS's 128-row M-tile wastes work under 128.
15095        // The hand-roll applies the per-tensor macro-scale as a POST-matmul MULTIPLY (scale_inplace(y,
15096        // scale)); CUTLASS's epilogue does D = alpha * (A@B^T), so alpha == scale reproduces it exactly
15097        // (NOT 1/scale — the plan sketch had this inverted; the kernel_check arm gates it). scale==1.0
15098        // for the common no-macro-scale case.
15099        #[cfg(memra_cutlass)]
15100        if m >= 128 && std::env::var("MEMRA_FP4_CUTLASS").is_ok() {
15101            if let GpuTensor::Quant {
15102                bytes,
15103                qtype,
15104                scale,
15105                row_bytes,
15106                cutlass,
15107                ..
15108            } = w
15109            {
15110                if *qtype == QT_NVFP4 && in_f % 64 == 0 {
15111                    if let Some(cw) = cutlass {
15112                        // Resident fast path: load-time-repacked B + swizzled SFB (no per-call repack).
15113                        let y = self.cutlass_fp4_gemm(
15114                            &cw.b_packed,
15115                            &cw.sfb_swizzled,
15116                            x,
15117                            *scale,
15118                            m,
15119                            out_f,
15120                            in_f,
15121                        )?;
15122                        return Ok(Some(y));
15123                    } else if std::env::var("MEMRA_FP4_CUTLASS_OTF").is_ok() {
15124                        // On-the-fly repack (MEMRA_FP4_CUTLASS_OTF): de-interleave + swizzle the B operand
15125                        // from raw bytes per prefill call. No resident doubling of the NVFP4 weight VRAM
15126                        // (the load-time repack ~doubles it) — needed for models that don't fit the
15127                        // resident path (e.g. the 27B on 24GB). Slower (per-call repack) but argmax-exact.
15128                        let (b_packed, sfb_sw) =
15129                            self.build_cutlass_weight(bytes, out_f, in_f, *row_bytes)?;
15130                        let y =
15131                            self.cutlass_fp4_gemm(&b_packed, &sfb_sw, x, *scale, m, out_f, in_f)?;
15132                        return Ok(Some(y));
15133                    }
15134                }
15135            }
15136        }
15137        if let GpuTensor::Quant {
15138            bytes,
15139            qtype,
15140            row_bytes,
15141            scale,
15142            rp,
15143            ..
15144        } = w
15145        {
15146            // A6: the hand-rolled W4A4 mxf4 GEMM reads 36B GGUF blocks — no rp port (MEMRA_FP4 is
15147            // an opt-in accuracy tradeoff); repacked tensors fall through to the int8 GEMM.
15148            if *qtype == QT_NVFP4 && in_f % 64 == 0 && !*rp {
15149                let y =
15150                    self.qmatvec_gemm_nvfp4_fp4(bytes, x, m, in_f, out_f, *row_bytes, *scale)?;
15151                return Ok(Some(y));
15152            }
15153        }
15154        Ok(None)
15155    }
15156
15157    /// rms_norm + fused fp16 twin (task #14): f32 output verbatim `rms_norm` + the fp16
15158    /// copy the f16-mirror GEMM group would otherwise produce with a standalone convert
15159    /// launch. BIT-IDENTICAL end-to-end (same reduction, same __float2half values).
15160    pub fn rms_norm_f16out(
15161        &self,
15162        x: &CudaSlice<f32>,
15163        w: &CudaSlice<f32>,
15164        dst: &mut CudaSlice<f32>,
15165        dst16: &mut CudaSlice<u8>,
15166        ncols: usize,
15167        nrows: usize,
15168        eps: f32,
15169    ) -> Result<(), Box<dyn std::error::Error>> {
15170        let f = self.func("rms_norm_f16out_f32");
15171        let cfg = LaunchConfig {
15172            grid_dim: (nrows as u32, 1, 1),
15173            block_dim: (rms_block(), 1, 1),
15174            shared_mem_bytes: 0,
15175        };
15176        let (nc, e) = (ncols as i32, eps);
15177        let __s_b = self.gpu.stream();
15178        let mut b = __s_b.launch_builder(&f);
15179        b.arg(x).arg(w).arg(dst).arg(dst16).arg(&nc).arg(&e);
15180        unsafe {
15181            b.launch(cfg)?;
15182        }
15183        Ok(())
15184    }
15185
15186    /// add+norm(+f16out) fusion for the prefill trunk (round 28; add_rms_norm precedent —
15187    /// bit-identical to add_f32 -> rms_norm_f16out). block_dim matches rms_norm_f16out's.
15188    #[allow(clippy::too_many_arguments)]
15189    pub fn add_rms_norm_f16out(
15190        &self,
15191        a: &CudaSlice<f32>,
15192        b: &CudaSlice<f32>,
15193        w: &CudaSlice<f32>,
15194        res: &mut CudaSlice<f32>,
15195        dst: &mut CudaSlice<f32>,
15196        dst16: &mut CudaSlice<u8>,
15197        ncols: usize,
15198        nrows: usize,
15199        eps: f32,
15200    ) -> Result<(), Box<dyn std::error::Error>> {
15201        let f = self.func("add_rms_norm_f16out_f32");
15202        let cfg = LaunchConfig {
15203            grid_dim: (nrows as u32, 1, 1),
15204            block_dim: (rms_block(), 1, 1),
15205            shared_mem_bytes: 0,
15206        };
15207        let (nc, e) = (ncols as i32, eps);
15208        let __s_lb = self.gpu.stream();
15209        let mut lb = __s_lb.launch_builder(&f);
15210        lb.arg(a)
15211            .arg(b)
15212            .arg(w)
15213            .arg(res)
15214            .arg(dst)
15215            .arg(dst16)
15216            .arg(&nc)
15217            .arg(&e);
15218        unsafe {
15219            lb.launch(cfg)?;
15220        }
15221        Ok(())
15222    }
15223
15224    /// matmul_group with a PRE-EMITTED fp16 activation (task #14: the producer norm fused
15225    /// the convert). Mirror-less members fall back to `matmul` on the f32 activation.
15226    pub fn matmul_group_xh(
15227        &self,
15228        ws: &[&crate::model::GpuTensor],
15229        x: &CudaSlice<f32>,
15230        xh: &CudaSlice<u8>,
15231        m: usize,
15232    ) -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
15233        let mut out = Vec::with_capacity(ws.len());
15234        let in_f = ws[0].in_features();
15235        for w in ws {
15236            if w.in_features() == in_f && m >= 16 && !self.verify_exact_on() {
15237                if let Some(y) = self.try_f16_gemm_pre(w, xh, m)? {
15238                    out.push(y);
15239                    continue;
15240                }
15241            }
15242            out.push(self.matmul(w, x, m)?);
15243        }
15244        Ok(out)
15245    }
15246
15247    /// task #14 pad-proofing: zero beta/g_log at rows >= len_d[0] (pads become identity
15248    /// GDN steps). Layouts [T, H].
15249    pub fn gdn_pad_mask(
15250        &self,
15251        beta: &mut CudaSlice<f32>,
15252        g_log: &mut CudaSlice<f32>,
15253        len_d: &CudaSlice<i32>,
15254        h: usize,
15255        t: usize,
15256    ) -> Result<(), Box<dyn std::error::Error>> {
15257        let f = self.func("gdn_pad_mask_f32");
15258        let cfg = LaunchConfig::for_num_elems((t * h) as u32);
15259        let (hi, ti) = (h as i32, t as i32);
15260        let __s_b = self.gpu.stream();
15261        let mut b = __s_b.launch_builder(&f);
15262        b.arg(beta).arg(g_log).arg(len_d).arg(&hi).arg(&ti);
15263        unsafe {
15264            b.launch(cfg)?;
15265        }
15266        Ok(())
15267    }
15268
15269    /// task #14 pad-proofing: dst[ncols] = src row (len_d[0]-1) — device-indexed last-row
15270    /// gather for the padded prime graph's h_seed/hlast.
15271    pub fn row_gather_dev(
15272        &self,
15273        src: &CudaSlice<f32>,
15274        dst: &mut CudaSlice<f32>,
15275        len_d: &CudaSlice<i32>,
15276        ncols: usize,
15277    ) -> Result<(), Box<dyn std::error::Error>> {
15278        let f = self.func("row_gather_dev_f32");
15279        let cfg = LaunchConfig::for_num_elems(ncols as u32);
15280        let nc = ncols as i32;
15281        let __s_b = self.gpu.stream();
15282        let mut b = __s_b.launch_builder(&f);
15283        b.arg(src).arg(dst).arg(len_d).arg(&nc);
15284        unsafe {
15285            b.launch(cfg)?;
15286        }
15287        Ok(())
15288    }
15289
15290    /// Grouped matmul: several weights consuming ONE activation (hybrid layers: the GDN
15291    /// 4-tuple wqkv/gate/beta/alpha, attention q/k/v, ffn gate/up). Semantics identical to
15292    /// calling `matmul` per weight; the f16-mirror arm converts the activation ONCE for the
15293    /// whole group instead of once per GEMM (the standalone converts were ~250 launches/prime
15294    /// of small-kernel gap fuel — nsys 2026-07-26). Any member without a mirror (or with a
15295    /// different in_f) falls back to its own `matmul` — behavior unchanged.
15296    pub fn matmul_group(
15297        &self,
15298        ws: &[&crate::model::GpuTensor],
15299        x: &CudaSlice<f32>,
15300        m: usize,
15301    ) -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
15302        use crate::model::GpuTensor;
15303        let mut out = Vec::with_capacity(ws.len());
15304        let any_mirror = ws
15305            .iter()
15306            .any(|w| matches!(w, GpuTensor::Quant { f16: Some(_), .. }));
15307        if m >= 16 && any_mirror && !self.verify_exact_on() {
15308            let in_f = ws[0].in_features();
15309            let xh = self.f16_act(x, m * in_f, in_f)?;
15310            for w in ws {
15311                if w.in_features() == in_f {
15312                    if let Some(y) = self.try_f16_gemm_pre(w, &xh, m)? {
15313                        out.push(y);
15314                        continue;
15315                    }
15316                }
15317                out.push(self.matmul(w, x, m)?);
15318            }
15319            return Ok(out);
15320        }
15321        for w in ws {
15322            out.push(self.matmul(w, x, m)?);
15323        }
15324        Ok(out)
15325    }
15326
15327    /// Cross-request grouped matmul (task #13): run ONE projection group over the
15328    /// CONCATENATION of several sequences' activations (m = sum of per-seq rows — the
15329    /// GEMM-batch win vLLM gets from continuous batching), then split each output back
15330    /// into per-seq buffers. Zero view plumbing: gather/scatter are stream-ordered D2D
15331    /// copies (~us at prime sizes). NUMERIC CONFIG NOTE: a GEMM at m=sum tiles K
15332    /// differently than per-seq GEMMs — argmax-gated like every prefill GEMM change.
15333    pub fn matmul_group_multi(
15334        &self,
15335        ws: &[&crate::model::GpuTensor],
15336        xs: &[&CudaSlice<f32>],
15337        ms: &[usize],
15338    ) -> Result<Vec<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
15339        assert_eq!(xs.len(), ms.len());
15340        let in_f = ws[0].in_features();
15341        let total: usize = ms.iter().sum();
15342        let mut xcat = self.uninit(total * in_f)?;
15343        let mut off = 0usize;
15344        for (x, &m) in xs.iter().zip(ms) {
15345            self.copy_into(&mut xcat, off * in_f, x, m * in_f)?;
15346            off += m;
15347        }
15348        let ys = self.matmul_group(ws, &xcat, total)?;
15349        let mut out: Vec<Vec<CudaSlice<f32>>> = (0..xs.len()).map(|_| Vec::new()).collect();
15350        for (w, y) in ws.iter().zip(ys) {
15351            let out_f = w.out_features();
15352            let mut off = 0usize;
15353            for (s, &m) in ms.iter().enumerate() {
15354                let mut ys_s = self.uninit(m * out_f)?;
15355                let src = y.slice(off * out_f..(off + m) * out_f);
15356                self.gpu.stream().memcpy_dtod(&src, &mut ys_s)?;
15357                out[s].push(ys_s);
15358                off += m;
15359            }
15360        }
15361        Ok(out)
15362    }
15363
15364    /// True if `w`'s qtype has a batched tensor-core GEMM kernel (the prefill T>1 root fix).
15365    /// Only the 4 daily-hot dtypes: Q8_0, Q4_K, Q6_K, NVFP4. NVFP4 needs in_f % 64 == 0.
15366    /// DEFAULT-ON (2026-06-28): measured pp512 9B-NVFP4 = 1413 tok/s WITH this GEMM vs 298 with the
15367    /// dp4a fallback (4.7x) AND MORE accurate (prefill logit maxdiff 0.159 vs dp4a 0.55, both argmax
15368    /// MATCH). The int8 tensor-core GEMM is unconditional (its historical MEMRA_GEMM opt-in gate
15369    /// shipped with Phase 0 — mma + smem swizzle + cp.async — and was removed). Prefill-only
15370    /// (m>=GEMM_M_THRESHOLD); m=1 decode keeps dp4a/MMVQ (this returns true but matmul only calls it
15371    /// at m>=threshold). Portable CUDA targets always use the correctness fallback; on sm_120a,
15372    /// MEMRA_NO_GEMM forces that same dp4a fallback (the bit-reference).
15373    pub fn gemm_supports(&self, w: &crate::model::GpuTensor) -> bool {
15374        use crate::model::GpuTensor;
15375        if !legacy_quant_gemm_allowed(
15376            cfg!(memra_portable_cuda),
15377            cfg!(memra_hopper_mma),
15378            std::env::var_os("MEMRA_NO_GEMM").is_some(),
15379        ) {
15380            return false;
15381        }
15382        match w {
15383            GpuTensor::Quant { qtype, .. } => {
15384                matches!(*qtype, QT_Q8_0 | QT_Q4_K | QT_Q6_K | QT_Q5_K | QT_Q4_0)
15385                    || (*qtype == QT_NVFP4 && w.in_features() % 64 == 0)
15386            }
15387            GpuTensor::Float { .. } | GpuTensor::FloatBf16 { .. } => false,
15388        }
15389    }
15390
15391    /// Batched tensor-core int8 GEMM with a PRE-QUANTIZED q8_1 activation (aq,ad). The prefill
15392    /// (T>1) root fix: decode each weight 32-block to int8 in shared memory ONCE per (row-tile,
15393    /// K-step) and reuse it across all BN tokens via mma.sync.m16n8k32.s8 — amortizing the weight
15394    /// read/decode N-fold (vs the dp4a matvec's per-token re-read). s32 accumulate is exact vs
15395    /// dp4a; only the final f32 block-scale rounding differs. Caller MUST have checked
15396    /// `gemm_supports(w)`. y[m,out] token-major. NVFP4 per-tensor macro-scale applied post.
15397    pub fn qmatvec_gemm(
15398        &self,
15399        w: &crate::model::GpuTensor,
15400        aq: &CudaSlice<i8>,
15401        ad: &CudaSlice<f32>,
15402        m: usize,
15403    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15404        use crate::model::GpuTensor;
15405        let in_f = w.in_features();
15406        let out_f = w.out_features();
15407        let (bytes, qtype, row_bytes, scale, rp) = match w {
15408            GpuTensor::Quant {
15409                bytes,
15410                qtype,
15411                row_bytes,
15412                scale,
15413                rp,
15414                ..
15415            } => (bytes, *qtype, *row_bytes, *scale, *rp),
15416            _ => unreachable!("gemm_supports guaranteed Quant"),
15417        };
15418        // wgmma arm (sm_90a, task 8): the m64n64k32 warpgroup kernel reads the rp4 split-plane
15419        // mirror AS-IS (qplane rows = its A operand, the half dplane its scales) and the same
15420        // (aq, ad) activation planes. Same numeric class as the mma kernel below (exact s32 per
15421        // 32-block, one f32 scale fold per block, ascending K) — argmax/tolerance gated like
15422        // every prefill GEMM, not bit-gated. MEMRA_WGMMA=0 restores the portable kernel.
15423        if cfg!(memra_hopper_mma) && qtype == QT_Q8_0 && out_f % 64 == 0 && wgmma_gemm_enabled() {
15424            if let GpuTensor::Quant { rp4: Some(m4), .. } = w {
15425                let mut y = self.qmatvec_gemm_q8_0_wgmma_raw(m4, aq, ad, m, in_f, out_f)?;
15426                if scale != 1.0 {
15427                    self.scale_inplace(&mut y, scale, m * out_f)?;
15428                }
15429                return Ok(y);
15430            }
15431        }
15432        let name = match qtype {
15433            QT_Q8_0 => "qmatvec_gemm_q8_0",
15434            QT_Q4_K => "qmatvec_gemm_q4_K",
15435            QT_Q4_0 => {
15436                if rp {
15437                    "qmatvec_gemm_q4_0_rp"
15438                } else {
15439                    "qmatvec_gemm_q4_0"
15440                }
15441            }
15442            QT_Q5_K => "qmatvec_gemm_q5_K",
15443            QT_Q6_K => "qmatvec_gemm_q6_K",
15444            QT_NVFP4 => {
15445                if rp {
15446                    "qmatvec_gemm_nvfp4_rp"
15447                } else {
15448                    "qmatvec_gemm_nvfp4"
15449                }
15450            }
15451            _ => unreachable!(),
15452        };
15453        let f = self.func(name);
15454        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output: skip memset
15455        // CTA tile MUST match the .cu per-kernel tile. MMQ-PORT: kernel1 (Q8_0/Q4_K/Q5_K) runs llama's
15456        // 128x128 SQUARE tile (K1_BM=128 x K1_BN=128, 8 warps); kernel2 (Q6_K/NVFP4) keeps 64x256, 4 warps
15457        // (the macro BM/BN in the .cu). Grid dims are selected by qtype so each launches its own tile.
15458        let is_k1 = matches!(qtype, QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q4_0);
15459        // TUNE SEAM: MEMRA_GEMM_K1_LAUNCH overrides kernel1's launch tile to match a -D-swept fatbin.
15460        let k1_tile = if is_k1 {
15461            k1_launch_override().unwrap_or((128, 128, 8))
15462        } else {
15463            (128, 128, 8)
15464        };
15465        let (bm, bn): (u32, u32) = if is_k1 {
15466            (k1_tile.0, k1_tile.1)
15467        } else {
15468            (64, 256)
15469        };
15470        let warps: u32 = if is_k1 {
15471            k1_tile.2
15472        } else {
15473            match qtype {
15474                QT_NVFP4 => 8,
15475                _ => 4,
15476            }
15477        };
15478        let cfg = LaunchConfig {
15479            grid_dim: ((out_f as u32 + bm - 1) / bm, (m as u32 + bn - 1) / bn, 1),
15480            block_dim: (32, warps, 1),
15481            shared_mem_bytes: 0,
15482        };
15483        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
15484        let __s_b = self.gpu.stream();
15485        let mut b = __s_b.launch_builder(&f);
15486        b.arg(bytes)
15487            .arg(aq)
15488            .arg(ad)
15489            .arg(&mut y)
15490            .arg(&inf)
15491            .arg(&outf)
15492            .arg(&mi)
15493            .arg(&rb);
15494        unsafe {
15495            b.launch(cfg)?;
15496        }
15497        if scale != 1.0 {
15498            self.scale_inplace(&mut y, scale, m * out_f)?;
15499        }
15500        Ok(y)
15501    }
15502
15503    /// Test entry: run the GEMM directly from raw weight bytes + qtype (no GpuTensor). Quantizes
15504    /// the f32 activation `x` to q8_1 internally then launches the tensor-core GEMM. NVFP4 per-tensor
15505    /// macro-scale is NOT applied here (caller passes it separately, like the dp4a path). Used by
15506    /// kernel_check for the bit-equivalence gate vs qmatvec_*_dp4a.
15507    pub fn qmatvec_gemm_raw(
15508        &self,
15509        bytes: &CudaSlice<u8>,
15510        x: &CudaSlice<f32>,
15511        m: usize,
15512        in_f: usize,
15513        out_f: usize,
15514        qtype: i32,
15515        row_bytes: usize,
15516    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15517        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15518        let name = match qtype {
15519            QT_Q8_0 => "qmatvec_gemm_q8_0",
15520            QT_Q4_K => "qmatvec_gemm_q4_K",
15521            QT_Q4_0 => "qmatvec_gemm_q4_0",
15522            QT_Q5_K => "qmatvec_gemm_q5_K",
15523            QT_Q6_K => "qmatvec_gemm_q6_K",
15524            QT_NVFP4 => "qmatvec_gemm_nvfp4",
15525            QT_NVFP4_RP => "qmatvec_gemm_nvfp4_rp",
15526            _ => panic!("qmatvec_gemm_raw: qtype {qtype} has no GEMM kernel"),
15527        };
15528        let f = self.func(name);
15529        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output: skip memset
15530        // MMQ-PORT: kernel1 (Q8_0/Q4_K/Q5_K) = llama 128x128 tile, 8 warps; kernel2 (Q6_K/NVFP4) = 64x256,
15531        // 4/8 warps. Grid tile per qtype (must match the .cu K1_BM/K1_BN vs BM/BN). KEEP IN SYNC w/ qmatvec_gemm.
15532        let is_k1 = matches!(qtype, QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q4_0);
15533        // TUNE SEAM: MEMRA_GEMM_K1_LAUNCH overrides kernel1's launch tile to match a -D-swept fatbin.
15534        let k1_tile = if is_k1 {
15535            k1_launch_override().unwrap_or((128, 128, 8))
15536        } else {
15537            (128, 128, 8)
15538        };
15539        let (bm, bn): (u32, u32) = if is_k1 {
15540            (k1_tile.0, k1_tile.1)
15541        } else {
15542            (64, 256)
15543        };
15544        let warps: u32 = if is_k1 {
15545            k1_tile.2
15546        } else {
15547            match qtype {
15548                QT_NVFP4 | QT_NVFP4_RP => 8,
15549                _ => 4,
15550            }
15551        };
15552        let cfg = LaunchConfig {
15553            grid_dim: ((out_f as u32 + bm - 1) / bm, (m as u32 + bn - 1) / bn, 1),
15554            block_dim: (32, warps, 1),
15555            shared_mem_bytes: 0,
15556        };
15557        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
15558        let __s_b = self.gpu.stream();
15559        let mut b = __s_b.launch_builder(&f);
15560        b.arg(bytes)
15561            .arg(&aq)
15562            .arg(&ad)
15563            .arg(&mut y)
15564            .arg(&inf)
15565            .arg(&outf)
15566            .arg(&mi)
15567            .arg(&rb);
15568        unsafe {
15569            b.launch(cfg)?;
15570        }
15571        Ok(y)
15572    }
15573
15574    /// H100 warpgroup GEMM raw entry (task 8): launch `qmatvec_gemm_q8_0_wgmma` on an rp4
15575    /// split-plane mirror + pre-quantized (aq, ad) activation planes. One warpgroup (128 thr)
15576    /// owns a 64x64 C tile; grid (out_f/64, ceil(m/64)). out_f % 64 == 0 REQUIRED (row loads
15577    /// and dplane scale reads are unguarded); the token edge is guarded in-kernel.
15578    /// Standalone harness verdict (tools/bench_q8_gemm_wgmma.cu, 4096x4096x512): rel 1.6e-05
15579    /// vs CPU ref, 179us vs the portable mma kernel's 688us (3.84x, unpipelined).
15580    pub fn qmatvec_gemm_q8_0_wgmma_raw(
15581        &self,
15582        rp4: &CudaSlice<u8>,
15583        aq: &CudaSlice<i8>,
15584        ad: &CudaSlice<f32>,
15585        m: usize,
15586        in_f: usize,
15587        out_f: usize,
15588    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15589        assert!(
15590            out_f % 64 == 0 && in_f % 32 == 0,
15591            "wgmma GEMM needs out_f%64==0, in_f%32==0"
15592        );
15593        let f = self.func("qmatvec_gemm_q8_0_wgmma");
15594        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output
15595        let cfg = LaunchConfig {
15596            grid_dim: ((out_f / 64) as u32, (m as u32).div_ceil(64), 1),
15597            block_dim: (128, 1, 1),
15598            shared_mem_bytes: 0,
15599        };
15600        let (inf, outf, mi) = (in_f as i32, out_f as i32, m as i32);
15601        let __s_b = self.gpu.stream();
15602        let mut b = __s_b.launch_builder(&f);
15603        b.arg(rp4)
15604            .arg(aq)
15605            .arg(ad)
15606            .arg(&mut y)
15607            .arg(&inf)
15608            .arg(&outf)
15609            .arg(&mi);
15610        unsafe {
15611            b.launch(cfg)?;
15612        }
15613        Ok(y)
15614    }
15615
15616    /// y[i] *= s. NVFP4 per-tensor macro-scale broadcast over the whole output.
15617    pub fn scale_inplace(
15618        &self,
15619        y: &mut CudaSlice<f32>,
15620        s: f32,
15621        n: usize,
15622    ) -> Result<(), Box<dyn std::error::Error>> {
15623        let f = self.func("scale_f32");
15624        let cfg = LaunchConfig::for_num_elems(n as u32);
15625        let (sf, ni) = (s, n as i32);
15626        let __s_b = self.gpu.stream();
15627        let mut b = __s_b.launch_builder(&f);
15628        b.arg(y).arg(&sf).arg(&ni);
15629        unsafe {
15630            b.launch(cfg)?;
15631        }
15632        Ok(())
15633    }
15634
15635    /// MEMRA_FULL_PREC dequant-on-use: expand a bf16-resident weight (`GpuTensor::FloatBf16`, raw
15636    /// bf16 bytes) to a transient f32 scratch of `n` elements, which then feeds the existing f32
15637    /// cuBLASLt GEMV. The scratch is freed when the caller drops it, so peak VRAM = resident bf16
15638    /// weights + ONE (largest) weight's f32 expansion + activations. SLOW IS FINE (research mode).
15639    pub fn bf16_to_f32(
15640        &self,
15641        data: &cudarc::driver::CudaView<'_, u8>,
15642        n: usize,
15643    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15644        let mut out = self.alloc_uninit::<f32>(n)?;
15645        let f = self.func("bf16_to_f32");
15646        let cfg = LaunchConfig::for_num_elems(n as u32);
15647        let ni = n as i32;
15648        let __s_b = self.gpu.stream();
15649        let mut b = __s_b.launch_builder(&f);
15650        b.arg(data).arg(&mut out).arg(&ni);
15651        unsafe {
15652            b.launch(cfg)?;
15653        }
15654        Ok(out)
15655    }
15656
15657    /// Chunked bf16 linear (MEMRA_FULL_PREC): y[m,out] = x @ W_bf16^T with the f32 dequant scratch
15658    /// bounded to CHUNK_ROWS rows (256MB at in_f=4096) instead of the whole weight — the 4GB
15659    /// lm_head expansion OOM'd the 24GB budget. Row-chunking partitions OUTPUT rows; each row's
15660    /// dot is computed by the identical kernel on identical bytes, so per-(token,row) results are
15661    /// bit-identical to the unchunked form. `exact` selects linear_decode_exact (per-column m=1
15662    /// calls, the spec-verify contract) vs plain linear.
15663    fn linear_bf16_chunked(
15664        &self,
15665        x: &CudaSlice<f32>,
15666        data: &CudaSlice<u8>,
15667        m: usize,
15668        in_f: usize,
15669        out_f: usize,
15670        exact: bool,
15671    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15672        const CHUNK_BYTES: usize = 256 << 20;
15673        let chunk_rows = (CHUNK_BYTES / (in_f * 4)).max(1).min(out_f);
15674        if chunk_rows >= out_f {
15675            let wf32 = self.bf16_to_f32(&data.slice(0..in_f * out_f * 2), in_f * out_f)?;
15676            return if exact {
15677                self.linear_decode_exact(x, &wf32, m, in_f, out_f)
15678            } else {
15679                self.linear(x, &wf32, m, in_f, out_f)
15680            };
15681        }
15682        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
15683        let mut r0 = 0usize;
15684        while r0 < out_f {
15685            let rows = chunk_rows.min(out_f - r0);
15686            let wslice = data.slice(r0 * in_f * 2..(r0 + rows) * in_f * 2);
15687            let wf32 = self.bf16_to_f32(&wslice, in_f * rows)?;
15688            let yc = if exact {
15689                self.linear_decode_exact(x, &wf32, m, in_f, rows)?
15690            } else {
15691                self.linear(x, &wf32, m, in_f, rows)?
15692            };
15693            // scatter [m, rows] into y[m, out_f] at column offset r0 (m is tiny in decode/verify)
15694            for mi in 0..m {
15695                let src = yc.slice(mi * rows..(mi + 1) * rows);
15696                let mut dst = y.slice_mut(mi * out_f + r0..mi * out_f + r0 + rows);
15697                self.gpu.stream().memcpy_dtod(&src, &mut dst)?;
15698            }
15699            r0 += rows;
15700        }
15701        Ok(y)
15702    }
15703
15704    /// On-device linear: y[m,out] = x[m,in] @ W[out,in]^T, weights row-major [out,in] (ggml).
15705    /// cuBLASLt col-major mapping (see memra_runtime::Gpu::linear_f32 for the derivation).
15706    /// DECODE-EXACT float linear: per-column m=1 cuBLASLt calls. cuBLASLt's reduction split is
15707    /// n-dependent (lt_ndep probe: m=1 vs m=2 col0 differs every bit), so spec-verify batches
15708    /// must not batch float matmuls the T=1 decode chain runs at m=1. Used by the small-t MoE
15709    /// router/shexp sites and matmul_decode_exact's Float arm.
15710    pub fn linear_decode_exact(
15711        &self,
15712        x: &CudaSlice<f32>,
15713        w: &CudaSlice<f32>,
15714        m_tokens: usize,
15715        in_f: usize,
15716        out_f: usize,
15717    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15718        if m_tokens == 1 {
15719            return self.linear(x, w, 1, in_f, out_f);
15720        }
15721        let xv = self.view(x, m_tokens * in_f);
15722        let mut y = self.alloc_uninit::<f32>(m_tokens * out_f)?;
15723        for t in 0..m_tokens {
15724            let row = xv.slice(t * in_f..(t + 1) * in_f);
15725            let mut xr = self.alloc_uninit::<f32>(in_f)?;
15726            self.copy_view_into(&mut xr, 0, &row, in_f)?;
15727            let yr = self.linear(&xr, w, 1, in_f, out_f)?;
15728            self.copy_into(&mut y, t * out_f, &yr, out_f)?;
15729        }
15730        Ok(y)
15731    }
15732
15733    pub fn linear(
15734        &self,
15735        x: &CudaSlice<f32>,
15736        w: &CudaSlice<f32>,
15737        m_tokens: usize,
15738        in_f: usize,
15739        out_f: usize,
15740    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15741        use cudarc::cublaslt::{Matmul, MatmulConfig};
15742        let mut c = self.alloc_uninit::<f32>(m_tokens * out_f)?; // cuBLASLt beta=0: C fully written
15743        let cfg = MatmulConfig {
15744            transa: true,
15745            transb: false,
15746            transc: false,
15747            m: out_f as u64,
15748            n: m_tokens as u64,
15749            k: in_f as u64,
15750            alpha: 1.0,
15751            lda: in_f as i64,
15752            ldb: in_f as i64,
15753            beta: 0.0,
15754            ldc: out_f as i64,
15755            stride_a: None,
15756            stride_b: None,
15757            stride_c: None,
15758            stride_bias: None,
15759            batch_size: None,
15760        };
15761        unsafe {
15762            self.gpu.blas.matmul(cfg, w, x, &mut c, None, None)?;
15763        }
15764        Ok(c)
15765    }
15766
15767    /// Naive SDPA. Q:[head_dim,n_head,T], K/V:[head_dim,n_head_kv,T_kv] -> O:[head_dim,n_head,T].
15768    pub fn sdpa_naive(
15769        &self,
15770        q: &CudaSlice<f32>,
15771        k: &CudaSlice<f32>,
15772        v: &CudaSlice<f32>,
15773        o: &mut CudaSlice<f32>,
15774        head_dim: usize,
15775        n_head: usize,
15776        n_head_kv: usize,
15777        t: usize,
15778        t_kv: usize,
15779        scale: f32,
15780        causal: bool,
15781    ) -> Result<(), Box<dyn std::error::Error>> {
15782        let f = self.func("sdpa_naive_f32");
15783        let cfg = LaunchConfig {
15784            grid_dim: (n_head as u32, t as u32, 1),
15785            block_dim: (128, 1, 1),
15786            shared_mem_bytes: (t_kv * 4) as u32,
15787        };
15788        let (hd, nh, nhkv, ti, tkvi, cz) = (
15789            head_dim as i32,
15790            n_head as i32,
15791            n_head_kv as i32,
15792            t as i32,
15793            t_kv as i32,
15794            causal as i32,
15795        );
15796        let __s_b = self.gpu.stream();
15797        let mut b = __s_b.launch_builder(&f);
15798        b.arg(q)
15799            .arg(k)
15800            .arg(v)
15801            .arg(o)
15802            .arg(&hd)
15803            .arg(&nh)
15804            .arg(&nhkv)
15805            .arg(&ti)
15806            .arg(&tkvi)
15807            .arg(&scale)
15808            .arg(&cz);
15809        unsafe {
15810            b.launch(cfg)?;
15811        }
15812        Ok(())
15813    }
15814
15815    /// Island twin (lane/gemma-vision): causal + sliding-window attention with
15816    /// bidirectional image islands. `span_id` labels each absolute kv position
15817    /// (-1 text, >=0 island id); same-island keys are visible unconditionally,
15818    /// reproducing the reference's non-causal image batch. window 0 = no window.
15819    #[allow(clippy::too_many_arguments)]
15820    pub fn sdpa_naive_island(
15821        &self,
15822        q: &CudaSlice<f32>,
15823        k: &CudaSlice<f32>,
15824        v: &CudaSlice<f32>,
15825        o: &mut CudaSlice<f32>,
15826        span_id: &CudaSlice<i32>,
15827        head_dim: usize,
15828        n_head: usize,
15829        n_head_kv: usize,
15830        t: usize,
15831        t_kv: usize,
15832        scale: f32,
15833        window: usize,
15834    ) -> Result<(), Box<dyn std::error::Error>> {
15835        let f = self.func("sdpa_naive_island_f32");
15836        let cfg = LaunchConfig {
15837            grid_dim: (n_head as u32, t as u32, 1),
15838            block_dim: (128, 1, 1),
15839            shared_mem_bytes: (t_kv * 4) as u32,
15840        };
15841        let (hd, nh, nhkv, ti, tkvi, wi) = (
15842            head_dim as i32,
15843            n_head as i32,
15844            n_head_kv as i32,
15845            t as i32,
15846            t_kv as i32,
15847            window as i32,
15848        );
15849        let __s_b = self.gpu.stream();
15850        let mut b = __s_b.launch_builder(&f);
15851        b.arg(q)
15852            .arg(k)
15853            .arg(v)
15854            .arg(o)
15855            .arg(span_id)
15856            .arg(&hd)
15857            .arg(&nh)
15858            .arg(&nhkv)
15859            .arg(&ti)
15860            .arg(&tkvi)
15861            .arg(&scale)
15862            .arg(&wi);
15863        unsafe {
15864            b.launch(cfg)?;
15865        }
15866        Ok(())
15867    }
15868
15869    /// Windowed sdpa_naive twin (gemma4 R6): masks keys older than q_pos-(window-1).
15870    #[allow(clippy::too_many_arguments)]
15871    pub fn sdpa_naive_w(
15872        &self,
15873        q: &CudaSlice<f32>,
15874        k: &CudaSlice<f32>,
15875        v: &CudaSlice<f32>,
15876        o: &mut CudaSlice<f32>,
15877        head_dim: usize,
15878        n_head: usize,
15879        n_head_kv: usize,
15880        t: usize,
15881        t_kv: usize,
15882        scale: f32,
15883        causal: bool,
15884        window: usize,
15885    ) -> Result<(), Box<dyn std::error::Error>> {
15886        let f = self.func("sdpa_naive_w_f32");
15887        let cfg = LaunchConfig {
15888            grid_dim: (n_head as u32, t as u32, 1),
15889            block_dim: (128, 1, 1),
15890            shared_mem_bytes: (t_kv * 4) as u32,
15891        };
15892        let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
15893            head_dim as i32,
15894            n_head as i32,
15895            n_head_kv as i32,
15896            t as i32,
15897            t_kv as i32,
15898            causal as i32,
15899            window as i32,
15900        );
15901        let __s_b = self.gpu.stream();
15902        let mut b = __s_b.launch_builder(&f);
15903        b.arg(q)
15904            .arg(k)
15905            .arg(v)
15906            .arg(o)
15907            .arg(&hd)
15908            .arg(&nh)
15909            .arg(&nhkv)
15910            .arg(&ti)
15911            .arg(&tkvi)
15912            .arg(&scale)
15913            .arg(&cz)
15914            .arg(&wi);
15915        unsafe {
15916            b.launch(cfg)?;
15917        }
15918        Ok(())
15919    }
15920
15921    /// Lo-clipped windowed sdpa_naive twin (lane/dflash2-longctx, DFLASH2-EVAL §10.6(c)).
15922    /// Same mask law as `sdpa_naive_w`, but keys below every query's window floor are never
15923    /// read: kv_lo = max(0, (t_kv - t) + 1 - window) — the oldest key visible to the OLDEST
15924    /// query row (q_pos = t_kv - t). Dynamic shared memory shrinks from t_kv*4 bytes (which
15925    /// blows the 48KB launch bound at ~12k rows — the B2 ctx crash) to (t_kv - kv_lo)*4 =
15926    /// (window - 1 + t)*4, and the key scan drops from O(t_kv) to O(window + t). Output is
15927    /// byte-identical to `sdpa_naive_w` (masked keys contribute exact zeros to same-order
15928    /// reductions; kernel_check `sdpa_naive_w_lo` pins it). window == 0 (no window) keeps
15929    /// kv_lo = 0 and is then shape-identical to the legacy kernel, including its bound.
15930    #[allow(clippy::too_many_arguments)]
15931    pub fn sdpa_naive_w_lo(
15932        &self,
15933        q: &CudaSlice<f32>,
15934        k: &CudaSlice<f32>,
15935        v: &CudaSlice<f32>,
15936        o: &mut CudaSlice<f32>,
15937        head_dim: usize,
15938        n_head: usize,
15939        n_head_kv: usize,
15940        t: usize,
15941        t_kv: usize,
15942        scale: f32,
15943        causal: bool,
15944        window: usize,
15945    ) -> Result<(), Box<dyn std::error::Error>> {
15946        let kv_lo = if window > 0 {
15947            (t_kv - t + 1).saturating_sub(window)
15948        } else {
15949            0
15950        };
15951        let smem = (t_kv - kv_lo) * 4;
15952        if smem > 48 * 1024 {
15953            return Err(format!(
15954                "sdpa_naive_w_lo: window {window} + T {t} rows need {smem} bytes of dynamic \
15955                 shared memory (> 48KB launch bound) — this kernel clips the OLD side only; \
15956                 a window this wide needs the multi-pass long-ctx kernel"
15957            )
15958            .into());
15959        }
15960        let f = self.func("sdpa_naive_w_lo_f32");
15961        let cfg = LaunchConfig {
15962            grid_dim: (n_head as u32, t as u32, 1),
15963            block_dim: (128, 1, 1),
15964            shared_mem_bytes: smem as u32,
15965        };
15966        let (hd, nh, nhkv, ti, tkvi, cz, wi, lo) = (
15967            head_dim as i32,
15968            n_head as i32,
15969            n_head_kv as i32,
15970            t as i32,
15971            t_kv as i32,
15972            causal as i32,
15973            window as i32,
15974            kv_lo as i32,
15975        );
15976        let __s_b = self.gpu.stream();
15977        let mut b = __s_b.launch_builder(&f);
15978        b.arg(q)
15979            .arg(k)
15980            .arg(v)
15981            .arg(o)
15982            .arg(&hd)
15983            .arg(&nh)
15984            .arg(&nhkv)
15985            .arg(&ti)
15986            .arg(&tkvi)
15987            .arg(&scale)
15988            .arg(&cz)
15989            .arg(&wi)
15990            .arg(&lo);
15991        unsafe {
15992            b.launch(cfg)?;
15993        }
15994        Ok(())
15995    }
15996
15997    /// SDPA where K/V are CudaViews into a resident KV cache (decode hot path, no host round-trip).
15998    pub fn sdpa_naive_view(
15999        &self,
16000        q: &CudaSlice<f32>,
16001        k: &cudarc::driver::CudaView<f32>,
16002        v: &cudarc::driver::CudaView<f32>,
16003        o: &mut CudaSlice<f32>,
16004        head_dim: usize,
16005        n_head: usize,
16006        n_head_kv: usize,
16007        t: usize,
16008        t_kv: usize,
16009        scale: f32,
16010        causal: bool,
16011    ) -> Result<(), Box<dyn std::error::Error>> {
16012        let f = self.func("sdpa_naive_f32");
16013        let cfg = LaunchConfig {
16014            grid_dim: (n_head as u32, t as u32, 1),
16015            block_dim: (128, 1, 1),
16016            shared_mem_bytes: (t_kv * 4) as u32,
16017        };
16018        let (hd, nh, nhkv, ti, tkvi, cz) = (
16019            head_dim as i32,
16020            n_head as i32,
16021            n_head_kv as i32,
16022            t as i32,
16023            t_kv as i32,
16024            causal as i32,
16025        );
16026        let __s_b = self.gpu.stream();
16027        let mut b = __s_b.launch_builder(&f);
16028        b.arg(q)
16029            .arg(k)
16030            .arg(v)
16031            .arg(o)
16032            .arg(&hd)
16033            .arg(&nh)
16034            .arg(&nhkv)
16035            .arg(&ti)
16036            .arg(&tkvi)
16037            .arg(&scale)
16038            .arg(&cz);
16039        unsafe {
16040            b.launch(cfg)?;
16041        }
16042        Ok(())
16043    }
16044
16045    /// Correctness fallback for quantized resident K/V views. Dequantizes K and V once into f32
16046    /// workspaces, then calls `sdpa_naive`. This is an explicit API: the optimized prefill view
16047    /// dispatch remains unchanged, so callers can use it as a reference or compatibility path.
16048    /// Dequant a quantized KV view into caller-owned f32 buffers (one grid-stride launch).
16049    /// `g` picks the kf8vf8-module stamp for e4m3 caches (same flag contract as fa_decode/
16050    /// fa_prefill_view). Used by the E4B shared-KV prefill arms (2026-07-31) to feed the
16051    /// f32 fa_prefill_w / fa_prefill_hd512 twins from the target layer's quantized rows.
16052    #[allow(clippy::too_many_arguments)]
16053    pub fn fa_dequant_kv_view_f32(
16054        &self,
16055        k: &cudarc::driver::CudaView<u8>,
16056        v: &cudarc::driver::CudaView<u8>,
16057        kf: &mut CudaSlice<f32>,
16058        vf: &mut CudaSlice<f32>,
16059        kv_dim_k: usize,
16060        kv_dim_v: usize,
16061        t_kv: usize,
16062        k_tok_bytes: usize,
16063        v_tok_bytes: usize,
16064        g: bool,
16065    ) -> Result<(), Box<dyn std::error::Error>> {
16066        let f = if g {
16067            self.func_g("fa_dequant_kv_ws_f32")
16068        } else {
16069            self.func("fa_dequant_kv_ws_f32")
16070        };
16071        let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
16072        let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
16073        let cfg = LaunchConfig {
16074            grid_dim: (nblk.max(1), 1, 1),
16075            block_dim: (256, 1, 1),
16076            shared_mem_bytes: 0,
16077        };
16078        let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
16079        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
16080        let __s_b = self.gpu.stream();
16081        let mut b = __s_b.launch_builder(&f);
16082        b.arg(k)
16083            .arg(v)
16084            .arg(&mut *kf)
16085            .arg(&mut *vf)
16086            .arg(&kdk)
16087            .arg(&kdv)
16088            .arg(&tkvi)
16089            .arg(&ktb)
16090            .arg(&vtb);
16091        unsafe {
16092            b.launch(cfg)?;
16093        }
16094        Ok(())
16095    }
16096
16097    #[allow(clippy::too_many_arguments)]
16098    pub fn sdpa_naive_quantized_view(
16099        &self,
16100        q: &CudaSlice<f32>,
16101        k: &cudarc::driver::CudaView<u8>,
16102        v: &cudarc::driver::CudaView<u8>,
16103        o: &mut CudaSlice<f32>,
16104        head_dim: usize,
16105        n_head: usize,
16106        n_head_kv: usize,
16107        t: usize,
16108        t_kv: usize,
16109        scale: f32,
16110        causal: bool,
16111        k_tok_bytes: usize,
16112        v_tok_bytes: usize,
16113    ) -> Result<(), Box<dyn std::error::Error>> {
16114        let kv_dim = n_head_kv * head_dim;
16115        let mut kf = self.uninit(t_kv * kv_dim)?;
16116        let mut vf = self.uninit(t_kv * kv_dim)?;
16117        let f = self.func("fa_dequant_kv_ws_f32");
16118        let total = (2 * t_kv * kv_dim) as u64;
16119        let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
16120        let cfg = LaunchConfig {
16121            grid_dim: (nblk.max(1), 1, 1),
16122            block_dim: (256, 1, 1),
16123            shared_mem_bytes: 0,
16124        };
16125        let (kv_dim_i, t_kv_i) = (kv_dim as i32, t_kv as i32);
16126        let (k_tok_bytes_i, v_tok_bytes_i) = (k_tok_bytes as i64, v_tok_bytes as i64);
16127        let __s_b = self.gpu.stream();
16128        let mut b = __s_b.launch_builder(&f);
16129        b.arg(k)
16130            .arg(v)
16131            .arg(&mut kf)
16132            .arg(&mut vf)
16133            .arg(&kv_dim_i)
16134            .arg(&kv_dim_i)
16135            .arg(&t_kv_i)
16136            .arg(&k_tok_bytes_i)
16137            .arg(&v_tok_bytes_i);
16138        unsafe { b.launch(cfg)? };
16139        self.sdpa_naive(
16140            q, &kf, &vf, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
16141        )
16142    }
16143
16144    /// WINDOWED twin of `sdpa_naive_quantized_view` (step35 SWA prefill): dequant the KV byte
16145    /// view into f32 workspaces with the SAME `fa_dequant_kv_ws_f32` launch, then run
16146    /// `sdpa_naive_w` instead of `sdpa_naive`. `window == 0` is the unwindowed form (the kernel
16147    /// treats a non-positive window as "no window mask"), so this is a strict superset of the
16148    /// unwindowed function above and produces bit-identical output at window == 0.
16149    ///
16150    /// Why this exists: EVERY windowed FlashAttention stamp in flash_attn.cu is head_dim-256
16151    /// only (`fa_prefill_w_f32` == `fa_prefill_f32_body<256>`, and the quantized-view windowed
16152    /// twins likewise), while step35 is head_dim 128. Its SWA layers therefore have no windowed
16153    /// FA path and take this f32 floor in v0 — same cache bytes, same numeric class as the
16154    /// unwindowed quantized-view fallback, so the chunk-invariance contract holds on both.
16155    #[allow(clippy::too_many_arguments)]
16156    pub fn sdpa_naive_w_quantized_view(
16157        &self,
16158        q: &CudaSlice<f32>,
16159        k: &cudarc::driver::CudaView<u8>,
16160        v: &cudarc::driver::CudaView<u8>,
16161        o: &mut CudaSlice<f32>,
16162        head_dim: usize,
16163        n_head: usize,
16164        n_head_kv: usize,
16165        t: usize,
16166        t_kv: usize,
16167        scale: f32,
16168        causal: bool,
16169        window: usize,
16170        k_tok_bytes: usize,
16171        v_tok_bytes: usize,
16172    ) -> Result<(), Box<dyn std::error::Error>> {
16173        let kv_dim = n_head_kv * head_dim;
16174        let mut kf = self.uninit(t_kv * kv_dim)?;
16175        let mut vf = self.uninit(t_kv * kv_dim)?;
16176        let f = self.func("fa_dequant_kv_ws_f32");
16177        let total = (2 * t_kv * kv_dim) as u64;
16178        let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
16179        let cfg = LaunchConfig {
16180            grid_dim: (nblk.max(1), 1, 1),
16181            block_dim: (256, 1, 1),
16182            shared_mem_bytes: 0,
16183        };
16184        let (kv_dim_i, t_kv_i) = (kv_dim as i32, t_kv as i32);
16185        let (k_tok_bytes_i, v_tok_bytes_i) = (k_tok_bytes as i64, v_tok_bytes as i64);
16186        let __s_b = self.gpu.stream();
16187        let mut b = __s_b.launch_builder(&f);
16188        b.arg(k)
16189            .arg(v)
16190            .arg(&mut kf)
16191            .arg(&mut vf)
16192            .arg(&kv_dim_i)
16193            .arg(&kv_dim_i)
16194            .arg(&t_kv_i)
16195            .arg(&k_tok_bytes_i)
16196            .arg(&v_tok_bytes_i);
16197        unsafe { b.launch(cfg)? };
16198        self.sdpa_naive_w(
16199            q, &kf, &vf, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal, window,
16200        )
16201    }
16202
16203    /// Hand-written FlashAttention prefill (sm_120, FA-2 online softmax on validated mma.sync,
16204    /// head_dim 256 or 128 (template-stamped twins), GQA, causal). Replaces sdpa_naive for T>1.
16205    /// Q/K/V/O [head_dim, n_head(_kv), T].
16206    pub fn fa_prefill(
16207        &self,
16208        q: &CudaSlice<f32>,
16209        k: &CudaSlice<f32>,
16210        v: &CudaSlice<f32>,
16211        o: &mut CudaSlice<f32>,
16212        head_dim: usize,
16213        n_head: usize,
16214        n_head_kv: usize,
16215        t: usize,
16216        t_kv: usize,
16217        scale: f32,
16218        causal: bool,
16219    ) -> Result<(), Box<dyn std::error::Error>> {
16220        if portable_mma_gated() {
16221            return self.sdpa_naive(
16222                q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
16223            );
16224        }
16225        // FA3 v10 arm (task #20, OPT-IN MEMRA_FA3=1 — harness-proven 883us vs the shipped
16226        // kernel's 993us at T=2048): TMA-swizzled wgmma FA, fresh causal hd256 only.
16227        // NEW NUMERIC CONFIG (GDN-mma precedent): online softmax / bf16-P class — the
16228        // run-gen argmax + greedy-stream batteries arbitrate; not bit-paired.
16229        // PROMOTED default-ON hopper (2026-07-27): 3-seed 2048-prime -> 128-decode
16230        // streams MATCH vs mma, full battery green, lane interleaved 5/5 (+2.4%).
16231        // MEMRA_FA3=0 reverts; kernel-check pins the mma config regardless.
16232        let fa3_on = head_dim == 256
16233            && causal
16234            && t == t_kv
16235            && match std::env::var("MEMRA_FA3").as_deref() {
16236                Ok("0") => false,
16237                Ok("1") => true,
16238                _ => cfg!(memra_hopper_mma),
16239            };
16240        if fa3_on {
16241            let n = t * n_head * head_dim;
16242            let nkv = t * n_head_kv * head_dim;
16243            let mut q16 = self.alloc_u8_uninit(n * 2)?;
16244            let mut k16 = self.alloc_u8_uninit(nkv * 2)?;
16245            let mut v16 = self.alloc_u8_uninit(nkv * 2)?;
16246            self.f32_to_bf16_into(q, &mut q16, n)?;
16247            self.f32_to_bf16_into(k, &mut k16, nkv)?;
16248            self.f32_to_bf16_into(v, &mut v16, nkv)?;
16249            let rc = {
16250                use cudarc::driver::{DevicePtr, DevicePtrMut};
16251                let stream = self.gpu.stream();
16252                let (qp, _g1) = q16.device_ptr(&stream);
16253                let (kp, _g2) = k16.device_ptr(&stream);
16254                let (vp, _g3) = v16.device_ptr(&stream);
16255                let (op, _g4) = o.device_ptr_mut(&stream);
16256                unsafe {
16257                    memra_fa3_prefill(
16258                        qp as *const core::ffi::c_void,
16259                        kp as *const core::ffi::c_void,
16260                        vp as *const core::ffi::c_void,
16261                        op as *mut f32,
16262                        t as i32,
16263                        n_head as i32,
16264                        n_head_kv as i32,
16265                        head_dim as i32,
16266                        scale,
16267                        stream.cu_stream() as *mut core::ffi::c_void,
16268                    )
16269                }
16270            };
16271            if rc != 0 {
16272                return Err(format!("memra_fa3_prefill rc={rc}").into());
16273            }
16274            return Ok(());
16275        }
16276        // FLOOR PORT (P2+P0a+P0b+P1): 4 warps/CTA, BLOCK_Q=64 query rows, BK=32 KV tile,
16277        // Q-in-reg + register-O, grid.y=n_head_kv (4 Q-heads share staged K/V).
16278        // P1 plain arm (MEMRA_FA_P1=1 opt-in until the qwen battery): the engine-study body
16279        // (FA2 schedule + boundary split + swizzle) on the non-windowed lane. bf16 pre-convert.
16280        static FA_P1: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16281        let fa_p1 = *FA_P1.get_or_init(|| std::env::var("MEMRA_FA_P1").as_deref() == Ok("1"));
16282        if fa_p1 && head_dim == 256 && !std::env::var("MEMRA_FA_FLOOR").is_ok() {
16283            const BLOCK_Q: usize = 64;
16284            const BKX: usize = 32;
16285            let f = self.func("fa_prefill_bf16_p1");
16286            let shmem = (2 * (2 * BKX * head_dim + BLOCK_Q * BKX)
16287                + 4 * (BLOCK_Q * BKX + 2 * BLOCK_Q)) as u32;
16288            use cudarc::driver::sys::CUfunction_attribute_enum as A;
16289            f.set_attribute(
16290                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16291                shmem as i32,
16292            )?;
16293            let cfg = LaunchConfig {
16294                grid_dim: (
16295                    (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
16296                    n_head as u32,
16297                    1,
16298                ),
16299                block_dim: (32, 4, 1),
16300                shared_mem_bytes: shmem,
16301            };
16302            let (hd, nh, nhkv, ti, tkvi, cz) = (
16303                head_dim as i32,
16304                n_head as i32,
16305                n_head_kv as i32,
16306                t as i32,
16307                t_kv as i32,
16308                causal as i32,
16309            );
16310            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
16311            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
16312            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
16313            let __s_b = self.gpu.stream();
16314            let mut b = __s_b.launch_builder(&f);
16315            b.arg(&qb)
16316                .arg(&kb)
16317                .arg(&vb)
16318                .arg(o)
16319                .arg(&hd)
16320                .arg(&nh)
16321                .arg(&nhkv)
16322                .arg(&ti)
16323                .arg(&tkvi)
16324                .arg(&scale)
16325                .arg(&cz);
16326            unsafe {
16327                b.launch(cfg)?;
16328            }
16329            return Ok(());
16330        }
16331        // Edge 5a (DEFAULT): fa_prefill_f32_pp — register-resident softmax (no sSw smem
16332        // round-trip), the FA3 softmax-GEMM overlap variant. ncu (pp512): short_scoreboard
16333        // 4.32->3.47, wait 1.99->1.45, per-call ~577us->~440us (1.31x) at flat 12.1% warps /
16334        // 255 regs / 2 CTAs (occupancy preserved). Bit-safe: 9B+27B argmax MATCH, rel 2.55e-3
16335        // vs floor 3.03e-3. MEMRA_FA_FLOOR reverts to the serialized-softmax floor kernel.
16336        const BK: usize = 32;
16337        // W2 lane (MEMRA_FA_PP_W2=1, ncu 2026-07-26): 2-warp/32-row CTA tile doubles grid.x —
16338        // bit-identical per-row math, pure coverage trade for the 6.25%-occupancy starvation.
16339        let w2 = std::env::var("MEMRA_FA_PP_W2").as_deref() == Ok("1");
16340        let (block_q, warps, w2_sfx): (usize, u32, &str) =
16341            if w2 { (32, 2, "_w2") } else { (64, 4, "") };
16342        // hd128 twins (2026-07-07): the prefill kernels are template-stamped at 256 (original
16343        // names, dispatch unchanged) and 128 (`_hd128`, the MiniMax-M3 class). Callers gate
16344        // other head_dims to sdpa_naive before reaching here.
16345        let hd_sfx = fa_hd_suffix(head_dim)?;
16346        let floor = std::env::var("MEMRA_FA_FLOOR").is_ok();
16347        // BF16-KV staging lane (2026-07-26, default ON): the kernel converts K/V to bf16
16348        // during staging anyway — pre-converting to bf16 mirrors is BIT-IDENTICAL (same
16349        // __float2bfloat16 values into the same mma) and turns the 67%-of-stalls scalar
16350        // staging into int4 vector copies. MEMRA_FA_BF16KV=0 reverts.
16351        let bf16kv = !floor && !w2 && std::env::var("MEMRA_FA_BF16KV").as_deref() != Ok("0");
16352        let (kb16, vb16) = if bf16kv {
16353            let n = t_kv * n_head_kv * head_dim;
16354            let mut kb = self.alloc_u8_uninit(n * 2)?;
16355            let mut vb = self.alloc_u8_uninit(n * 2)?;
16356            let fcv = self.func("f32_to_bf16_bulk");
16357            let ni = n as i64;
16358            let cfgc = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
16359            let __s_b = self.gpu.stream();
16360            let mut b = __s_b.launch_builder(&fcv);
16361            b.arg(k).arg(&mut kb).arg(&ni);
16362            unsafe {
16363                b.launch(cfgc)?;
16364            }
16365            let __s_b = self.gpu.stream();
16366            let mut b = __s_b.launch_builder(&fcv);
16367            b.arg(v).arg(&mut vb).arg(&ni);
16368            unsafe {
16369                b.launch(cfgc)?;
16370            }
16371            (Some(kb), Some(vb))
16372        } else {
16373            (None, None)
16374        };
16375        let f = self.func(&if bf16kv {
16376            format!("fa_prefill_bf16kv_pp{hd_sfx}")
16377        } else {
16378            format!(
16379                "fa_prefill_f32{}{}{hd_sfx}",
16380                if floor { "" } else { "_pp" },
16381                if floor { "" } else { w2_sfx }
16382            )
16383        });
16384        // persistent smem: bf16*(KV_STAGES*(sK + sV) + sP) + f32*(sS + sM + sL);
16385        // the bf16kv ring doubles the K/V stages (KV_STAGES=2).
16386        let kv_stages = if bf16kv { 2 } else { 1 };
16387        let shmem = (2 * (kv_stages * 2 * BK * head_dim + block_q * BK)
16388            + 4 * (block_q * BK + 2 * block_q)) as u32;
16389        use cudarc::driver::sys::CUfunction_attribute_enum as A;
16390        f.set_attribute(
16391            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16392            shmem as i32,
16393        )?;
16394        let cfg = LaunchConfig {
16395            grid_dim: (
16396                (t as u32 + block_q as u32 - 1) / block_q as u32,
16397                n_head as u32,
16398                1,
16399            ),
16400            block_dim: (32, warps, 1),
16401            shared_mem_bytes: shmem,
16402        };
16403        let (hd, nh, nhkv, ti, tkvi, cz) = (
16404            head_dim as i32,
16405            n_head as i32,
16406            n_head_kv as i32,
16407            t as i32,
16408            t_kv as i32,
16409            causal as i32,
16410        );
16411        let __s_b = self.gpu.stream();
16412        let mut b = __s_b.launch_builder(&f);
16413        b.arg(q);
16414        match (&kb16, &vb16) {
16415            (Some(kb), Some(vb)) => {
16416                b.arg(kb).arg(vb);
16417            }
16418            _ => {
16419                b.arg(k).arg(v);
16420            }
16421        }
16422        b.arg(o)
16423            .arg(&hd)
16424            .arg(&nh)
16425            .arg(&nhkv)
16426            .arg(&ti)
16427            .arg(&tkvi)
16428            .arg(&scale)
16429            .arg(&cz);
16430        unsafe {
16431            b.launch(cfg)?;
16432        }
16433        Ok(())
16434    }
16435
16436    /// Windowed FA prefill (gemma4 SWA layers past the sliding window, hd256): fa_prefill's
16437    /// exact dispatch (pp default, MEMRA_FA_FLOOR seam) with the sliding-window mask + tile
16438    /// skip in-kernel. Replaces the O(T*T_kv) scalar sdpa_naive_w on the prime path.
16439    #[allow(clippy::too_many_arguments)]
16440    pub fn fa_prefill_w(
16441        &self,
16442        q: &CudaSlice<f32>,
16443        k: &CudaSlice<f32>,
16444        v: &CudaSlice<f32>,
16445        o: &mut CudaSlice<f32>,
16446        head_dim: usize,
16447        n_head: usize,
16448        n_head_kv: usize,
16449        t: usize,
16450        t_kv: usize,
16451        scale: f32,
16452        causal: bool,
16453        window: usize,
16454    ) -> Result<(), Box<dyn std::error::Error>> {
16455        // sm_90a rides the mma twins (portable_mma_gated, 2026-07-31 — the raw
16456        // portable_cuda gate was stale-conservative on Hopper; fa_prefill already flipped).
16457        if portable_mma_gated() {
16458            return self.sdpa_naive_w(
16459                q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal, window,
16460            );
16461        }
16462        // Default: bf16-prestaged twin (same treatment as hd512 — Q/K/V pre-converted once,
16463        // int4 stage copies; bit-identical, kernel_check-gated). MEMRA_FAW_STAGE=f32 reverts;
16464        // MEMRA_FA_FLOOR keeps the f32 floor stamp untouched.
16465        static FAW_F32: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16466        let faw_f32 =
16467            *FAW_F32.get_or_init(|| std::env::var("MEMRA_FAW_STAGE").as_deref() == Ok("f32"));
16468        let floor = std::env::var("MEMRA_FA_FLOOR").is_ok();
16469        self.fa_prefill_w_arm(
16470            q,
16471            k,
16472            v,
16473            o,
16474            head_dim,
16475            n_head,
16476            n_head_kv,
16477            t,
16478            t_kv,
16479            scale,
16480            causal,
16481            window,
16482            floor || faw_f32,
16483            floor,
16484        )
16485    }
16486
16487    /// Windowed FA prefill with PRE-CONVERTED bf16 operands (producer-emitted; 31B glue lane).
16488    /// Launches the P1 stamp directly — callers guarantee qb/kb/vb hold the exact bf16 of q/k/v.
16489    #[allow(clippy::too_many_arguments)]
16490    pub fn fa_prefill_w_pre(
16491        &self,
16492        qb: &CudaSlice<u8>,
16493        kb: &CudaSlice<u8>,
16494        vb: &CudaSlice<u8>,
16495        o: &mut CudaSlice<f32>,
16496        head_dim: usize,
16497        n_head: usize,
16498        n_head_kv: usize,
16499        t: usize,
16500        t_kv: usize,
16501        scale: f32,
16502        causal: bool,
16503        window: usize,
16504        v_f16: bool,
16505    ) -> Result<(), Box<dyn std::error::Error>> {
16506        const BLOCK_Q: usize = 64;
16507        const BK: usize = 32;
16508        debug_assert_eq!(head_dim, 256);
16509        let hp = fa_f16pv_on() && faw_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
16510        debug_assert!(!v_f16 || hp, "f16 V emitted but the SWA hp arm is off");
16511        if hp {
16512            const BLOCK_QH: usize = 32;
16513            // V bytes must be f16 for the h2 stamp; producer normally emits f16 (v_f16),
16514            // else re-encode through the pooled scratch (stream-ordered reuse).
16515            let mut vguard = self.fa_vf16_scratch.lock().unwrap();
16516            let vh: &CudaSlice<u8> = if v_f16 {
16517                vb
16518            } else {
16519                let n = t_kv * n_head_kv * head_dim;
16520                if vguard.as_ref().map(|b| b.len() < n * 2).unwrap_or(true) {
16521                    *vguard = Some(self.alloc_uninit::<u8>(n * 2)?);
16522                }
16523                self.bf16_to_f16_into(vb, n, vguard.as_mut().unwrap())?;
16524                vguard.as_ref().unwrap()
16525            };
16526            let f = self.func("fa_prefill_w_bf16_p1h2");
16527            let shmem = (2 * (2 * BK * head_dim + 2 * BLOCK_QH * BK) + 4 * (2 * BLOCK_QH)) as u32;
16528            use cudarc::driver::sys::CUfunction_attribute_enum as A;
16529            f.set_attribute(
16530                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16531                shmem as i32,
16532            )?;
16533            let cfg = LaunchConfig {
16534                grid_dim: ((t as u32).div_ceil(BLOCK_QH as u32), (n_head / 2) as u32, 1),
16535                block_dim: (32, 4, 1),
16536                shared_mem_bytes: shmem,
16537            };
16538            let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
16539                head_dim as i32,
16540                n_head as i32,
16541                n_head_kv as i32,
16542                t as i32,
16543                t_kv as i32,
16544                causal as i32,
16545                window as i32,
16546            );
16547            let __s_b = self.gpu.stream();
16548            let mut b = __s_b.launch_builder(&f);
16549            b.arg(qb)
16550                .arg(kb)
16551                .arg(vh)
16552                .arg(o)
16553                .arg(&hd)
16554                .arg(&nh)
16555                .arg(&nhkv)
16556                .arg(&ti)
16557                .arg(&tkvi)
16558                .arg(&scale)
16559                .arg(&cz)
16560                .arg(&wi);
16561            unsafe {
16562                b.launch(cfg)?;
16563            }
16564            return Ok(());
16565        }
16566        let f = self.func("fa_prefill_w_bf16_p1");
16567        let shmem =
16568            (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
16569        use cudarc::driver::sys::CUfunction_attribute_enum as A;
16570        f.set_attribute(
16571            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16572            shmem as i32,
16573        )?;
16574        let cfg = LaunchConfig {
16575            grid_dim: (
16576                (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
16577                n_head as u32,
16578                1,
16579            ),
16580            block_dim: (32, 4, 1),
16581            shared_mem_bytes: shmem,
16582        };
16583        let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
16584            head_dim as i32,
16585            n_head as i32,
16586            n_head_kv as i32,
16587            t as i32,
16588            t_kv as i32,
16589            causal as i32,
16590            window as i32,
16591        );
16592        let __s_b = self.gpu.stream();
16593        let mut b = __s_b.launch_builder(&f);
16594        b.arg(qb)
16595            .arg(kb)
16596            .arg(vb)
16597            .arg(o)
16598            .arg(&hd)
16599            .arg(&nh)
16600            .arg(&nhkv)
16601            .arg(&ti)
16602            .arg(&tkvi)
16603            .arg(&scale)
16604            .arg(&cz)
16605            .arg(&wi);
16606        unsafe {
16607            b.launch(cfg)?;
16608        }
16609        Ok(())
16610    }
16611
16612    /// Windowed FA prefill with the stage arm FORCED — the kernel_check bit-identity entry.
16613    #[allow(clippy::too_many_arguments)]
16614    pub fn fa_prefill_w_arm(
16615        &self,
16616        q: &CudaSlice<f32>,
16617        k: &CudaSlice<f32>,
16618        v: &CudaSlice<f32>,
16619        o: &mut CudaSlice<f32>,
16620        head_dim: usize,
16621        n_head: usize,
16622        n_head_kv: usize,
16623        t: usize,
16624        t_kv: usize,
16625        scale: f32,
16626        causal: bool,
16627        window: usize,
16628        f32_stage: bool,
16629        floor: bool,
16630    ) -> Result<(), Box<dyn std::error::Error>> {
16631        const BLOCK_Q: usize = 64;
16632        const BK: usize = 32;
16633        debug_assert_eq!(head_dim, 256, "fa_prefill_w is stamped hd256 only");
16634        // P1 (2026-07-22 engine study): per-head Br=64 stamp with the FA2 schedule (V-copy
16635        // over GEMM0, next-K over softmax+GEMM1) + boundary/interior mask split. FP order
16636        // preserved -> bit-identical (gated). MEMRA_FAW_P1=0 reverts to the g4/o2 arms.
16637        static P1_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16638        let p1 = !floor
16639            && !f32_stage
16640            && *P1_ON.get_or_init(|| {
16641                std::env::var("MEMRA_FAW_P1")
16642                    .map(|v| v != "0")
16643                    .unwrap_or(true)
16644            });
16645        let hp =
16646            p1 && fa_f16pv_on() && faw_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
16647        if hp {
16648            const BLOCK_QH: usize = 32;
16649            let f = self.func("fa_prefill_w_bf16_p1h2");
16650            let shmem = (2 * (2 * BK * head_dim + 2 * BLOCK_QH * BK) + 4 * (2 * BLOCK_QH)) as u32;
16651            use cudarc::driver::sys::CUfunction_attribute_enum as A;
16652            f.set_attribute(
16653                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16654                shmem as i32,
16655            )?;
16656            let cfg = LaunchConfig {
16657                grid_dim: ((t as u32).div_ceil(BLOCK_QH as u32), (n_head / 2) as u32, 1),
16658                block_dim: (32, 4, 1),
16659                shared_mem_bytes: shmem,
16660            };
16661            let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
16662                head_dim as i32,
16663                n_head as i32,
16664                n_head_kv as i32,
16665                t as i32,
16666                t_kv as i32,
16667                causal as i32,
16668                window as i32,
16669            );
16670            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
16671            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
16672            let vh = self.f32_to_f16(v, t_kv * n_head_kv * head_dim)?;
16673            let __s_b = self.gpu.stream();
16674            let mut b = __s_b.launch_builder(&f);
16675            b.arg(&qb)
16676                .arg(&kb)
16677                .arg(&vh)
16678                .arg(o)
16679                .arg(&hd)
16680                .arg(&nh)
16681                .arg(&nhkv)
16682                .arg(&ti)
16683                .arg(&tkvi)
16684                .arg(&scale)
16685                .arg(&cz)
16686                .arg(&wi);
16687            unsafe {
16688                b.launch(cfg)?;
16689            }
16690            return Ok(());
16691        }
16692        if p1 {
16693            let f = self.func("fa_prefill_w_bf16_p1");
16694            let shmem =
16695                (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
16696            use cudarc::driver::sys::CUfunction_attribute_enum as A;
16697            f.set_attribute(
16698                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16699                shmem as i32,
16700            )?;
16701            let cfg = LaunchConfig {
16702                grid_dim: (
16703                    (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
16704                    n_head as u32,
16705                    1,
16706                ),
16707                block_dim: (32, 4, 1),
16708                shared_mem_bytes: shmem,
16709            };
16710            let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
16711                head_dim as i32,
16712                n_head as i32,
16713                n_head_kv as i32,
16714                t as i32,
16715                t_kv as i32,
16716                causal as i32,
16717                window as i32,
16718            );
16719            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
16720            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
16721            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
16722            let __s_b = self.gpu.stream();
16723            let mut b = __s_b.launch_builder(&f);
16724            b.arg(&qb)
16725                .arg(&kb)
16726                .arg(&vb)
16727                .arg(o)
16728                .arg(&hd)
16729                .arg(&nh)
16730                .arg(&nhkv)
16731                .arg(&ti)
16732                .arg(&tkvi)
16733                .arg(&scale)
16734                .arg(&cz)
16735                .arg(&wi);
16736            unsafe {
16737                b.launch(cfg)?;
16738            }
16739            return Ok(());
16740        }
16741        // MQA head-grouping (MEMRA_FAW_G4=0 reverts): 4 heads/CTA share the staged K/V —
16742        // per-(head,row) FP chain identical to the per-head stamp -> bit-identical (gated).
16743        static G4_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16744        let g4 = !floor
16745            && !f32_stage
16746            && n_head_kv == 1
16747            && n_head % 4 == 0
16748            && *G4_ON.get_or_init(|| {
16749                std::env::var("MEMRA_FAW_G4")
16750                    .map(|v| v != "0")
16751                    .unwrap_or(true)
16752            });
16753        if g4 {
16754            const SP_M: usize = 16;
16755            // Occupancy-2 twin (MEMRA_FAW_O2=0 reverts): one shared K/V buffer inside the dead
16756            // Q-stage region -> ~36.5KB smem, 2 CTA/SM (the llama hd256 mechanism). Bit-identical.
16757            static O2_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16758            let o2 = *O2_ON.get_or_init(|| {
16759                std::env::var("MEMRA_FAW_O2")
16760                    .map(|v| v != "0")
16761                    .unwrap_or(true)
16762            });
16763            let f = self.func(if o2 {
16764                "fa_prefill_w_bf16_g4o2"
16765            } else {
16766                "fa_prefill_w_bf16_g4"
16767            });
16768            let shmem = if o2 {
16769                (2 * (4 * SP_M * head_dim + 4 * SP_M * BK) + 4 * (4 * SP_M)) as u32
16770            } else {
16771                (2 * (2 * BK * head_dim + 4 * SP_M * head_dim + 4 * SP_M * BK) + 4 * (4 * SP_M))
16772                    as u32
16773            };
16774            use cudarc::driver::sys::CUfunction_attribute_enum as A;
16775            f.set_attribute(
16776                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16777                shmem as i32,
16778            )?;
16779            let cfg = LaunchConfig {
16780                grid_dim: ((t as u32).div_ceil(SP_M as u32), (n_head / 4) as u32, 1),
16781                block_dim: (32, 4, 1),
16782                shared_mem_bytes: shmem,
16783            };
16784            let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
16785                head_dim as i32,
16786                n_head as i32,
16787                n_head_kv as i32,
16788                t as i32,
16789                t_kv as i32,
16790                causal as i32,
16791                window as i32,
16792            );
16793            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
16794            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
16795            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
16796            let __s_b = self.gpu.stream();
16797            let mut b = __s_b.launch_builder(&f);
16798            b.arg(&qb)
16799                .arg(&kb)
16800                .arg(&vb)
16801                .arg(o)
16802                .arg(&hd)
16803                .arg(&nh)
16804                .arg(&nhkv)
16805                .arg(&ti)
16806                .arg(&tkvi)
16807                .arg(&scale)
16808                .arg(&cz)
16809                .arg(&wi);
16810            unsafe {
16811                b.launch(cfg)?;
16812            }
16813            return Ok(());
16814        }
16815        let f = self.func(if floor {
16816            "fa_prefill_w_f32"
16817        } else if f32_stage {
16818            "fa_prefill_w_f32_pp"
16819        } else {
16820            "fa_prefill_w_bf16_pp"
16821        });
16822        let shmem =
16823            (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
16824        use cudarc::driver::sys::CUfunction_attribute_enum as A;
16825        f.set_attribute(
16826            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16827            shmem as i32,
16828        )?;
16829        let cfg = LaunchConfig {
16830            grid_dim: (
16831                (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
16832                n_head as u32,
16833                1,
16834            ),
16835            block_dim: (32, 4, 1),
16836            shared_mem_bytes: shmem,
16837        };
16838        let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
16839            head_dim as i32,
16840            n_head as i32,
16841            n_head_kv as i32,
16842            t as i32,
16843            t_kv as i32,
16844            causal as i32,
16845            window as i32,
16846        );
16847        if f32_stage {
16848            let __s_b = self.gpu.stream();
16849            let mut b = __s_b.launch_builder(&f);
16850            b.arg(q)
16851                .arg(k)
16852                .arg(v)
16853                .arg(o)
16854                .arg(&hd)
16855                .arg(&nh)
16856                .arg(&nhkv)
16857                .arg(&ti)
16858                .arg(&tkvi)
16859                .arg(&scale)
16860                .arg(&cz)
16861                .arg(&wi);
16862            unsafe {
16863                b.launch(cfg)?;
16864            }
16865        } else {
16866            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
16867            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
16868            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
16869            let __s_b = self.gpu.stream();
16870            let mut b = __s_b.launch_builder(&f);
16871            b.arg(&qb)
16872                .arg(&kb)
16873                .arg(&vb)
16874                .arg(o)
16875                .arg(&hd)
16876                .arg(&nh)
16877                .arg(&nhkv)
16878                .arg(&ti)
16879                .arg(&tkvi)
16880                .arg(&scale)
16881                .arg(&cz)
16882                .arg(&wi);
16883            unsafe {
16884                b.launch(cfg)?;
16885            }
16886        }
16887        Ok(())
16888    }
16889
16890    /// hd512 FA prefill (gemma4 GLOBAL layers): BLOCK_Q=32 x 2 warps, Q staged in smem,
16891    /// grid.z = 2 O-halves (each CTA computes the full 512-dim scores, accumulates half the
16892    /// V dims). Replaces the scalar sdpa_naive on the prime path's globals.
16893    #[allow(clippy::too_many_arguments)]
16894    pub fn fa_prefill_hd512(
16895        &self,
16896        q: &CudaSlice<f32>,
16897        k: &CudaSlice<f32>,
16898        v: &CudaSlice<f32>,
16899        o: &mut CudaSlice<f32>,
16900        head_dim: usize,
16901        n_head: usize,
16902        n_head_kv: usize,
16903        t: usize,
16904        t_kv: usize,
16905        scale: f32,
16906        causal: bool,
16907    ) -> Result<(), Box<dyn std::error::Error>> {
16908        // sm_90a rides the mma twins (portable_mma_gated, 2026-07-31 — same flip as _w).
16909        if portable_mma_gated() {
16910            return self.sdpa_naive(
16911                q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
16912            );
16913        }
16914        // Default: pre-convert Q/K/V to bf16 once and stage int4 (8 bf16/copy) — at 1 CTA/SM the
16915        // synchronous stage serializes with compute and MQA re-stages the same K/V per head CTA;
16916        // pre-converting halves staged bytes and cuts stage instructions 8x. BIT-IDENTICAL to the
16917        // f32-staged kernel (the converter applies the same __float2bfloat16 the stage applied;
16918        // kernel_check gates the identity). MEMRA_FA512_STAGE=f32 = rollback to the f32 kernel.
16919        static F32_STAGE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16920        let f32_stage =
16921            *F32_STAGE.get_or_init(|| std::env::var("MEMRA_FA512_STAGE").as_deref() == Ok("f32"));
16922        // Single-pass arm (MEMRA_FA512_SP=0 reverts to the z=2 bf16 kernel): GEMM0 split-K across
16923        // the 2 warps instead of recomputed per O-half CTA — the 2026-07-22 kernel-diff excess.
16924        // Own numeric config (partial-sum order) — battery-gated.
16925        static SP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16926        let sp = !f32_stage
16927            && *SP_ON.get_or_init(|| {
16928                std::env::var("MEMRA_FA512_SP")
16929                    .map(|v| v != "0")
16930                    .unwrap_or(true)
16931            });
16932        self.fa_prefill_hd512_arm(
16933            q,
16934            k,
16935            v,
16936            o,
16937            head_dim,
16938            n_head,
16939            n_head_kv,
16940            t,
16941            t_kv,
16942            scale,
16943            causal,
16944            f32_stage,
16945            sp,
16946            sp && fa_f16pv_on(),
16947        )
16948    }
16949
16950    /// hd512 single-pass FA with PRE-CONVERTED bf16 operands (producer-emitted).
16951    #[allow(clippy::too_many_arguments)]
16952    pub fn fa_prefill_hd512_pre(
16953        &self,
16954        qb: &CudaSlice<u8>,
16955        kb: &CudaSlice<u8>,
16956        vb: &CudaSlice<u8>,
16957        o: &mut CudaSlice<f32>,
16958        head_dim: usize,
16959        n_head: usize,
16960        n_head_kv: usize,
16961        t: usize,
16962        t_kv: usize,
16963        scale: f32,
16964        causal: bool,
16965        v_f16: bool,
16966    ) -> Result<(), Box<dyn std::error::Error>> {
16967        debug_assert_eq!(head_dim, 512);
16968        const SP_M: usize = 16;
16969        const BKS: usize = 32;
16970        // f16-P/V door (MEMRA_FA_F16PV=1): P and the P@V accumulation in f16 (llama's fa=1 VKQ
16971        // class); KQ/softmax/rescale-band/final-normalize stay f32. Own numeric config —
16972        // battery-gated. V bytes must be f16 for the sp16 kernel (stage/ldmatrix are typeless).
16973        let f16pv = fa_f16pv_on();
16974        let nw = if f16pv { fa512_wide_warps() } else { 2 };
16975        let hp = f16pv && fa512_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
16976        debug_assert!(!v_f16 || f16pv, "f16 V emitted without the door on");
16977        let mut vguard = self.fa_vf16_scratch.lock().unwrap();
16978        let vref: &CudaSlice<u8> = if f16pv && !v_f16 {
16979            // Fallback re-encode (producer emitted bf16); the emit lane normally hands f16.
16980            let n = t_kv * n_head_kv * head_dim;
16981            let need = n * 2;
16982            if vguard.as_ref().map(|b| b.len() < need).unwrap_or(true) {
16983                *vguard = Some(self.alloc_uninit::<u8>(need)?);
16984            }
16985            let dst = vguard.as_mut().unwrap();
16986            self.bf16_to_f16_into(vb, n, dst)?;
16987            vguard.as_ref().unwrap()
16988        } else {
16989            vb
16990        };
16991        let f = self.func(if hp {
16992            "fa_prefill_bf16_hd512_sp16h2"
16993        } else {
16994            match (f16pv, nw) {
16995                (true, 4) => "fa_prefill_bf16_hd512_sp16w4",
16996                (true, _) => "fa_prefill_bf16_hd512_sp16",
16997                _ => "fa_prefill_bf16_hd512_sp",
16998            }
16999        });
17000        let (nwarp, npart) = if hp {
17001            (4usize, 4usize)
17002        } else if nw > 2 {
17003            (nw, nw)
17004        } else {
17005            (2, 1)
17006        };
17007        // h2 drops sQ (Q register-resident) and doubles sP/sS/sL for the head pair.
17008        let shmem = if hp {
17009            (2 * (2 * BKS * head_dim + 2 * SP_M * BKS) + 4 * (2 * npart * SP_M * BKS + 2 * SP_M))
17010                as u32
17011        } else {
17012            (2 * (SP_M * head_dim + 2 * BKS * head_dim + SP_M * BKS)
17013                + 4 * (npart * SP_M * BKS + SP_M)) as u32
17014        };
17015        use cudarc::driver::sys::CUfunction_attribute_enum as A;
17016        f.set_attribute(
17017            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
17018            shmem as i32,
17019        )?;
17020        let grid_y = if hp {
17021            (n_head / 2) as u32
17022        } else {
17023            n_head as u32
17024        };
17025        let cfg = LaunchConfig {
17026            grid_dim: ((t as u32).div_ceil(SP_M as u32), grid_y, 1),
17027            block_dim: (32, nwarp as u32, 1),
17028            shared_mem_bytes: shmem,
17029        };
17030        let (hd, nh, nhkv, ti, tkvi, cz) = (
17031            head_dim as i32,
17032            n_head as i32,
17033            n_head_kv as i32,
17034            t as i32,
17035            t_kv as i32,
17036            causal as i32,
17037        );
17038        let __s_b = self.gpu.stream();
17039        let mut b = __s_b.launch_builder(&f);
17040        b.arg(qb)
17041            .arg(kb)
17042            .arg(vref)
17043            .arg(o)
17044            .arg(&hd)
17045            .arg(&nh)
17046            .arg(&nhkv)
17047            .arg(&ti)
17048            .arg(&tkvi)
17049            .arg(&scale)
17050            .arg(&cz);
17051        unsafe {
17052            b.launch(cfg)?;
17053        }
17054        Ok(())
17055    }
17056
17057    /// hd512 FA prefill with the stage/sp arms FORCED — the kernel_check gate entry
17058    /// (`fa_prefill_hd512` picks the arms from MEMRA_FA512_STAGE / MEMRA_FA512_SP).
17059    #[allow(clippy::too_many_arguments)]
17060    pub fn fa_prefill_hd512_arm(
17061        &self,
17062        q: &CudaSlice<f32>,
17063        k: &CudaSlice<f32>,
17064        v: &CudaSlice<f32>,
17065        o: &mut CudaSlice<f32>,
17066        head_dim: usize,
17067        n_head: usize,
17068        n_head_kv: usize,
17069        t: usize,
17070        t_kv: usize,
17071        scale: f32,
17072        causal: bool,
17073        f32_stage: bool,
17074        sp: bool,
17075        f16pv: bool,
17076    ) -> Result<(), Box<dyn std::error::Error>> {
17077        debug_assert_eq!(head_dim, 512, "fa_prefill_hd512 is hd512 only");
17078        if sp && !f32_stage {
17079            // Single-pass: 16 q-rows/CTA, 2 warps, grid (ceil(T/16), n_head, 1).
17080            // smem: sQ[16][512] + sK[32][512] + sV[32][512] + sP[16][32] (bf16) + sS[16][32]+sL f32.
17081            // f16pv: sp16 kernel — f16 P + f16 P@V accum, V operand encoded f16.
17082            const SP_M: usize = 16;
17083            const BKS: usize = 32;
17084            let nw = if f16pv { fa512_wide_warps() } else { 2 };
17085            let hp = f16pv && fa512_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
17086            let f = self.func(if hp {
17087                "fa_prefill_bf16_hd512_sp16h2"
17088            } else {
17089                match (f16pv, nw) {
17090                    (true, 4) => "fa_prefill_bf16_hd512_sp16w4",
17091                    (true, _) => "fa_prefill_bf16_hd512_sp16",
17092                    _ => "fa_prefill_bf16_hd512_sp",
17093                }
17094            });
17095            let (nwarp, npart) = if hp {
17096                (4usize, 4usize)
17097            } else if nw > 2 {
17098                (nw, nw)
17099            } else {
17100                (2, 1)
17101            };
17102            let shmem = if hp {
17103                (2 * (2 * BKS * head_dim + 2 * SP_M * BKS)
17104                    + 4 * (2 * npart * SP_M * BKS + 2 * SP_M)) as u32
17105            } else {
17106                (2 * (SP_M * head_dim + 2 * BKS * head_dim + SP_M * BKS)
17107                    + 4 * (npart * SP_M * BKS + SP_M)) as u32
17108            };
17109            use cudarc::driver::sys::CUfunction_attribute_enum as A;
17110            f.set_attribute(
17111                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
17112                shmem as i32,
17113            )?;
17114            let grid_y = if hp {
17115                (n_head / 2) as u32
17116            } else {
17117                n_head as u32
17118            };
17119            let cfg = LaunchConfig {
17120                grid_dim: ((t as u32).div_ceil(SP_M as u32), grid_y, 1),
17121                block_dim: (32, nwarp as u32, 1),
17122                shared_mem_bytes: shmem,
17123            };
17124            let (hd, nh, nhkv, ti, tkvi, cz) = (
17125                head_dim as i32,
17126                n_head as i32,
17127                n_head_kv as i32,
17128                t as i32,
17129                t_kv as i32,
17130                causal as i32,
17131            );
17132            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
17133            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
17134            let vb = if f16pv {
17135                self.f32_to_f16(v, t_kv * n_head_kv * head_dim)?
17136            } else {
17137                self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?
17138            };
17139            let __s_b = self.gpu.stream();
17140            let mut b = __s_b.launch_builder(&f);
17141            b.arg(&qb)
17142                .arg(&kb)
17143                .arg(&vb)
17144                .arg(o)
17145                .arg(&hd)
17146                .arg(&nh)
17147                .arg(&nhkv)
17148                .arg(&ti)
17149                .arg(&tkvi)
17150                .arg(&scale)
17151                .arg(&cz);
17152            unsafe {
17153                b.launch(cfg)?;
17154            }
17155            return Ok(());
17156        }
17157        const BLOCK_Q: usize = 32;
17158        const BK: usize = 32;
17159        const HALF: usize = 256;
17160        let f = self.func(if f32_stage {
17161            "fa_prefill_f32_hd512"
17162        } else {
17163            "fa_prefill_bf16_hd512"
17164        });
17165        // sQ[32][512] + sK[BK][512] + sV[BK][256] + sP[32][BK] (bf16) + sL[32] f32
17166        let shmem = (2 * (BLOCK_Q * head_dim + BK * head_dim + BK * HALF + BLOCK_Q * BK)
17167            + 4 * BLOCK_Q) as u32;
17168        use cudarc::driver::sys::CUfunction_attribute_enum as A;
17169        f.set_attribute(
17170            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
17171            shmem as i32,
17172        )?;
17173        let cfg = LaunchConfig {
17174            grid_dim: (
17175                (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
17176                n_head as u32,
17177                2,
17178            ),
17179            block_dim: (32, 2, 1),
17180            shared_mem_bytes: shmem,
17181        };
17182        let (hd, nh, nhkv, ti, tkvi, cz) = (
17183            head_dim as i32,
17184            n_head as i32,
17185            n_head_kv as i32,
17186            t as i32,
17187            t_kv as i32,
17188            causal as i32,
17189        );
17190        if f32_stage {
17191            let __s_b = self.gpu.stream();
17192            let mut b = __s_b.launch_builder(&f);
17193            b.arg(q)
17194                .arg(k)
17195                .arg(v)
17196                .arg(o)
17197                .arg(&hd)
17198                .arg(&nh)
17199                .arg(&nhkv)
17200                .arg(&ti)
17201                .arg(&tkvi)
17202                .arg(&scale)
17203                .arg(&cz);
17204            unsafe {
17205                b.launch(cfg)?;
17206            }
17207        } else {
17208            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
17209            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
17210            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
17211            let __s_b = self.gpu.stream();
17212            let mut b = __s_b.launch_builder(&f);
17213            b.arg(&qb)
17214                .arg(&kb)
17215                .arg(&vb)
17216                .arg(o)
17217                .arg(&hd)
17218                .arg(&nh)
17219                .arg(&nhkv)
17220                .arg(&ti)
17221                .arg(&tkvi)
17222                .arg(&scale)
17223                .arg(&cz);
17224            unsafe {
17225                b.launch(cfg)?;
17226            }
17227        }
17228        Ok(())
17229    }
17230
17231    /// rope_neox2 with bf16 EMIT (31B glue lane): identical rope math/stores plus the post-rope
17232    /// values written as bf16 — the FA q/k operands come from this launch (bit-identical to the
17233    /// separate f32_to_bf16 the FA entries would run).
17234    #[allow(clippy::too_many_arguments)]
17235    pub fn rope_neox2_bf16e(
17236        &self,
17237        q: &mut CudaSlice<f32>,
17238        k: &mut CudaSlice<f32>,
17239        qb: &mut CudaSlice<u8>,
17240        kb: &mut CudaSlice<u8>,
17241        pos: &CudaSlice<i32>,
17242        head_dim: usize,
17243        n_dims: usize,
17244        nh_q: usize,
17245        nh_k: usize,
17246        n_tokens: usize,
17247        base: f32,
17248        freq_scale: f32,
17249        ff: Option<&CudaSlice<f32>>,
17250    ) -> Result<(), Box<dyn std::error::Error>> {
17251        let f = self.func("rope_neox2_bf16e_f32");
17252        let rows = ((nh_q + nh_k) * n_tokens) as u32;
17253        let cfg = LaunchConfig {
17254            grid_dim: (rows, 1, 1),
17255            block_dim: ((head_dim / 2) as u32, 1, 1),
17256            shared_mem_bytes: 0,
17257        };
17258        let theta_scale = base.powf(-2.0 / n_dims as f32);
17259        let (hd, nd, nhq, nhk, nt) = (
17260            head_dim as i32,
17261            n_dims as i32,
17262            nh_q as i32,
17263            nh_k as i32,
17264            n_tokens as i32,
17265        );
17266        let __s_b = self.gpu.stream();
17267        let mut b = __s_b.launch_builder(&f);
17268        match ff {
17269            Some(t) => {
17270                b.arg(&mut *q)
17271                    .arg(&mut *k)
17272                    .arg(&mut *qb)
17273                    .arg(&mut *kb)
17274                    .arg(pos)
17275                    .arg(&hd)
17276                    .arg(&nd)
17277                    .arg(&nhq)
17278                    .arg(&nhk)
17279                    .arg(&nt)
17280                    .arg(&theta_scale)
17281                    .arg(&freq_scale)
17282                    .arg(t);
17283                unsafe {
17284                    b.launch(cfg)?;
17285                }
17286            }
17287            None => {
17288                let null: u64 = 0;
17289                b.arg(&mut *q)
17290                    .arg(&mut *k)
17291                    .arg(&mut *qb)
17292                    .arg(&mut *kb)
17293                    .arg(pos)
17294                    .arg(&hd)
17295                    .arg(&nd)
17296                    .arg(&nhq)
17297                    .arg(&nhk)
17298                    .arg(&nt)
17299                    .arg(&theta_scale)
17300                    .arg(&freq_scale)
17301                    .arg(&null);
17302                unsafe {
17303                    b.launch(cfg)?;
17304                }
17305            }
17306        }
17307        Ok(())
17308    }
17309
17310    /// Flat f32 -> bf16 conversion into a fresh scratch buffer (2 bytes/elem). `n % 4 == 0`
17311    /// (float4 in, 4x bf16 out). Feeds the bf16-staged hd512 FA prefill.
17312    pub fn f32_to_bf16(
17313        &self,
17314        x: &CudaSlice<f32>,
17315        n: usize,
17316    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
17317        assert!(n % 4 == 0, "f32_to_bf16 requires n % 4 == 0, got {n}");
17318        let mut y = self.alloc_uninit::<u8>(n * 2)?;
17319        let f = self.func("f32_to_bf16_flat");
17320        let n_i = n as i64;
17321        let cfg = LaunchConfig {
17322            grid_dim: (((n / 4) as u32).div_ceil(256), 1, 1),
17323            block_dim: (256, 1, 1),
17324            shared_mem_bytes: 0,
17325        };
17326        let __s_b = self.gpu.stream();
17327        let mut b = __s_b.launch_builder(&f);
17328        b.arg(x).arg(&mut y).arg(&n_i);
17329        unsafe {
17330            b.launch(cfg)?;
17331        }
17332        Ok(y)
17333    }
17334
17335    pub fn f32_to_f16(
17336        &self,
17337        x: &CudaSlice<f32>,
17338        n: usize,
17339    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
17340        assert!(n % 4 == 0, "f32_to_f16 requires n % 4 == 0, got {n}");
17341        let mut y = self.alloc_uninit::<u8>(n * 2)?;
17342        let f = self.func("f32_to_f16_flat");
17343        let n_i = n as i64;
17344        let cfg = LaunchConfig {
17345            grid_dim: (((n / 4) as u32).div_ceil(256), 1, 1),
17346            block_dim: (256, 1, 1),
17347            shared_mem_bytes: 0,
17348        };
17349        let __s_b = self.gpu.stream();
17350        let mut b = __s_b.launch_builder(&f);
17351        b.arg(x).arg(&mut y).arg(&n_i);
17352        unsafe {
17353            b.launch(cfg)?;
17354        }
17355        Ok(y)
17356    }
17357
17358    /// bf16 bytes -> f16 bytes, n elements (the f16-P/V door's V re-encode on the emit lane).
17359    pub fn bf16_to_f16(
17360        &self,
17361        xb: &CudaSlice<u8>,
17362        n: usize,
17363    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
17364        let mut y = self.alloc_uninit::<u8>(n * 2)?;
17365        self.bf16_to_f16_into(xb, n, &mut y)?;
17366        Ok(y)
17367    }
17368
17369    /// Same conversion into a caller-owned (pooled) buffer; `y.len() >= n*2`.
17370    pub fn bf16_to_f16_into(
17371        &self,
17372        xb: &CudaSlice<u8>,
17373        n: usize,
17374        y: &mut CudaSlice<u8>,
17375    ) -> Result<(), Box<dyn std::error::Error>> {
17376        assert!(n % 2 == 0, "bf16_to_f16 requires n % 2 == 0, got {n}");
17377        assert!(y.len() >= n * 2);
17378        let f = self.func("bf16_to_f16_flat");
17379        let n2 = (n / 2) as i64;
17380        let cfg = LaunchConfig {
17381            grid_dim: (((n / 2) as u32).div_ceil(256), 1, 1),
17382            block_dim: (256, 1, 1),
17383            shared_mem_bytes: 0,
17384        };
17385        let __s_b = self.gpu.stream();
17386        let mut b = __s_b.launch_builder(&f);
17387        b.arg(xb).arg(y).arg(&n2);
17388        unsafe {
17389            b.launch(cfg)?;
17390        }
17391        Ok(())
17392    }
17393
17394    /// task #18 (attn side): varlen FA — bf16 K/V mirrors (2 launches) + ONE
17395    /// fa_prefill_bf16kv launch for every fresh sequence. Same per-block math as the
17396    /// per-seq path (bit-gateable). Caller guarantees: fresh causal (T_kv == T),
17397    /// head_dim in {256, 128}, bf16kv lane on.
17398    #[allow(clippy::too_many_arguments)]
17399    pub fn fa_prefill_vl8(
17400        &self,
17401        seqs: &[FaSeqVl],
17402        head_dim: usize,
17403        n_head: usize,
17404        n_head_kv: usize,
17405        scale: f32,
17406    ) -> Result<(), Box<dyn std::error::Error>> {
17407        const BK: usize = 32;
17408        let b = seqs.len();
17409        assert!(b >= 1 && b <= 8);
17410        let mut packed = [FaSeqVl::default(); 8];
17411        packed[..b].copy_from_slice(seqs);
17412        let v = FaVl8(packed);
17413        let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
17414        let ept = (n_head_kv * head_dim) as i32;
17415        {
17416            let f = self.func("fa_mirror_vl");
17417            let max_n = (max_t as i64) * ept as i64;
17418            let blocks = ((max_n as u32).div_ceil(4)).div_ceil(256);
17419            for which in 0..2i32 {
17420                let cfg = LaunchConfig {
17421                    grid_dim: (blocks, 1, b as u32),
17422                    block_dim: (256, 1, 1),
17423                    shared_mem_bytes: 0,
17424                };
17425                let __s_lb = self.gpu.stream();
17426                let mut lb = __s_lb.launch_builder(&f);
17427                lb.arg(&v).arg(&ept).arg(&which);
17428                unsafe {
17429                    lb.launch(cfg)?;
17430                }
17431            }
17432        }
17433        let hd_sfx = fa_hd_suffix(head_dim)?;
17434        let f = self.func(&format!("fa_prefill_bf16kv_vl{hd_sfx}"));
17435        let block_q = 64usize;
17436        let kv_stages = 2usize;
17437        let shmem = (2 * (kv_stages * 2 * BK * head_dim + block_q * BK)
17438            + 4 * (block_q * BK + 2 * block_q)) as u32;
17439        use cudarc::driver::sys::CUfunction_attribute_enum as A;
17440        f.set_attribute(
17441            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
17442            shmem as i32,
17443        )?;
17444        let cfg = LaunchConfig {
17445            grid_dim: (max_t.div_ceil(block_q as u32), n_head as u32, b as u32),
17446            block_dim: (32, 4, 1),
17447            shared_mem_bytes: shmem,
17448        };
17449        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
17450        let __s_lb = self.gpu.stream();
17451        let mut lb = __s_lb.launch_builder(&f);
17452        lb.arg(&v).arg(&hd).arg(&nh).arg(&nhkv).arg(&scale);
17453        unsafe {
17454            lb.launch(cfg)?;
17455        }
17456        Ok(())
17457    }
17458
17459    /// task #18 (attn pre-FA): varlen split + QK-norm + RoPE + KV-append — FOUR launches
17460    /// for every fresh sequence (was 6 x B, plus the q/k/v split copies which the view
17461    /// inputs remove entirely). Fresh-only (append at t0=0, RoPE pos = token index).
17462    #[allow(clippy::too_many_arguments)]
17463    pub fn attn_pre_vl8(
17464        &self,
17465        seqs: &[AttnPreVl],
17466        wq: &CudaSlice<f32>,
17467        wk: &CudaSlice<f32>,
17468        head_dim: usize,
17469        rope_dims: usize,
17470        n_head: usize,
17471        n_head_kv: usize,
17472        eps: f32,
17473        freq_base: f32,
17474        freq_scale: f32,
17475        kv_dim_k: usize,
17476        kv_dim_v: usize,
17477        k_tok_bytes: usize,
17478        v_tok_bytes: usize,
17479    ) -> Result<(), Box<dyn std::error::Error>> {
17480        let b = seqs.len();
17481        assert!(b >= 1 && b <= 8);
17482        let mut packed = [AttnPreVl::default(); 8];
17483        packed[..b].copy_from_slice(seqs);
17484        let v = AttnPreVl8(packed);
17485        let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
17486        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
17487        {
17488            let f = self.func("q_gate_split_vl");
17489            let n = max_t * (n_head * head_dim) as u32;
17490            let cfg = LaunchConfig {
17491                grid_dim: (n.div_ceil(256), 1, b as u32),
17492                block_dim: (256, 1, 1),
17493                shared_mem_bytes: 0,
17494            };
17495            let __s_lb = self.gpu.stream();
17496            let mut lb = __s_lb.launch_builder(&f);
17497            lb.arg(&v).arg(&hd).arg(&nh);
17498            unsafe {
17499                lb.launch(cfg)?;
17500            }
17501        }
17502        {
17503            let f = self.func("attn_rms_vl");
17504            let cfg = LaunchConfig {
17505                grid_dim: (max_t * n_head as u32, 2, b as u32),
17506                block_dim: (rms_block(), 1, 1),
17507                shared_mem_bytes: 0,
17508            };
17509            let __s_lb = self.gpu.stream();
17510            let mut lb = __s_lb.launch_builder(&f);
17511            lb.arg(&v)
17512                .arg(wq)
17513                .arg(wk)
17514                .arg(&hd)
17515                .arg(&nh)
17516                .arg(&nhkv)
17517                .arg(&eps);
17518            unsafe {
17519                lb.launch(cfg)?;
17520            }
17521        }
17522        {
17523            let f = self.func("attn_rope_vl");
17524            let theta_scale = freq_base.powf(-2.0 / rope_dims as f32);
17525            let nd = rope_dims as i32;
17526            let cfg = LaunchConfig {
17527                grid_dim: (max_t * n_head as u32, 2, b as u32),
17528                block_dim: ((head_dim / 2) as u32, 1, 1),
17529                shared_mem_bytes: 0,
17530            };
17531            let __s_lb = self.gpu.stream();
17532            let mut lb = __s_lb.launch_builder(&f);
17533            lb.arg(&v)
17534                .arg(&hd)
17535                .arg(&nd)
17536                .arg(&nh)
17537                .arg(&nhkv)
17538                .arg(&theta_scale)
17539                .arg(&freq_scale);
17540            unsafe {
17541                lb.launch(cfg)?;
17542            }
17543        }
17544        {
17545            let f = self.func("append_kv_vl");
17546            let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
17547            let cfg = LaunchConfig {
17548                grid_dim: (nblk, max_t, b as u32),
17549                block_dim: (32, 1, 1),
17550                shared_mem_bytes: 0,
17551            };
17552            let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
17553            let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
17554            let __s_lb = self.gpu.stream();
17555            let mut lb = __s_lb.launch_builder(&f);
17556            lb.arg(&v).arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
17557            unsafe {
17558                lb.launch(cfg)?;
17559            }
17560        }
17561        Ok(())
17562    }
17563
17564    /// FA prefill where K/V are QUANTIZED CudaViews into the resident byte KV cache (the T=K verify
17565    /// path, MTP-PLAN §D.3). Uses `fa_prefill_q` (inline-dequant during stage-to-smem). The view's
17566    /// base+offset pointer is honored; the kernel reads [0..t_kv*tok_bytes). Q is the T fresh query
17567    /// rows; t = T, t_kv = cache len. k_tok_bytes/v_tok_bytes are the per-token byte strides.
17568    pub fn fa_prefill_view(
17569        &self,
17570        q: &CudaSlice<f32>,
17571        k: &cudarc::driver::CudaView<u8>,
17572        v: &cudarc::driver::CudaView<u8>,
17573        o: &mut CudaSlice<f32>,
17574        head_dim: usize,
17575        n_head: usize,
17576        n_head_kv: usize,
17577        t: usize,
17578        t_kv: usize,
17579        scale: f32,
17580        causal: bool,
17581        k_tok_bytes: usize,
17582        v_tok_bytes: usize,
17583        g: bool,
17584    ) -> Result<(), Box<dyn std::error::Error>> {
17585        if portable_mma_gated() {
17586            return self.sdpa_naive_quantized_view(
17587                q,
17588                k,
17589                v,
17590                o,
17591                head_dim,
17592                n_head,
17593                n_head_kv,
17594                t,
17595                t_kv,
17596                scale,
17597                causal,
17598                k_tok_bytes,
17599                v_tok_bytes,
17600            );
17601        }
17602        const BLOCK_Q: usize = 64;
17603        const BK: usize = 32;
17604        // g = e4m3 cache: the kernel parses via DQ_K_ELEM/DQ_V_ELEM (format macros) — the
17605        // kf8vf8-module stamp reads fp8 with the identical MMA/softmax/PV body.
17606        let name = format!("fa_prefill_q{}", fa_hd_suffix(head_dim)?);
17607        let f = if g {
17608            self.func_g(&name)
17609        } else {
17610            self.func(&name)
17611        };
17612        let shmem =
17613            (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
17614        use cudarc::driver::sys::CUfunction_attribute_enum as A;
17615        f.set_attribute(
17616            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
17617            shmem as i32,
17618        )?;
17619        let cfg = LaunchConfig {
17620            grid_dim: (
17621                (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
17622                n_head as u32,
17623                1,
17624            ),
17625            block_dim: (32, 4, 1),
17626            shared_mem_bytes: shmem,
17627        };
17628        let (hd, nh, nhkv, ti, tkvi, cz) = (
17629            head_dim as i32,
17630            n_head as i32,
17631            n_head_kv as i32,
17632            t as i32,
17633            t_kv as i32,
17634            causal as i32,
17635        );
17636        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
17637        let __s_b = self.gpu.stream();
17638        let mut b = __s_b.launch_builder(&f);
17639        b.arg(q)
17640            .arg(k)
17641            .arg(v)
17642            .arg(o)
17643            .arg(&hd)
17644            .arg(&nh)
17645            .arg(&nhkv)
17646            .arg(&ti)
17647            .arg(&tkvi)
17648            .arg(&scale)
17649            .arg(&cz)
17650            .arg(&ktb)
17651            .arg(&vtb);
17652        unsafe {
17653            b.launch(cfg)?;
17654        }
17655        Ok(())
17656    }
17657
17658    /// ARC B (2026-07-05): dequant-once chunk-prime FA. Same contract as `fa_prefill_view`, but
17659    /// instead of every (q-block, head) CTA re-dequanting the whole quantized KV stream inline
17660    /// (T/64 x n_head redundant at chunk prime — 30.5% of the 32k prime wall), dequant the full
17661    /// [t_kv, kv_dim] K and V ONCE into a resident bf16 workspace (fa_dequant_kv_ws_bf16), then
17662    /// run `fa_prefill_qw` (the bf16-workspace twin) over it. EXACT: the workspace holds the same
17663    /// __float2bfloat16(dq_*_elem(...)) values fa_prefill_q stages to smem, and the twin's MMA/
17664    /// softmax/PV code is byte-identical -> bit-identical O (kernel_check pins bitdiff=0).
17665    /// The workspace allocation is REUSED across layers/chunks (grown to the largest shape);
17666    /// contents are rewritten per call. MEMRA_PRIME_DEQW=0 falls back to fa_prefill_view (callers gate).
17667    #[allow(clippy::too_many_arguments)]
17668    pub fn fa_prefill_view_ws(
17669        &self,
17670        q: &CudaSlice<f32>,
17671        k: &cudarc::driver::CudaView<u8>,
17672        v: &cudarc::driver::CudaView<u8>,
17673        o: &mut CudaSlice<f32>,
17674        head_dim: usize,
17675        n_head: usize,
17676        n_head_kv: usize,
17677        t: usize,
17678        t_kv: usize,
17679        scale: f32,
17680        causal: bool,
17681        k_tok_bytes: usize,
17682        v_tok_bytes: usize,
17683        g: bool,
17684    ) -> Result<(), Box<dyn std::error::Error>> {
17685        if portable_mma_gated() {
17686            return self.sdpa_naive_quantized_view(
17687                q,
17688                k,
17689                v,
17690                o,
17691                head_dim,
17692                n_head,
17693                n_head_kv,
17694                t,
17695                t_kv,
17696                scale,
17697                causal,
17698                k_tok_bytes,
17699                v_tok_bytes,
17700            );
17701        }
17702        const BLOCK_Q: usize = 64;
17703        const BK: usize = 32;
17704        let kv_dim_k = n_head_kv * head_dim;
17705        let kv_dim_v = n_head_kv * head_dim;
17706        let k_ws_bytes = t_kv * kv_dim_k * 2; // bf16
17707        let v_ws_bytes = t_kv * kv_dim_v * 2;
17708        // Lock held across BOTH launches: enqueue-only (µs), all compute serializes on gpu.stream.
17709        let mut guard = self.prime_deqw_ws.lock().unwrap();
17710        let need_grow = match guard.as_ref() {
17711            Some((kw, vw)) => kw.len() < k_ws_bytes || vw.len() < v_ws_bytes,
17712            None => true,
17713        };
17714        if need_grow {
17715            let grow = |cur: usize, need: usize| if cur >= need { cur } else { need };
17716            let (ck, cv) = guard
17717                .as_ref()
17718                .map(|(a, b)| (a.len(), b.len()))
17719                .unwrap_or((0, 0));
17720            *guard = Some((
17721                self.alloc_u8(grow(ck, k_ws_bytes))?,
17722                self.alloc_u8(grow(cv, v_ws_bytes))?,
17723            ));
17724        }
17725        let (kw, vw) = guard.as_mut().unwrap();
17726        // pass 1: dequant K+V once into the bf16 workspace (grid-stride, 1 thread/elem)
17727        {
17728            // only THIS pass parses KV bytes — pass 2 reads the bf16 workspace (format-free).
17729            let f = if g {
17730                self.func_g("fa_dequant_kv_ws_bf16")
17731            } else {
17732                self.func("fa_dequant_kv_ws_bf16")
17733            };
17734            let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
17735            let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
17736            let cfg = LaunchConfig {
17737                grid_dim: (nblk.max(1), 1, 1),
17738                block_dim: (256, 1, 1),
17739                shared_mem_bytes: 0,
17740            };
17741            let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
17742            let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
17743            let __s_b = self.gpu.stream();
17744            let mut b = __s_b.launch_builder(&f);
17745            b.arg(k)
17746                .arg(v)
17747                .arg(&mut *kw)
17748                .arg(&mut *vw)
17749                .arg(&kdk)
17750                .arg(&kdv)
17751                .arg(&tkvi)
17752                .arg(&ktb)
17753                .arg(&vtb);
17754            unsafe {
17755                b.launch(cfg)?;
17756            }
17757        }
17758        // pass 2: the bf16-workspace prefill twin (same tile sizes/loop structure as fa_prefill_q).
17759        // DEFAULT: cp.async double-buffered staging twin (fa_prefill_qw_db, +32KB smem for the
17760        // second K/V tile pair, 1 CTA/SM): overlaps tile n+1's L2->smem copy with tile n's MMA.
17761        // Bit-identical output (staging is a pure byte copy; kernel_check pins bitdiff=0 under
17762        // both twins). A/B (27B g7e, N=3): 32k prime 17.10->16.51s, 16k 9.09->8.65s — the copy
17763        // latency hides behind the MMA pipe and beats the 2-CTA/SM occupancy of the sync twin.
17764        // MEMRA_PRIME_DEQW_DB=0 falls back to the single-buffer twin.
17765        let db = std::env::var("MEMRA_PRIME_DEQW_DB")
17766            .map(|v| v != "0")
17767            .unwrap_or(true);
17768        {
17769            let hd_sfx = fa_hd_suffix(head_dim)?;
17770            let f = self.func(&format!(
17771                "fa_prefill_qw{}{hd_sfx}",
17772                if db { "_db" } else { "" }
17773            ));
17774            let shmem = if db {
17775                // 4x KV tile buffers (bf16) + sP (bf16) + sL (f32)
17776                (2 * (4 * BK * head_dim + BLOCK_Q * BK) + 4 * BLOCK_Q) as u32
17777            } else {
17778                (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32
17779            };
17780            use cudarc::driver::sys::CUfunction_attribute_enum as A;
17781            f.set_attribute(
17782                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
17783                shmem as i32,
17784            )?;
17785            let cfg = LaunchConfig {
17786                grid_dim: (
17787                    (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
17788                    n_head as u32,
17789                    1,
17790                ),
17791                block_dim: (32, 4, 1),
17792                shared_mem_bytes: shmem,
17793            };
17794            let (hd, nh, nhkv, ti, tkvi, cz) = (
17795                head_dim as i32,
17796                n_head as i32,
17797                n_head_kv as i32,
17798                t as i32,
17799                t_kv as i32,
17800                causal as i32,
17801            );
17802            let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
17803            let __s_b = self.gpu.stream();
17804            let mut b = __s_b.launch_builder(&f);
17805            b.arg(q)
17806                .arg(&*kw)
17807                .arg(&*vw)
17808                .arg(o)
17809                .arg(&hd)
17810                .arg(&nh)
17811                .arg(&nhkv)
17812                .arg(&ti)
17813                .arg(&tkvi)
17814                .arg(&scale)
17815                .arg(&cz)
17816                .arg(&kdk)
17817                .arg(&kdv);
17818            unsafe {
17819                b.launch(cfg)?;
17820            }
17821        }
17822        Ok(())
17823    }
17824
17825    /// WINDOWED `fa_prefill_view_ws` twin at head_dim 128 (lane/pp-prefill 2026-08-07):
17826    /// step35's SWA prefill (win=512, 33 of 45 layers) previously had NO windowed FA prefill
17827    /// stamp — every windowed twin was hd256-only — and took `sdpa_naive_w_quantized_view`,
17828    /// the f32 floor, at 565 ms/layer on a pp4096 where the hd128 FA family does the harder
17829    /// causal-4096 in 3.3 ms (41% of the whole prime; research/pp-prefill-20260807 anatomy).
17830    /// Same two-pass shape as the unwindowed function: dequant K/V ONCE into the resident
17831    /// bf16 workspace, then the windowed qw kernel (`fa_prefill_qw_db_w_hd128`, cp.async
17832    /// double-buffered; MEMRA_PRIME_DEQW_DB=0 selects the single-buffer twin). The window
17833    /// mask is `fa_prefill_f32_body`'s exact predicate; `window == 0` is bit-identical to
17834    /// `fa_prefill_view_ws` by construction (default-arg body). NEW NUMERIC CLASS vs the
17835    /// f32 floor on SWA rows (bf16 MMA online-softmax vs f32 serial softmax) — adoption is
17836    /// gated by the full battery, and the class must change UNIFORMLY for a whole request
17837    /// (kernel selection keys on seq_end, never per chunk — the chunkfix law).
17838    /// hd128-only deliberately: the only windowed-prefill consumer at another head_dim is
17839    /// gemma4 (hd256), which already has `fa_prefill_w_f32`.
17840    #[allow(clippy::too_many_arguments)]
17841    pub fn fa_prefill_view_ws_w_hd128(
17842        &self,
17843        q: &CudaSlice<f32>,
17844        k: &cudarc::driver::CudaView<u8>,
17845        v: &cudarc::driver::CudaView<u8>,
17846        o: &mut CudaSlice<f32>,
17847        head_dim: usize,
17848        n_head: usize,
17849        n_head_kv: usize,
17850        t: usize,
17851        t_kv: usize,
17852        scale: f32,
17853        causal: bool,
17854        window: usize,
17855        k_tok_bytes: usize,
17856        v_tok_bytes: usize,
17857    ) -> Result<(), Box<dyn std::error::Error>> {
17858        assert_eq!(
17859            head_dim, 128,
17860            "fa_prefill_view_ws_w_hd128: only the hd128 twin is stamped"
17861        );
17862        if portable_mma_gated() {
17863            return self.sdpa_naive_w_quantized_view(
17864                q,
17865                k,
17866                v,
17867                o,
17868                head_dim,
17869                n_head,
17870                n_head_kv,
17871                t,
17872                t_kv,
17873                scale,
17874                causal,
17875                window,
17876                k_tok_bytes,
17877                v_tok_bytes,
17878            );
17879        }
17880        const BLOCK_Q: usize = 64;
17881        const BK: usize = 32;
17882        let kv_dim_k = n_head_kv * head_dim;
17883        let kv_dim_v = n_head_kv * head_dim;
17884        let k_ws_bytes = t_kv * kv_dim_k * 2; // bf16
17885        let v_ws_bytes = t_kv * kv_dim_v * 2;
17886        let mut guard = self.prime_deqw_ws.lock().unwrap();
17887        let need_grow = match guard.as_ref() {
17888            Some((kw, vw)) => kw.len() < k_ws_bytes || vw.len() < v_ws_bytes,
17889            None => true,
17890        };
17891        if need_grow {
17892            let grow = |cur: usize, need: usize| if cur >= need { cur } else { need };
17893            let (ck, cv) = guard
17894                .as_ref()
17895                .map(|(a, b)| (a.len(), b.len()))
17896                .unwrap_or((0, 0));
17897            *guard = Some((
17898                self.alloc_u8(grow(ck, k_ws_bytes))?,
17899                self.alloc_u8(grow(cv, v_ws_bytes))?,
17900            ));
17901        }
17902        let (kw, vw) = guard.as_mut().unwrap();
17903        // pass 1: dequant K+V once into the bf16 workspace (identical to fa_prefill_view_ws —
17904        // the workspace bytes are the SAME __float2bfloat16(dq(...)) values either way).
17905        {
17906            let f = self.func("fa_dequant_kv_ws_bf16");
17907            let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
17908            let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
17909            let cfg = LaunchConfig {
17910                grid_dim: (nblk.max(1), 1, 1),
17911                block_dim: (256, 1, 1),
17912                shared_mem_bytes: 0,
17913            };
17914            let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
17915            let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
17916            let __s_b = self.gpu.stream();
17917            let mut b = __s_b.launch_builder(&f);
17918            b.arg(k)
17919                .arg(v)
17920                .arg(&mut *kw)
17921                .arg(&mut *vw)
17922                .arg(&kdk)
17923                .arg(&kdv)
17924                .arg(&tkvi)
17925                .arg(&ktb)
17926                .arg(&vtb);
17927            unsafe {
17928                b.launch(cfg)?;
17929            }
17930        }
17931        // pass 2: the WINDOWED qw twin (db default, same as the unwindowed wrapper).
17932        let db = std::env::var("MEMRA_PRIME_DEQW_DB")
17933            .map(|v| v != "0")
17934            .unwrap_or(true);
17935        {
17936            let f = self.func(if db {
17937                "fa_prefill_qw_db_w_hd128"
17938            } else {
17939                "fa_prefill_qw_w_hd128"
17940            });
17941            let shmem = if db {
17942                (2 * (4 * BK * head_dim + BLOCK_Q * BK) + 4 * BLOCK_Q) as u32
17943            } else {
17944                (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32
17945            };
17946            use cudarc::driver::sys::CUfunction_attribute_enum as A;
17947            f.set_attribute(
17948                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
17949                shmem as i32,
17950            )?;
17951            let cfg = LaunchConfig {
17952                grid_dim: (
17953                    (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
17954                    n_head as u32,
17955                    1,
17956                ),
17957                block_dim: (32, 4, 1),
17958                shared_mem_bytes: shmem,
17959            };
17960            let (hd, nh, nhkv, ti, tkvi, cz) = (
17961                head_dim as i32,
17962                n_head as i32,
17963                n_head_kv as i32,
17964                t as i32,
17965                t_kv as i32,
17966                causal as i32,
17967            );
17968            let (kdk, kdv, wnd) = (kv_dim_k as i32, kv_dim_v as i32, window as i32);
17969            let __s_b = self.gpu.stream();
17970            let mut b = __s_b.launch_builder(&f);
17971            b.arg(q)
17972                .arg(&*kw)
17973                .arg(&*vw)
17974                .arg(o)
17975                .arg(&hd)
17976                .arg(&nh)
17977                .arg(&nhkv)
17978                .arg(&ti)
17979                .arg(&tkvi)
17980                .arg(&scale)
17981                .arg(&cz)
17982                .arg(&kdk)
17983                .arg(&kdv)
17984                .arg(&wnd);
17985            unsafe {
17986                b.launch(cfg)?;
17987            }
17988        }
17989        Ok(())
17990    }
17991
17992    /// FA decode (T=1 split-K) over the resident QUANTIZED KV cache (q8_0 K / q5_1 V) as u8 views.
17993    /// Replaces sdpa_naive_view for decode; inline-dequants per element. k_tok_bytes/v_tok_bytes are
17994    /// the per-token byte strides (differ: q8_0=34*nblk, q5_1=24*nblk per token).
17995    pub fn fa_decode(
17996        &self,
17997        q: &CudaSlice<f32>,
17998        k: &cudarc::driver::CudaView<u8>,
17999        v: &cudarc::driver::CudaView<u8>,
18000        o: &mut CudaSlice<f32>,
18001        head_dim: usize,
18002        n_head: usize,
18003        n_head_kv: usize,
18004        t_kv: usize,
18005        scale: f32,
18006        k_tok_bytes: usize,
18007        v_tok_bytes: usize,
18008    ) -> Result<(), Box<dyn std::error::Error>> {
18009        self.fa_decode_kvmod(
18010            q,
18011            k,
18012            v,
18013            o,
18014            head_dim,
18015            n_head,
18016            n_head_kv,
18017            t_kv,
18018            scale,
18019            k_tok_bytes,
18020            v_tok_bytes,
18021            false,
18022        )
18023    }
18024
18025    /// `fa_decode` with an explicit fp8-module flag (`g`): gemma windowed layers under
18026    /// MEMRA_GEMMA_WKV read an e4m3 cache — every kernel must come from the kf8vf8 module
18027    /// and the v4 lane (q8_0-hardcoded staging) is excluded.
18028    #[allow(clippy::too_many_arguments)]
18029    /// UNIFIED scalar decode launch (fa_decode_f32, nullable-ctr): ONE symbol for host-len
18030    /// (kvmod eager) and device-len (graph/stream) callers — the textually-identical f32_dc
18031    /// twin compiled apart and its ULP drift flipped 31B verify argmaxes (2026-07-12).
18032    #[allow(clippy::too_many_arguments)]
18033    #[allow(clippy::too_many_arguments)]
18034    fn fa_decode_scalar_unified(
18035        &self,
18036        q: &cudarc::driver::CudaView<f32>,
18037        k: &cudarc::driver::CudaView<u8>,
18038        v: &cudarc::driver::CudaView<u8>,
18039        o: &mut cudarc::driver::CudaViewMut<f32>,
18040        head_dim: usize,
18041        n_head: usize,
18042        n_head_kv: usize,
18043        t_kv_host: usize,
18044        t_kv_dev: Option<&CudaSlice<i32>>,
18045        scale: f32,
18046        n_splits: usize,
18047        split_keys: usize,
18048        k_tok_bytes: usize,
18049        v_tok_bytes: usize,
18050        g: bool,
18051        part_o: &mut CudaSlice<f32>,
18052        part_m: &mut CudaSlice<f32>,
18053        part_l: &mut CudaSlice<f32>,
18054        q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
18055    ) -> Result<(), Box<dyn std::error::Error>> {
18056        let f = if g {
18057            self.func_g("fa_decode_f32")
18058        } else {
18059            self.fa_func("fa_decode_f32", head_dim)
18060        };
18061        let cfg = LaunchConfig {
18062            grid_dim: (n_head as u32, n_splits as u32, 1),
18063            block_dim: (head_dim as u32, 1, 1),
18064            shared_mem_bytes: (4 * (head_dim + 32)) as u32,
18065        };
18066        let (hd, nh, nhkv, nsp) = (
18067            head_dim as i32,
18068            n_head as i32,
18069            n_head_kv as i32,
18070            n_splits as i32,
18071        );
18072        let (ktb, vtb, tkvi, ski) = (
18073            k_tok_bytes as i64,
18074            v_tok_bytes as i64,
18075            t_kv_host as i32,
18076            split_keys as i32,
18077        );
18078        let __s_b = self.gpu.stream();
18079        let mut b = __s_b.launch_builder(&f);
18080        match t_kv_dev {
18081            Some(d) => {
18082                b.arg(q)
18083                    .arg(k)
18084                    .arg(v)
18085                    .arg(&mut *part_o)
18086                    .arg(&mut *part_m)
18087                    .arg(&mut *part_l)
18088                    .arg(&hd)
18089                    .arg(&nh)
18090                    .arg(&nhkv)
18091                    .arg(&tkvi)
18092                    .arg(d)
18093                    .arg(&scale)
18094                    .arg(&nsp)
18095                    .arg(&ski)
18096                    .arg(&ktb)
18097                    .arg(&vtb);
18098                unsafe {
18099                    b.launch(cfg)?;
18100                }
18101            }
18102            None => {
18103                let null: u64 = 0;
18104                b.arg(q)
18105                    .arg(k)
18106                    .arg(v)
18107                    .arg(&mut *part_o)
18108                    .arg(&mut *part_m)
18109                    .arg(&mut *part_l)
18110                    .arg(&hd)
18111                    .arg(&nh)
18112                    .arg(&nhkv)
18113                    .arg(&tkvi)
18114                    .arg(&null)
18115                    .arg(&scale)
18116                    .arg(&nsp)
18117                    .arg(&ski)
18118                    .arg(&ktb)
18119                    .arg(&vtb);
18120                unsafe {
18121                    b.launch(cfg)?;
18122                }
18123            }
18124        }
18125        let cfg2 = LaunchConfig {
18126            grid_dim: (n_head as u32, 1, 1),
18127            block_dim: (head_dim as u32, 1, 1),
18128            shared_mem_bytes: 0,
18129        };
18130        if let Some((oq, od)) = q8_out {
18131            // wave-5b: q8-emitting combine — the wo matmul_pre consumes the pair directly.
18132            let fc = if g {
18133                self.func_g("fa_decode_combine_q8_1")
18134            } else {
18135                self.fa_func("fa_decode_combine_q8_1", head_dim)
18136            };
18137            let __s_b2 = self.gpu.stream();
18138            let mut b2 = __s_b2.launch_builder(&fc);
18139            b2.arg(&*part_o)
18140                .arg(&*part_m)
18141                .arg(&*part_l)
18142                .arg(oq)
18143                .arg(od)
18144                .arg(&hd)
18145                .arg(&nh)
18146                .arg(&nsp);
18147            unsafe {
18148                b2.launch(cfg2)?;
18149            }
18150            return Ok(());
18151        }
18152        let fc = if g {
18153            self.func_g("fa_decode_combine_f32")
18154        } else {
18155            self.fa_func("fa_decode_combine_f32", head_dim)
18156        };
18157        let __s_b2 = self.gpu.stream();
18158        let mut b2 = __s_b2.launch_builder(&fc);
18159        b2.arg(&*part_o)
18160            .arg(&*part_m)
18161            .arg(&*part_l)
18162            .arg(o)
18163            .arg(&hd)
18164            .arg(&nh)
18165            .arg(&nsp);
18166        unsafe {
18167            b2.launch(cfg2)?;
18168        }
18169        Ok(())
18170    }
18171
18172    pub fn fa_decode_kvmod(
18173        &self,
18174        q: &CudaSlice<f32>,
18175        k: &cudarc::driver::CudaView<u8>,
18176        v: &cudarc::driver::CudaView<u8>,
18177        o: &mut CudaSlice<f32>,
18178        head_dim: usize,
18179        n_head: usize,
18180        n_head_kv: usize,
18181        t_kv: usize,
18182        scale: f32,
18183        k_tok_bytes: usize,
18184        v_tok_bytes: usize,
18185        g: bool,
18186    ) -> Result<(), Box<dyn std::error::Error>> {
18187        let q_view = q.as_view();
18188        let mut o_view = o.as_view_mut();
18189        self.fa_decode_kvmod_view(
18190            &q_view,
18191            k,
18192            v,
18193            &mut o_view,
18194            head_dim,
18195            n_head,
18196            n_head_kv,
18197            t_kv,
18198            scale,
18199            k_tok_bytes,
18200            v_tok_bytes,
18201            g,
18202        )
18203    }
18204
18205    /// Row-view entry into `fa_decode_kvmod`. The kernel sees the selected Q/output rows as its
18206    /// base pointers, so the launch geometry and arithmetic are identical to the owned-slice entry.
18207    /// Batched fallback callers use this to avoid materializing rows around an otherwise unchanged
18208    /// per-session KV view and FA launch.
18209    #[allow(clippy::too_many_arguments)]
18210    pub fn fa_decode_kvmod_view(
18211        &self,
18212        q: &cudarc::driver::CudaView<f32>,
18213        k: &cudarc::driver::CudaView<u8>,
18214        v: &cudarc::driver::CudaView<u8>,
18215        o: &mut cudarc::driver::CudaViewMut<f32>,
18216        head_dim: usize,
18217        n_head: usize,
18218        n_head_kv: usize,
18219        t_kv: usize,
18220        scale: f32,
18221        k_tok_bytes: usize,
18222        v_tok_bytes: usize,
18223        g: bool,
18224    ) -> Result<(), Box<dyn std::error::Error>> {
18225        // PERF-4: the warp-per-token vec path replaces the scalar element-per-thread fa_decode_f32 —
18226        // warp-per-token fa_decode_vec_q (grid=(n_head_kv,n_splits), block=(32,gqa_ratio)).
18227        // The block dequants each KV tile ONCE into smem (bf16) and broadcasts to all gqa Q-head
18228        // warps -> each KV byte leaves HBM/L2 ~1x/group (vs 4x). ARGS identical; func/grid/block/
18229        // smem/n_splits differ. fa_decode_f32 stays the bit-reference fallback. Combine is shared.
18230        //
18231        // SPLIT-K: the scalar path has grid.x=n_head (32) blocks; the vec path only has
18232        // grid.x=n_head_kv (8). To avoid starving the GPU at mid ctx, the vec path splits MORE
18233        // aggressively (64 keys/split vs 256) so grid.y rises and 8*n_splits fills the SMs.
18234        // At VERY short ctx (t_kv<96) even 1 split can't fill the GPU from 8 KV heads, so the
18235        // broadcast can't beat the scalar path's 4x-more-blocks latency hiding — fall back to
18236        // scalar there (measured crossover: vec 0.68x at t_kv=64, 1.23x at t_kv=96, 2.2x at 256).
18237        // DEFAULT-ON (2026-06-28): clean clock-locked sweep proved vec beats scalar at every
18238        // t_kv>=96 and the gain WIDENS with ctx (graph decode: +9.5% @128, +11.6% @512, +11.8%
18239        // @2048) — the KV-byte-broadcast (4x fewer HBM reads/group) compounds as attention grows.
18240        // MEMRA_NO_FA_VEC forces the scalar bit-reference. Below FA_VEC_MIN_TKV the scalar path's
18241        // 4x-more-blocks (grid.x=n_head=32 vs n_head_kv=8) hides latency better, so keep scalar there.
18242        // g + no-v4: the g-module REGISTER twin mis-decodes the gemma windowed shape
18243        // (root-cause open, jsonl) — only reachable by forcing v4 off (MEMRA_FA_V4_MAX);
18244        // fall to the exact scalar there instead of the broken register arm.
18245        let mut fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && t_kv >= fa_vec_min_tkv();
18246        // hd256 under g: v4-or-scalar. hd128/other under g ride the REGISTER g-lane (the
18247        // dq_K_lane/dq_V_lane macros are format-aware; v3 is format-gated off, v2/smem
18248        // arms are excluded under g). Mirrored in kvmod / fa_decode_dc / fa_geom_eager.
18249        if g && head_dim == 256 && !fa_v4_at(t_kv) {
18250            fa_vec = false;
18251        }
18252        let sp = fa_split_keys(t_kv, n_head_kv);
18253        let n_splits = if fa_vec {
18254            ((t_kv + sp - 1) / sp).max(1)
18255        } else {
18256            ((t_kv + 255) / 256).max(1)
18257        };
18258        let o_len = n_head * n_splits * head_dim;
18259        let ml_len = n_head * n_splits;
18260        let mut part_guard = self.fa_part_pool.lock().unwrap();
18261        if part_guard
18262            .as_ref()
18263            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
18264            .unwrap_or(true)
18265        {
18266            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
18267            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
18268            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
18269            // later live allocations land at those addresses, and the next graph REPLAY writes
18270            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
18271            // output corruption began the burst after the trunk's t_kv growth first realloc'd
18272            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
18273            // the baked addresses alive (single-stream: eager writes the new buffers, replays
18274            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
18275            // (total retired < final size).
18276            let old = part_guard.take();
18277            let (co, cm) = old
18278                .as_ref()
18279                .map(|pp| (pp.0.len(), pp.1.len()))
18280                .unwrap_or((0, 0));
18281            if let Some(old) = old {
18282                self.fa_part_retired.lock().unwrap().push(old);
18283            }
18284            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
18285                eprintln!(
18286                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
18287                    co, o_len, cm, ml_len
18288                );
18289            }
18290            *part_guard = Some((
18291                self.alloc_uninit::<f32>(o_len.max(2 * co))?,
18292                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
18293                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
18294            ));
18295        }
18296        let pg = part_guard.as_mut().unwrap();
18297        self.gpu
18298            .stream()
18299            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
18300        self.gpu
18301            .stream()
18302            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
18303        self.gpu
18304            .stream()
18305            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
18306        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
18307        let (part_o, part_m, part_l) = (&mut *part_o, &mut *part_m, &mut *part_l);
18308        let (hd, nh, nhkv, tkvi, nsp) = (
18309            head_dim as i32,
18310            n_head as i32,
18311            n_head_kv as i32,
18312            t_kv as i32,
18313            n_splits as i32,
18314        );
18315        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
18316        // The vec kernel holds head_dim/32 register accumulators (FA_DEC_MAX_DPL=8 -> head_dim<=256).
18317        // All shipped models use head_dim=256; fall back to scalar for anything wider rather than
18318        // silently truncating the accumulator.
18319        let fa_vec = fa_vec && head_dim <= 512 && head_dim % 32 == 0;
18320        // hd-512 vec crossover (MEMRA_FA512_MIN, default 512): the DPL16 twin wins at depth
18321        // (82.5 -> vec at 1736) but the scalar's more-blocks latency hiding wins at tiny t_kv
18322        // (the same scalar-floor physics as hd256's old 96 floor; short-ctx plain regressed
18323        // 178.4 -> 173.7 when 512 rode vec unconditionally).
18324        let fa512_min = fa512_min_tkv();
18325        // FA-DEEP pick (bit-identical twins, see fa_deep_at): default module only — the
18326        // g-module keeps the v4 pick (its class is not the depth-decay class).
18327        let deep = fa_vec
18328            && head_dim == 256
18329            && fa_v4_at(t_kv)
18330            && !g
18331            && fa_deep_at(t_kv)
18332            && !matches!(fa_v4_mode(), "noB3" | "stage");
18333        let (f, cfg) = if fa_vec && head_dim == 512 && t_kv >= fa512_min {
18334            // gemma4 globals (hd 512): the DPL16 register twin (fa_decode_vec_q body with a
18335            // 16-slot accumulator ceiling). Scalar fallback measured 82.5us/layer at 1736 ctx.
18336            let gqa = (n_head / n_head_kv).max(1) as u32;
18337            let fv = self.fa_func("fa_decode_vec_q_dpl16", head_dim);
18338            (
18339                fv,
18340                LaunchConfig {
18341                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
18342                    block_dim: (32, gqa, 1),
18343                    shared_mem_bytes: 0,
18344                },
18345            )
18346        } else if fa_vec && head_dim <= 256 {
18347            let gqa = (n_head / n_head_kv).max(1) as u32;
18348            // DEEP-CTX smem twin (2026-07-05): the register-dequant path's GQA reuse rides L2,
18349            // which holds to ~8k ctx but dies at 40k (layer KV ~37MB) — the 4 GQA warps then
18350            // re-read every KV byte from DRAM (4x traffic). Above MEMRA_FA_SMEM_TKV (default
18351            // 1024 — the 2026-07-05 crossover re-sweep on real prompts: p3 spec 73.8->79.2 at
18352            // 2048, flat down to 512, p2 +5%, p1/9B unchanged; the ARC-A probe's synthetic
18353            // 2.1x smem-at-all-depths pointed here; 0=never) dispatch the smem-broadcast twin:
18354            // dequant each tile ONCE per block.
18355            // Bit-identical per (token,split): same bf16 round-trip, same accumulation order,
18356            // same partial layout -> same combine. Short/mid ctx keeps the register path (it won
18357            // there by 12x — latency, not bandwidth, rules small KV).
18358            static SMEM_TKV: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
18359            let smem_tkv = *SMEM_TKV.get_or_init(|| {
18360                std::env::var("MEMRA_FA_SMEM_TKV")
18361                    .ok()
18362                    .and_then(|v| v.parse().ok())
18363                    .unwrap_or_else(|| {
18364                        FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
18365                    })
18366            });
18367            if fa_v4_at(t_kv) && head_dim == 256 {
18368                // FA v4 lane (2026-07-10): key-per-lane score phase, zero shuffles per key.
18369                // NEW NUMERIC CONFIG (chunk-serial per-key dot) — battery-arbitrated.
18370                // g (fp8-windowed): the v4 staging is format-aware (2026-07-12) — kf8vf8 module.
18371                let v4name = match fa_v4_mode() {
18372                    "noB3" => "fa_decode_vec_q_v4_noB3", // phase probe (WRONG OUTPUT)
18373                    "stage" => "fa_decode_vec_q_v4_stage", // phase probe (WRONG OUTPUT)
18374                    _ if deep => "fa_decode_vec_q_v4_deep",
18375                    _ => "fa_decode_vec_q_v4",
18376                };
18377                let fv = if g {
18378                    self.func_g(v4name)
18379                } else {
18380                    self.func(v4name)
18381                };
18382                // fa_v4_smem (deep: fa_v4_deep_smem, +640B row pads) + sV (g: raw e4m3 sV
18383                // tile = 1B/elem — half the smem, 3->5 blocks/SM)
18384                let shmem = (if deep { 12160 } else { 11520 }
18385                    + 32 * head_dim * if g { 1 } else { 2 }) as u32;
18386                use cudarc::driver::sys::CUfunction_attribute_enum as A;
18387                fv.set_attribute(
18388                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18389                    shmem as i32,
18390                )?;
18391                (
18392                    fv,
18393                    LaunchConfig {
18394                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
18395                        block_dim: (32, gqa, 1),
18396                        shared_mem_bytes: shmem,
18397                    },
18398                )
18399            } else if fa_v3_active(head_dim) {
18400                // FA v3 lane: dp4a-K hybrid (register-quantized Q, raw q8_0 K, staged-V kept).
18401                // smem = sV only (half of v2's).
18402                let fv = if g {
18403                    self.func_g("fa_decode_vec_q_v3")
18404                } else {
18405                    self.func("fa_decode_vec_q_v3")
18406                };
18407                let shmem = (32 * head_dim * 2) as u32; // sV bf16 [FA_DEC_TILE=32][hd]
18408                (
18409                    fv,
18410                    LaunchConfig {
18411                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
18412                        block_dim: (32, gqa, 1),
18413                        shared_mem_bytes: shmem,
18414                    },
18415                )
18416            } else if fa_v2_on() {
18417                // FAVENDOR lane: llama fattn-vec tile-batched softmax + wide-load staging on
18418                // OUR smem KV broadcast. Replaces BOTH per-key twins when on; same grid/block/
18419                // partials; same 32KB sK+sV tile as the smem twin.
18420                let fv = if g {
18421                    self.func_g("fa_decode_vec_q_v2")
18422                } else {
18423                    self.func("fa_decode_vec_q_v2")
18424                };
18425                let shmem = (2 * 32 * head_dim * 2) as u32; // sK+sV bf16 [FA_DEC_TILE=32][hd]
18426                (
18427                    fv,
18428                    LaunchConfig {
18429                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
18430                        block_dim: (32, gqa, 1),
18431                        shared_mem_bytes: shmem,
18432                    },
18433                )
18434            } else if smem_tkv > 0 && t_kv >= smem_tkv && !g && !(head_dim == 512 && Self::gkv_on())
18435            {
18436                // (fp8 exclusions: the smem twin's V-stage is q5_1-hardcoded — neither the wkv
18437                // windowed layers (g) nor the gkv globals (hd512) may be forced onto it via
18438                // MEMRA_FA_SMEM_TKV; they fall through to the format-clean register/scalar arms.)
18439                let fv = if g {
18440                    self.func_g("fa_decode_vec_q_smem")
18441                } else {
18442                    self.func("fa_decode_vec_q_smem")
18443                };
18444                let shmem = (2 * 32 * head_dim * 2) as u32; // sK+sV bf16 [FA_DEC_TILE=32][hd]
18445                use cudarc::driver::sys::CUfunction_attribute_enum as A;
18446                fv.set_attribute(
18447                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18448                    shmem as i32,
18449                )?;
18450                (
18451                    fv,
18452                    LaunchConfig {
18453                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
18454                        block_dim: (32, gqa, 1),
18455                        shared_mem_bytes: shmem,
18456                    },
18457                )
18458            } else {
18459                // REGISTER-DEQUANT kernel (2026-07-03): per-warp direct q8_0/q5_1 register
18460                // dequant, zero dynamic shared memory.
18461                let fv = if g {
18462                    self.func_g("fa_decode_vec_q")
18463                } else {
18464                    self.func("fa_decode_vec_q")
18465                };
18466                (
18467                    fv,
18468                    LaunchConfig {
18469                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
18470                        block_dim: (32, gqa, 1),
18471                        shared_mem_bytes: 0,
18472                    },
18473                )
18474            }
18475        } else {
18476            // UNIFIED scalar (nullable-ctr symbol shared with graph/stream callers). The
18477            // split ladder value rides along so ns_eff reproduces THIS n_splits in-kernel.
18478            return self.fa_decode_scalar_unified(
18479                q,
18480                k,
18481                v,
18482                o,
18483                head_dim,
18484                n_head,
18485                n_head_kv,
18486                t_kv,
18487                None,
18488                scale,
18489                n_splits,
18490                if fa_vec { sp } else { 256 },
18491                k_tok_bytes,
18492                v_tok_bytes,
18493                g,
18494                part_o,
18495                part_m,
18496                part_l,
18497                None,
18498            );
18499        };
18500        let __s_b = self.gpu.stream();
18501        let mut b = __s_b.launch_builder(&f);
18502        b.arg(q)
18503            .arg(k)
18504            .arg(v)
18505            .arg(&mut *part_o)
18506            .arg(&mut *part_m)
18507            .arg(&mut *part_l)
18508            .arg(&hd)
18509            .arg(&nh)
18510            .arg(&nhkv)
18511            .arg(&tkvi)
18512            .arg(&scale)
18513            .arg(&nsp)
18514            .arg(&ktb)
18515            .arg(&vtb);
18516        unsafe {
18517            b.launch(cfg)?;
18518        }
18519        // (combine re-tile refuted in the fa-deep lane — flat/worse both shapes; the v4
18520        // combine stays for all arms. Receipts research/fa-decode-deep-20260802/.)
18521        let (fc, cfg2) = (
18522            if g {
18523                self.func_g("fa_decode_combine_f32")
18524            } else {
18525                self.fa_func("fa_decode_combine_f32", head_dim)
18526            },
18527            LaunchConfig {
18528                grid_dim: (n_head as u32, 1, 1),
18529                block_dim: (head_dim as u32, 1, 1),
18530                shared_mem_bytes: 0,
18531            },
18532        );
18533        let __s_b2 = self.gpu.stream();
18534        let mut b2 = __s_b2.launch_builder(&fc);
18535        b2.arg(&*part_o)
18536            .arg(&*part_m)
18537            .arg(&*part_l)
18538            .arg(o)
18539            .arg(&hd)
18540            .arg(&nh)
18541            .arg(&nsp);
18542        unsafe {
18543            b2.launch(cfg2)?;
18544        }
18545        Ok(())
18546    }
18547
18548    /// BATCHED-TICK increment 2: ONE fa_decode launch covering ALL B sequences of the
18549    /// batched decode step (blockIdx.z = sequence). Per-seq K/V cache bases ride a device
18550    /// pointer table (`kv_ptrs`, [2B] interleaved k0,v0,...); per-seq key bounds ride the
18551    /// tick's position table (`pos_seq`, T_kv = pos+1). v4-lane only: the CALLER
18552    /// (decode_batch) gates every row through `fa_seqs_eligible` AND one `fa_split_keys`
18553    /// rung (`split_keys`), so each sequence's split partition, key walk and combine order
18554    /// reproduce its per-seq eager v4 program exactly (kernel-check pins seqs-vs-loop bit
18555    /// identity; decode-batch-gate strict pins the whole tick vs decode_step_h).
18556    /// q is the stacked [B, n_head, head_dim] tick buffer read in place (no per-seq q
18557    /// copies); o is written [B, n_head, head_dim] in place (no per-seq a copies).
18558    #[allow(clippy::too_many_arguments)]
18559    pub fn fa_decode_batch_seqs_v4(
18560        &self,
18561        q: &CudaSlice<f32>,
18562        kv_ptrs: &cudarc::driver::CudaView<u64>,
18563        pos_seq: &CudaSlice<i32>,
18564        o: &mut CudaSlice<f32>,
18565        head_dim: usize,
18566        n_head: usize,
18567        n_head_kv: usize,
18568        b_n: usize,
18569        t_kv_max: usize,
18570        scale: f32,
18571        split_keys: usize,
18572        k_tok_bytes: usize,
18573        v_tok_bytes: usize,
18574    ) -> Result<(), Box<dyn std::error::Error>> {
18575        debug_assert!(head_dim == 256, "seqs twin is v4-stamped (hd256 only)");
18576        let n_splits_max = (t_kv_max + split_keys - 1) / split_keys;
18577        let o_len = b_n * n_head * n_splits_max * head_dim;
18578        let ml_len = b_n * n_head * n_splits_max;
18579        let mut part_guard = self.fa_part_pool.lock().unwrap();
18580        if part_guard
18581            .as_ref()
18582            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
18583            .unwrap_or(true)
18584        {
18585            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
18586            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
18587            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
18588            // later live allocations land at those addresses, and the next graph REPLAY writes
18589            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
18590            // output corruption began the burst after the trunk's t_kv growth first realloc'd
18591            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
18592            // the baked addresses alive (single-stream: eager writes the new buffers, replays
18593            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
18594            // (total retired < final size).
18595            let old = part_guard.take();
18596            let (co, cm) = old
18597                .as_ref()
18598                .map(|pp| (pp.0.len(), pp.1.len()))
18599                .unwrap_or((0, 0));
18600            if let Some(old) = old {
18601                self.fa_part_retired.lock().unwrap().push(old);
18602            }
18603            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
18604                eprintln!(
18605                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
18606                    co, o_len, cm, ml_len
18607                );
18608            }
18609            *part_guard = Some((
18610                self.alloc_uninit::<f32>(o_len.max(2 * co))?,
18611                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
18612                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
18613            ));
18614        }
18615        let pg = part_guard.as_mut().unwrap();
18616        self.gpu
18617            .stream()
18618            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
18619        self.gpu
18620            .stream()
18621            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
18622        self.gpu
18623            .stream()
18624            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
18625        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
18626        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
18627        let (nspm, spk) = (n_splits_max as i32, split_keys as i32);
18628        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
18629        let gqa = (n_head / n_head_kv).max(1) as u32;
18630        let f = self.func("fa_decode_vec_q_seqs_v4");
18631        // fa_v4_smem (11520B) + sV bf16 tile — the v4 eager arm's sizing on the default module.
18632        let shmem = (11520 + 32 * head_dim * 2) as u32;
18633        use cudarc::driver::sys::CUfunction_attribute_enum as A;
18634        f.set_attribute(
18635            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18636            shmem as i32,
18637        )?;
18638        let cfg = LaunchConfig {
18639            grid_dim: (n_head_kv as u32, n_splits_max as u32, b_n as u32),
18640            block_dim: (32, gqa, 1),
18641            shared_mem_bytes: shmem,
18642        };
18643        {
18644            let __s_b = self.gpu.stream();
18645            let mut b = __s_b.launch_builder(&f);
18646            b.arg(q)
18647                .arg(kv_ptrs)
18648                .arg(pos_seq)
18649                .arg(&mut *part_o)
18650                .arg(&mut *part_m)
18651                .arg(&mut *part_l)
18652                .arg(&hd)
18653                .arg(&nh)
18654                .arg(&nhkv)
18655                .arg(&scale)
18656                .arg(&nspm)
18657                .arg(&spk)
18658                .arg(&ktb)
18659                .arg(&vtb);
18660            unsafe {
18661                b.launch(cfg)?;
18662            }
18663        }
18664        let fc = self.func("fa_decode_combine_seqs");
18665        let cfg2 = LaunchConfig {
18666            grid_dim: (n_head as u32, b_n as u32, 1),
18667            block_dim: (head_dim as u32, 1, 1),
18668            shared_mem_bytes: 0,
18669        };
18670        let __s_b2 = self.gpu.stream();
18671        let mut b2 = __s_b2.launch_builder(&fc);
18672        b2.arg(&*part_o)
18673            .arg(&*part_m)
18674            .arg(&*part_l)
18675            .arg(o)
18676            .arg(&hd)
18677            .arg(&nh)
18678            .arg(pos_seq)
18679            .arg(&nspm)
18680            .arg(&spk);
18681        unsafe {
18682            b2.launch(cfg2)?;
18683        }
18684        Ok(())
18685    }
18686
18687    /// BATCHED-TICK increment 2: z-batched decode KV append — one launch appends this
18688    /// step's B rows, each into ITS OWN sequence cache at slot pos_seq[z], through the same
18689    /// [2B] interleaved pointer table the seqs FA reads. Each (block, z) warp executes the
18690    /// per-token appender's exact warp program on row z of the stacked [B, kv_dim] k/v —
18691    /// written cache bytes are BIT-IDENTICAL to the B per-seq calls it replaces
18692    /// (kernel-check pins the bytes). Default flash module only (callers exclude fp8-KV).
18693    #[allow(clippy::too_many_arguments)]
18694    pub fn append_kv_quantized_seqs(
18695        &self,
18696        k_rows: &CudaSlice<f32>,
18697        v_rows: &CudaSlice<f32>,
18698        kv_ptrs: &cudarc::driver::CudaView<u64>,
18699        pos_seq: &CudaSlice<i32>,
18700        b_n: usize,
18701        kv_dim_k: usize,
18702        kv_dim_v: usize,
18703        k_tok_bytes: usize,
18704        v_tok_bytes: usize,
18705    ) -> Result<(), Box<dyn std::error::Error>> {
18706        let f = self.func("append_quantize_kv_q8_0_q5_1_seqs");
18707        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
18708        let cfg = LaunchConfig {
18709            grid_dim: (nblk, b_n as u32, 1),
18710            block_dim: (32, 1, 1),
18711            shared_mem_bytes: 0,
18712        };
18713        let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
18714        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
18715        let __s_b = self.gpu.stream();
18716        let mut b = __s_b.launch_builder(&f);
18717        b.arg(k_rows)
18718            .arg(v_rows)
18719            .arg(kv_ptrs)
18720            .arg(pos_seq)
18721            .arg(&kdk)
18722            .arg(&kdv)
18723            .arg(&ktb)
18724            .arg(&vtb);
18725        unsafe {
18726            b.launch(cfg)?;
18727        }
18728        Ok(())
18729    }
18730
18731    /// True iff the MULTI-ROW verify FA (`fa_decode_rows`) is usable for a verify batch whose
18732    /// FIRST row attends `base_len + 1` keys: every row must take the SAME kernel eager decode
18733    /// would (the vec path) — mirrors fa_decode's gate exactly (MEMRA_NO_FA_VEC + FA_VEC_MIN_TKV +
18734    /// head_dim), evaluated at the MINIMUM row bound so no row could have picked scalar.
18735    /// MEMRA_FA_ROWS_OFF=1 is the A/B + fallback seam (per-row loop).
18736    pub fn fa_rows_eligible(&self, base_len: usize, head_dim: usize) -> bool {
18737        std::env::var("MEMRA_NO_FA_VEC").is_err()
18738            && std::env::var("MEMRA_FA_ROWS_OFF").is_err()
18739            && base_len + 1 >= fa_vec_min_tkv()
18740            && head_dim <= 256
18741            && head_dim % 32 == 0
18742    }
18743
18744    /// MULTI-ROW verify FA: run fa_decode_vec_q's EXACT per-row program for T causal query rows
18745    /// (row r attends keys [0..base_len+r+1)) in ONE kernel launch with grid.z = row, plus ONE
18746    /// row-batched combine. Replaces the T separate (fa_decode + combine) launches of the spec
18747    /// verify — same per-row split partition (n_splits_r = ceil(t_kv_r/split_keys), the
18748    /// fa_split_keys formula), same key-walk order, same reduce shapes => bit-identical outputs
18749    /// per row (kernel-check pins rows-vs-loop byte identity; run-spec is the end gate).
18750    /// Caller must have checked `fa_rows_eligible(base_len, head_dim)`.
18751    /// q is the verify's token-major [T, n_head, head_dim] stack; o is written [T, n_head, head_dim].
18752    #[allow(clippy::too_many_arguments)]
18753    pub fn fa_decode_rows(
18754        &self,
18755        q: &CudaSlice<f32>,
18756        k: &cudarc::driver::CudaView<u8>,
18757        v: &cudarc::driver::CudaView<u8>,
18758        o: &mut CudaSlice<f32>,
18759        head_dim: usize,
18760        n_head: usize,
18761        n_head_kv: usize,
18762        base_len: usize,
18763        t: usize,
18764        scale: f32,
18765        k_tok_bytes: usize,
18766        v_tok_bytes: usize,
18767        // hd512 dpl16 twin is DEVICE-LEN (graph arc): base_dev/plus feed the
18768        // kernel; host base_len keeps sizing the splits/partials. hd256 twins
18769        // keep the host arg. None is a bug for hd512 (asserted below).
18770        base_dev: Option<(&CudaSlice<i32>, i32)>,
18771        // K and V planes hold the same values (gemma globals, wv:=wk): pick
18772        // the _kv twin — V plane never read, value rides the q8_0 key dq.
18773        kv_shared: bool,
18774        // this layer's cache is e4m3 (gemma windowed under wkv): resolve the
18775        // hd256 rows kernel from the kf8vf8 module. PER-CALL — a global env
18776        // check here hijacked qwen/kernel-check hd256 rows (8 FAILs, 230ebbe).
18777        g: bool,
18778        // t=1 decode arm only: emit (int8, per-32 scales) from the dc combine
18779        // (hd512 path) — the standalone quantize launch folds away.
18780        mut q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
18781    ) -> Result<(), Box<dyn std::error::Error>> {
18782        debug_assert!(base_len + 1 >= fa_vec_min_tkv() && head_dim <= 512 && head_dim % 32 == 0);
18783        let t_kv_max = base_len + t; // LAST row's key bound
18784        let mut sp = fa_split_keys(t_kv_max, n_head_kv); // env/default — same value every row
18785        // hd512 split override (MEMRA_FA_SP512, 2026-07-11): gemma globals have n_head_kv=2 so
18786        // the grid is (2 x n_splits) — at depth ~29 splits = 58 blocks on 82 SMs (half idle,
18787        // rows_dpl16 8x off its byte floor). EVERY gemma hd512 caller shares THIS wrapper
18788        // (parity law), so the partition is freely tunable — verify and decode move together.
18789        if head_dim == 512 {
18790            static SP512: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
18791            // default 16 (2026-07-11 depth sweep, N=2: plain 155.4->156.5, depth spec
18792            // 236.9->250.4; 12/24/32 all worse). hd512 exists only on gemma globals.
18793            let v = *SP512.get_or_init(|| {
18794                std::env::var("MEMRA_FA_SP512")
18795                    .ok()
18796                    .and_then(|x| x.parse().ok())
18797                    .unwrap_or(0)
18798            });
18799            sp = if v >= 8 {
18800                v
18801            } else {
18802                FA_SP512_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
18803            };
18804        }
18805        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
18806        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
18807        let gqa = (n_head / n_head_kv).max(1) as u32;
18808        // LADDER-RUNG STRADDLE FIX (issue #10, 2026-07-13, g7e-proven): one sp for every row
18809        // diverges from eager decode when a split-ladder rung falls INSIDE the batch — row r's
18810        // eager twin used fa_split_keys(t_kv_r), the batch used fa_split_keys(t_kv_max), and
18811        // the different partition changes the combine's FP order (greedy tie flips at depth;
18812        // MEMRA_FA_SPLIT=64 pin -> PASS on the exact g7e failing config). Fix: group
18813        // consecutive rows by their OWN ladder value and launch once per group — each row then
18814        // executes the exact per-row program eager ran. Rungs land once per doubling, so this
18815        // is 1 launch in the common case and 2 on a crossing round. hd512 keeps one group (its
18816        // sp override is t_kv-independent by construction).
18817        let mut groups: Vec<(usize, usize, usize)> = Vec::new(); // (row0, t_g, sp_g)
18818        if head_dim == 512 || fa_split_keys(base_len + 1, n_head_kv) == sp {
18819            groups.push((0, t, sp));
18820        } else {
18821            let mut r0 = 0usize;
18822            while r0 < t {
18823                let sp_g = fa_split_keys(base_len + r0 + 1, n_head_kv);
18824                let mut r1 = r0 + 1;
18825                while r1 < t && fa_split_keys(base_len + r1 + 1, n_head_kv) == sp_g {
18826                    r1 += 1;
18827                }
18828                groups.push((r0, r1 - r0, sp_g));
18829                r0 = r1;
18830            }
18831        }
18832        // Deep-ctx smem twin for the VERIFY rows (2026-07-05): same threshold + rationale as
18833        // fa_decode's dispatch — at 40k the register path's GQA L2-reuse premise is dead and the
18834        // verify multiplies the 4x DRAM re-read by T rows. Bit-identical per (row,token,split).
18835        static SMEM_TKV_R: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
18836        let smem_tkv = *SMEM_TKV_R.get_or_init(|| {
18837            std::env::var("MEMRA_FA_SMEM_TKV")
18838                .ok()
18839                .and_then(|v| v.parse().ok())
18840                .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
18841        });
18842        let v4 = fa_v4_at(base_len + t) && head_dim == 256;
18843        let v3 = fa_v3_active(head_dim);
18844        let smem_rows =
18845            head_dim <= 256 && !v3 && !fa_v2_on() && smem_tkv > 0 && t_kv_max >= smem_tkv;
18846        // kv_shared twin RETIRED (2026-07-11 depth run-gen gate): the wv:=wk premise fails
18847        // POST-cache — cached K is k-normed+roped, cached V is not; the twin fed roped keys
18848        // in as values. Verify/decode/stream gates were blind (both sides shared the wrong
18849        // symbol — the parity law's blind spot); only prefill-vs-decode at depth caught it.
18850        let _ = kv_shared;
18851        // i2 twin: 2-key interleaved walk (MEMRA_FA_I2=0 reverts). i4 probed NEGATIVE
18852        // (157.3 vs 161.2 depth plain — register pressure past i2's sweet spot; jsonl).
18853        let i2 = head_dim == 512 && std::env::var("MEMRA_FA_I2").as_deref() != Ok("0");
18854        // v4-hd512 (MEMRA_FA_V512=1 opt-in, 2026-07-14): the v4 key-per-lane recipe on the
18855        // globals lane (depth profile: i2 ~4.6x off its byte floor — the v3-class
18856        // reduce-per-key latency signature). NEW NUMERIC CONFIG shared by every hd512
18857        // caller (decode+verify flip together); run-gen argmax + acceptance arbitrate.
18858        // T-BATCHED hd512 (DEFAULT ON 2026-07-14, MEMRA_FA_TB512=0 seam): one block per
18859        // (kv_head, split) stages its tile once and loops the rows over it — kills the
18860        // x t DRAM re-read of the full-ctx globals (depth cell +1.4%, plain flat, N=3
18861        // interleaved). FIXED absolute partition = NEW NUMERIC for the combine order,
18862        // shared by every hd512 caller through this wrapper (decode+verify flip together;
18863        // depth stream identical, acceptance unshifted, spec 256/256 x3 models).
18864        // Requires sp <= 32 (single staged tile; acc reused per row). The z-form v4_512
18865        // sibling (in-kernel dp4a port alone) probed FLAT — hd512 was DRAM-re-read-bound,
18866        // not unpack-bound; jsonl 2026-07-14.
18867        static TB512: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
18868        // gqa <= 16 = fa_v4_smem_512's q-array capacity; past it fall to the register twins.
18869        let tb512 = head_dim == 512
18870            && sp <= 32
18871            && n_head / n_head_kv.max(1) <= 16
18872            && *TB512.get_or_init(|| std::env::var("MEMRA_FA_TB512").as_deref() != Ok("0"));
18873        let fname = if tb512 {
18874            "fa_decode_vec_q_rows_v4_512_tb"
18875        } else if i2 {
18876            "fa_decode_vec_q_rows_dpl16_i2"
18877        } else if head_dim == 512 {
18878            "fa_decode_vec_q_rows_dpl16"
18879        }
18880        // gemma globals (parity law)
18881        else if v4 {
18882            "fa_decode_vec_q_rows_v4"
18883        } else if v3 {
18884            "fa_decode_vec_q_rows_v3"
18885        } else if fa_v2_on() {
18886            "fa_decode_vec_q_rows_v2"
18887        } else if smem_rows {
18888            "fa_decode_vec_q_rows_smem"
18889        } else {
18890            "fa_decode_vec_q_rows"
18891        };
18892        let f = if head_dim == 512 {
18893            self.fa_func(fname, head_dim)
18894        } else if g {
18895            // FP8-WINDOWED: hd256 rows over an e4m3 cache — kf8vf8 module, SAME symbol
18896            // choice as decode's kvmod dispatch (parity law: excluding v4 here paired
18897            // g-module rows against decode's g-module v4 — different programs, short-VG
18898            // maxdiff 2.0 / spec stream 0/128, 2026-07-12). rows_v4 is format-aware
18899            // since fda9790; only the smem twin stays excluded (V-stage q5_1-only).
18900            // hd128 (qwen fp8-KV) lands on the base/register rows via fname — the
18901            // dq macros are format-aware.
18902            self.func_g(if smem_rows {
18903                "fa_decode_vec_q_rows"
18904            } else {
18905                fname
18906            })
18907        } else {
18908            self.func(fname)
18909        };
18910        let shmem = if tb512 {
18911            // fa_v4_smem_512 (q 9KB gqa<=16 + k tile 18KB) + sV 32*512 (e4m3 module halves it)
18912            let gk = Self::gkv_on();
18913            let sh =
18914                (8192 + 1024 + 32 * 512 + 32 * 64 + 32 * head_dim * if gk { 1 } else { 2 }) as u32;
18915            use cudarc::driver::sys::CUfunction_attribute_enum as A;
18916            f.set_attribute(
18917                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18918                sh as i32,
18919            )?;
18920            sh
18921        } else if v4 || v3 || smem_rows || fa_v2_on() {
18922            // v4: fa_v4_smem (11.5KB) + sV; v3 stages sV only; v2/smem twins stage sK+sV.
18923            let sh = (if v4 {
18924                11520 + 32 * head_dim * if g { 1 } else { 2 }
18925            } else if v3 {
18926                32 * head_dim * 2
18927            } else {
18928                2 * 32 * head_dim * 2
18929            }) as u32;
18930            use cudarc::driver::sys::CUfunction_attribute_enum as A;
18931            f.set_attribute(
18932                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18933                sh as i32,
18934            )?;
18935            sh
18936        } else {
18937            0
18938        };
18939        // Per-GROUP launches (single group in the common case — identical to the pre-fix
18940        // single launch there): each group gets its own partials (the rows kernel indexes
18941        // partials by its LOCAL grid.z row) and q/o row-offset views.
18942        for &(r0, t_g, sp_g) in &groups {
18943            let n_splits_g = (base_len + r0 + t_g).div_ceil(sp_g);
18944            let (nspm, spk) = (n_splits_g as i32, sp_g as i32);
18945            let base_i = (base_len + r0) as i32;
18946            let o_len = t_g * n_head * n_splits_g * head_dim;
18947            let ml_len = t_g * n_head * n_splits_g;
18948            let mut part_guard = self.fa_part_pool.lock().unwrap();
18949            if part_guard
18950                .as_ref()
18951                .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
18952                .unwrap_or(true)
18953            {
18954                // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
18955                // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
18956                // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
18957                // later live allocations land at those addresses, and the next graph REPLAY writes
18958                // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
18959                // output corruption began the burst after the trunk's t_kv growth first realloc'd
18960                // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
18961                // the baked addresses alive (single-stream: eager writes the new buffers, replays
18962                // touch only the old — never concurrently). Doubling growth bounds retired VRAM
18963                // (total retired < final size).
18964                let old = part_guard.take();
18965                let (co, cm) = old
18966                    .as_ref()
18967                    .map(|pp| (pp.0.len(), pp.1.len()))
18968                    .unwrap_or((0, 0));
18969                if let Some(old) = old {
18970                    self.fa_part_retired.lock().unwrap().push(old);
18971                }
18972                if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
18973                    eprintln!(
18974                        "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
18975                        co, o_len, cm, ml_len
18976                    );
18977                }
18978                *part_guard = Some((
18979                    self.alloc_uninit::<f32>(o_len.max(2 * co))?,
18980                    self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
18981                    self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
18982                ));
18983            }
18984            let pg = part_guard.as_mut().unwrap();
18985            self.gpu
18986                .stream()
18987                .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
18988            self.gpu
18989                .stream()
18990                .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
18991            self.gpu
18992                .stream()
18993                .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
18994            let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
18995            let (part_o, part_m, part_l) = (&mut *part_o, &mut *part_m, &mut *part_l);
18996            let qv = self.view(q, t * n_head * head_dim);
18997            let q_g = qv.slice(r0 * n_head * head_dim..(r0 + t_g) * n_head * head_dim);
18998            let cfg = LaunchConfig {
18999                grid_dim: (n_head_kv as u32, n_splits_g as u32, t_g as u32),
19000                block_dim: (32, gqa, 1),
19001                shared_mem_bytes: shmem,
19002            };
19003            {
19004                let __s_b = self.gpu.stream();
19005                let mut b = __s_b.launch_builder(&f);
19006                if tb512 {
19007                    // rows-inner launch: grid.z dropped, the kernel loops n_rows itself.
19008                    let (bd, plus) =
19009                        base_dev.expect("hd512 rows twin requires a device base counter");
19010                    let plus_g = plus + r0 as i32;
19011                    let nr = t_g as i32;
19012                    if Self::pdl_on() && Self::pdl_wb_on() {
19013                        // wave-B2b: flavor mirrors fa_func(fname, 512) = gkv.
19014                        use cudarc::driver::{DevicePtr, DevicePtrMut};
19015                        let s = &self.gpu.stream();
19016                        let (pq, _b0) = q_g.device_ptr(s);
19017                        let (pk, _b1) = k.device_ptr(s);
19018                        let (pv, _b2) = v.device_ptr(s);
19019                        let (po, _b3) = part_o.device_ptr_mut(s);
19020                        let (pm, _b4) = part_m.device_ptr_mut(s);
19021                        let (pl, _b5) = part_l.device_ptr_mut(s);
19022                        let (pb, _b6) = bd.device_ptr(s);
19023                        let mut ps = [
19024                            &pq as *const _ as *mut std::ffi::c_void,
19025                            &pk as *const _ as *mut _,
19026                            &pv as *const _ as *mut _,
19027                            &po as *const _ as *mut _,
19028                            &pm as *const _ as *mut _,
19029                            &pl as *const _ as *mut _,
19030                            &hd as *const _ as *mut _,
19031                            &nh as *const _ as *mut _,
19032                            &nhkv as *const _ as *mut _,
19033                            &pb as *const _ as *mut _,
19034                            &plus_g as *const _ as *mut _,
19035                            &scale as *const _ as *mut _,
19036                            &nspm as *const _ as *mut _,
19037                            &spk as *const _ as *mut _,
19038                            &ktb as *const _ as *mut _,
19039                            &vtb as *const _ as *mut _,
19040                            &nr as *const _ as *mut _,
19041                        ];
19042                        unsafe {
19043                            self.launch_pdl_flash(
19044                                Self::gkv_on(),
19045                                "fa_decode_vec_q_rows_v4_512_tb",
19046                                (n_head_kv as u32, n_splits_g as u32, 1),
19047                                (32, gqa, 1),
19048                                shmem,
19049                                &mut ps,
19050                            )?;
19051                        }
19052                    } else {
19053                        let cfg_tb = LaunchConfig {
19054                            grid_dim: (n_head_kv as u32, n_splits_g as u32, 1),
19055                            block_dim: (32, gqa, 1),
19056                            shared_mem_bytes: shmem,
19057                        };
19058                        b.arg(&q_g)
19059                            .arg(k)
19060                            .arg(v)
19061                            .arg(&mut *part_o)
19062                            .arg(&mut *part_m)
19063                            .arg(&mut *part_l)
19064                            .arg(&hd)
19065                            .arg(&nh)
19066                            .arg(&nhkv)
19067                            .arg(bd)
19068                            .arg(&plus_g)
19069                            .arg(&scale)
19070                            .arg(&nspm)
19071                            .arg(&spk)
19072                            .arg(&ktb)
19073                            .arg(&vtb)
19074                            .arg(&nr);
19075                        unsafe {
19076                            b.launch(cfg_tb)?;
19077                        }
19078                    }
19079                } else if head_dim == 512 {
19080                    let (bd, plus) =
19081                        base_dev.expect("hd512 rows twin requires a device base counter");
19082                    let plus_g = plus + r0 as i32;
19083                    b.arg(&q_g)
19084                        .arg(k)
19085                        .arg(v)
19086                        .arg(&mut *part_o)
19087                        .arg(&mut *part_m)
19088                        .arg(&mut *part_l)
19089                        .arg(&hd)
19090                        .arg(&nh)
19091                        .arg(&nhkv)
19092                        .arg(bd)
19093                        .arg(&plus_g)
19094                        .arg(&scale)
19095                        .arg(&nspm)
19096                        .arg(&spk)
19097                        .arg(&ktb)
19098                        .arg(&vtb);
19099                    unsafe {
19100                        b.launch(cfg)?;
19101                    }
19102                } else {
19103                    b.arg(&q_g)
19104                        .arg(k)
19105                        .arg(v)
19106                        .arg(&mut *part_o)
19107                        .arg(&mut *part_m)
19108                        .arg(&mut *part_l)
19109                        .arg(&hd)
19110                        .arg(&nh)
19111                        .arg(&nhkv)
19112                        .arg(&base_i)
19113                        .arg(&scale)
19114                        .arg(&nspm)
19115                        .arg(&spk)
19116                        .arg(&ktb)
19117                        .arg(&vtb);
19118                    unsafe {
19119                        b.launch(cfg)?;
19120                    }
19121                }
19122            }
19123            let cfg2 = LaunchConfig {
19124                grid_dim: (n_head as u32, t_g as u32, 1),
19125                block_dim: (head_dim as u32, 1, 1),
19126                shared_mem_bytes: 0,
19127            };
19128            let mut o_g = o.slice_mut(r0 * n_head * head_dim..(r0 + t_g) * n_head * head_dim);
19129            if head_dim == 512 {
19130                // device-len combine (shared by verify/eager/graph — parity by symbol): the
19131                // per-row n_splits derives from the SAME counter the rows kernel read.
19132                let (bd, plus) = base_dev.unwrap();
19133                let plus_g = plus + r0 as i32;
19134                if let Some((oq, od)) = q8_out.as_mut() {
19135                    // wave-5b port (2026-07-23, t=1 decode only): q8-emitting dc combine.
19136                    debug_assert!(t == 1, "rows q8 emit is a t=1 decode arm");
19137                    if Self::pdl_on() && Self::pdl_wb_on() {
19138                        // wave-B2: flavor mirrors fa_func (hd512 + gkv → kf8vf8).
19139                        use cudarc::driver::{DevicePtr, DevicePtrMut};
19140                        let s = &self.gpu.stream();
19141                        let (po, _g0) = part_o.device_ptr(s);
19142                        let (pm, _g1) = part_m.device_ptr(s);
19143                        let (pl, _g2) = part_l.device_ptr(s);
19144                        let (pq, _g3) = oq.device_ptr_mut(s);
19145                        let (pd, _g4) = od.device_ptr_mut(s);
19146                        let (pb, _g5) = bd.device_ptr(s);
19147                        let mut ps = [
19148                            &po as *const _ as *mut std::ffi::c_void,
19149                            &pm as *const _ as *mut _,
19150                            &pl as *const _ as *mut _,
19151                            &pq as *const _ as *mut _,
19152                            &pd as *const _ as *mut _,
19153                            &hd as *const _ as *mut _,
19154                            &nh as *const _ as *mut _,
19155                            &pb as *const _ as *mut _,
19156                            &plus_g as *const _ as *mut _,
19157                            &nspm as *const _ as *mut _,
19158                            &spk as *const _ as *mut _,
19159                        ];
19160                        unsafe {
19161                            self.launch_pdl_flash(
19162                                Self::gkv_on(),
19163                                "fa_decode_combine_rows_dc_q8_1",
19164                                cfg2.grid_dim,
19165                                cfg2.block_dim,
19166                                0,
19167                                &mut ps,
19168                            )?;
19169                        }
19170                        continue;
19171                    }
19172                    let fc = self.fa_func("fa_decode_combine_rows_dc_q8_1", head_dim);
19173                    let __s_b2 = self.gpu.stream();
19174                    let mut b2 = __s_b2.launch_builder(&fc);
19175                    b2.arg(&*part_o)
19176                        .arg(&*part_m)
19177                        .arg(&*part_l)
19178                        .arg(&mut **oq)
19179                        .arg(&mut **od)
19180                        .arg(&hd)
19181                        .arg(&nh)
19182                        .arg(bd)
19183                        .arg(&plus_g)
19184                        .arg(&nspm)
19185                        .arg(&spk);
19186                    unsafe {
19187                        b2.launch(cfg2)?;
19188                    }
19189                    continue;
19190                }
19191                let fc = self.fa_func("fa_decode_combine_rows_dc", head_dim);
19192                let __s_b2 = self.gpu.stream();
19193                let mut b2 = __s_b2.launch_builder(&fc);
19194                b2.arg(&*part_o)
19195                    .arg(&*part_m)
19196                    .arg(&*part_l)
19197                    .arg(&mut o_g)
19198                    .arg(&hd)
19199                    .arg(&nh)
19200                    .arg(bd)
19201                    .arg(&plus_g)
19202                    .arg(&nspm)
19203                    .arg(&spk);
19204                unsafe {
19205                    b2.launch(cfg2)?;
19206                }
19207            } else {
19208                // q8 emit is wired for the hd512 dc-combine arm only — a Some here would
19209                // leave the caller's pair unwritten (consumer would read garbage).
19210                assert!(
19211                    q8_out.is_none(),
19212                    "rows q8 emit requires the hd512 dc combine"
19213                );
19214                let fc = self.func("fa_decode_combine_rows");
19215                let __s_b2 = self.gpu.stream();
19216                let mut b2 = __s_b2.launch_builder(&fc);
19217                b2.arg(&*part_o)
19218                    .arg(&*part_m)
19219                    .arg(&*part_l)
19220                    .arg(&mut o_g)
19221                    .arg(&hd)
19222                    .arg(&nh)
19223                    .arg(&base_i)
19224                    .arg(&nspm)
19225                    .arg(&spk);
19226                unsafe {
19227                    b2.launch(cfg2)?;
19228                }
19229            }
19230        }
19231        Ok(())
19232    }
19233
19234    /// WINDOWED verify rows (gemma R6 deep-ctx): every row attends exactly `window` keys —
19235    /// bit-identical per row to the T=1 decode's fa_decode over the window VIEW. Caller gates
19236    /// base_len + 1 >= window (no under-window rows) and head_dim == 256 (v4 stamp).
19237    #[allow(clippy::too_many_arguments)]
19238    pub fn fa_decode_rows_w(
19239        &self,
19240        q: &CudaSlice<f32>,
19241        k: &cudarc::driver::CudaView<u8>,
19242        v: &cudarc::driver::CudaView<u8>,
19243        o: &mut CudaSlice<f32>,
19244        head_dim: usize,
19245        n_head: usize,
19246        n_head_kv: usize,
19247        base_dev: &CudaSlice<i32>,
19248        base_plus: i32,
19249        t: usize,
19250        scale: f32,
19251        window: usize,
19252        k_tok_bytes: usize,
19253        v_tok_bytes: usize,
19254        q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
19255    ) -> Result<(), Box<dyn std::error::Error>> {
19256        // DEVICE-LEN (graph arc step 1, 2026-07-11): the causal base rides an i32 counter
19257        // (kernel T_kv = dev[0] + base_plus + r + 1) so depth graphs can replay with len
19258        // advancing on-device. dc paths pass kvl.len_d with plus=-1; verify/eager sync the
19259        // counter with one async set_i32_one first. Partials/splits size from `window` (host).
19260        debug_assert!(head_dim == 256);
19261        // windowed split (MEMRA_FA_SPW, default 32 — re-swept 2026-07-12 under the raw-e4m3 sV
19262        // occupancy ceiling (4 blocks/SM): t=1 decode is GRID-limited (win/sp splits x nkv
19263        // blocks), so smaller splits fill the ceiling — 1.7k 174.4/174.0 vs 48's 170.7/170.3,
19264        // 4.9k 159.8 vs 157.4 (N=2 interleaved, stable window). Spec serving prefers 64
19265        // (verify t=K+1 fills the grid via grid.z=t; depth K=7 281.3 vs 249.3 at 32) — set
19266        // MEMRA_FA_SPW=64 there, same config law as MEMRA_GEMMA_GKV=0. MUST be one value for
19267        // ALL widths: a t-keyed probe broke decode-vs-verify combine order (stream 9/128).
19268        let sp = {
19269            static SPW: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
19270            let v = *SPW.get_or_init(|| {
19271                std::env::var("MEMRA_FA_SPW")
19272                    .ok()
19273                    .and_then(|x| x.parse().ok())
19274                    .unwrap_or(0)
19275            });
19276            if v >= 8 {
19277                v
19278            } else {
19279                FA_SPW_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
19280            }
19281        };
19282        let n_splits_max = (window + sp - 1) / sp;
19283        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
19284        let (nspm, spk, wini) = (n_splits_max as i32, sp as i32, window as i32);
19285        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
19286        let gqa = (n_head / n_head_kv).max(1) as u32;
19287        let o_len = t * n_head * n_splits_max * head_dim;
19288        let ml_len = t * n_head * n_splits_max;
19289        let mut part_guard = self.fa_part_pool.lock().unwrap();
19290        if part_guard
19291            .as_ref()
19292            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
19293            .unwrap_or(true)
19294        {
19295            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
19296            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
19297            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
19298            // later live allocations land at those addresses, and the next graph REPLAY writes
19299            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
19300            // output corruption began the burst after the trunk's t_kv growth first realloc'd
19301            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
19302            // the baked addresses alive (single-stream: eager writes the new buffers, replays
19303            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
19304            // (total retired < final size).
19305            let old = part_guard.take();
19306            let (co, cm) = old
19307                .as_ref()
19308                .map(|pp| (pp.0.len(), pp.1.len()))
19309                .unwrap_or((0, 0));
19310            if let Some(old) = old {
19311                self.fa_part_retired.lock().unwrap().push(old);
19312            }
19313            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
19314                eprintln!(
19315                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
19316                    co, o_len, cm, ml_len
19317                );
19318            }
19319            *part_guard = Some((
19320                self.alloc_uninit::<f32>(o_len.max(2 * co))?,
19321                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
19322                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
19323            ));
19324        }
19325        let pg = part_guard.as_mut().unwrap();
19326        self.gpu
19327            .stream()
19328            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
19329        self.gpu
19330            .stream()
19331            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
19332        self.gpu
19333            .stream()
19334            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
19335        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
19336        // Lane pick: decode AND verify both land here in the windowed regime (parity law —
19337        // hybrid_forward verify_attn), so the pick only needs internal consistency, not
19338        // clone-of-decode bit fidelity (SASS-proven impossible for textually identical
19339        // kernels, jsonl 2026-07-10). v4 under the threshold; smem twin at/above the smem
19340        // floor (deep-ctx broadcast win); register twin between.
19341        static SMEM_TKV_W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
19342        let smem_tkv = *SMEM_TKV_W.get_or_init(|| {
19343            std::env::var("MEMRA_FA_SMEM_TKV")
19344                .ok()
19345                .and_then(|v| v.parse().ok())
19346                .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
19347        });
19348        // MULTI-ROW v4: resurrected 2026-07-14 (the '33 tok/s collapse' was a paired-map
19349        // partial-write bug, not the mechanism) and falsified HONESTLY at gqa 2: bit-exact
19350        // but −1.7% on the 31B depth cell — the sp helper warp already hides staging
19351        // in-block, and mr trades L2-cheap redundant bytes for serialized per-warp gqa
19352        // score/B3 chains. Arm deleted; jsonl row 2026-07-14 is the record.
19353        use cudarc::driver::sys::CUfunction_attribute_enum as A;
19354        // FP8-WINDOWED (wkv): the v4 family is format-aware (2026-07-12 KFMT/VFMT staging
19355        // arms) — wkv rides the SAME lane logic, resolved from the kf8vf8 module. One symbol
19356        // per (lane, format-module) keeps parity structural; the old register-i2 detour
19357        // (-33%) is retired.
19358        let wg = Self::wkv_on();
19359        // STAGING-PARALLEL v4 (MEMRA_FA_SPW2, default ON at gqa==1): warp 1 = staging helper
19360        // (v4 is 61% staging); score phases identical to v4_w. Same symbol all t.
19361        let sp2 =
19362            gqa <= 4 && fa_v4_at(window) && std::env::var("MEMRA_FA_SPW2").as_deref() != Ok("0");
19363        if sp2 {
19364            let sh = (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32;
19365            if Self::pdl_on() && Self::pdl_wb_on() {
19366                // wave-B2b: flavor mirrors wg.
19367                use cudarc::driver::{DevicePtr, DevicePtrMut};
19368                let s = &self.gpu.stream();
19369                let (pq, _b0) = q.device_ptr(s);
19370                let (pk, _b1) = k.device_ptr(s);
19371                let (pv, _b2) = v.device_ptr(s);
19372                let (po, _b3) = part_o.device_ptr_mut(s);
19373                let (pm, _b4) = part_m.device_ptr_mut(s);
19374                let (pl, _b5) = part_l.device_ptr_mut(s);
19375                let (pb, _b6) = base_dev.device_ptr(s);
19376                let mut ps = [
19377                    &pq as *const _ as *mut std::ffi::c_void,
19378                    &pk as *const _ as *mut _,
19379                    &pv as *const _ as *mut _,
19380                    &po as *const _ as *mut _,
19381                    &pm as *const _ as *mut _,
19382                    &pl as *const _ as *mut _,
19383                    &hd as *const _ as *mut _,
19384                    &nh as *const _ as *mut _,
19385                    &nhkv as *const _ as *mut _,
19386                    &pb as *const _ as *mut _,
19387                    &base_plus as *const _ as *mut _,
19388                    &scale as *const _ as *mut _,
19389                    &nspm as *const _ as *mut _,
19390                    &spk as *const _ as *mut _,
19391                    &ktb as *const _ as *mut _,
19392                    &vtb as *const _ as *mut _,
19393                    &wini as *const _ as *mut _,
19394                ];
19395                unsafe {
19396                    self.launch_pdl_flash(
19397                        wg,
19398                        "fa_decode_vec_q_rows_v4_w_sp",
19399                        (n_head_kv as u32, n_splits_max as u32, t as u32),
19400                        (32, gqa + 1, 1),
19401                        sh,
19402                        &mut ps,
19403                    )?;
19404                }
19405            } else {
19406                let f = if wg {
19407                    self.func_g("fa_decode_vec_q_rows_v4_w_sp")
19408                } else {
19409                    self.func("fa_decode_vec_q_rows_v4_w_sp")
19410                };
19411                f.set_attribute(
19412                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19413                    sh as i32,
19414                )?;
19415                let cfg = LaunchConfig {
19416                    grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
19417                    block_dim: (32, gqa + 1, 1),
19418                    shared_mem_bytes: sh,
19419                };
19420                let __s_b = self.gpu.stream();
19421                let mut b = __s_b.launch_builder(&f);
19422                b.arg(q)
19423                    .arg(k)
19424                    .arg(v)
19425                    .arg(&mut *part_o)
19426                    .arg(&mut *part_m)
19427                    .arg(&mut *part_l)
19428                    .arg(&hd)
19429                    .arg(&nh)
19430                    .arg(&nhkv)
19431                    .arg(base_dev)
19432                    .arg(&base_plus)
19433                    .arg(&scale)
19434                    .arg(&nspm)
19435                    .arg(&spk)
19436                    .arg(&ktb)
19437                    .arg(&vtb)
19438                    .arg(&wini);
19439                unsafe {
19440                    b.launch(cfg)?;
19441                }
19442            }
19443        } else {
19444            if fa_v4_at(window) && Self::pdl_on() && Self::pdl_wb_on() {
19445                // wave-B2b: the v4_w pick only (smem/reg twins stay builder-launched).
19446                let sh = (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32;
19447                use cudarc::driver::{DevicePtr, DevicePtrMut};
19448                let s = &self.gpu.stream();
19449                let (pq, _b0) = q.device_ptr(s);
19450                let (pk, _b1) = k.device_ptr(s);
19451                let (pv, _b2) = v.device_ptr(s);
19452                let (po, _b3) = part_o.device_ptr_mut(s);
19453                let (pm, _b4) = part_m.device_ptr_mut(s);
19454                let (pl, _b5) = part_l.device_ptr_mut(s);
19455                let (pb, _b6) = base_dev.device_ptr(s);
19456                let mut ps = [
19457                    &pq as *const _ as *mut std::ffi::c_void,
19458                    &pk as *const _ as *mut _,
19459                    &pv as *const _ as *mut _,
19460                    &po as *const _ as *mut _,
19461                    &pm as *const _ as *mut _,
19462                    &pl as *const _ as *mut _,
19463                    &hd as *const _ as *mut _,
19464                    &nh as *const _ as *mut _,
19465                    &nhkv as *const _ as *mut _,
19466                    &pb as *const _ as *mut _,
19467                    &base_plus as *const _ as *mut _,
19468                    &scale as *const _ as *mut _,
19469                    &nspm as *const _ as *mut _,
19470                    &spk as *const _ as *mut _,
19471                    &ktb as *const _ as *mut _,
19472                    &vtb as *const _ as *mut _,
19473                    &wini as *const _ as *mut _,
19474                ];
19475                unsafe {
19476                    self.launch_pdl_flash(
19477                        wg,
19478                        "fa_decode_vec_q_rows_v4_w",
19479                        (n_head_kv as u32, n_splits_max as u32, t as u32),
19480                        (32, gqa, 1),
19481                        sh,
19482                        &mut ps,
19483                    )?;
19484                }
19485            } else {
19486                let pick = |name: &str| {
19487                    if wg {
19488                        self.func_g(name)
19489                    } else {
19490                        self.func(name)
19491                    }
19492                };
19493                let (f, sh) = if fa_v4_at(window) {
19494                    let f = pick("fa_decode_vec_q_rows_v4_w");
19495                    (f, (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32)
19496                } else if smem_tkv > 0 && window >= smem_tkv {
19497                    // NOTE: the smem twin's V-stage is still q5_1-hardcoded — unreachable under wkv
19498                    // at the gemma window (v4 covers it); revisit if the smem floor ever drops.
19499                    (
19500                        pick("fa_decode_vec_q_rows_smem_w"),
19501                        (2 * 32 * head_dim * 2) as u32,
19502                    )
19503                } else {
19504                    (pick("fa_decode_vec_q_rows_reg_w"), 0u32)
19505                };
19506                f.set_attribute(
19507                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19508                    sh as i32,
19509                )?;
19510                let cfg = LaunchConfig {
19511                    grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
19512                    block_dim: (32, gqa, 1),
19513                    shared_mem_bytes: sh,
19514                };
19515                let __s_b = self.gpu.stream();
19516                let mut b = __s_b.launch_builder(&f);
19517                b.arg(q)
19518                    .arg(k)
19519                    .arg(v)
19520                    .arg(&mut *part_o)
19521                    .arg(&mut *part_m)
19522                    .arg(&mut *part_l)
19523                    .arg(&hd)
19524                    .arg(&nh)
19525                    .arg(&nhkv)
19526                    .arg(base_dev)
19527                    .arg(&base_plus)
19528                    .arg(&scale)
19529                    .arg(&nspm)
19530                    .arg(&spk)
19531                    .arg(&ktb)
19532                    .arg(&vtb)
19533                    .arg(&wini);
19534                unsafe {
19535                    b.launch(cfg)?;
19536                }
19537            }
19538        }
19539        let cfg2 = LaunchConfig {
19540            grid_dim: (n_head as u32, t as u32, 1),
19541            block_dim: (head_dim as u32, 1, 1),
19542            shared_mem_bytes: 0,
19543        };
19544        if let Some((oq, od)) = q8_out {
19545            // wave-5b port (2026-07-23): q8-emitting combine — the t=1 decode's wo matvec
19546            // consumes the pair directly; the standalone quantize launch folds away.
19547            if Self::pdl_on() && Self::pdl_wb_on() {
19548                // wave-B2: flavor mirrors the builder's wg choice.
19549                use cudarc::driver::{DevicePtr, DevicePtrMut};
19550                let s = &self.gpu.stream();
19551                let (po, _g0) = part_o.device_ptr(s);
19552                let (pm, _g1) = part_m.device_ptr(s);
19553                let (pl, _g2) = part_l.device_ptr(s);
19554                let (pq, _g3) = oq.device_ptr_mut(s);
19555                let (pd, _g4) = od.device_ptr_mut(s);
19556                let mut ps = [
19557                    &po as *const _ as *mut std::ffi::c_void,
19558                    &pm as *const _ as *mut _,
19559                    &pl as *const _ as *mut _,
19560                    &pq as *const _ as *mut _,
19561                    &pd as *const _ as *mut _,
19562                    &hd as *const _ as *mut _,
19563                    &nh as *const _ as *mut _,
19564                    &nspm as *const _ as *mut _,
19565                    &spk as *const _ as *mut _,
19566                    &wini as *const _ as *mut _,
19567                ];
19568                unsafe {
19569                    self.launch_pdl_flash(
19570                        wg,
19571                        "fa_decode_combine_rows_w_q8_1",
19572                        cfg2.grid_dim,
19573                        cfg2.block_dim,
19574                        0,
19575                        &mut ps,
19576                    )?;
19577                }
19578                return Ok(());
19579            }
19580            let fc = if wg {
19581                self.func_g("fa_decode_combine_rows_w_q8_1")
19582            } else {
19583                self.func("fa_decode_combine_rows_w_q8_1")
19584            };
19585            let __s_b2 = self.gpu.stream();
19586            let mut b2 = __s_b2.launch_builder(&fc);
19587            b2.arg(&*part_o)
19588                .arg(&*part_m)
19589                .arg(&*part_l)
19590                .arg(oq)
19591                .arg(od)
19592                .arg(&hd)
19593                .arg(&nh)
19594                .arg(&nspm)
19595                .arg(&spk)
19596                .arg(&wini);
19597            unsafe {
19598                b2.launch(cfg2)?;
19599            }
19600            return Ok(());
19601        }
19602        let fc = if wg {
19603            self.func_g("fa_decode_combine_rows_w")
19604        } else {
19605            self.func("fa_decode_combine_rows_w")
19606        };
19607        let __s_b2 = self.gpu.stream();
19608        let mut b2 = __s_b2.launch_builder(&fc);
19609        b2.arg(&*part_o)
19610            .arg(&*part_m)
19611            .arg(&*part_l)
19612            .arg(o)
19613            .arg(&hd)
19614            .arg(&nh)
19615            .arg(&nspm)
19616            .arg(&spk)
19617            .arg(&wini);
19618        unsafe {
19619            b2.launch(cfg2)?;
19620        }
19621        Ok(())
19622    }
19623
19624    /// ROUND-STREAM stage (c): fa rows with the causal base from a device counter. Two lanes:
19625    /// v3 (qwen stream, fa_v3_active) and v4 (gemma hd256 burst — rows_v4_dc, g-module aware);
19626    /// `t_kv_upper` sizes splits/partials — the same one-sp-for-all-rows approximation class
19627    /// the host rows path already uses (battery-arbitrated); actual per-row bounds derive
19628    /// in-kernel from the counter (+ base_plus, v4 lane only — v3's kernel has no plus arg).
19629    #[allow(clippy::too_many_arguments)]
19630    pub fn fa_decode_rows_dc(
19631        &self,
19632        q: &CudaSlice<f32>,
19633        k: &cudarc::driver::CudaView<u8>,
19634        v: &cudarc::driver::CudaView<u8>,
19635        o: &mut CudaSlice<f32>,
19636        head_dim: usize,
19637        n_head: usize,
19638        n_head_kv: usize,
19639        base_dev: &CudaSlice<i32>,
19640        t_kv_upper: usize,
19641        t: usize,
19642        scale: f32,
19643        k_tok_bytes: usize,
19644        v_tok_bytes: usize,
19645        base_plus: i32,
19646        g: bool,
19647    ) -> Result<(), Box<dyn std::error::Error>> {
19648        let v4 = head_dim == 256 && fa_v4_at(t_kv_upper);
19649        assert!(
19650            v4 || fa_v3_active(head_dim),
19651            "stream fa rows requires the v3 or v4 lane"
19652        );
19653        assert!(v4 || base_plus == 0, "v3_dc kernel takes no plus arg");
19654        if v4 {
19655            let sp = fa_split_keys(t_kv_upper, n_head_kv);
19656            let n_splits_max = (t_kv_upper + sp - 1) / sp;
19657            let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
19658            let (nspm, spk) = (n_splits_max as i32, sp as i32);
19659            let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
19660            let gqa = (n_head / n_head_kv).max(1) as u32;
19661            let o_len = t * n_head * n_splits_max * head_dim;
19662            let ml_len = t * n_head * n_splits_max;
19663            let mut part_guard = self.fa_part_pool.lock().unwrap();
19664            if part_guard
19665                .as_ref()
19666                .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
19667                .unwrap_or(true)
19668            {
19669                // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
19670                // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
19671                // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
19672                // later live allocations land at those addresses, and the next graph REPLAY writes
19673                // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
19674                // output corruption began the burst after the trunk's t_kv growth first realloc'd
19675                // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
19676                // the baked addresses alive (single-stream: eager writes the new buffers, replays
19677                // touch only the old — never concurrently). Doubling growth bounds retired VRAM
19678                // (total retired < final size).
19679                let old = part_guard.take();
19680                let (co, cm) = old
19681                    .as_ref()
19682                    .map(|pp| (pp.0.len(), pp.1.len()))
19683                    .unwrap_or((0, 0));
19684                if let Some(old) = old {
19685                    self.fa_part_retired.lock().unwrap().push(old);
19686                }
19687                if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
19688                    eprintln!(
19689                        "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
19690                        co, o_len, cm, ml_len
19691                    );
19692                }
19693                *part_guard = Some((
19694                    self.alloc_uninit::<f32>(o_len.max(2 * co))?,
19695                    self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
19696                    self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
19697                ));
19698            }
19699            let pg = part_guard.as_mut().unwrap();
19700            self.gpu
19701                .stream()
19702                .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
19703            self.gpu
19704                .stream()
19705                .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
19706            self.gpu
19707                .stream()
19708                .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
19709            let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
19710            let f = if g {
19711                self.func_g("fa_decode_vec_q_rows_v4_dc")
19712            } else {
19713                self.func("fa_decode_vec_q_rows_v4_dc")
19714            };
19715            let sh = (11520 + 32 * head_dim * if g { 1 } else { 2 }) as u32;
19716            use cudarc::driver::sys::CUfunction_attribute_enum as A;
19717            f.set_attribute(
19718                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19719                sh as i32,
19720            )?;
19721            let cfg = LaunchConfig {
19722                grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
19723                block_dim: (32, gqa, 1),
19724                shared_mem_bytes: sh,
19725            };
19726            let __s_b = self.gpu.stream();
19727            let mut b = __s_b.launch_builder(&f);
19728            b.arg(q)
19729                .arg(k)
19730                .arg(v)
19731                .arg(&mut *part_o)
19732                .arg(&mut *part_m)
19733                .arg(&mut *part_l)
19734                .arg(&hd)
19735                .arg(&nh)
19736                .arg(&nhkv)
19737                .arg(base_dev)
19738                .arg(&base_plus)
19739                .arg(&scale)
19740                .arg(&nspm)
19741                .arg(&spk)
19742                .arg(&ktb)
19743                .arg(&vtb);
19744            unsafe {
19745                b.launch(cfg)?;
19746            }
19747            let fc = self.func("fa_decode_combine_rows_dc");
19748            let cfg2 = LaunchConfig {
19749                grid_dim: (n_head as u32, t as u32, 1),
19750                block_dim: (head_dim as u32, 1, 1),
19751                shared_mem_bytes: 0,
19752            };
19753            let __s_b2 = self.gpu.stream();
19754            let mut b2 = __s_b2.launch_builder(&fc);
19755            b2.arg(&*part_o)
19756                .arg(&*part_m)
19757                .arg(&*part_l)
19758                .arg(o)
19759                .arg(&hd)
19760                .arg(&nh)
19761                .arg(base_dev)
19762                .arg(&base_plus)
19763                .arg(&nspm)
19764                .arg(&spk);
19765            unsafe {
19766                b2.launch(cfg2)?;
19767            }
19768            return Ok(());
19769        }
19770        let sp = fa_split_keys(t_kv_upper, n_head_kv);
19771        let n_splits_max = (t_kv_upper + sp - 1) / sp;
19772        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
19773        let (nspm, spk) = (n_splits_max as i32, sp as i32);
19774        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
19775        let gqa = (n_head / n_head_kv).max(1) as u32;
19776        let o_len = t * n_head * n_splits_max * head_dim;
19777        let ml_len = t * n_head * n_splits_max;
19778        let mut part_guard = self.fa_part_pool.lock().unwrap();
19779        if part_guard
19780            .as_ref()
19781            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
19782            .unwrap_or(true)
19783        {
19784            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
19785            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
19786            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
19787            // later live allocations land at those addresses, and the next graph REPLAY writes
19788            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
19789            // output corruption began the burst after the trunk's t_kv growth first realloc'd
19790            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
19791            // the baked addresses alive (single-stream: eager writes the new buffers, replays
19792            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
19793            // (total retired < final size).
19794            let old = part_guard.take();
19795            let (co, cm) = old
19796                .as_ref()
19797                .map(|pp| (pp.0.len(), pp.1.len()))
19798                .unwrap_or((0, 0));
19799            if let Some(old) = old {
19800                self.fa_part_retired.lock().unwrap().push(old);
19801            }
19802            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
19803                eprintln!(
19804                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
19805                    co, o_len, cm, ml_len
19806                );
19807            }
19808            *part_guard = Some((
19809                self.alloc_uninit::<f32>(o_len.max(2 * co))?,
19810                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
19811                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
19812            ));
19813        }
19814        let pg = part_guard.as_mut().unwrap();
19815        self.gpu
19816            .stream()
19817            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
19818        self.gpu
19819            .stream()
19820            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
19821        self.gpu
19822            .stream()
19823            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
19824        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
19825        let f = self.func("fa_decode_vec_q_rows_v3_dc");
19826        let sh = (32 * head_dim * 2) as u32;
19827        use cudarc::driver::sys::CUfunction_attribute_enum as A;
19828        f.set_attribute(
19829            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19830            sh as i32,
19831        )?;
19832        let cfg = LaunchConfig {
19833            grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
19834            block_dim: (32, gqa, 1),
19835            shared_mem_bytes: sh,
19836        };
19837        let __s_b = self.gpu.stream();
19838        let mut b = __s_b.launch_builder(&f);
19839        b.arg(q)
19840            .arg(k)
19841            .arg(v)
19842            .arg(&mut *part_o)
19843            .arg(&mut *part_m)
19844            .arg(&mut *part_l)
19845            .arg(&hd)
19846            .arg(&nh)
19847            .arg(&nhkv)
19848            .arg(base_dev)
19849            .arg(&scale)
19850            .arg(&nspm)
19851            .arg(&spk)
19852            .arg(&ktb)
19853            .arg(&vtb);
19854        unsafe {
19855            b.launch(cfg)?;
19856        }
19857        let fc = self.func("fa_decode_combine_rows_dc");
19858        let cfg2 = LaunchConfig {
19859            grid_dim: (n_head as u32, t as u32, 1),
19860            block_dim: (head_dim as u32, 1, 1),
19861            shared_mem_bytes: 0,
19862        };
19863        let plus0 = 0i32;
19864        let __s_b2 = self.gpu.stream();
19865        let mut b2 = __s_b2.launch_builder(&fc);
19866        b2.arg(&*part_o)
19867            .arg(&*part_m)
19868            .arg(&*part_l)
19869            .arg(o)
19870            .arg(&hd)
19871            .arg(&nh)
19872            .arg(base_dev)
19873            .arg(&plus0)
19874            .arg(&nspm)
19875            .arg(&spk);
19876        unsafe {
19877            b2.launch(cfg2)?;
19878        }
19879        Ok(())
19880    }
19881
19882    /// Device-counter variant of `fa_decode` (CUDA-GRAPH-PLAN Phase 2). The sequence length is read
19883    /// from `t_kv_dev[0]` (resident device i32[1]) for the attention loop bound + per-split key range;
19884    /// the GRID `n_splits` is sized for `bucket_max` (the bucket's max t_kv — baked at capture time).
19885    /// Empty splits (key range beyond the actual t_kv) write an empty partial (m=NEG_INF) so the
19886    /// shared combine skips them -> bit-correct for ANY actual t_kv <= bucket_max.
19887    ///
19888    /// BIT-IDENTITY (the gate): pass `bucket_max == actual_t_kv` and this reproduces `fa_decode`
19889    /// EXACTLY (same n_splits, same per, same split boundaries, same combine) while reading t_kv from
19890    /// device. Bucketing (bucket_max > t_kv) is for the future captured path and changes split
19891    /// grouping (different but mathematically-equal log-sum-exp merge).
19892    pub fn fa_decode_dc(
19893        &self,
19894        q: &CudaSlice<f32>,
19895        k: &cudarc::driver::CudaView<u8>,
19896        v: &cudarc::driver::CudaView<u8>,
19897        o: &mut CudaSlice<f32>,
19898        head_dim: usize,
19899        n_head: usize,
19900        n_head_kv: usize,
19901        t_kv_dev: &CudaSlice<i32>,
19902        bucket_max: usize,
19903        scale: f32,
19904        k_tok_bytes: usize,
19905        v_tok_bytes: usize,
19906        g: bool,
19907    ) -> Result<(), Box<dyn std::error::Error>> {
19908        self.fa_decode_dc_q8(
19909            q,
19910            k,
19911            v,
19912            o,
19913            head_dim,
19914            n_head,
19915            n_head_kv,
19916            t_kv_dev,
19917            bucket_max,
19918            scale,
19919            k_tok_bytes,
19920            v_tok_bytes,
19921            g,
19922            None,
19923        )
19924    }
19925
19926    /// `fa_decode_dc` with an optional q8_1 sink (wave 5b): when `q8_out` is given the
19927    /// combine emits (int8, per-32 scales) for the wo matmul_pre and skips the f32 O write.
19928    #[allow(clippy::too_many_arguments)]
19929    pub fn fa_decode_dc_q8(
19930        &self,
19931        q: &CudaSlice<f32>,
19932        k: &cudarc::driver::CudaView<u8>,
19933        v: &cudarc::driver::CudaView<u8>,
19934        o: &mut CudaSlice<f32>,
19935        head_dim: usize,
19936        n_head: usize,
19937        n_head_kv: usize,
19938        t_kv_dev: &CudaSlice<i32>,
19939        bucket_max: usize,
19940        scale: f32,
19941        k_tok_bytes: usize,
19942        v_tok_bytes: usize,
19943        g: bool,
19944        q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
19945    ) -> Result<(), Box<dyn std::error::Error>> {
19946        // The fa_vec gate + n_splits are sized from bucket_max (host, fixed at capture). The kernel
19947        // reads the ACTUAL t_kv from t_kv_dev for the per-split bound. DEFAULT-ON to MATCH the eager
19948        // `fa_decode` gate above — graph capture must mirror eager's kernel choice or the graph-vs-eager
19949        // bit-identity gate breaks. MEMRA_NO_FA_VEC forces scalar on BOTH paths in lockstep.
19950        // `g` = this layer's cache is e4m3 (gemma windowed under wkv) — every pick below must
19951        // mirror fa_decode_kvmod's g-routing or the graph diverges from eager (short/mid 1/96,
19952        // 2026-07-12).
19953        let mut fa_vec =
19954            std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
19955        if g && head_dim == 256 && !fa_v4_at(bucket_max) {
19956            fa_vec = false;
19957        } // mirror kvmod/geom
19958        let sp = fa_split_keys(bucket_max, n_head_kv);
19959        let n_splits = if fa_vec {
19960            ((bucket_max + sp - 1) / sp).max(1)
19961        } else {
19962            ((bucket_max + 255) / 256).max(1)
19963        };
19964        let o_len = n_head * n_splits * head_dim;
19965        let ml_len = n_head * n_splits;
19966        let mut part_guard = self.fa_part_pool.lock().unwrap();
19967        if part_guard
19968            .as_ref()
19969            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
19970            .unwrap_or(true)
19971        {
19972            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
19973            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
19974            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
19975            // later live allocations land at those addresses, and the next graph REPLAY writes
19976            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
19977            // output corruption began the burst after the trunk's t_kv growth first realloc'd
19978            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
19979            // the baked addresses alive (single-stream: eager writes the new buffers, replays
19980            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
19981            // (total retired < final size).
19982            let old = part_guard.take();
19983            let (co, cm) = old
19984                .as_ref()
19985                .map(|pp| (pp.0.len(), pp.1.len()))
19986                .unwrap_or((0, 0));
19987            if let Some(old) = old {
19988                self.fa_part_retired.lock().unwrap().push(old);
19989            }
19990            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
19991                eprintln!(
19992                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
19993                    co, o_len, cm, ml_len
19994                );
19995            }
19996            *part_guard = Some((
19997                self.alloc_uninit::<f32>(o_len.max(2 * co))?,
19998                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
19999                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
20000            ));
20001        }
20002        let pg = part_guard.as_mut().unwrap();
20003        self.gpu
20004            .stream()
20005            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
20006        self.gpu
20007            .stream()
20008            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
20009        self.gpu
20010            .stream()
20011            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
20012        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
20013        let (hd, nh, nhkv, nsp) = (
20014            head_dim as i32,
20015            n_head as i32,
20016            n_head_kv as i32,
20017            n_splits as i32,
20018        );
20019        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
20020        let fa_vec = fa_vec && head_dim <= 512 && head_dim % 32 == 0;
20021        // FA-DEEP pick keyed on bucket_max (the fa_v4_at precedent) — bit-identical twins,
20022        // so a threshold falling between t_kv and bucket_max cannot diverge eager-vs-graph.
20023        let deep = fa_vec
20024            && head_dim == 256
20025            && fa_v4_at(bucket_max)
20026            && !g
20027            && fa_deep_at(bucket_max)
20028            && !matches!(fa_v4_mode(), "noB3" | "stage");
20029        let (f, cfg) = if fa_vec
20030            && head_dim == 512
20031            && bucket_max >= {
20032                static FA512_MIN_DC: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
20033                *FA512_MIN_DC.get_or_init(|| {
20034                    std::env::var("MEMRA_FA512_MIN")
20035                        .ok()
20036                        .and_then(|v| v.parse().ok())
20037                        .unwrap_or(512)
20038                })
20039            } {
20040            // gemma globals dc twin (mirror the eager dpl16 pick incl the crossover floor).
20041            let gqa = (n_head / n_head_kv).max(1) as u32;
20042            (
20043                self.fa_func("fa_decode_vec_q_dpl16_dc", head_dim),
20044                LaunchConfig {
20045                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
20046                    block_dim: (32, gqa, 1),
20047                    shared_mem_bytes: 0,
20048                },
20049            )
20050        } else if fa_vec && head_dim == 512 {
20051            // under the 512 floor eager runs scalar — the SAME unified symbol, ctr non-null;
20052            // ns_eff in-kernel reproduces eager's ceil(t_kv/sp) partition for the LIVE len.
20053            let q_view = q.as_view();
20054            let mut o_view = o.as_view_mut();
20055            return self.fa_decode_scalar_unified(
20056                &q_view,
20057                k,
20058                v,
20059                &mut o_view,
20060                head_dim,
20061                n_head,
20062                n_head_kv,
20063                0,
20064                Some(t_kv_dev),
20065                scale,
20066                n_splits,
20067                sp,
20068                k_tok_bytes,
20069                v_tok_bytes,
20070                g,
20071                &mut *part_o,
20072                &mut *part_m,
20073                &mut *part_l,
20074                q8_out,
20075            );
20076        } else if fa_vec && head_dim == 256 && fa_v4_at(bucket_max) {
20077            // gemma/qwen v4 dc twin (eager default lane) — capture must mirror eager's pick,
20078            // incl the g-module route + raw-e4m3 sV sizing.
20079            let gqa = (n_head / n_head_kv).max(1) as u32;
20080            let fv = if g {
20081                self.func_g("fa_decode_vec_q_v4_dc")
20082            } else if deep {
20083                self.func("fa_decode_vec_q_v4_deep_dc")
20084            } else {
20085                self.func("fa_decode_vec_q_v4_dc")
20086            };
20087            let shmem =
20088                (if deep { 12160 } else { 11520 } + 32 * head_dim * if g { 1 } else { 2 }) as u32;
20089            use cudarc::driver::sys::CUfunction_attribute_enum as A;
20090            fv.set_attribute(
20091                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20092                shmem as i32,
20093            )?;
20094            (
20095                fv,
20096                LaunchConfig {
20097                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
20098                    block_dim: (32, gqa, 1),
20099                    shared_mem_bytes: shmem,
20100                },
20101            )
20102        } else if fa_vec && fa_v3_active(head_dim) {
20103            // FA v3 lane _dc twin: the captured graph must run the SAME walk body as eager
20104            // under MEMRA_FA_V3=1 (eager, rows-verify and graph switch together).
20105            let gqa = (n_head / n_head_kv).max(1) as u32;
20106            let fv = if g {
20107                self.func_g("fa_decode_vec_q_v3_dc")
20108            } else {
20109                self.func("fa_decode_vec_q_v3_dc")
20110            };
20111            let shmem = (32 * head_dim * 2) as u32; // sV bf16 [FA_DEC_TILE=32][hd]
20112            (
20113                fv,
20114                LaunchConfig {
20115                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
20116                    block_dim: (32, gqa, 1),
20117                    shared_mem_bytes: shmem,
20118                },
20119            )
20120        } else if fa_vec && fa_v2_on() {
20121            // FAVENDOR lane: v2 _dc twin — the captured graph must run the SAME walk body as
20122            // eager under MEMRA_FA_V2=1 or graph_decode_gate's bit-identity breaks (the flag is
20123            // a numeric config; eager, rows-verify and graph all switch together).
20124            let gqa = (n_head / n_head_kv).max(1) as u32;
20125            let fv = if g {
20126                self.func_g("fa_decode_vec_q_v2_dc")
20127            } else {
20128                self.func("fa_decode_vec_q_v2_dc")
20129            };
20130            let shmem = (2 * 32 * head_dim * 2) as u32; // sK+sV bf16 [FA_DEC_TILE=32][hd]
20131            (
20132                fv,
20133                LaunchConfig {
20134                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
20135                    block_dim: (32, gqa, 1),
20136                    shared_mem_bytes: shmem,
20137                },
20138            )
20139        } else if fa_vec {
20140            let gqa = (n_head / n_head_kv).max(1) as u32;
20141            // REGISTER-DEQUANT twin: zero dynamic smem (see fa_decode above).
20142            let fv = if g {
20143                self.func_g("fa_decode_vec_q_dc")
20144            } else {
20145                self.func("fa_decode_vec_q_dc")
20146            };
20147            (
20148                fv,
20149                LaunchConfig {
20150                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
20151                    block_dim: (32, gqa, 1),
20152                    shared_mem_bytes: 0,
20153                },
20154            )
20155        } else {
20156            let q_view = q.as_view();
20157            let mut o_view = o.as_view_mut();
20158            return self.fa_decode_scalar_unified(
20159                &q_view,
20160                k,
20161                v,
20162                &mut o_view,
20163                head_dim,
20164                n_head,
20165                n_head_kv,
20166                0,
20167                Some(t_kv_dev),
20168                scale,
20169                n_splits,
20170                if fa_vec { sp } else { 256 },
20171                k_tok_bytes,
20172                v_tok_bytes,
20173                g,
20174                &mut *part_o,
20175                &mut *part_m,
20176                &mut *part_l,
20177                q8_out,
20178            );
20179        };
20180        let ski = sp as i32; // one-partition law: the twins derive ns_eff from (T_kv, ski)
20181        let __s_b = self.gpu.stream();
20182        let mut b = __s_b.launch_builder(&f);
20183        b.arg(q)
20184            .arg(k)
20185            .arg(v)
20186            .arg(&mut *part_o)
20187            .arg(&mut *part_m)
20188            .arg(&mut *part_l)
20189            .arg(&hd)
20190            .arg(&nh)
20191            .arg(&nhkv)
20192            .arg(t_kv_dev)
20193            .arg(&scale)
20194            .arg(&nsp)
20195            .arg(&ski)
20196            .arg(&ktb)
20197            .arg(&vtb);
20198        unsafe {
20199            b.launch(cfg)?;
20200        }
20201        let cfg2 = LaunchConfig {
20202            grid_dim: (n_head as u32, 1, 1),
20203            block_dim: (head_dim as u32, 1, 1),
20204            shared_mem_bytes: 0,
20205        };
20206        if let Some((oq, od)) = q8_out {
20207            let fc = if g {
20208                self.func_g("fa_decode_combine_q8_1")
20209            } else {
20210                self.fa_func("fa_decode_combine_q8_1", head_dim)
20211            };
20212            let __s_b2 = self.gpu.stream();
20213            let mut b2 = __s_b2.launch_builder(&fc);
20214            b2.arg(&*part_o)
20215                .arg(&*part_m)
20216                .arg(&*part_l)
20217                .arg(oq)
20218                .arg(od)
20219                .arg(&hd)
20220                .arg(&nh)
20221                .arg(&nsp);
20222            unsafe {
20223                b2.launch(cfg2)?;
20224            }
20225            return Ok(());
20226        }
20227        let fc = if g {
20228            self.func_g("fa_decode_combine_f32")
20229        } else {
20230            self.fa_func("fa_decode_combine_f32", head_dim)
20231        };
20232        let __s_b2 = self.gpu.stream();
20233        let mut b2 = __s_b2.launch_builder(&fc);
20234        b2.arg(&*part_o)
20235            .arg(&*part_m)
20236            .arg(&*part_l)
20237            .arg(o)
20238            .arg(&hd)
20239            .arg(&nh)
20240            .arg(&nsp);
20241        unsafe {
20242            b2.launch(cfg2)?;
20243        }
20244        Ok(())
20245    }
20246
20247    /// EAGER fa_decode geometry for a given actual `t_kv` (CUDA-GRAPH-PLAN §3.3 bucketing). Returns
20248    /// `(fa_vec, n_splits)` EXACTLY as `fa_decode` computes them so the graph-capture path can key its
20249    /// bucket on the same `(kernel, n_splits)` pair and pass a `bucket_max` that reproduces eager's
20250    /// n_splits bit-for-bit. (Per = ceil(t_kv/n_splits) is then recomputed from the DEVICE t_kv inside
20251    /// the kernel and matches eager when n_splits matches — the bit-identity contract.)
20252    pub fn fa_geom_eager(
20253        &self,
20254        t_kv: usize,
20255        head_dim: usize,
20256        n_head_kv: usize,
20257        g: bool,
20258    ) -> (bool, usize) {
20259        // MUST mirror `fa_decode` / `fa_decode_dc` (default-ON 2026-06-28). This is the bucket-key
20260        // source: if it disagrees with the actual kernel pick, the graph captures the wrong path and
20261        // replay diverges from eager. All three sites read MEMRA_NO_FA_VEC in lockstep.
20262        let fa_ok = std::env::var("MEMRA_NO_FA_VEC").is_err() && t_kv >= fa_vec_min_tkv();
20263        // hd512 dpl16 vec lane (gemma globals, 2026-07-11 graph-arc fix): the original key
20264        // hardcoded vec = hd<=256, so for hd512 it bucketed by the SCALAR 256-key splits while
20265        // the dpl16/rows_dpl16 kernels split by the ladder — n_splits changed WITHIN a bucket
20266        // (mid-ctx graph mismatch at pos 19 + partials OOB at longer runs). Mirror the real
20267        // fa_decode dispatch: vec512 above the fa512 floor, vec256 as before.
20268        let vec512 = fa_ok && head_dim == 512 && t_kv >= fa512_min_tkv();
20269        let mut fa_vec = vec512 || (fa_ok && head_dim <= 256 && head_dim % 32 == 0);
20270        // g (fp8-windowed): mirror kvmod's clamp — only the v4 lane parses e4m3 in the vec
20271        // family; everything else falls to the g-module scalar.
20272        // hd256 under g: v4-or-scalar. hd128/other under g ride the REGISTER g-lane (the
20273        // dq_K_lane/dq_V_lane macros are format-aware; v3 is format-gated off, v2/smem
20274        // arms are excluded under g). Mirrored in kvmod / fa_decode_dc / fa_geom_eager.
20275        if g && head_dim == 256 && !fa_v4_at(t_kv) {
20276            fa_vec = false;
20277        }
20278        let sp = fa_split_keys(t_kv, n_head_kv);
20279        let n_splits = if fa_vec {
20280            ((t_kv + sp - 1) / sp).max(1)
20281        } else {
20282            ((t_kv + 255) / 256).max(1)
20283        };
20284        (fa_vec, n_splits)
20285    }
20286
20287    /// `bucket_max` (host t_kv to feed `fa_decode_dc` / `full_attn_decode_dc`) that makes the _dc
20288    /// kernel pick the SAME (fa_vec, n_splits) as eager would for actual `t_kv`. Because the dc
20289    /// launcher derives both from `bucket_max` via the same formulas, we just hand it `t_kv` itself:
20290    /// the n_splits is then identical, and the per-split boundaries (computed from the DEVICE t_kv in
20291    /// the kernel) match eager exactly. The bucket KEY (for the graph HashMap) is `(fa_vec, n_splits)`.
20292    pub fn fa_bucket_key(
20293        &self,
20294        t_kv: usize,
20295        head_dim: usize,
20296        n_head_kv: usize,
20297        g: bool,
20298    ) -> (bool, usize) {
20299        self.fa_geom_eager(t_kv, head_dim, n_head_kv, g)
20300    }
20301
20302    /// CUDA-graph capture wrapper (CUDA-GRAPH-PLAN §3.2, llama.cpp warmup pattern). Runs `step`
20303    /// inline TWICE (warmup — lets the caching allocator settle to stable pointers and any one-time
20304    /// kernel attribute/JIT happen outside capture), then captures a THIRD invocation on the Engine's
20305    /// decode stream (RELAXED mode) and instantiates it into a replayable `CudaGraph`. The closure
20306    /// must enqueue ONLY device work on `e.stream()` (no dtoh / no synchronize / no host branch on
20307    /// device data) — every per-step varying scalar must come from a device counter. Returns the
20308    /// instantiated graph; `CudaGraph::launch()` replays the whole step in one dispatch.
20309    /// `capture_graph` with CAPTURE-RETAIN: every Engine allocation made during the warmups
20310    /// and the capture is kept alive in the returned keeper — hold it as long as the graph
20311    /// replays (transients returning to the pool get reused by unrelated work and corrupt
20312    /// replays; the draft-graph root cause). Model-generic, next capture reuses it.
20313    pub fn capture_graph_retained<F>(
20314        &self,
20315        step: F,
20316    ) -> Result<
20317        (
20318            cudarc::driver::CudaGraph,
20319            Vec<Box<dyn std::any::Any + Send>>,
20320        ),
20321        Box<dyn std::error::Error>,
20322    >
20323    where
20324        F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
20325    {
20326        use cudarc::driver::sys::CUgraphInstantiate_flags;
20327        self.capture_graph_retained_flags(
20328            CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
20329            step,
20330        )
20331    }
20332
20333    /// Retained capture with an explicit instantiate flag. ALLOC-FREE captured graphs
20334    /// (zero mem nodes — the gemma slotted door) should pass UPLOAD instead of
20335    /// AUTO_FREE_ON_LAUNCH: the auto-free flag's launch-time mem-pool scan was measured at
20336    /// ~0.25us/node (205us on the 826-node step) even with nothing to free.
20337    pub fn capture_graph_retained_flags<F>(
20338        &self,
20339        flags: cudarc::driver::sys::CUgraphInstantiate_flags,
20340        mut step: F,
20341    ) -> Result<
20342        (
20343            cudarc::driver::CudaGraph,
20344            Vec<Box<dyn std::any::Any + Send>>,
20345        ),
20346        Box<dyn std::error::Error>,
20347    >
20348    where
20349        F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
20350    {
20351        use cudarc::driver::sys::CUstreamCaptureMode;
20352        // KEEP scope = WARMUPS ONLY (2026-07-13): keep_if_capturing retains via
20353        // CudaSlice::clone, which is a device ALLOC + D2D COPY on the stream — clones made
20354        // while the capture region is open become dead copy NODES replayed every launch
20355        // (E4B: 1440 copies = 0.74ms/token, the whole graph-vs-eager regression). The
20356        // warmup runs allocate the same transient sequence at the same pool addresses, so
20357        // retaining the warmup clones preserves the draft-graph fix without polluting the
20358        // captured graph.
20359        self.capture_keep.lock().unwrap().clear();
20360        let was_tracking = self.gpu.ctx.is_event_tracking();
20361        if was_tracking {
20362            unsafe {
20363                self.gpu.ctx.disable_event_tracking();
20364            }
20365        }
20366        let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
20367            self.capture_keep_on
20368                .store(true, std::sync::atomic::Ordering::Relaxed);
20369            let w = (|| {
20370                step(self)?;
20371                step(self)
20372            })();
20373            self.capture_keep_on
20374                .store(false, std::sync::atomic::Ordering::Relaxed);
20375            w?;
20376            self.gpu.stream().synchronize()?;
20377            self.gpu
20378                .stream()
20379                .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
20380            let r = step(self);
20381            let g = self.gpu.stream().end_capture(flags);
20382            r?;
20383            let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
20384            graph.upload()?;
20385            Ok(graph)
20386        };
20387        let result = run();
20388        self.capture_keep_on
20389            .store(false, std::sync::atomic::Ordering::Relaxed);
20390        if was_tracking {
20391            unsafe {
20392                self.gpu.ctx.enable_event_tracking();
20393            }
20394        }
20395        let keeper = std::mem::take(&mut *self.capture_keep.lock().unwrap());
20396        Ok((result?, keeper))
20397    }
20398
20399    pub fn capture_graph<F>(
20400        &self,
20401        mut step: F,
20402    ) -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>>
20403    where
20404        F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
20405    {
20406        use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
20407        // EVENT TRACKING OFF for capture. The Engine creates a 2nd stream (copy_stream) so cudarc is in
20408        // multi-stream mode and, by default, records a CudaEvent per CudaSlice alloc/use to serialize
20409        // cross-stream access. Those per-buffer event waits issue stream ops that are NOT permitted
20410        // inside a capture region (CUDA_ERROR_STREAM_CAPTURE_UNSUPPORTED). The captured decode step is
20411        // strictly SINGLE-STREAM (every kernel on gpu.stream), so this synchronization is unnecessary
20412        // here — disable it for the whole warmup+capture, re-enable after. SAFETY: the decode-dc path
20413        // touches only gpu.stream; no buffer crosses to copy_stream during capture.
20414        let was_tracking = self.gpu.ctx.is_event_tracking();
20415        if was_tracking {
20416            unsafe {
20417                self.gpu.ctx.disable_event_tracking();
20418            }
20419        }
20420        // Q1 PROBE (MEMRA_GRAPH_IFLAG): the generic capture body's cuMemAllocAsync nodes are
20421        // EXACTLY BALANCED by in-graph free nodes (measured census q27: 1589 ALLOC / 1589
20422        // FREE), so AUTO_FREE_ON_LAUNCH has nothing to reclaim at launch — it only pays its
20423        // per-node launch-time mem-pool scan. `upload` / `none` select the alternatives to
20424        // measure that scan's real cost on the generic path. Diagnostic door only; the
20425        // default stays AUTO_FREE until a measured A/B justifies moving it.
20426        let iflag = {
20427            static F: std::sync::OnceLock<CUgraphInstantiate_flags> = std::sync::OnceLock::new();
20428            *F.get_or_init(|| match std::env::var("MEMRA_GRAPH_IFLAG").as_deref() {
20429                // UPLOAD = the gemma slotted door's zero-mem-node choice; PRIORITY = the flag
20430                // hybrid_forward.rs:5935 actually ships (both drop the auto-free launch scan).
20431                Ok("upload") => CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_UPLOAD,
20432                Ok("priority") => {
20433                    CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY
20434                }
20435                _ => CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
20436            })
20437        };
20438        // MEMRA_GRAPH_CAPTIME=1 (Q1 lane): phase-resolved capture cost. Recapture is paid at
20439        // every kernel-class crossing, so it — not steady-state decode — is the quantity a
20440        // mem-node reduction could plausibly shrink. Only `instantiate` (cuStreamEndCapture +
20441        // cuGraphInstantiateWithFlags) and `upload` scale with node count; the warmups are
20442        // eager step executions and are node-count-invariant. Printing the split bounds the
20443        // refactor's ceiling instead of assuming it.
20444        let ct = {
20445            static T: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20446            *T.get_or_init(|| std::env::var("MEMRA_GRAPH_CAPTIME").as_deref() == Ok("1"))
20447        };
20448        // MEMRA_GRAPH_WARMUPS (Q1 lane; DEFAULT 1 since lane/graph-warmups 2026-08-05): the
20449        // phase split showed the eager warmups are 80% of recapture cost (q27 27.4 of 34.4 ms
20450        // pod / 42% of 52.6 ms 5090) — 3x larger than the ENTIRE mem-node ceiling the audit
20451        // chased, and node-count-invariant, so no capture-body refactor could touch it.
20452        // Warmup 2's theorized job was async-pool ADDRESS STABILITY: warmup 1's allocs may
20453        // grow/map the pool, warmup 2 re-walks the same sequence over the freed blocks so the
20454        // captured third run bakes settled addresses. That hazard is the #68 stale-baked-
20455        // address class — which the engine now guards STRUCTURALLY rather than by re-walking:
20456        // in-body transients are captured as BALANCED in-graph alloc/free node pairs (census
20457        // 1589/1589 — replays allocate for themselves; no baked transient pointers), every
20458        // externally-referenced buffer is stable-pointer by design (fa_part_pool retires-on-
20459        // grow and never frees, resident counters/scratch, cache set in place), and the
20460        // draft-graph path additionally rides capture_graph_retained (capture_keep holds all
20461        // warmup+capture allocs alive). One warmup therefore suffices for kernel-attr
20462        // settling and pool mapping. Arbitrated adversarially, not by taste:
20463        // graph-warmup-stress (pool-growth cycles large<->small x10, overlap arm, forced
20464        // recaptures over freed blocks — bit-identity vs eager + canary teeth) is GREEN at
20465        // warmups=1 on the deployment rig, plus graph-decode-gate 256-step bit-identity,
20466        // graph-session-gate, run-spec K=1..8 (receipts research/graph-warmups-5090-20260805/
20467        // + the pod's research/graph-allocfree-20260805/). Measured: recapture -38..-42% q27 /
20468        // -41% q9, decode +~1%, capture+prime -13ms. MEMRA_GRAPH_WARMUPS=2 = the rollback
20469        // seam; tools/graph-warmup-stress-gate.sh = the gate any regression re-runs.
20470        let warmups = {
20471            static W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
20472            *W.get_or_init(|| {
20473                std::env::var("MEMRA_GRAPH_WARMUPS")
20474                    .ok()
20475                    .and_then(|v| v.parse().ok())
20476                    .filter(|n| *n >= 1)
20477                    .unwrap_or(1)
20478            })
20479        };
20480        let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
20481            let t_w = std::time::Instant::now();
20482            // warmup: inline runs (no capture) so allocator pointers + kernel attrs are stable.
20483            for _ in 0..warmups {
20484                step(self)?;
20485            }
20486            self.gpu.stream().synchronize()?;
20487            let ms_warm = t_w.elapsed().as_secs_f64() * 1e3;
20488            // capture the third run.
20489            let t_c = std::time::Instant::now();
20490            self.gpu
20491                .stream()
20492                .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
20493            // If the body errors mid-capture, end the capture before propagating so the stream isn't
20494            // left in a capturing state.
20495            let r = step(self);
20496            let ms_body = t_c.elapsed().as_secs_f64() * 1e3;
20497            let t_i = std::time::Instant::now();
20498            let g = self.gpu.stream().end_capture(iflag);
20499            let ms_inst = t_i.elapsed().as_secs_f64() * 1e3;
20500            r?;
20501            let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
20502            let t_u = std::time::Instant::now();
20503            graph.upload()?;
20504            if ct {
20505                println!(
20506                    "[graph-captime] warmup2x {ms_warm:.2} ms  capture-body {ms_body:.2} ms  \
20507                          instantiate {ms_inst:.2} ms  upload {:.2} ms",
20508                    t_u.elapsed().as_secs_f64() * 1e3
20509                );
20510            }
20511            Ok(graph)
20512        };
20513        let result = run();
20514        if was_tracking {
20515            unsafe {
20516                self.gpu.ctx.enable_event_tracking();
20517            }
20518        }
20519        result
20520    }
20521
20522    /// gdn_scan variant where state_in/out are CudaViews (resident SSM state, in-place per step).
20523    pub fn gdn_scan_s128_view(
20524        &self,
20525        q: &CudaSlice<f32>,
20526        k: &CudaSlice<f32>,
20527        v: &CudaSlice<f32>,
20528        g: &CudaSlice<f32>,
20529        beta: &CudaSlice<f32>,
20530        state_in: &cudarc::driver::CudaView<f32>,
20531        state_out: &mut cudarc::driver::CudaViewMut<f32>,
20532        o: &mut CudaSlice<f32>,
20533        n_head: usize,
20534        t: usize,
20535        scale: f32,
20536    ) -> Result<(), Box<dyn std::error::Error>> {
20537        let f = self.func("gdn_scan_s128");
20538        const S_V: u32 = 128;
20539        const WARP: u32 = 32;
20540        const COLS: u32 = 4;
20541        let cfg = LaunchConfig {
20542            grid_dim: (n_head as u32, 1, S_V / COLS),
20543            block_dim: (WARP, COLS, 1),
20544            shared_mem_bytes: 0,
20545        };
20546        let (h, ti) = (n_head as i32, t as i32);
20547        let __s_b = self.gpu.stream();
20548        let mut b = __s_b.launch_builder(&f);
20549        b.arg(q)
20550            .arg(k)
20551            .arg(v)
20552            .arg(g)
20553            .arg(beta)
20554            .arg(state_in)
20555            .arg(state_out)
20556            .arg(o)
20557            .arg(&h)
20558            .arg(&ti)
20559            .arg(&scale);
20560        unsafe {
20561            b.launch(cfg)?;
20562        }
20563        Ok(())
20564    }
20565
20566    /// conv1d where the input is a CudaView (resident conv state assembled in place).
20567    pub fn ssm_conv1d_view(
20568        &self,
20569        x: &cudarc::driver::CudaView<f32>,
20570        w: &CudaSlice<f32>,
20571        y: &mut CudaSlice<f32>,
20572        conv_dim: usize,
20573        t: usize,
20574        d_conv: usize,
20575        silu: bool,
20576    ) -> Result<(), Box<dyn std::error::Error>> {
20577        let f = self.func("ssm_conv1d_silu_f32");
20578        // grid.x = channel, grid.y = T-tiles (block 256 strides over T) — parallel over both axes.
20579        let cfg = LaunchConfig {
20580            grid_dim: (conv_dim as u32, ((t as u32 + 255) / 256).max(1), 1),
20581            block_dim: (256, 1, 1),
20582            shared_mem_bytes: 0,
20583        };
20584        let (cd, ti, dc, s) = (conv_dim as i32, t as i32, d_conv as i32, silu as i32);
20585        let __s_b = self.gpu.stream();
20586        let mut b = __s_b.launch_builder(&f);
20587        b.arg(x).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc).arg(&s);
20588        unsafe {
20589            b.launch(cfg)?;
20590        }
20591        Ok(())
20592    }
20593
20594    /// Depthwise causal conv1d + optional SiLU.
20595    /// x:[conv_dim, T+d_conv-1] channel-major (first d_conv-1 cols = carried state),
20596    /// w:[d_conv, conv_dim] kernel-major, y:[conv_dim, T] channel-major.
20597    /// FUSED prefill conv (token-major input, zero left-state): replaces
20598    /// transpose + zeros + conv_left_pad + ssm_conv1d with ONE launch reading the matmul output
20599    /// directly. Output channel-major [conv_dim, T], SiLU applied. BIT-IDENTICAL accumulation.
20600    pub fn ssm_conv1d_tm(
20601        &self,
20602        qkv_tm: &CudaSlice<f32>,
20603        w: &CudaSlice<f32>,
20604        y: &mut CudaSlice<f32>,
20605        conv_dim: usize,
20606        t: usize,
20607        d_conv: usize,
20608    ) -> Result<(), Box<dyn std::error::Error>> {
20609        let f = self.func("ssm_conv1d_tm_f32");
20610        let cfg = LaunchConfig {
20611            grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
20612            block_dim: (256, 1, 1),
20613            shared_mem_bytes: 0,
20614        };
20615        let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
20616        let __s_b = self.gpu.stream();
20617        let mut b = __s_b.launch_builder(&f);
20618        b.arg(qkv_tm).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc);
20619        unsafe {
20620            b.launch(cfg)?;
20621        }
20622        Ok(())
20623    }
20624
20625    /// BATCHED verify conv (T>1, carried state): window reads the resident conv ring for
20626    /// negative rows; separate ring-update launch afterwards. BIT-IDENTICAL per value to the
20627    /// T=1 chain. T >= pad rides the pure input-column ring update (unchanged legacy path);
20628    /// T < pad (the MEMRA_SPEC_M2 t=2 verify arm) needs old-ring sources for the roll — the
20629    /// update kernel would race reading the ring it rewrites, so that arm clones the ring
20630    /// (dtod) and rolls via ssm_conv_ring_rebuild (PURE COPIES: the ring stores raw input
20631    /// columns; the final ring == what T sequential decode ring rolls leave).
20632    pub fn ssm_conv1d_tm_state(
20633        &self,
20634        qkv_tm: &CudaSlice<f32>,
20635        conv_state: &mut CudaSlice<f32>,
20636        w: &CudaSlice<f32>,
20637        y: &mut CudaSlice<f32>,
20638        conv_dim: usize,
20639        t: usize,
20640        d_conv: usize,
20641    ) -> Result<(), Box<dyn std::error::Error>> {
20642        self.ssm_conv1d_tm_state_pad(qkv_tm, conv_state, w, y, conv_dim, t, d_conv, None)
20643    }
20644
20645    /// task #14: `pad_len` = device true length for PADDED prime graphs — the ring update
20646    /// reads rows [len-pad, len) instead of the pad tail. None = the classic host-T path.
20647    #[allow(clippy::too_many_arguments)]
20648    pub fn ssm_conv1d_tm_state_pad(
20649        &self,
20650        qkv_tm: &CudaSlice<f32>,
20651        conv_state: &mut CudaSlice<f32>,
20652        w: &CudaSlice<f32>,
20653        y: &mut CudaSlice<f32>,
20654        conv_dim: usize,
20655        t: usize,
20656        d_conv: usize,
20657        pad_len: Option<&CudaSlice<i32>>,
20658    ) -> Result<(), Box<dyn std::error::Error>> {
20659        assert!(t >= 1, "ssm_conv1d_tm_state requires T >= 1");
20660        // clone BEFORE the window kernel is issued is not required (stream-ordered: the dtod and
20661        // the window kernel both read the pre-roll ring; the roll launches after both) — but
20662        // cloning first keeps the ordering trivially correct under any future stream split.
20663        let ring_old = if t < d_conv - 1 {
20664            Some(self.clone_dtod(conv_state)?)
20665        } else {
20666            None
20667        };
20668        {
20669            let f = self.func("ssm_conv1d_tm_state_f32");
20670            let cfg = LaunchConfig {
20671                grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
20672                block_dim: (256, 1, 1),
20673                shared_mem_bytes: 0,
20674            };
20675            let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
20676            let __s_b = self.gpu.stream();
20677            let mut b = __s_b.launch_builder(&f);
20678            b.arg(qkv_tm)
20679                .arg(&*conv_state)
20680                .arg(w)
20681                .arg(y)
20682                .arg(&cd)
20683                .arg(&ti)
20684                .arg(&dc);
20685            unsafe {
20686                b.launch(cfg)?;
20687            }
20688        }
20689        match (ring_old, pad_len) {
20690            (None, Some(len_d)) => {
20691                let f = self.func("ssm_conv_ring_update_dev_f32");
20692                let n = conv_dim * (d_conv - 1);
20693                let cfg = LaunchConfig::for_num_elems(n as u32);
20694                let (cd, dc) = (conv_dim as i32, d_conv as i32);
20695                let __s_b = self.gpu.stream();
20696                let mut b = __s_b.launch_builder(&f);
20697                b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
20698                unsafe {
20699                    b.launch(cfg)?;
20700                }
20701            }
20702            (None, None) => {
20703                let f = self.func("ssm_conv_ring_update_f32");
20704                let n = conv_dim * (d_conv - 1);
20705                let cfg = LaunchConfig::for_num_elems(n as u32);
20706                let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
20707                let __s_b = self.gpu.stream();
20708                let mut b = __s_b.launch_builder(&f);
20709                b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
20710                unsafe {
20711                    b.launch(cfg)?;
20712                }
20713            }
20714            (Some(old), _) => {
20715                self.ssm_conv_ring_rebuild(qkv_tm, &old, conv_state, conv_dim, t, d_conv)?
20716            }
20717        }
20718        Ok(())
20719    }
20720
20721    /// qkv-view twin (task #16): batched prime reads the concat GEMM output directly.
20722    pub fn ssm_conv1d_tm_state_pad_v(
20723        &self,
20724        qkv_tm: &cudarc::driver::CudaView<f32>,
20725        conv_state: &mut CudaSlice<f32>,
20726        w: &CudaSlice<f32>,
20727        y: &mut CudaSlice<f32>,
20728        conv_dim: usize,
20729        t: usize,
20730        d_conv: usize,
20731        pad_len: Option<&CudaSlice<i32>>,
20732    ) -> Result<(), Box<dyn std::error::Error>> {
20733        assert!(t >= 1, "ssm_conv1d_tm_state requires T >= 1");
20734        // clone BEFORE the window kernel is issued is not required (stream-ordered: the dtod and
20735        // the window kernel both read the pre-roll ring; the roll launches after both) — but
20736        // cloning first keeps the ordering trivially correct under any future stream split.
20737        let ring_old = if t < d_conv - 1 {
20738            Some(self.clone_dtod(conv_state)?)
20739        } else {
20740            None
20741        };
20742        {
20743            let f = self.func("ssm_conv1d_tm_state_f32");
20744            let cfg = LaunchConfig {
20745                grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
20746                block_dim: (256, 1, 1),
20747                shared_mem_bytes: 0,
20748            };
20749            let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
20750            let __s_b = self.gpu.stream();
20751            let mut b = __s_b.launch_builder(&f);
20752            b.arg(qkv_tm)
20753                .arg(&*conv_state)
20754                .arg(w)
20755                .arg(y)
20756                .arg(&cd)
20757                .arg(&ti)
20758                .arg(&dc);
20759            unsafe {
20760                b.launch(cfg)?;
20761            }
20762        }
20763        match (ring_old, pad_len) {
20764            (None, Some(len_d)) => {
20765                let f = self.func("ssm_conv_ring_update_dev_f32");
20766                let n = conv_dim * (d_conv - 1);
20767                let cfg = LaunchConfig::for_num_elems(n as u32);
20768                let (cd, dc) = (conv_dim as i32, d_conv as i32);
20769                let __s_b = self.gpu.stream();
20770                let mut b = __s_b.launch_builder(&f);
20771                b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
20772                unsafe {
20773                    b.launch(cfg)?;
20774                }
20775            }
20776            (None, None) => {
20777                let f = self.func("ssm_conv_ring_update_f32");
20778                let n = conv_dim * (d_conv - 1);
20779                let cfg = LaunchConfig::for_num_elems(n as u32);
20780                let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
20781                let __s_b = self.gpu.stream();
20782                let mut b = __s_b.launch_builder(&f);
20783                b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
20784                unsafe {
20785                    b.launch(cfg)?;
20786                }
20787            }
20788            (Some(_), _) => unreachable!(
20789                "ssm_conv1d_tm_state_pad_v: T < d_conv-1 has no view path (PRIME_MIN_T gates it)"
20790            ),
20791        }
20792        Ok(())
20793    }
20794
20795    /// PREFIX conv-ring rebuild (spec REPLAY-FREE partial accept): overwrite the resident ring
20796    /// with the state a T=1 chain holds after only the FIRST `tc` columns of `qkv_tm` — the last
20797    /// `pad` entries of [ring_old | cols 0..tc-1]. PURE COPIES (the ring stores raw inputs; no
20798    /// arithmetic, cannot perturb FP order). `ring_old` = the pre-round snapshot ring.
20799    pub fn ssm_conv_ring_rebuild(
20800        &self,
20801        qkv_tm: &CudaSlice<f32>,
20802        ring_old: &CudaSlice<f32>,
20803        conv_state: &mut CudaSlice<f32>,
20804        conv_dim: usize,
20805        tc: usize,
20806        d_conv: usize,
20807    ) -> Result<(), Box<dyn std::error::Error>> {
20808        let f = self.func("ssm_conv_ring_rebuild_f32");
20809        let n = conv_dim * (d_conv - 1);
20810        let cfg = LaunchConfig::for_num_elems(n as u32);
20811        let (cd, ti, dc) = (conv_dim as i32, tc as i32, d_conv as i32);
20812        let __s_b = self.gpu.stream();
20813        let mut b = __s_b.launch_builder(&f);
20814        b.arg(qkv_tm)
20815            .arg(ring_old)
20816            .arg(conv_state)
20817            .arg(&cd)
20818            .arg(&ti)
20819            .arg(&dc);
20820        unsafe {
20821            b.launch(cfg)?;
20822        }
20823        Ok(())
20824    }
20825
20826    /// FUSED decode GDN prep (T=1): repack + q/k L2-norm + beta sigmoid + g_log in one launch.
20827    /// Replaces 5 tiny serialized kernels on the decode critical path. L2 reduce runs as a 32-lane
20828    /// warp tree (vs l2_norm_f32's 256-thread two-level tree) — same math, different FP sum order;
20829    /// the argmax + run-spec gates are the authority.
20830    #[allow(clippy::too_many_arguments)]
20831    pub fn gdn_prep_decode(
20832        &self,
20833        conv_out: &CudaSlice<f32>,
20834        beta_raw: &CudaSlice<f32>,
20835        alpha: &CudaSlice<f32>,
20836        dt_bias: &CudaSlice<f32>,
20837        a: &CudaSlice<f32>,
20838        q_l2: &mut CudaSlice<f32>,
20839        k_l2: &mut CudaSlice<f32>,
20840        v_g: &mut CudaSlice<f32>,
20841        beta: &mut CudaSlice<f32>,
20842        g_log: &mut CudaSlice<f32>,
20843        d_state: usize,
20844        num_v: usize,
20845        num_k: usize,
20846        key_dim: usize,
20847        eps: f32,
20848    ) -> Result<(), Box<dyn std::error::Error>> {
20849        let f = self.func("gdn_prep_decode_f32");
20850        let cfg = LaunchConfig {
20851            grid_dim: (num_v as u32, 1, 1),
20852            block_dim: (32, 4, 1),
20853            shared_mem_bytes: 0,
20854        };
20855        let (ds, nv, nk, kd) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
20856        let __s_b = self.gpu.stream();
20857        let mut b = __s_b.launch_builder(&f);
20858        b.arg(conv_out)
20859            .arg(beta_raw)
20860            .arg(alpha)
20861            .arg(dt_bias)
20862            .arg(a)
20863            .arg(q_l2)
20864            .arg(k_l2)
20865            .arg(v_g)
20866            .arg(beta)
20867            .arg(g_log)
20868            .arg(&ds)
20869            .arg(&nv)
20870            .arg(&nk)
20871            .arg(&kd)
20872            .arg(&eps);
20873        unsafe {
20874            b.launch(cfg)?;
20875        }
20876        Ok(())
20877    }
20878
20879    /// FUSED prefill conv + GDN repack: token-major qkv -> q_g/k_g/v_g in ONE launch (no conv_out
20880    /// materialization, no qkv_to_gdn_repack pass). BIT-IDENTICAL values; scatter matches
20881    /// qkv_to_gdn_repack's modulo head-repeat mapping exactly.
20882    #[allow(clippy::too_many_arguments)]
20883    pub fn ssm_conv1d_gdn(
20884        &self,
20885        qkv_tm: &CudaSlice<f32>,
20886        w: &CudaSlice<f32>,
20887        q_g: &mut CudaSlice<f32>,
20888        k_g: &mut CudaSlice<f32>,
20889        v_g: &mut CudaSlice<f32>,
20890        conv_dim: usize,
20891        t: usize,
20892        d_conv: usize,
20893        d_state: usize,
20894        num_v: usize,
20895        num_k: usize,
20896        key_dim: usize,
20897    ) -> Result<(), Box<dyn std::error::Error>> {
20898        let f = self.func("ssm_conv1d_gdn_f32");
20899        let cfg = LaunchConfig {
20900            grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
20901            block_dim: (256, 1, 1),
20902            shared_mem_bytes: 0,
20903        };
20904        let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
20905        let (ds, nv, nk, kd) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
20906        let __s_b = self.gpu.stream();
20907        let mut b = __s_b.launch_builder(&f);
20908        b.arg(qkv_tm)
20909            .arg(w)
20910            .arg(q_g)
20911            .arg(k_g)
20912            .arg(v_g)
20913            .arg(&cd)
20914            .arg(&ti)
20915            .arg(&dc)
20916            .arg(&ds)
20917            .arg(&nv)
20918            .arg(&nk)
20919            .arg(&kd);
20920        unsafe {
20921            b.launch(cfg)?;
20922        }
20923        Ok(())
20924    }
20925
20926    pub fn ssm_conv1d(
20927        &self,
20928        x: &CudaSlice<f32>,
20929        w: &CudaSlice<f32>,
20930        y: &mut CudaSlice<f32>,
20931        conv_dim: usize,
20932        t: usize,
20933        d_conv: usize,
20934        silu: bool,
20935    ) -> Result<(), Box<dyn std::error::Error>> {
20936        let f = self.func("ssm_conv1d_silu_f32");
20937        let cfg = LaunchConfig {
20938            grid_dim: (conv_dim as u32, ((t as u32 + 255) / 256).max(1), 1),
20939            block_dim: (256, 1, 1),
20940            shared_mem_bytes: 0,
20941        };
20942        let (cd, ti, dc, s) = (conv_dim as i32, t as i32, d_conv as i32, silu as i32);
20943        let __s_b = self.gpu.stream();
20944        let mut b = __s_b.launch_builder(&f);
20945        b.arg(x).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc).arg(&s);
20946        unsafe {
20947            b.launch(cfg)?;
20948        }
20949        Ok(())
20950    }
20951
20952    /// Gated DeltaNet scan, S_v=128. q,k,v:[128,H,T]; g,beta:[H,T]; state:[128,128,H] transposed;
20953    /// o:[128,H,T]. Single sequence.
20954    pub fn gdn_scan_s128(
20955        &self,
20956        q: &CudaSlice<f32>,
20957        k: &CudaSlice<f32>,
20958        v: &CudaSlice<f32>,
20959        g: &CudaSlice<f32>,
20960        beta: &CudaSlice<f32>,
20961        state_in: &CudaSlice<f32>,
20962        state_out: &mut CudaSlice<f32>,
20963        o: &mut CudaSlice<f32>,
20964        n_head: usize,
20965        t: usize,
20966        scale: f32,
20967    ) -> Result<(), Box<dyn std::error::Error>> {
20968        let f = self.func("gdn_scan_s128");
20969        const S_V: u32 = 128;
20970        const WARP: u32 = 32;
20971        const COLS_PER_BLOCK: u32 = 4;
20972        let cfg = LaunchConfig {
20973            grid_dim: (n_head as u32, 1, S_V / COLS_PER_BLOCK),
20974            block_dim: (WARP, COLS_PER_BLOCK, 1),
20975            shared_mem_bytes: 0,
20976        };
20977        let (h, ti) = (n_head as i32, t as i32);
20978        let __s_b = self.gpu.stream();
20979        let mut b = __s_b.launch_builder(&f);
20980        b.arg(q)
20981            .arg(k)
20982            .arg(v)
20983            .arg(g)
20984            .arg(beta)
20985            .arg(state_in)
20986            .arg(state_out)
20987            .arg(o)
20988            .arg(&h)
20989            .arg(&ti)
20990            .arg(&scale);
20991        unsafe {
20992            b.launch(cfg)?;
20993        }
20994        Ok(())
20995    }
20996
20997    // ==== B2' batched decode state ops (decode_batch.rs) ====
20998    // Per-seq state pointers ride device u64 arrays (views into the per-step pointer table).
20999    // Bodies are the single-seq kernels per sequence — bit-identical per row.
21000
21001    #[allow(clippy::too_many_arguments)]
21002    pub fn ssm_conv1d_fused_decode_b(
21003        &self,
21004        qkv_cols: &CudaSlice<f32>,
21005        conv_state_ptrs: &cudarc::driver::CudaView<u64>,
21006        w: &CudaSlice<f32>,
21007        conv_outs: &mut CudaSlice<f32>,
21008        conv_dim: usize,
21009        d_conv: usize,
21010        b_n: usize,
21011    ) -> Result<(), Box<dyn std::error::Error>> {
21012        let f = self.func("ssm_conv1d_fused_decode_b_f32");
21013        let cfg = LaunchConfig {
21014            grid_dim: (((conv_dim + 255) / 256) as u32, 1, b_n as u32),
21015            block_dim: (256, 1, 1),
21016            shared_mem_bytes: 0,
21017        };
21018        let (cd, dc) = (conv_dim as i32, d_conv as i32);
21019        let __s_b = self.gpu.stream();
21020        let mut b = __s_b.launch_builder(&f);
21021        b.arg(qkv_cols)
21022            .arg(conv_state_ptrs)
21023            .arg(w)
21024            .arg(conv_outs)
21025            .arg(&cd)
21026            .arg(&dc);
21027        unsafe {
21028            b.launch(cfg)?;
21029        }
21030        Ok(())
21031    }
21032
21033    #[allow(clippy::too_many_arguments)]
21034    pub fn gdn_prep_decode_b(
21035        &self,
21036        conv_outs: &CudaSlice<f32>,
21037        beta_raws: &CudaSlice<f32>,
21038        alphas: &CudaSlice<f32>,
21039        dt_bias: &CudaSlice<f32>,
21040        a: &CudaSlice<f32>,
21041        q_l2: &mut CudaSlice<f32>,
21042        k_l2: &mut CudaSlice<f32>,
21043        v_g: &mut CudaSlice<f32>,
21044        beta: &mut CudaSlice<f32>,
21045        g_log: &mut CudaSlice<f32>,
21046        d_state: usize,
21047        num_v: usize,
21048        num_k: usize,
21049        key_dim: usize,
21050        eps: f32,
21051        conv_dim: usize,
21052        b_n: usize,
21053    ) -> Result<(), Box<dyn std::error::Error>> {
21054        let f = self.func("gdn_prep_decode_b_f32");
21055        let cfg = LaunchConfig {
21056            grid_dim: (num_v as u32, 1, b_n as u32),
21057            block_dim: (32, 4, 1),
21058            shared_mem_bytes: 0,
21059        };
21060        let (ds, nv, nk, kd, cd) = (
21061            d_state as i32,
21062            num_v as i32,
21063            num_k as i32,
21064            key_dim as i32,
21065            conv_dim as i32,
21066        );
21067        let __s_b = self.gpu.stream();
21068        let mut b = __s_b.launch_builder(&f);
21069        b.arg(conv_outs)
21070            .arg(beta_raws)
21071            .arg(alphas)
21072            .arg(dt_bias)
21073            .arg(a)
21074            .arg(q_l2)
21075            .arg(k_l2)
21076            .arg(v_g)
21077            .arg(beta)
21078            .arg(g_log)
21079            .arg(&ds)
21080            .arg(&nv)
21081            .arg(&nk)
21082            .arg(&kd)
21083            .arg(&eps)
21084            .arg(&cd);
21085        unsafe {
21086            b.launch(cfg)?;
21087        }
21088        Ok(())
21089    }
21090
21091    #[allow(clippy::too_many_arguments)]
21092    pub fn gdn_scan_s128_batched(
21093        &self,
21094        q: &CudaSlice<f32>,
21095        k: &CudaSlice<f32>,
21096        v: &CudaSlice<f32>,
21097        g: &CudaSlice<f32>,
21098        beta: &CudaSlice<f32>,
21099        state_in_ptrs: &cudarc::driver::CudaView<u64>,
21100        state_out_ptrs: &cudarc::driver::CudaView<u64>,
21101        o: &mut CudaSlice<f32>,
21102        n_head: usize,
21103        b_n: usize,
21104        scale: f32,
21105    ) -> Result<(), Box<dyn std::error::Error>> {
21106        let f = self.func("gdn_scan_s128_b");
21107        const S_V: u32 = 128;
21108        const WARP: u32 = 32;
21109        const COLS_PER_BLOCK: u32 = 4;
21110        let cfg = LaunchConfig {
21111            grid_dim: (n_head as u32, b_n as u32, S_V / COLS_PER_BLOCK),
21112            block_dim: (WARP, COLS_PER_BLOCK, 1),
21113            shared_mem_bytes: 0,
21114        };
21115        let h = n_head as i32;
21116        let __s_b = self.gpu.stream();
21117        let mut b = __s_b.launch_builder(&f);
21118        b.arg(q)
21119            .arg(k)
21120            .arg(v)
21121            .arg(g)
21122            .arg(beta)
21123            .arg(state_in_ptrs)
21124            .arg(state_out_ptrs)
21125            .arg(o)
21126            .arg(&h)
21127            .arg(&scale);
21128        unsafe {
21129            b.launch(cfg)?;
21130        }
21131        Ok(())
21132    }
21133
21134    /// VIEW twins of the three GDN decode state wrappers (lane/verify-launchslim): identical
21135    /// launches, row args as views into the caller's packed [T, ...] buffers — the t-parallel
21136    /// verify's per-row loop passes slices instead of paying an arithmetic-free dtod per row
21137    /// per kernel (48 layers x T rows x 4 copies/round on the money path). Same kernels, same
21138    /// numeric class; only the pointer arithmetic moved host-side.
21139    #[allow(clippy::too_many_arguments)]
21140    pub fn ssm_conv1d_fused_decode_b_view(
21141        &self,
21142        qkv_cols: &cudarc::driver::CudaView<f32>,
21143        conv_state_ptrs: &cudarc::driver::CudaView<u64>,
21144        w: &CudaSlice<f32>,
21145        conv_outs: &mut CudaSlice<f32>,
21146        conv_dim: usize,
21147        d_conv: usize,
21148        b_n: usize,
21149    ) -> Result<(), Box<dyn std::error::Error>> {
21150        let f = self.func("ssm_conv1d_fused_decode_b_f32");
21151        let cfg = LaunchConfig {
21152            grid_dim: (((conv_dim + 255) / 256) as u32, 1, b_n as u32),
21153            block_dim: (256, 1, 1),
21154            shared_mem_bytes: 0,
21155        };
21156        let (cd, dc) = (conv_dim as i32, d_conv as i32);
21157        let __s_b = self.gpu.stream();
21158        let mut b = __s_b.launch_builder(&f);
21159        b.arg(qkv_cols)
21160            .arg(conv_state_ptrs)
21161            .arg(w)
21162            .arg(conv_outs)
21163            .arg(&cd)
21164            .arg(&dc);
21165        unsafe {
21166            b.launch(cfg)?;
21167        }
21168        Ok(())
21169    }
21170
21171    #[allow(clippy::too_many_arguments)]
21172    pub fn gdn_prep_decode_b_view(
21173        &self,
21174        conv_outs: &CudaSlice<f32>,
21175        beta_raws: &cudarc::driver::CudaView<f32>,
21176        alphas: &cudarc::driver::CudaView<f32>,
21177        dt_bias: &CudaSlice<f32>,
21178        a: &CudaSlice<f32>,
21179        q_l2: &mut CudaSlice<f32>,
21180        k_l2: &mut CudaSlice<f32>,
21181        v_g: &mut CudaSlice<f32>,
21182        beta: &mut CudaSlice<f32>,
21183        g_log: &mut CudaSlice<f32>,
21184        d_state: usize,
21185        num_v: usize,
21186        num_k: usize,
21187        key_dim: usize,
21188        eps: f32,
21189        conv_dim: usize,
21190        b_n: usize,
21191    ) -> Result<(), Box<dyn std::error::Error>> {
21192        let f = self.func("gdn_prep_decode_b_f32");
21193        let cfg = LaunchConfig {
21194            grid_dim: (num_v as u32, 1, b_n as u32),
21195            block_dim: (32, 4, 1),
21196            shared_mem_bytes: 0,
21197        };
21198        let (ds, nv, nk, kd, cd) = (
21199            d_state as i32,
21200            num_v as i32,
21201            num_k as i32,
21202            key_dim as i32,
21203            conv_dim as i32,
21204        );
21205        let __s_b = self.gpu.stream();
21206        let mut b = __s_b.launch_builder(&f);
21207        b.arg(conv_outs)
21208            .arg(beta_raws)
21209            .arg(alphas)
21210            .arg(dt_bias)
21211            .arg(a)
21212            .arg(q_l2)
21213            .arg(k_l2)
21214            .arg(v_g)
21215            .arg(beta)
21216            .arg(g_log)
21217            .arg(&ds)
21218            .arg(&nv)
21219            .arg(&nk)
21220            .arg(&kd)
21221            .arg(&eps)
21222            .arg(&cd);
21223        unsafe {
21224            b.launch(cfg)?;
21225        }
21226        Ok(())
21227    }
21228
21229    #[allow(clippy::too_many_arguments)]
21230    pub fn gdn_scan_s128_batched_view(
21231        &self,
21232        q: &CudaSlice<f32>,
21233        k: &CudaSlice<f32>,
21234        v: &CudaSlice<f32>,
21235        g: &CudaSlice<f32>,
21236        beta: &CudaSlice<f32>,
21237        state_in_ptrs: &cudarc::driver::CudaView<u64>,
21238        state_out_ptrs: &cudarc::driver::CudaView<u64>,
21239        o: &mut cudarc::driver::CudaViewMut<f32>,
21240        n_head: usize,
21241        b_n: usize,
21242        scale: f32,
21243    ) -> Result<(), Box<dyn std::error::Error>> {
21244        let f = self.func("gdn_scan_s128_b");
21245        const S_V: u32 = 128;
21246        const WARP: u32 = 32;
21247        const COLS_PER_BLOCK: u32 = 4;
21248        let cfg = LaunchConfig {
21249            grid_dim: (n_head as u32, b_n as u32, S_V / COLS_PER_BLOCK),
21250            block_dim: (WARP, COLS_PER_BLOCK, 1),
21251            shared_mem_bytes: 0,
21252        };
21253        let h = n_head as i32;
21254        let __s_b = self.gpu.stream();
21255        let mut b = __s_b.launch_builder(&f);
21256        b.arg(q)
21257            .arg(k)
21258            .arg(v)
21259            .arg(g)
21260            .arg(beta)
21261            .arg(state_in_ptrs)
21262            .arg(state_out_ptrs)
21263            .arg(o)
21264            .arg(&h)
21265            .arg(&scale);
21266        unsafe {
21267            b.launch(cfg)?;
21268        }
21269        Ok(())
21270    }
21271
21272    /// A4 seam: chunked WY GDN prefill. DEFAULT ON (`MEMRA_GDN_CHUNKED=0` = rollback to the
21273    /// sequential scan). Flipped 2026-07-04 with the full battery green: kernel-check ALL
21274    /// GREEN x {9B, 27B} incl the f64-truth chunk gates; run-gen argmax 82==82 both models
21275    /// on AND off (24/24 sweep runs); run-spec K={1,2,3,4,6,8} PASS x {9B synth, 9B text,
21276    /// 27B p2, 27B p3}; e2e first-16-token agreement 6/6 (full-256 drifts at index 47-125
21277    /// on 5/6 prompts — accepted cache-state-FP class, batched-prime precedent).
21278    /// PREFILL-ONLY: decode + spec verify never route here (decode==verify dispatch
21279    /// identity law); prime_cache/forward/forward_last are the only callers.
21280    pub fn gdn_chunked_enabled() -> bool {
21281        static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
21282        *E.get_or_init(|| {
21283            std::env::var("MEMRA_GDN_CHUNKED")
21284                .map(|v| v != "0")
21285                .unwrap_or(true)
21286        })
21287    }
21288
21289    /// A4 chunk size (MEMRA_GDN_CHUNK, default 32 — the sweep winner: the O(T*C) chunk
21290    /// matrices grow with C while the sequential state pass is C-flat, so smaller chunks
21291    /// win; C=32/64 also get the register-history solve template). Clamped to multiples
21292    /// of 32 in [32, 128] (kernel row mappings require it).
21293    pub fn gdn_chunk_size() -> usize {
21294        static C: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
21295        *C.get_or_init(|| {
21296            let c: usize = std::env::var("MEMRA_GDN_CHUNK")
21297                .ok()
21298                .and_then(|v| v.parse().ok())
21299                .unwrap_or(32);
21300            c.clamp(32, 128) / 32 * 32
21301        })
21302    }
21303
21304    /// A4: chunked WY / blockwise-inverse GDN prefill (see cu/hybrid.cu K1-K5 header for the
21305    /// math). Same contract as `gdn_scan_s128` (layouts, state ping-pong) but chunk-parallel:
21306    /// NOT bit-identical to the sequential scan (chunked FP accumulation order); run-gen
21307    /// argmax + run-spec batteries are the accuracy authority. PREFILL callers only.
21308    #[allow(clippy::too_many_arguments)]
21309    /// task #18: K1-K3 of the chunked WY scan (shared by the per-seq path and the
21310    /// batched-prime varlen path). Returns (gcum, P, U, W); `A` is K3-internal.
21311    #[allow(clippy::too_many_arguments, clippy::type_complexity)]
21312    #[allow(clippy::too_many_arguments)]
21313    pub fn gdn_chunk_k123(
21314        &self,
21315        q: &CudaSlice<f32>,
21316        k: &CudaSlice<f32>,
21317        v: &CudaSlice<f32>,
21318        g: &CudaSlice<f32>,
21319        beta: &CudaSlice<f32>,
21320        wb16: Option<&mut CudaSlice<u8>>,
21321        n_head: usize,
21322        t: usize,
21323        c: usize,
21324        hk: usize,
21325        k2w: Option<(&CudaSlice<u8>, &CudaSlice<u8>, &mut CudaSlice<u8>)>,
21326    ) -> Result<
21327        (
21328            CudaSlice<f32>,
21329            CudaSlice<f32>,
21330            CudaSlice<f32>,
21331            CudaSlice<f32>,
21332        ),
21333        Box<dyn std::error::Error>,
21334    > {
21335        const D: usize = 128;
21336        let h = n_head;
21337        let nc = (t + c - 1) / c;
21338        let (hi, ti, ci) = (h as i32, t as i32, c as i32);
21339        let mut gcum = self.uninit(t * h)?;
21340        let mut a = self.uninit(nc * h * c * c)?;
21341        let mut p = self.uninit(nc * h * c * c)?;
21342        let mut u = self.uninit(nc * h * c * D)?;
21343        let mut w = self.uninit(nc * h * c * D)?;
21344        {
21345            // K1
21346            let f = self.func("gdn_chunk_cumgate_f32");
21347            let cfg = LaunchConfig {
21348                grid_dim: (nc as u32, h as u32, 1),
21349                block_dim: (32, 1, 1),
21350                shared_mem_bytes: 0,
21351            };
21352            let __s_b = self.gpu.stream();
21353            let mut b = __s_b.launch_builder(&f);
21354            b.arg(g).arg(&mut gcum).arg(&hi).arg(&ti).arg(&ci);
21355            unsafe {
21356                b.launch(cfg)?;
21357            }
21358        }
21359        if let Some((qb, kb, pb)) = k2w {
21360            // K2-wgmma (MEMRA_GDN_WGMMA path, c==32): A + pre-masked Pb16 in one kernel;
21361            // the P f32 buffer stays UNWRITTEN (its only wgmma-path consumer is Pb16).
21362            assert!(c == 32, "gdn_k2_wgmma is a C==32 tile");
21363            let f = self.func("gdn_k2_wgmma");
21364            let cfg = LaunchConfig {
21365                grid_dim: (nc as u32, h as u32, 1),
21366                block_dim: (128, 1, 1),
21367                shared_mem_bytes: 0,
21368            };
21369            let hki = hk as i32;
21370            let __s_b = self.gpu.stream();
21371            let mut b = __s_b.launch_builder(&f);
21372            b.arg(qb)
21373                .arg(kb)
21374                .arg(&gcum)
21375                .arg(beta)
21376                .arg(&mut a)
21377                .arg(&mut *pb)
21378                .arg(&hi)
21379                .arg(&ti)
21380                .arg(&ci)
21381                .arg(&hki);
21382            unsafe {
21383                b.launch(cfg)?;
21384            }
21385        } else if c <= 64 && !portable_mma_gated() {
21386            // K2 register-tiled (2x2 outputs/thread, whole-chunk smem k tile)
21387            let f = self.func("gdn_chunk_attn_f32");
21388            let jt = ((c + 31) / 32) as u32;
21389            let cfg = LaunchConfig {
21390                grid_dim: (nc as u32, h as u32, jt),
21391                block_dim: (256, 1, 1),
21392                shared_mem_bytes: 0,
21393            };
21394            let hki = hk as i32;
21395            let __s_b = self.gpu.stream();
21396            let mut b = __s_b.launch_builder(&f);
21397            b.arg(q)
21398                .arg(k)
21399                .arg(&gcum)
21400                .arg(beta)
21401                .arg(&mut a)
21402                .arg(&mut p)
21403                .arg(&hi)
21404                .arg(&ti)
21405                .arg(&ci)
21406                .arg(&hki);
21407            unsafe {
21408                b.launch(cfg)?;
21409            }
21410        } else {
21411            // K2 generic (C = 128, or the portable target's low-smem fallback)
21412            assert!(
21413                hk == h,
21414                "generic K2 is broadcast-only (de-broadcast rides C==32)"
21415            );
21416            let f = self.func("gdn_chunk_attn_g_f32");
21417            let cfg = LaunchConfig {
21418                grid_dim: (nc as u32, h as u32, 1),
21419                block_dim: (32, 8, 1),
21420                shared_mem_bytes: 0,
21421            };
21422            let __s_b = self.gpu.stream();
21423            let mut b = __s_b.launch_builder(&f);
21424            b.arg(q)
21425                .arg(k)
21426                .arg(&gcum)
21427                .arg(beta)
21428                .arg(&mut a)
21429                .arg(&mut p)
21430                .arg(&hi)
21431                .arg(&ti)
21432                .arg(&ci);
21433            unsafe {
21434                b.launch(cfg)?;
21435            }
21436        }
21437        {
21438            // K3 (register-history templates for C=32/64; local-memory generic otherwise)
21439            let cfg = LaunchConfig {
21440                grid_dim: (nc as u32, h as u32, 1),
21441                block_dim: (256, 1, 1),
21442                shared_mem_bytes: 0,
21443            };
21444            match c {
21445                32 | 64 => {
21446                    let f = self.func(if c == 32 {
21447                        "gdn_chunk_solve32_f32"
21448                    } else {
21449                        "gdn_chunk_solve64_f32"
21450                    });
21451                    // mirror-fold: W's bf16 twin emitted on store (0 = skip)
21452                    let wb: u64 = match wb16 {
21453                        Some(d) => self.addr_u8(d),
21454                        None => 0,
21455                    };
21456                    let hki = hk as i32;
21457                    let __s_b = self.gpu.stream();
21458                    let mut b = __s_b.launch_builder(&f);
21459                    b.arg(v)
21460                        .arg(k)
21461                        .arg(&a)
21462                        .arg(&gcum)
21463                        .arg(&mut u)
21464                        .arg(&mut w)
21465                        .arg(&wb)
21466                        .arg(&hi)
21467                        .arg(&ti)
21468                        .arg(&hki);
21469                    unsafe {
21470                        b.launch(cfg)?;
21471                    }
21472                }
21473                _ => {
21474                    assert!(hk == h, "generic K3 is broadcast-only");
21475                    let f = self.func("gdn_chunk_solve_f32");
21476                    let __s_b = self.gpu.stream();
21477                    let mut b = __s_b.launch_builder(&f);
21478                    b.arg(v)
21479                        .arg(k)
21480                        .arg(&a)
21481                        .arg(&gcum)
21482                        .arg(&mut u)
21483                        .arg(&mut w)
21484                        .arg(&hi)
21485                        .arg(&ti)
21486                        .arg(&ci);
21487                    unsafe {
21488                        b.launch(cfg)?;
21489                    }
21490                }
21491            }
21492        }
21493        Ok((gcum, p, u, w))
21494    }
21495
21496    /// task #21 de-broadcast seam: q/k stored at num_k distinct GQA heads instead of
21497    /// the num_v broadcast. MEMRA_GDN_DB=0 reverts. Only the chunked prefill path
21498    /// consumes the compact layout (hk plumbed; hk == H reproduces broadcast exactly).
21499    pub fn gdn_db_on() -> bool {
21500        std::env::var("MEMRA_GDN_DB").as_deref() != Ok("0")
21501    }
21502
21503    /// Whether the K4/K5 mma pair serves at chunk size `c` (mirrors gdn_scan_chunked's
21504    /// seam read — env re-read per call ON PURPOSE, kernel-check pins both configs).
21505    /// DEFAULT ON for sm_120a builds too (lane/moeprime-nvfp4-direct, 2026-08-21): the pair
21506    /// was qualified on 90a only and left env-opt-in elsewhere; measured on Blackwell it
21507    /// wins on BOTH rigs — one RTX PRO 6000 (ornith15 pp14715 12,036 -> 12,751/12,957,
21508    /// +6-8%, both orders) and the local 5090 (q38-27b pp6435 1,397/1,429 -> 1,427/1,446,
21509    /// both orders) — with kernel-check/run-gen/margin-gate/run-spec green under the flag.
21510    /// bf16 HMMA (m16n8k16) is sm_80-class PTX; only the wgmma nest stays Hopper-gated.
21511    pub fn gdn_mma_enabled(&self, c: usize) -> bool {
21512        !portable_mma_gated()
21513            && c == 32
21514            && match std::env::var("MEMRA_GDN_MMA").as_deref() {
21515                Ok("1") => true,
21516                Ok("0") => false,
21517                _ => gdn_mma_default_on(),
21518            }
21519    }
21520
21521    /// task #22: whether the fused K4+K5 (+K2) wgmma path serves (nested inside the
21522    /// mma config; same per-call env read discipline). HARD-gated to the Hopper build:
21523    /// the wgmma asm bodies exist only at __CUDA_ARCH__ == 900 (MEMRA_K45_REAL,
21524    /// wgmma_common.cuh) — on every other arch the kernel compiles EMPTY, so an env
21525    /// force would silently produce garbage. Required since the sm_120a mma default
21526    /// flip made MEMRA_GDN_WGMMA=1 alone reach this branch there.
21527    pub fn gdn_wgmma_on(&self, c: usize) -> bool {
21528        cfg!(memra_hopper_mma)
21529            && self.gdn_mma_enabled(c)
21530            && std::env::var("MEMRA_GDN_WGMMA").as_deref() != Ok("0")
21531    }
21532
21533    /// task #18 conv-fuse: carried-ring conv + SiLU + GDN repack in ONE pass (the
21534    /// conv_out intermediate and its transposed re-read disappear — 11.8ms of the
21535    /// T=2048 prime). Ring update stays the separate follow-up launch (pad-aware).
21536    /// BIT-IDENTICAL values to ssm_conv1d_tm_state_pad + qkv_to_gdn_repack.
21537    #[allow(clippy::too_many_arguments)]
21538    pub fn ssm_conv1d_gdn_state_pad(
21539        &self,
21540        qkv_tm: &cudarc::driver::CudaView<f32>,
21541        conv_state: &mut CudaSlice<f32>,
21542        w: &CudaSlice<f32>,
21543        q_g: &mut CudaSlice<f32>,
21544        k_g: &mut CudaSlice<f32>,
21545        v_g: &mut CudaSlice<f32>,
21546        conv_dim: usize,
21547        t: usize,
21548        d_conv: usize,
21549        d_state: usize,
21550        num_v: usize,
21551        num_k: usize,
21552        key_dim: usize,
21553        hk: usize,
21554        pad_len: Option<&CudaSlice<i32>>,
21555    ) -> Result<(), Box<dyn std::error::Error>> {
21556        assert!(
21557            t >= d_conv - 1,
21558            "fused state conv requires T >= pad (PRIME_MIN_T gates)"
21559        );
21560        {
21561            let f = self.func("ssm_conv1d_gdn_state_f32");
21562            let cfg = LaunchConfig {
21563                grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
21564                block_dim: (256, 1, 1),
21565                shared_mem_bytes: 0,
21566            };
21567            let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
21568            let (ds, nv, nk, kd, hki) = (
21569                d_state as i32,
21570                num_v as i32,
21571                num_k as i32,
21572                key_dim as i32,
21573                hk as i32,
21574            );
21575            let __s_b = self.gpu.stream();
21576            let mut b = __s_b.launch_builder(&f);
21577            b.arg(qkv_tm)
21578                .arg(&*conv_state)
21579                .arg(w)
21580                .arg(q_g)
21581                .arg(k_g)
21582                .arg(v_g)
21583                .arg(&cd)
21584                .arg(&ti)
21585                .arg(&dc)
21586                .arg(&ds)
21587                .arg(&nv)
21588                .arg(&nk)
21589                .arg(&kd)
21590                .arg(&hki);
21591            unsafe {
21592                b.launch(cfg)?;
21593            }
21594        }
21595        match pad_len {
21596            Some(len_d) => {
21597                let f = self.func("ssm_conv_ring_update_dev_f32");
21598                let n = conv_dim * (d_conv - 1);
21599                let cfg = LaunchConfig::for_num_elems(n as u32);
21600                let (cd, dc) = (conv_dim as i32, d_conv as i32);
21601                let __s_b = self.gpu.stream();
21602                let mut b = __s_b.launch_builder(&f);
21603                b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
21604                unsafe {
21605                    b.launch(cfg)?;
21606                }
21607            }
21608            None => {
21609                let f = self.func("ssm_conv_ring_update_f32");
21610                let n = conv_dim * (d_conv - 1);
21611                let cfg = LaunchConfig::for_num_elems(n as u32);
21612                let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
21613                let __s_b = self.gpu.stream();
21614                let mut b = __s_b.launch_builder(&f);
21615                b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
21616                unsafe {
21617                    b.launch(cfg)?;
21618                }
21619            }
21620        }
21621        Ok(())
21622    }
21623
21624    /// task #18 increment 2: allocate ONE sequence's chunk buffers (no launches) —
21625    /// K1-K5 all run varlen afterwards. `a`/`w` become struct members so the varlen
21626    /// K2/K3 can write them.
21627    pub fn gdn_chunk_alloc(
21628        &self,
21629        n_head: usize,
21630        t: usize,
21631        c: usize,
21632        hk: usize,
21633    ) -> Result<GdnChunkBufs, Box<dyn std::error::Error>> {
21634        const D: usize = 128;
21635        assert!(
21636            c == 32,
21637            "gdn_chunk_alloc: varlen chain is the C==32 mma pair"
21638        );
21639        let h = n_head;
21640        let nc = (t + c - 1) / c;
21641        Ok(GdnChunkBufs {
21642            gcum: self.uninit(t * h)?,
21643            a: self.uninit(nc * h * c * c)?,
21644            p: self.uninit(nc * h * c * c)?,
21645            u: self.uninit(nc * h * c * D)?,
21646            w: self.uninit(nc * h * c * D)?,
21647            kb16: self.alloc_u8_uninit(t * hk * D * 2)?,
21648            wb16: self.alloc_u8_uninit(nc * h * c * D * 2)?,
21649            y16: self.alloc_u8_uninit(nc * h * c * D * 2)?,
21650            ssnap16: self.alloc_u8_uninit(nc * h * D * D * 2)?,
21651            qb16: self.alloc_u8_uninit(t * hk * D * 2)?,
21652            pb16: self.alloc_u8_uninit(nc * h * c * c * 2)?,
21653            o: self.uninit(D * h * t)?,
21654            t,
21655            nc,
21656        })
21657    }
21658
21659    /// view-source twin of f32_to_bf16 (the batched FA3 v mirror reads a concat view).
21660    pub fn f32_to_bf16_v(
21661        &self,
21662        x: &cudarc::driver::CudaView<f32>,
21663        dst: &mut CudaSlice<u8>,
21664        n: usize,
21665    ) -> Result<(), Box<dyn std::error::Error>> {
21666        let f = self.func("f32_to_bf16_bulk");
21667        let ni = n as i64;
21668        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
21669        let __s_b = self.gpu.stream();
21670        let mut b = __s_b.launch_builder(&f);
21671        b.arg(x).arg(dst).arg(&ni);
21672        unsafe {
21673            b.launch(cfg)?;
21674        }
21675        Ok(())
21676    }
21677
21678    /// f32 -> bf16 bulk mirror into a caller buffer (the K4/K5 operand mirrors).
21679    pub fn f32_to_bf16_into(
21680        &self,
21681        x: &CudaSlice<f32>,
21682        dst: &mut CudaSlice<u8>,
21683        n: usize,
21684    ) -> Result<(), Box<dyn std::error::Error>> {
21685        let f = self.func("f32_to_bf16_bulk");
21686        let ni = n as i64;
21687        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
21688        let __s_b = self.gpu.stream();
21689        let mut b = __s_b.launch_builder(&f);
21690        b.arg(x).arg(dst).arg(&ni);
21691        unsafe {
21692            b.launch(cfg)?;
21693        }
21694        Ok(())
21695    }
21696
21697    /// task #18 increment 2: varlen K1+K2+K3 — three launches run every sequence's
21698    /// cumgate/attn/solve (per-block math identical to the per-seq kernels).
21699    pub fn gdn_chunk_k123_vl8(
21700        &self,
21701        seqs: &[GdnSeqVl],
21702        n_head: usize,
21703        hk: usize,
21704        wq: Option<&GdnWVl8>,
21705    ) -> Result<(), Box<dyn std::error::Error>> {
21706        let b = seqs.len();
21707        assert!(b >= 1 && b <= 8, "gdn_chunk_k123_vl8: 1..=8 sequences");
21708        let mut packed = [GdnSeqVl::default(); 8];
21709        packed[..b].copy_from_slice(seqs);
21710        let v = GdnVl8(packed);
21711        let (hi, ci) = (n_head as i32, 32i32);
21712        let max_nc = seqs.iter().map(|a| a.nc).max().unwrap() as u32;
21713        {
21714            let f = self.func("gdn_chunk_cumgate_vl");
21715            let cfg = LaunchConfig {
21716                grid_dim: (max_nc, n_head as u32, b as u32),
21717                block_dim: (32, 1, 1),
21718                shared_mem_bytes: 0,
21719            };
21720            let __s_lb = self.gpu.stream();
21721            let mut lb = __s_lb.launch_builder(&f);
21722            lb.arg(&v).arg(&hi).arg(&ci);
21723            unsafe {
21724                lb.launch(cfg)?;
21725            }
21726        }
21727        let hki = hk as i32;
21728        if let Some(w) = wq {
21729            // K2-wgmma vl twin (writes A + pre-masked Pb16)
21730            let f = self.func("gdn_k2_wgmma_vl");
21731            let cfg = LaunchConfig {
21732                grid_dim: (max_nc, n_head as u32, b as u32),
21733                block_dim: (128, 1, 1),
21734                shared_mem_bytes: 0,
21735            };
21736            let __s_lb = self.gpu.stream();
21737            let mut lb = __s_lb.launch_builder(&f);
21738            lb.arg(&v).arg(w).arg(&hi).arg(&ci).arg(&hki);
21739            unsafe {
21740                lb.launch(cfg)?;
21741            }
21742        } else {
21743            let f = self.func("gdn_chunk_attn_vl");
21744            let cfg = LaunchConfig {
21745                grid_dim: (max_nc, n_head as u32, b as u32),
21746                block_dim: (256, 1, 1),
21747                shared_mem_bytes: 0,
21748            };
21749            let __s_lb = self.gpu.stream();
21750            let mut lb = __s_lb.launch_builder(&f);
21751            lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
21752            unsafe {
21753                lb.launch(cfg)?;
21754            }
21755        }
21756        {
21757            let f = self.func("gdn_chunk_solve32_vl");
21758            let cfg = LaunchConfig {
21759                grid_dim: (max_nc, n_head as u32, b as u32),
21760                block_dim: (256, 1, 1),
21761                shared_mem_bytes: 0,
21762            };
21763            let __s_lb = self.gpu.stream();
21764            let mut lb = __s_lb.launch_builder(&f);
21765            lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
21766            unsafe {
21767                lb.launch(cfg)?;
21768            }
21769        }
21770        Ok(())
21771    }
21772
21773    /// task #18 increment 3: varlen PREP chain — conv(+ring) / repack / fused-l2 /
21774    /// fused gate-prep, 5 launches for every sequence (per-element math identical
21775    /// to the per-seq kernels; l2/gate fusions write disjoint outputs).
21776    #[allow(clippy::too_many_arguments)]
21777    pub fn gdn_prep_vl8(
21778        &self,
21779        seqs: &[GdnPrepVl],
21780        conv_w: &CudaSlice<f32>,
21781        dt_bias: &CudaSlice<f32>,
21782        a: &CudaSlice<f32>,
21783        conv_dim: usize,
21784        d_conv: usize,
21785        d_state: usize,
21786        num_v: usize,
21787        num_k: usize,
21788        key_dim: usize,
21789        hk: usize,
21790        eps: f32,
21791    ) -> Result<(), Box<dyn std::error::Error>> {
21792        let b = seqs.len();
21793        assert!(b >= 1 && b <= 8);
21794        let mut packed = [GdnPrepVl::default(); 8];
21795        packed[..b].copy_from_slice(seqs);
21796        let v = GdnPrepVl8(packed);
21797        let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
21798        let (cdi, dci) = (conv_dim as i32, d_conv as i32);
21799        let conv_fuse = std::env::var("MEMRA_CONV_FUSE").as_deref() != Ok("0");
21800        assert!(
21801            conv_fuse || hk == num_v,
21802            "de-broadcast requires the fused conv"
21803        );
21804        if conv_fuse {
21805            let f = self.func("ssm_conv1d_gdn_state_vl");
21806            let cfg = LaunchConfig {
21807                grid_dim: ((conv_dim as u32).div_ceil(256), max_t, b as u32),
21808                block_dim: (256, 1, 1),
21809                shared_mem_bytes: 0,
21810            };
21811            let (dsi, nvi, nki, kdi, hki) = (
21812                d_state as i32,
21813                num_v as i32,
21814                num_k as i32,
21815                key_dim as i32,
21816                hk as i32,
21817            );
21818            let __s_lb = self.gpu.stream();
21819            let mut lb = __s_lb.launch_builder(&f);
21820            lb.arg(&v)
21821                .arg(conv_w)
21822                .arg(&cdi)
21823                .arg(&dci)
21824                .arg(&dsi)
21825                .arg(&nvi)
21826                .arg(&nki)
21827                .arg(&kdi)
21828                .arg(&hki);
21829            unsafe {
21830                lb.launch(cfg)?;
21831            }
21832        } else {
21833            let f = self.func("ssm_conv1d_tm_state_vl");
21834            let cfg = LaunchConfig {
21835                grid_dim: ((conv_dim as u32).div_ceil(256), max_t, b as u32),
21836                block_dim: (256, 1, 1),
21837                shared_mem_bytes: 0,
21838            };
21839            let __s_lb = self.gpu.stream();
21840            let mut lb = __s_lb.launch_builder(&f);
21841            lb.arg(&v).arg(conv_w).arg(&cdi).arg(&dci);
21842            unsafe {
21843                lb.launch(cfg)?;
21844            }
21845        }
21846        {
21847            let f = self.func("ssm_conv_ring_update_vl");
21848            let n = (conv_dim * (d_conv - 1)) as u32;
21849            let cfg = LaunchConfig {
21850                grid_dim: (n.div_ceil(256), 1, b as u32),
21851                block_dim: (256, 1, 1),
21852                shared_mem_bytes: 0,
21853            };
21854            let __s_lb = self.gpu.stream();
21855            let mut lb = __s_lb.launch_builder(&f);
21856            lb.arg(&v).arg(&cdi).arg(&dci);
21857            unsafe {
21858                lb.launch(cfg)?;
21859            }
21860        }
21861        if !conv_fuse {
21862            let f = self.func("qkv_to_gdn_repack_vl");
21863            let n = max_t * (num_v * d_state) as u32;
21864            let cfg = LaunchConfig {
21865                grid_dim: (n.div_ceil(256), 1, b as u32),
21866                block_dim: (256, 1, 1),
21867                shared_mem_bytes: 0,
21868            };
21869            let (dsi, nvi, nki, kdi) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
21870            let __s_lb = self.gpu.stream();
21871            let mut lb = __s_lb.launch_builder(&f);
21872            lb.arg(&v).arg(&dsi).arg(&nvi).arg(&nki).arg(&kdi);
21873            unsafe {
21874                lb.launch(cfg)?;
21875            }
21876        }
21877        if Self::l2_v2_on(d_state) {
21878            let f = self.func("gdn_l2_v2_vl");
21879            let cfg = LaunchConfig {
21880                grid_dim: ((max_t * hk as u32).div_ceil(8), 2, b as u32),
21881                block_dim: (256, 1, 1),
21882                shared_mem_bytes: 0,
21883            };
21884            let (dsi, nvi) = (d_state as i32, hk as i32);
21885            let __s_lb = self.gpu.stream();
21886            let mut lb = __s_lb.launch_builder(&f);
21887            lb.arg(&v).arg(&dsi).arg(&nvi).arg(&eps);
21888            unsafe {
21889                lb.launch(cfg)?;
21890            }
21891        } else {
21892            let f = self.func("gdn_l2_vl");
21893            let cfg = LaunchConfig {
21894                grid_dim: (max_t * hk as u32, 2, b as u32),
21895                block_dim: (256, 1, 1),
21896                shared_mem_bytes: 0,
21897            };
21898            let (dsi, nvi) = (d_state as i32, hk as i32);
21899            let __s_lb = self.gpu.stream();
21900            let mut lb = __s_lb.launch_builder(&f);
21901            lb.arg(&v).arg(&dsi).arg(&nvi).arg(&eps);
21902            unsafe {
21903                lb.launch(cfg)?;
21904            }
21905        }
21906        {
21907            let f = self.func("gdn_gate_prep_vl");
21908            let n = max_t * num_v as u32;
21909            let cfg = LaunchConfig {
21910                grid_dim: (n.div_ceil(256), 1, b as u32),
21911                block_dim: (256, 1, 1),
21912                shared_mem_bytes: 0,
21913            };
21914            let nvi = num_v as i32;
21915            let __s_lb = self.gpu.stream();
21916            let mut lb = __s_lb.launch_builder(&f);
21917            lb.arg(&v).arg(dt_bias).arg(a).arg(&nvi);
21918            unsafe {
21919                lb.launch(cfg)?;
21920            }
21921        }
21922        Ok(())
21923    }
21924
21925    /// varlen bf16 mirrors over the gdnseq_t table (which: 0 = k_l2 -> kb16, 1 = w -> wb16).
21926    pub fn gdn_mirror_vl8(
21927        &self,
21928        seqs: &[GdnSeqVl],
21929        n_head: usize,
21930        which: i32,
21931        hk: usize,
21932    ) -> Result<(), Box<dyn std::error::Error>> {
21933        let b = seqs.len();
21934        assert!(b >= 1 && b <= 8);
21935        let mut packed = [GdnSeqVl::default(); 8];
21936        packed[..b].copy_from_slice(seqs);
21937        let v = GdnVl8(packed);
21938        let ept = (if which == 0 { hk } else { n_head } * 128) as i32;
21939        let max_n = seqs
21940            .iter()
21941            .map(|s| {
21942                if which == 0 {
21943                    s.t as i64 * ept as i64
21944                } else {
21945                    s.nc as i64 * ept as i64 * 32
21946                }
21947            })
21948            .max()
21949            .unwrap();
21950        let f = self.func("gdn_mirror_vl");
21951        let blocks = ((max_n as u32).div_ceil(4)).div_ceil(256);
21952        let cfg = LaunchConfig {
21953            grid_dim: (blocks, 1, b as u32),
21954            block_dim: (256, 1, 1),
21955            shared_mem_bytes: 0,
21956        };
21957        let __s_lb = self.gpu.stream();
21958        let mut lb = __s_lb.launch_builder(&f);
21959        lb.arg(&v).arg(&ept).arg(&which);
21960        unsafe {
21961            lb.launch(cfg)?;
21962        }
21963        Ok(())
21964    }
21965
21966    /// varlen gated-norm tail (+f16out) — one launch replaces B gated_rmsnorm calls.
21967    pub fn gdn_tail_vl8(
21968        &self,
21969        seqs: &[GdnPrepVl],
21970        norm_w: &CudaSlice<f32>,
21971        d_state: usize,
21972        num_v: usize,
21973        eps: f32,
21974    ) -> Result<(), Box<dyn std::error::Error>> {
21975        let b = seqs.len();
21976        assert!(b >= 1 && b <= 8);
21977        let mut packed = [GdnPrepVl::default(); 8];
21978        packed[..b].copy_from_slice(seqs);
21979        let v = GdnPrepVl8(packed);
21980        let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
21981        let f = self.func("gated_rmsnorm_f16out_vl");
21982        // block_dim MUST match gated_rmsnorm's (128): the reduction tree order pins the scale
21983        let cfg = LaunchConfig {
21984            grid_dim: (max_t * num_v as u32, 1, b as u32),
21985            block_dim: (128, 1, 1),
21986            shared_mem_bytes: 0,
21987        };
21988        let (dsi, nvi) = (d_state as i32, num_v as i32);
21989        let __s_lb = self.gpu.stream();
21990        let mut lb = __s_lb.launch_builder(&f);
21991        lb.arg(&v).arg(norm_w).arg(&dsi).arg(&nvi).arg(&eps);
21992        unsafe {
21993            lb.launch(cfg)?;
21994        }
21995        Ok(())
21996    }
21997
21998    /// Raw device address helpers for the varlen by-value arg struct (single-stream
21999    /// launches; every buffer outlives the call — the f16 FFI discipline).
22000    pub fn addr_f32(&self, x: &CudaSlice<f32>) -> u64 {
22001        use cudarc::driver::DevicePtr;
22002        let s = self.gpu.stream();
22003        let (p, _g) = x.device_ptr(&s);
22004        p as u64
22005    }
22006    pub fn addr_f32_mut(&self, x: &mut CudaSlice<f32>) -> u64 {
22007        use cudarc::driver::DevicePtrMut;
22008        let s = self.gpu.stream();
22009        let (p, _g) = x.device_ptr_mut(&s);
22010        p as u64
22011    }
22012    pub fn addr_f32v(&self, x: &cudarc::driver::CudaView<f32>) -> u64 {
22013        use cudarc::driver::DevicePtr;
22014        let s = self.gpu.stream();
22015        let (p, _g) = x.device_ptr(&s);
22016        p as u64
22017    }
22018    pub fn addr_u8(&self, x: &CudaSlice<u8>) -> u64 {
22019        use cudarc::driver::DevicePtr;
22020        let s = self.gpu.stream();
22021        let (p, _g) = x.device_ptr(&s);
22022        p as u64
22023    }
22024
22025    /// task #18: the varlen K4+K5 pair — TWO launches run every sequence's state pass
22026    /// and output pass (grid gains a seq dim; per-block math identical to the per-seq
22027    /// launches, so this is strictly bit-gateable against them).
22028    pub fn gdn_chunk_vl8(
22029        &self,
22030        seqs: &[GdnSeqVl],
22031        n_head: usize,
22032        scale: f32,
22033        hk: usize,
22034        wq: Option<&GdnWVl8>,
22035    ) -> Result<(), Box<dyn std::error::Error>> {
22036        const NSPLIT: u32 = 4;
22037        let b = seqs.len();
22038        assert!(b >= 1 && b <= 8, "gdn_chunk_vl8: 1..=8 sequences");
22039        let mut packed = [GdnSeqVl::default(); 8];
22040        packed[..b].copy_from_slice(seqs);
22041        let v = GdnVl8(packed);
22042        let (hi, ci) = (n_head as i32, 32i32);
22043        let max_nc = seqs.iter().map(|a| a.nc).max().unwrap() as u32;
22044        let hki = hk as i32;
22045        if let Some(w) = wq {
22046            // K4+K5 fused wgmma vl twin: one launch, Y/Ssnap never materialized.
22047            let f = self.func("gdn_k45_wgmma_vl");
22048            let cfg = LaunchConfig {
22049                grid_dim: (n_head as u32, NSPLIT, b as u32),
22050                block_dim: (256, 1, 1),
22051                shared_mem_bytes: 0,
22052            };
22053            let __s_lb = self.gpu.stream();
22054            let mut lb = __s_lb.launch_builder(&f);
22055            lb.arg(&v).arg(w).arg(&scale).arg(&hi).arg(&ci).arg(&hki);
22056            unsafe {
22057                lb.launch(cfg)?;
22058            }
22059            let _ = max_nc;
22060            return Ok(());
22061        }
22062        {
22063            let f = self.func("gdn_chunk_state_mma_vl");
22064            let cfg = LaunchConfig {
22065                grid_dim: (n_head as u32, NSPLIT, b as u32),
22066                block_dim: (256, 1, 1),
22067                shared_mem_bytes: 0,
22068            };
22069            let __s_lb = self.gpu.stream();
22070            let mut lb = __s_lb.launch_builder(&f);
22071            lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
22072            unsafe {
22073                lb.launch(cfg)?;
22074            }
22075        }
22076        {
22077            let f = self.func("gdn_chunk_output_mma_vl");
22078            let cfg = LaunchConfig {
22079                grid_dim: (max_nc, n_head as u32, b as u32),
22080                block_dim: (256, 1, 1),
22081                shared_mem_bytes: 0,
22082            };
22083            let __s_lb = self.gpu.stream();
22084            let mut lb = __s_lb.launch_builder(&f);
22085            lb.arg(&v).arg(&hi).arg(&ci).arg(&scale).arg(&hki);
22086            unsafe {
22087                lb.launch(cfg)?;
22088            }
22089        }
22090        Ok(())
22091    }
22092    pub fn gdn_scan_chunked(
22093        &self,
22094        q: &CudaSlice<f32>,
22095        k: &CudaSlice<f32>,
22096        v: &CudaSlice<f32>,
22097        g: &CudaSlice<f32>,
22098        beta: &CudaSlice<f32>,
22099        kb16_pre: Option<&CudaSlice<u8>>,
22100        qb16_pre: Option<&CudaSlice<u8>>,
22101        state_in: &CudaSlice<f32>,
22102        state_out: &mut CudaSlice<f32>,
22103        o: &mut CudaSlice<f32>,
22104        n_head: usize,
22105        t: usize,
22106        scale: f32,
22107        c: usize,
22108        hk: usize,
22109    ) -> Result<(), Box<dyn std::error::Error>> {
22110        const D: usize = 128;
22111        const NSPLIT: u32 = 4;
22112        assert!(c >= 1 && c <= 128, "gdn_scan_chunked: C must be in 1..=128");
22113        let h = n_head;
22114        let nc = (t + c - 1) / c;
22115        let (hi, ti, ci) = (h as i32, t as i32, c as i32);
22116        // mirror-fold (round 27): on the mma path W's bf16 twin is emitted by K3's store
22117        // (wb16 pre-allocated and threaded through k123) and k's by the producer l2 when
22118        // the caller hands `kb16_pre` — both standalone mirror passes disappear.
22119        // default mirrors gdn_mma_enabled (incl. the sm_120a-build ON, lane/moeprime-nvfp4-direct)
22120        let gdn_mma_pre = !portable_mma_gated()
22121            && c == 32
22122            && match std::env::var("MEMRA_GDN_MMA").as_deref() {
22123                Ok("1") => true,
22124                Ok("0") => false,
22125                _ => gdn_mma_default_on(),
22126            };
22127        let mut wb16_pre: Option<CudaSlice<u8>> = if gdn_mma_pre {
22128            Some(self.alloc_u8_uninit(nc * h * c * D * 2)?)
22129        } else {
22130            None
22131        };
22132        // K2-wgmma pre-work (MEMRA_GDN_WGMMA): the kb16/qb16 mirrors hoist ABOVE K123 so
22133        // K2 rides them via cp.async; K2 writes the pre-masked Pb16 directly (the
22134        // gdn_p_bf16_masked pass and the in-branch mirror builds disappear).
22135        // mirrors gdn_wgmma_on: hard Hopper-build gate (empty wgmma bodies elsewhere)
22136        let gdn_wgmma_pre = cfg!(memra_hopper_mma)
22137            && gdn_mma_pre
22138            && std::env::var("MEMRA_GDN_WGMMA").as_deref() != Ok("0");
22139        let nk = t * hk * D;
22140        let mut kb16_local: Option<CudaSlice<u8>> = None;
22141        if gdn_mma_pre && kb16_pre.is_none() {
22142            let mut kb = self.alloc_u8_uninit(nk * 2)?;
22143            let f = self.func("f32_to_bf16_bulk");
22144            let n2 = nk as i64;
22145            let cfg2 = LaunchConfig::for_num_elems((nk as u32).div_ceil(4));
22146            let __s_b = self.gpu.stream();
22147            let mut b = __s_b.launch_builder(&f);
22148            b.arg(k).arg(&mut kb).arg(&n2);
22149            unsafe {
22150                b.launch(cfg2)?;
22151            }
22152            kb16_local = Some(kb);
22153        }
22154        let kb16_ref0: Option<&CudaSlice<u8>> = kb16_local.as_ref().or(kb16_pre);
22155        if let Some(kb) = kb16_pre {
22156            assert!(kb.len() >= nk * 2, "kb16_pre too small");
22157        }
22158        let mut qb16: Option<CudaSlice<u8>> = None;
22159        let mut pb16: Option<CudaSlice<u8>> = None;
22160        if gdn_wgmma_pre {
22161            // mirror-fold (round 35): prep's l2 v2 emits qb16 in-epilogue (kb16 pattern);
22162            // the standalone bulk cvt only serves callers without the prep mirror.
22163            if qb16_pre.is_none() {
22164                let mut qb = self.alloc_u8_uninit(nk * 2)?;
22165                let f = self.func("f32_to_bf16_bulk");
22166                let n2 = nk as i64;
22167                let cfg2 = LaunchConfig::for_num_elems((nk as u32).div_ceil(4));
22168                let __s_b = self.gpu.stream();
22169                let mut b = __s_b.launch_builder(&f);
22170                b.arg(q).arg(&mut qb).arg(&n2);
22171                unsafe {
22172                    b.launch(cfg2)?;
22173                }
22174                qb16 = Some(qb);
22175            } else if let Some(qb) = qb16_pre {
22176                assert!(qb.len() >= nk * 2, "qb16_pre too small");
22177            }
22178            pb16 = Some(self.alloc_u8_uninit(nc * h * c * c * 2)?);
22179        }
22180        let qb16_ref0: Option<&CudaSlice<u8>> = qb16.as_ref().or(qb16_pre);
22181        let k2w = if gdn_wgmma_pre {
22182            Some((
22183                *qb16_ref0.as_ref().unwrap(),
22184                *kb16_ref0.as_ref().unwrap(),
22185                pb16.as_mut().unwrap(),
22186            ))
22187        } else {
22188            None
22189        };
22190        let (gcum, p, u, w) =
22191            self.gdn_chunk_k123(q, k, v, g, beta, wb16_pre.as_mut(), n_head, t, c, hk, k2w)?;
22192        let _ = &w;
22193        let mut y = self.uninit(nc * h * c * D)?;
22194        let mut ssnap = self.uninit(nc * h * D * D)?; // chunk-start state snapshots (K5 phase 1)
22195        // K4-MMA seam (MEMRA_GDN_MMA; harness verdict 1.75x — tools/bench_gdn_k4.cu, ledger
22196        // 2026-07-26): M in mma accumulator fragments, bf16 W/k mirrors through a cp.async
22197        // ring. C==32 only (the kernel's tile). PROMOTED default-ON on the Hopper lane
22198        // after the STATE-CARRY battery (2026-07-26): 2048-token prime (64 in-kernel state
22199        // carries) -> 256 greedy decode tokens IDENTICAL to f32 on 3 seeds, AND chunked-
22200        // continuation prime (MEMRA_PRIME_CHUNK=512, 4 cross-call carries via cache.recur)
22201        // IDENTICAL on 2 seeds; plus argmax MATCH, pp512 +3.5% (17286), oracle out
22202        // mean_rel ~1e-4. kernel-check pins BOTH configs (f32 tight band forced =0; mma
22203        // band 8e-2/8e-1 vs f64 truth). =0 reverts; portable stays f32. NOT read via
22204        // OnceLock ON PURPOSE: kernel-check toggles the env per call to pin both forms.
22205        // default mirrors gdn_mma_enabled (incl. the sm_120a-build ON) — all three read
22206        // sites must agree or the pre-work arms while the scan takes the scalar route.
22207        let gdn_mma = !portable_mma_gated()
22208            && c == 32
22209            && match std::env::var("MEMRA_GDN_MMA").as_deref() {
22210                Ok("1") => true,
22211                Ok("0") => false,
22212                _ => gdn_mma_default_on(),
22213            };
22214        if gdn_mma {
22215            let wb16 = wb16_pre
22216                .take()
22217                .expect("mma path pre-allocates wb16 (K3 store fold)");
22218            let kb16_ref: &CudaSlice<u8> = kb16_ref0.expect("mma path pre-builds kb16 above K123");
22219            // K4+K5 FUSED wgmma seam (MEMRA_GDN_WGMMA, task #22; harness verdict
22220            // tools/bench_gdn_wgmma.cu v5, ledger 1f08b997: in-band Y 1.07e-2 / state
22221            // 1.03e-2 / O 1.08e-2, 91.3us vs 70.4 K4-only at H=32 T=512). K5's output
22222            // pass runs inside the persistent-M kernel; Y and Ssnap are never
22223            // materialized. New numeric class (gk folds into k^T instead of ys) —
22224            // explicit opt-in until the state-carry battery promotes it. Env read per
22225            // call (kernel-check pins configs by toggling env, GDN_MMA precedent).
22226            // PROMOTED default-ON hopper (2026-07-27): full battery green — harness
22227            // in-band, argmax gate PASS, 3-seed greedy IDENTICAL after ~2k prime,
22228            // chunked-continuation IDENTICAL, kernel-check + decode-batch gates green,
22229            // official prefill lane +0.74% interleaved x5 (5/5 rounds). =0 reverts.
22230            if gdn_wgmma_pre {
22231                // qb16/pb16 pre-built above K123 (K2-wgmma wrote the masked Pb16).
22232                let qb16 = qb16_ref0.unwrap();
22233                let pb16 = pb16.as_ref().unwrap();
22234                {
22235                    let f = self.func("gdn_k45_wgmma");
22236                    let cfg = LaunchConfig {
22237                        grid_dim: (h as u32, 4, 1),
22238                        block_dim: (256, 1, 1),
22239                        shared_mem_bytes: 0,
22240                    };
22241                    let hki = hk as i32;
22242                    let __s_b = self.gpu.stream();
22243                    let mut b = __s_b.launch_builder(&f);
22244                    b.arg(kb16_ref)
22245                        .arg(&gcum)
22246                        .arg(beta)
22247                        .arg(&u)
22248                        .arg(&wb16)
22249                        .arg(qb16)
22250                        .arg(pb16)
22251                        .arg(o)
22252                        .arg(&scale)
22253                        .arg(state_in)
22254                        .arg(&mut *state_out)
22255                        .arg(&hi)
22256                        .arg(&ti)
22257                        .arg(&ci)
22258                        .arg(&hki);
22259                    unsafe {
22260                        b.launch(cfg)?;
22261                    }
22262                }
22263                return Ok(());
22264            }
22265            // COUPLED PAIR: K4-mma writes Y and Ssnap as bf16 (their only consumer is
22266            // K5-mma, which rounds to bf16 regardless — identical numerics, half the
22267            // traffic; harness K5 63.0 -> 35.3us). Fresh bf16 buffers replace the f32 ones.
22268            let mut y16 = self.alloc_u8_uninit(nc * h * c * D * 2)?;
22269            let mut ssnap16 = self.alloc_u8_uninit(nc * h * D * D * 2)?;
22270            {
22271                let f = self.func("gdn_chunk_state_mma");
22272                let cfg = LaunchConfig {
22273                    grid_dim: (h as u32, NSPLIT, 1),
22274                    block_dim: (256, 1, 1),
22275                    shared_mem_bytes: 0,
22276                };
22277                let hki = hk as i32;
22278                let __s_b = self.gpu.stream();
22279                let mut b = __s_b.launch_builder(&f);
22280                b.arg(kb16_ref)
22281                    .arg(&gcum)
22282                    .arg(beta)
22283                    .arg(&u)
22284                    .arg(&wb16)
22285                    .arg(&mut y16)
22286                    .arg(&mut ssnap16)
22287                    .arg(state_in)
22288                    .arg(&mut *state_out)
22289                    .arg(&hi)
22290                    .arg(&ti)
22291                    .arg(&ci)
22292                    .arg(&hki);
22293                unsafe {
22294                    b.launch(cfg)?;
22295                }
22296            }
22297            {
22298                // K5-mma (bf16 St/Y consumers)
22299                let f = self.func("gdn_chunk_output_mma");
22300                let jt = ((c + 31) / 32) as u32;
22301                let cfg = LaunchConfig {
22302                    grid_dim: (nc as u32, h as u32, jt),
22303                    block_dim: (256, 1, 1),
22304                    shared_mem_bytes: 0,
22305                };
22306                let hki = hk as i32;
22307                let __s_b = self.gpu.stream();
22308                let mut b = __s_b.launch_builder(&f);
22309                b.arg(q)
22310                    .arg(&gcum)
22311                    .arg(&p)
22312                    .arg(&y16)
22313                    .arg(&ssnap16)
22314                    .arg(o)
22315                    .arg(&hi)
22316                    .arg(&ti)
22317                    .arg(&ci)
22318                    .arg(&scale)
22319                    .arg(&hki);
22320                unsafe {
22321                    b.launch(cfg)?;
22322                }
22323            }
22324            return Ok(());
22325        }
22326        {
22327            // K4 (sequential over chunks inside; blocks col-partition the state)
22328            let f = self.func("gdn_chunk_state_f32");
22329            let cfg = LaunchConfig {
22330                grid_dim: (h as u32, NSPLIT, 1),
22331                block_dim: (256, 1, 1),
22332                shared_mem_bytes: 0,
22333            };
22334            let __s_b = self.gpu.stream();
22335            let mut b = __s_b.launch_builder(&f);
22336            b.arg(k)
22337                .arg(&gcum)
22338                .arg(beta)
22339                .arg(&u)
22340                .arg(&w)
22341                .arg(&mut y)
22342                .arg(&mut ssnap)
22343                .arg(state_in)
22344                .arg(&mut *state_out)
22345                .arg(&hi)
22346                .arg(&ti)
22347                .arg(&ci);
22348            unsafe {
22349                b.launch(cfg)?;
22350            }
22351        }
22352        {
22353            // K5 (j-blocked: grid.z = 32-row output blocks per chunk; writes o fully)
22354            let f = self.func("gdn_chunk_output_f32");
22355            let jt = ((c + 31) / 32) as u32;
22356            let cfg = LaunchConfig {
22357                grid_dim: (nc as u32, h as u32, jt),
22358                block_dim: (256, 1, 1),
22359                shared_mem_bytes: 0,
22360            };
22361            let __s_b = self.gpu.stream();
22362            let mut b = __s_b.launch_builder(&f);
22363            b.arg(q)
22364                .arg(&gcum)
22365                .arg(&p)
22366                .arg(&y)
22367                .arg(&ssnap)
22368                .arg(o)
22369                .arg(&hi)
22370                .arg(&ti)
22371                .arg(&ci)
22372                .arg(&scale);
22373            unsafe {
22374                b.launch(cfg)?;
22375            }
22376        }
22377        Ok(())
22378    }
22379
22380    /// PREFILL GDN scan dispatch (the A4 seam): chunked WY form when enabled and T is in the
22381    /// batched-prefill regime, else the sequential scan. Callers: hybrid_forward::linear_attn
22382    /// (forward/forward_last) + linear_attn_prime (prime_cache). Decode (T=1) and the spec
22383    /// verify call `gdn_scan_s128` DIRECTLY — the decode==verify dispatch identity is untouched.
22384    ///
22385    /// MEMRA_GDN_DIFF=1: numerical-oracle mode — runs BOTH forms on the same inputs, prints the
22386    /// per-call (== per-layer, in call order) output/state error distribution, and keeps the
22387    /// SEQUENTIAL results so the run stays on the shipped path (stage-1 prototype evidence).
22388    #[allow(clippy::too_many_arguments)]
22389    #[allow(clippy::too_many_arguments)]
22390    pub fn gdn_scan_prefill(
22391        &self,
22392        q: &CudaSlice<f32>,
22393        k: &CudaSlice<f32>,
22394        v: &CudaSlice<f32>,
22395        g: &CudaSlice<f32>,
22396        beta: &CudaSlice<f32>,
22397        kb16_pre: Option<&CudaSlice<u8>>,
22398        qb16_pre: Option<&CudaSlice<u8>>,
22399        state_in: &CudaSlice<f32>,
22400        state_out: &mut CudaSlice<f32>,
22401        o: &mut CudaSlice<f32>,
22402        n_head: usize,
22403        t: usize,
22404        scale: f32,
22405        hk: usize,
22406    ) -> Result<(), Box<dyn std::error::Error>> {
22407        if std::env::var("MEMRA_GDN_DIFF").is_ok() && t >= 16 {
22408            assert!(hk == n_head, "GDN_DIFF oracle is broadcast-only");
22409            return self.gdn_scan_diff(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale);
22410        }
22411        if Self::gdn_chunked_enabled() && t >= 16 {
22412            self.gdn_scan_chunked(
22413                q,
22414                k,
22415                v,
22416                g,
22417                beta,
22418                kb16_pre,
22419                qb16_pre,
22420                state_in,
22421                state_out,
22422                o,
22423                n_head,
22424                t,
22425                scale,
22426                Self::gdn_chunk_size(),
22427                hk,
22428            )
22429        } else {
22430            assert!(
22431                hk == n_head,
22432                "s128 scan is broadcast-only (prep guarantees by predicate)"
22433            );
22434            self.gdn_scan_s128(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale)
22435        }
22436    }
22437
22438    /// Stage-1 oracle: run sequential AND chunked, report per-call error stats, keep sequential.
22439    #[allow(clippy::too_many_arguments)]
22440    fn gdn_scan_diff(
22441        &self,
22442        q: &CudaSlice<f32>,
22443        k: &CudaSlice<f32>,
22444        v: &CudaSlice<f32>,
22445        g: &CudaSlice<f32>,
22446        beta: &CudaSlice<f32>,
22447        state_in: &CudaSlice<f32>,
22448        state_out: &mut CudaSlice<f32>,
22449        o: &mut CudaSlice<f32>,
22450        n_head: usize,
22451        t: usize,
22452        scale: f32,
22453    ) -> Result<(), Box<dyn std::error::Error>> {
22454        static CALL: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
22455        let call = CALL.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
22456        let mut o_c = self.uninit(o.len())?;
22457        let mut st_c = self.uninit(state_out.len())?;
22458        self.gdn_scan_chunked(
22459            q,
22460            k,
22461            v,
22462            g,
22463            beta,
22464            None,
22465            None,
22466            state_in,
22467            &mut st_c,
22468            &mut o_c,
22469            n_head,
22470            t,
22471            scale,
22472            Self::gdn_chunk_size(),
22473            n_head,
22474        )?;
22475        self.gdn_scan_s128(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale)?;
22476        let (oh_s, oh_c) = (self.dtoh(o)?, self.dtoh(&o_c)?);
22477        let (sh_s, sh_c) = (self.dtoh(state_out)?, self.dtoh(&st_c)?);
22478        let stats = |a: &[f32], b: &[f32]| -> (f32, f32, f64) {
22479            let mut max_abs = 0f32;
22480            let mut max_rel = 0f32;
22481            let mut sum_rel = 0f64;
22482            for (x, y) in a.iter().zip(b) {
22483                let ad = (x - y).abs();
22484                let rel = ad / x.abs().max(y.abs()).max(1e-3);
22485                if ad > max_abs {
22486                    max_abs = ad;
22487                }
22488                if rel > max_rel {
22489                    max_rel = rel;
22490                }
22491                sum_rel += rel as f64;
22492            }
22493            (max_abs, max_rel, sum_rel / a.len() as f64)
22494        };
22495        let (o_ma, o_mr, o_mean) = stats(&oh_s, &oh_c);
22496        let (s_ma, s_mr, s_mean) = stats(&sh_s, &sh_c);
22497        println!(
22498            "[gdn-diff call {call:3} T={t} C={}] out: max_abs={o_ma:.3e} max_rel={o_mr:.3e} mean_rel={o_mean:.3e} | \
22499                  state: max_abs={s_ma:.3e} max_rel={s_mr:.3e} mean_rel={s_mean:.3e}",
22500            Self::gdn_chunk_size()
22501        );
22502        Ok(())
22503    }
22504
22505    /// softplus-based g_log: g_log[h,t] = a[h] * softplus(alpha[h,t] + dt_bias[h]). a pre-negated.
22506    pub fn gdn_glog(
22507        &self,
22508        alpha: &CudaSlice<f32>,
22509        dt_bias: &CudaSlice<f32>,
22510        a: &CudaSlice<f32>,
22511        g_log: &mut CudaSlice<f32>,
22512        n_head: usize,
22513        t: usize,
22514    ) -> Result<(), Box<dyn std::error::Error>> {
22515        let f = self.func("gdn_glog_f32");
22516        let cfg = LaunchConfig::for_num_elems((n_head * t) as u32);
22517        let (h, ti) = (n_head as i32, t as i32);
22518        let __s_b = self.gpu.stream();
22519        let mut b = __s_b.launch_builder(&f);
22520        b.arg(alpha).arg(dt_bias).arg(a).arg(g_log).arg(&h).arg(&ti);
22521        unsafe {
22522            b.launch(cfg)?;
22523        }
22524        Ok(())
22525    }
22526
22527    /// view twins (task #16): the batched prime's GDN core reads the CONCAT projection
22528    /// buffers at row offsets (CudaView) — same kernels, same values, no split copies.
22529    pub fn sigmoid_v(
22530        &self,
22531        x: &cudarc::driver::CudaView<f32>,
22532        y: &mut CudaSlice<f32>,
22533        n: usize,
22534    ) -> Result<(), Box<dyn std::error::Error>> {
22535        let f = self.func("sigmoid_f32");
22536        let cfg = LaunchConfig::for_num_elems(n as u32);
22537        let ni = n as i32;
22538        let __s_b = self.gpu.stream();
22539        let mut b = __s_b.launch_builder(&f);
22540        b.arg(x).arg(y).arg(&ni);
22541        unsafe {
22542            b.launch(cfg)?;
22543        }
22544        Ok(())
22545    }
22546
22547    pub fn gdn_glog_v(
22548        &self,
22549        alpha: &cudarc::driver::CudaView<f32>,
22550        dt_bias: &CudaSlice<f32>,
22551        a: &CudaSlice<f32>,
22552        g_log: &mut CudaSlice<f32>,
22553        n_head: usize,
22554        t: usize,
22555    ) -> Result<(), Box<dyn std::error::Error>> {
22556        let f = self.func("gdn_glog_f32");
22557        let cfg = LaunchConfig::for_num_elems((n_head * t) as u32);
22558        let (h, ti) = (n_head as i32, t as i32);
22559        let __s_b = self.gpu.stream();
22560        let mut b = __s_b.launch_builder(&f);
22561        b.arg(alpha).arg(dt_bias).arg(a).arg(g_log).arg(&h).arg(&ti);
22562        unsafe {
22563            b.launch(cfg)?;
22564        }
22565        Ok(())
22566    }
22567
22568    pub fn sigmoid(
22569        &self,
22570        x: &CudaSlice<f32>,
22571        y: &mut CudaSlice<f32>,
22572        n: usize,
22573    ) -> Result<(), Box<dyn std::error::Error>> {
22574        let f = self.func("sigmoid_f32");
22575        let cfg = LaunchConfig::for_num_elems(n as u32);
22576        let ni = n as i32;
22577        let __s_b = self.gpu.stream();
22578        let mut b = __s_b.launch_builder(&f);
22579        b.arg(x).arg(y).arg(&ni);
22580        unsafe {
22581            b.launch(cfg)?;
22582        }
22583        Ok(())
22584    }
22585
22586    /// attn out-gate fused epilogue (task #17): dst = a * sigmoid(g) + fp16 twin, one launch
22587    /// (replaces sigmoid + mul + convert). Bit-identical class.
22588    pub fn sig_mul_f16out(
22589        &self,
22590        a: &CudaSlice<f32>,
22591        g: &CudaSlice<f32>,
22592        dst: &mut CudaSlice<f32>,
22593        dst16: &mut CudaSlice<u8>,
22594        n: usize,
22595    ) -> Result<(), Box<dyn std::error::Error>> {
22596        let f = self.func("sig_mul_f16out_f32");
22597        let cfg = LaunchConfig::for_num_elems(n as u32);
22598        let ni = n as i32;
22599        let __s_b = self.gpu.stream();
22600        let mut b = __s_b.launch_builder(&f);
22601        b.arg(a).arg(g).arg(dst).arg(dst16).arg(&ni);
22602        unsafe {
22603            b.launch(cfg)?;
22604        }
22605        Ok(())
22606    }
22607
22608    /// step35 (Step-3.7-Flash) SEPARATE head-wise attention gate: one scalar per query head,
22609    /// broadcast over head_dim. `dst = a * sigmoid(g)` where `a`/`dst` are `[head_dim, n_head, T]`
22610    /// (the `q_gate_split` layout) and `g` is the PRE-sigmoid `attn_gate` projection output in
22611    /// token-major `[T, n_head]`. `dst16` is the optional fp16 operand for wo (None -> skipped).
22612    ///
22613    /// NOT interchangeable with `sig_mul_f16out`, which gates FULL WIDTH (qwen35 packs one gate
22614    /// value per (head, dim) element inside wq). Using this for that, or that for this, silently
22615    /// applies the wrong number of distinct gate values.
22616    #[allow(clippy::too_many_arguments)]
22617    pub fn attn_head_gate(
22618        &self,
22619        a: &CudaSlice<f32>,
22620        g: &CudaSlice<f32>,
22621        dst: &mut CudaSlice<f32>,
22622        dst16: Option<&mut CudaSlice<u8>>,
22623        head_dim: usize,
22624        n_head: usize,
22625        t: usize,
22626    ) -> Result<(), Box<dyn std::error::Error>> {
22627        let f = self.func("attn_head_gate_f32");
22628        let cfg = LaunchConfig::for_num_elems((head_dim * n_head * t) as u32);
22629        let (hd, nh, ti) = (head_dim as i32, n_head as i32, t as i32);
22630        // nullable device pointer by value (0 = skip), same convention as `l2_norm_pp`.
22631        let d16: u64 = match dst16 {
22632            Some(d) => self.addr_u8(d),
22633            None => 0,
22634        };
22635        let __s_b = self.gpu.stream();
22636        let mut b = __s_b.launch_builder(&f);
22637        b.arg(a)
22638            .arg(g)
22639            .arg(dst)
22640            .arg(&d16)
22641            .arg(&hd)
22642            .arg(&nh)
22643            .arg(&ti);
22644        unsafe {
22645            b.launch(cfg)?;
22646        }
22647        Ok(())
22648    }
22649
22650    /// step35 CLAMPED SwiGLU: `dst = min(silu(gate*gs), limit) * clamp(up*us, +-limit)`.
22651    /// Verbatim from llama.cpp `llama-graph.cpp:2146-2165` (routed, `swiglu_clamp_exp`) and
22652    /// `:1751-1770` (shared, `swiglu_clamp_shexp`), non-DEEPSEEK4 branch.
22653    ///
22654    /// This is NOT `swigluoai_mul_scaled`: that one clamps the gate BEFORE swish and multiplies by
22655    /// `(1 + clamp(up))`. Caller MUST check `limit > 1e-6` (upstream's eps gate) and use the plain
22656    /// `silu_mul_scaled` path otherwise — at limit=0 this kernel would clamp every positive
22657    /// activation to zero. On Step-3.7-Flash only layers 43 (7.0) and 44 (16.0) have a live limit.
22658    #[allow(clippy::too_many_arguments)]
22659    pub fn swiglu_clamped_mul_scaled(
22660        &self,
22661        gate: &CudaSlice<f32>,
22662        up: &CudaSlice<f32>,
22663        gs: f32,
22664        us: f32,
22665        limit: f32,
22666        dst: &mut CudaSlice<f32>,
22667        n: usize,
22668    ) -> Result<(), Box<dyn std::error::Error>> {
22669        debug_assert!(
22670            limit > 1e-6,
22671            "swiglu_clamped needs a live limit; use silu_mul_scaled"
22672        );
22673        let f = self.func("swiglu_clamped_mul_scaled_f32");
22674        let cfg = LaunchConfig::for_num_elems(n as u32);
22675        let ni = n as i32;
22676        let __s_b = self.gpu.stream();
22677        let mut b = __s_b.launch_builder(&f);
22678        b.arg(gate)
22679            .arg(up)
22680            .arg(&gs)
22681            .arg(&us)
22682            .arg(&limit)
22683            .arg(dst)
22684            .arg(&ni);
22685        unsafe {
22686            b.launch(cfg)?;
22687        }
22688        Ok(())
22689    }
22690
22691    /// gated RMSNorm: dst = RMSNorm(o, w[ncols]) * silu(z), per row of ncols. nrows blocks.
22692    pub fn gated_rmsnorm(
22693        &self,
22694        o: &CudaSlice<f32>,
22695        w: &CudaSlice<f32>,
22696        z: &CudaSlice<f32>,
22697        dst: &mut CudaSlice<f32>,
22698        ncols: usize,
22699        nrows: usize,
22700        eps: f32,
22701    ) -> Result<(), Box<dyn std::error::Error>> {
22702        let f = self.func("gated_rmsnorm_f32");
22703        let cfg = LaunchConfig {
22704            grid_dim: (nrows as u32, 1, 1),
22705            block_dim: (128, 1, 1),
22706            shared_mem_bytes: 0,
22707        };
22708        let (nc, e) = (ncols as i32, eps);
22709        let __s_b = self.gpu.stream();
22710        let mut b = __s_b.launch_builder(&f);
22711        b.arg(o).arg(w).arg(z).arg(dst).arg(&nc).arg(&e);
22712        unsafe {
22713            b.launch(cfg)?;
22714        }
22715        Ok(())
22716    }
22717
22718    /// f16out twin of `gated_rmsnorm` (task #17): epilogue also emits the fp16 operand for
22719    /// the ssm_out GEMM. Bit-identical class (same floats + the cvt kernel's __float2half).
22720    pub fn gated_rmsnorm_f16out(
22721        &self,
22722        o: &CudaSlice<f32>,
22723        w: &CudaSlice<f32>,
22724        z: &CudaSlice<f32>,
22725        dst: &mut CudaSlice<f32>,
22726        dst16: &mut CudaSlice<u8>,
22727        ncols: usize,
22728        nrows: usize,
22729        eps: f32,
22730    ) -> Result<(), Box<dyn std::error::Error>> {
22731        let f = self.func("gated_rmsnorm_f16out_f32");
22732        // block_dim MUST match gated_rmsnorm's (128): the reduction tree order pins the scale
22733        let cfg = LaunchConfig {
22734            grid_dim: (nrows as u32, 1, 1),
22735            block_dim: (128, 1, 1),
22736            shared_mem_bytes: 0,
22737        };
22738        let (nc, e) = (ncols as i32, eps);
22739        let __s_b = self.gpu.stream();
22740        let mut b = __s_b.launch_builder(&f);
22741        b.arg(o).arg(w).arg(z).arg(dst).arg(dst16).arg(&nc).arg(&e);
22742        unsafe {
22743            b.launch(cfg)?;
22744        }
22745        Ok(())
22746    }
22747
22748    /// add+RMSNorm emitting the f32 normed row AND its q8_1 quantization in one launch (the MoE
22749    /// layer input: z feeds the router matmul as f32, the expert dp4a as q8_1). BIT-IDENTICAL to
22750    /// add_rms_norm + quantize_q8_1. Returns (q, d) alongside the caller-provided res/z buffers.
22751    #[allow(clippy::too_many_arguments)]
22752    pub fn add_rms_norm_zq8(
22753        &self,
22754        a: &CudaSlice<f32>,
22755        b_in: &CudaSlice<f32>,
22756        w: &CudaSlice<f32>,
22757        res: &mut CudaSlice<f32>,
22758        z: &mut CudaSlice<f32>,
22759        ncols: usize,
22760        nrows: usize,
22761        eps: f32,
22762    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
22763        assert!(ncols % 32 == 0);
22764        let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
22765        let mut d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
22766        let f = self.func("add_rms_norm_zq8");
22767        let cfg = LaunchConfig {
22768            grid_dim: (nrows as u32, 1, 1),
22769            block_dim: (1024, 1, 1),
22770            shared_mem_bytes: 0,
22771        };
22772        let (nc, ep) = (ncols as i32, eps);
22773        let __s_b = self.gpu.stream();
22774        let mut b = __s_b.launch_builder(&f);
22775        b.arg(a)
22776            .arg(b_in)
22777            .arg(w)
22778            .arg(res)
22779            .arg(z)
22780            .arg(&mut q)
22781            .arg(&mut d)
22782            .arg(&nc)
22783            .arg(&ep);
22784        unsafe {
22785            b.launch(cfg)?;
22786        }
22787        Ok((q, d))
22788    }
22789
22790    /// gated RMSNorm emitting q8_1 directly (fused quantize epilogue) — the ssm_out matvec input.
22791    /// BIT-IDENTICAL bytes to gated_rmsnorm + quantize_q8_1 (ncols % 32 == 0; blocks never straddle
22792    /// rows). Saves one launch per linear-attn layer (36/token on the 9B).
22793    /// z-view twins of gated_rmsnorm(+f16out) — task #16 batched-prime split removal.
22794    pub fn gated_rmsnorm_zv(
22795        &self,
22796        o: &CudaSlice<f32>,
22797        w: &CudaSlice<f32>,
22798        z: &cudarc::driver::CudaView<f32>,
22799        dst: &mut CudaSlice<f32>,
22800        ncols: usize,
22801        nrows: usize,
22802        eps: f32,
22803    ) -> Result<(), Box<dyn std::error::Error>> {
22804        let f = self.func("gated_rmsnorm_f32");
22805        let cfg = LaunchConfig {
22806            grid_dim: (nrows as u32, 1, 1),
22807            block_dim: (128, 1, 1),
22808            shared_mem_bytes: 0,
22809        };
22810        let (nc, e) = (ncols as i32, eps);
22811        let __s_b = self.gpu.stream();
22812        let mut b = __s_b.launch_builder(&f);
22813        b.arg(o).arg(w).arg(z).arg(dst).arg(&nc).arg(&e);
22814        unsafe {
22815            b.launch(cfg)?;
22816        }
22817        Ok(())
22818    }
22819
22820    pub fn gated_rmsnorm_f16out_zv(
22821        &self,
22822        o: &CudaSlice<f32>,
22823        w: &CudaSlice<f32>,
22824        z: &cudarc::driver::CudaView<f32>,
22825        dst: &mut CudaSlice<f32>,
22826        dst16: &mut CudaSlice<u8>,
22827        ncols: usize,
22828        nrows: usize,
22829        eps: f32,
22830    ) -> Result<(), Box<dyn std::error::Error>> {
22831        let f = self.func("gated_rmsnorm_f16out_f32");
22832        // block_dim MUST match gated_rmsnorm's (128): the reduction tree order pins the scale
22833        let cfg = LaunchConfig {
22834            grid_dim: (nrows as u32, 1, 1),
22835            block_dim: (128, 1, 1),
22836            shared_mem_bytes: 0,
22837        };
22838        let (nc, e) = (ncols as i32, eps);
22839        let __s_b = self.gpu.stream();
22840        let mut b = __s_b.launch_builder(&f);
22841        b.arg(o).arg(w).arg(z).arg(dst).arg(dst16).arg(&nc).arg(&e);
22842        unsafe {
22843            b.launch(cfg)?;
22844        }
22845        Ok(())
22846    }
22847
22848    pub fn gated_rmsnorm_q8_1(
22849        &self,
22850        o: &CudaSlice<f32>,
22851        w: &CudaSlice<f32>,
22852        z: &CudaSlice<f32>,
22853        ncols: usize,
22854        nrows: usize,
22855        eps: f32,
22856    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
22857        assert!(ncols % 32 == 0);
22858        let f = self.func("gated_rmsnorm_q8_1");
22859        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
22860        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
22861        let cfg = LaunchConfig {
22862            grid_dim: (nrows as u32, 1, 1),
22863            block_dim: (128, 1, 1),
22864            shared_mem_bytes: 0,
22865        };
22866        let (nc, ep) = (ncols as i32, eps);
22867        let __s_b = self.gpu.stream();
22868        let mut b = __s_b.launch_builder(&f);
22869        b.arg(o)
22870            .arg(w)
22871            .arg(z)
22872            .arg(&mut out_q)
22873            .arg(&mut out_d)
22874            .arg(&nc)
22875            .arg(&ep);
22876        unsafe {
22877            b.launch(cfg)?;
22878        }
22879        Ok((out_q, out_d))
22880    }
22881
22882    /// transpose [rows,cols] row-major -> [cols,rows] row-major.
22883    pub fn transpose(
22884        &self,
22885        inp: &CudaSlice<f32>,
22886        rows: usize,
22887        cols: usize,
22888    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22889        let f = self.func("transpose_f32");
22890        let mut out = self.zeros(rows * cols)?;
22891        let cfg = LaunchConfig::for_num_elems((rows * cols) as u32);
22892        let (r, c) = (rows as i32, cols as i32);
22893        let __s_b = self.gpu.stream();
22894        let mut b = __s_b.launch_builder(&f);
22895        b.arg(inp).arg(&mut out).arg(&r).arg(&c);
22896        unsafe {
22897            b.launch(cfg)?;
22898        }
22899        Ok(out)
22900    }
22901
22902    /// repeat-interleave heads: in[head_dim,n_in,T] -> out[head_dim,n_out,T].
22903    pub fn repeat_heads(
22904        &self,
22905        inp: &CudaSlice<f32>,
22906        out: &mut CudaSlice<f32>,
22907        head_dim: usize,
22908        n_in: usize,
22909        n_out: usize,
22910        t: usize,
22911    ) -> Result<(), Box<dyn std::error::Error>> {
22912        let f = self.func("repeat_heads_f32");
22913        let cfg = LaunchConfig::for_num_elems((head_dim * n_out * t) as u32);
22914        let (hd, ni, no, ti) = (head_dim as i32, n_in as i32, n_out as i32, t as i32);
22915        let __s_b = self.gpu.stream();
22916        let mut b = __s_b.launch_builder(&f);
22917        b.arg(inp).arg(out).arg(&hd).arg(&ni).arg(&no).arg(&ti);
22918        unsafe {
22919            b.launch(cfg)?;
22920        }
22921        Ok(())
22922    }
22923
22924    /// q|gate split (on-device). qf:[T, n_head*2*head_dim] -> q_out,gate_out:[head_dim,n_head,T].
22925    /// Replaces the dtoh->host-double-loop->htod in full_attn / full_attn_decode.
22926    ///
22927    /// Bounds-checked against `qf` before the launch: the kernel reads `2*head_dim*n_head*t`
22928    /// floats, and running it on a `wq` output that carries no fused gate reads 2x off the end
22929    /// (silently, on the device). A layout mismatch is a typed `FusedQGateExtent` here instead.
22930    pub fn q_gate_split(
22931        &self,
22932        qf: &CudaSlice<f32>,
22933        q_out: &mut CudaSlice<f32>,
22934        gate_out: &mut CudaSlice<f32>,
22935        head_dim: usize,
22936        n_head: usize,
22937        t: usize,
22938    ) -> Result<(), Box<dyn std::error::Error>> {
22939        memra_gguf::config::check_fused_q_gate_extent(qf.len(), head_dim, n_head, t)?;
22940        let out_need = head_dim * n_head * t;
22941        if q_out.len() < out_need || gate_out.len() < out_need {
22942            return Err(format!(
22943                "q_gate_split destinations too small: need {out_need} each, have q={} gate={}",
22944                q_out.len(),
22945                gate_out.len()
22946            )
22947            .into());
22948        }
22949        let f = self.func("q_gate_split_f32");
22950        let cfg = LaunchConfig::for_num_elems((head_dim * n_head * t) as u32);
22951        let (hd, nh, ti) = (head_dim as i32, n_head as i32, t as i32);
22952        let __s_b = self.gpu.stream();
22953        let mut b = __s_b.launch_builder(&f);
22954        b.arg(qf)
22955            .arg(q_out)
22956            .arg(gate_out)
22957            .arg(&hd)
22958            .arg(&nh)
22959            .arg(&ti);
22960        unsafe {
22961            b.launch(cfg)?;
22962        }
22963        Ok(())
22964    }
22965
22966    /// qkv->GDN repack (on-device). conv_out:[conv_dim,T] channel-major ->
22967    /// q_g/k_g/v_g:[d_state,num_v,T] with q/k head-repeat kh = vh % num_k (validated modulo mapping).
22968    /// Replaces the dtoh->host-q/k/v-repack->3x-htod in linear_attn / linear_attn_decode.
22969    pub fn qkv_to_gdn_repack(
22970        &self,
22971        conv_out: &CudaSlice<f32>,
22972        q_g: &mut CudaSlice<f32>,
22973        k_g: &mut CudaSlice<f32>,
22974        v_g: &mut CudaSlice<f32>,
22975        d_state: usize,
22976        num_v: usize,
22977        num_k: usize,
22978        key_dim: usize,
22979        t: usize,
22980    ) -> Result<(), Box<dyn std::error::Error>> {
22981        let f = self.func("qkv_to_gdn_repack_f32");
22982        let cfg = LaunchConfig::for_num_elems((d_state * num_v * t) as u32);
22983        let (ds, nv, nk, kd, ti) = (
22984            d_state as i32,
22985            num_v as i32,
22986            num_k as i32,
22987            key_dim as i32,
22988            t as i32,
22989        );
22990        let __s_b = self.gpu.stream();
22991        let mut b = __s_b.launch_builder(&f);
22992        b.arg(conv_out)
22993            .arg(q_g)
22994            .arg(k_g)
22995            .arg(v_g)
22996            .arg(&ds)
22997            .arg(&nv)
22998            .arg(&nk)
22999            .arg(&kd)
23000            .arg(&ti);
23001        unsafe {
23002            b.launch(cfg)?;
23003        }
23004        Ok(())
23005    }
23006
23007    /// conv left zero-pad (prefill from zero state). src:[conv_dim,T] -> dst:[conv_dim,T+pad],
23008    /// cols 0..pad = 0, cols pad..pad+T = src. `dst` MUST be pre-zeroed. No dtoh/host-loop/htod.
23009    pub fn conv_left_pad(
23010        &self,
23011        src: &CudaSlice<f32>,
23012        dst: &mut CudaSlice<f32>,
23013        conv_dim: usize,
23014        t: usize,
23015        pad: usize,
23016    ) -> Result<(), Box<dyn std::error::Error>> {
23017        let f = self.func("conv_left_pad_f32");
23018        let cfg = LaunchConfig::for_num_elems((conv_dim * t) as u32);
23019        let (cd, ti, p) = (conv_dim as i32, t as i32, pad as i32);
23020        let __s_b = self.gpu.stream();
23021        let mut b = __s_b.launch_builder(&f);
23022        b.arg(src).arg(dst).arg(&cd).arg(&ti).arg(&p);
23023        unsafe {
23024            b.launch(cfg)?;
23025        }
23026        Ok(())
23027    }
23028
23029    /// conv-state assemble + ring roll (decode T=1). conv_state:[conv_dim,pad] (resident),
23030    /// qkv_col:[conv_dim] -> conv_in:[conv_dim,pad+1]; AND rolls conv_state (keep last pad cols).
23031    /// Replaces the dtoh->host-conv-ring-assemble->ring-update->htod in linear_attn_decode.
23032    pub fn conv_assemble_and_roll(
23033        &self,
23034        qkv_col: &CudaSlice<f32>,
23035        conv_state: &mut CudaSlice<f32>,
23036        conv_in: &mut CudaSlice<f32>,
23037        conv_dim: usize,
23038        pad: usize,
23039    ) -> Result<(), Box<dyn std::error::Error>> {
23040        let f = self.func("conv_assemble_and_roll_f32");
23041        let cfg = LaunchConfig::for_num_elems(conv_dim as u32);
23042        let (cd, p) = (conv_dim as i32, pad as i32);
23043        let __s_b = self.gpu.stream();
23044        let mut b = __s_b.launch_builder(&f);
23045        b.arg(qkv_col).arg(conv_state).arg(conv_in).arg(&cd).arg(&p);
23046        unsafe {
23047            b.launch(cfg)?;
23048        }
23049        Ok(())
23050    }
23051
23052    /// RANK3 LEVER (conv fuse, T=1 DECODE): fused conv_assemble_and_roll + ssm_conv1d_silu in ONE
23053    /// launch. Assembles the conv window [conv_state | qkv_col] in registers, computes the depthwise
23054    /// causal conv + SiLU into `conv_out`, and rolls the ring — never materializing conv_in to HBM.
23055    /// Replaces e.conv_assemble_and_roll(...) + e.ssm_conv1d(...). BIT-IDENTICAL to that two-kernel
23056    /// sequence (same 8-wide accumulation order, same SiLU). `conv_out` is [conv_dim] (T=1).
23057    pub fn ssm_conv1d_fused_decode(
23058        &self,
23059        qkv_col: &CudaSlice<f32>,
23060        conv_state: &mut CudaSlice<f32>,
23061        w: &CudaSlice<f32>,
23062        conv_out: &mut CudaSlice<f32>,
23063        conv_dim: usize,
23064        d_conv: usize,
23065    ) -> Result<(), Box<dyn std::error::Error>> {
23066        let f = self.func("ssm_conv1d_fused_decode_f32");
23067        let cfg = LaunchConfig::for_num_elems(conv_dim as u32);
23068        let (cd, dc) = (conv_dim as i32, d_conv as i32);
23069        let __s_b = self.gpu.stream();
23070        let mut b = __s_b.launch_builder(&f);
23071        b.arg(qkv_col)
23072            .arg(conv_state)
23073            .arg(w)
23074            .arg(conv_out)
23075            .arg(&cd)
23076            .arg(&dc);
23077        unsafe {
23078            b.launch(cfg)?;
23079        }
23080        Ok(())
23081    }
23082
23083    /// Copy a contiguous range [start, start+len) out of src into a fresh slice (device→device via host).
23084    /// Used for qkv split views. Small/rare; not perf-critical in Stage 1.
23085    pub fn slice_range(
23086        &self,
23087        src: &CudaSlice<f32>,
23088        start: usize,
23089        len: usize,
23090    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
23091        let host = self.gpu.stream().clone_dtoh(src)?;
23092        self.gpu.stream().synchronize()?;
23093        Ok(self.htod(&host[start..start + len])?)
23094    }
23095}
23096
23097#[cfg(test)]
23098mod target_dispatch_tests {
23099    use super::legacy_quant_gemm_allowed;
23100
23101    #[test]
23102    fn legacy_quant_gemm_arch_policy_honors_the_escape_hatch() {
23103        // sm_120a native lane
23104        assert!(legacy_quant_gemm_allowed(false, false, false));
23105        assert!(!legacy_quant_gemm_allowed(false, false, true));
23106        // pure portable lane (sm_89): gated
23107        assert!(!legacy_quant_gemm_allowed(true, false, false));
23108        assert!(!legacy_quant_gemm_allowed(true, false, true));
23109        // Hopper-MMA lane (sm_90a): portable build, int8-MMA GEMM re-admitted
23110        assert!(legacy_quant_gemm_allowed(true, true, false));
23111        assert!(!legacy_quant_gemm_allowed(true, true, true));
23112    }
23113
23114    #[cfg(all(memra_portable_cuda, not(memra_hopper_mma)))]
23115    #[test]
23116    fn portable_build_disables_legacy_quant_gemm_without_an_env_override() {
23117        assert!(!legacy_quant_gemm_allowed(
23118            cfg!(memra_portable_cuda),
23119            cfg!(memra_hopper_mma),
23120            false
23121        ));
23122    }
23123
23124    #[cfg(memra_hopper_mma)]
23125    #[test]
23126    fn hopper_mma_build_re_admits_legacy_quant_gemm() {
23127        assert!(legacy_quant_gemm_allowed(
23128            cfg!(memra_portable_cuda),
23129            cfg!(memra_hopper_mma),
23130            false
23131        ));
23132        assert!(super::portable_mma_gated() == false);
23133    }
23134}
23135
23136/// The memra-kv device seam (Phase D): the cache's 7 ops delegate to the engine's
23137/// inherent methods (inherent methods win name resolution, so no recursion).
23138impl memra_kv::KvDev for Engine {
23139    fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
23140        Engine::zeros(self, n)
23141    }
23142    fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
23143        Engine::uninit(self, n)
23144    }
23145    fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
23146        Engine::alloc_u8(self, n)
23147    }
23148    fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
23149        Engine::htod_i32(self, v)
23150    }
23151    fn clone_dtod(
23152        &self,
23153        src: &CudaSlice<f32>,
23154    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
23155        Engine::clone_dtod(self, src)
23156    }
23157    fn copy_into(
23158        &self,
23159        dst: &mut CudaSlice<f32>,
23160        off: usize,
23161        src: &CudaSlice<f32>,
23162        len: usize,
23163    ) -> Result<(), Box<dyn std::error::Error>> {
23164        Engine::copy_into(self, dst, off, src, len)
23165    }
23166    fn set_i32_one(
23167        &self,
23168        d: &mut CudaSlice<i32>,
23169        v: i32,
23170    ) -> Result<(), Box<dyn std::error::Error>> {
23171        Engine::set_i32_one(self, d, v)
23172    }
23173}
23174
23175#[cfg(test)]
23176mod fused_gate_bounds_tests {
23177    use super::*;
23178
23179    /// The fused `[q|gate]` split's read-site guard, on the device.
23180    ///
23181    /// `q_gate_split_f32` reads `2*head_dim*n_head*T` floats out of `qf`. A checkpoint whose gate
23182    /// is a SEPARATE tensor produces a `wq` output of exactly half that, so before 2026-08-19 the
23183    /// kernel launched and read 2x past the end of the allocation — an out-of-bounds DEVICE read:
23184    /// no panic, no error, just whatever memory follows. The guard turns it into a typed
23185    /// `FusedQGateExtent` before the launch.
23186    ///
23187    /// Catch demonstration for this test (guard temporarily removed, then restored):
23188    /// `compute-sanitizer --tool memcheck` on the half-width case reported invalid `__global__`
23189    /// reads of size 4 in `q_gate_split_f32`; with the guard in place the same run is clean and
23190    /// the call returns `Err`. Receipt in the lane report.
23191    #[test]
23192    #[ignore = "requires a CUDA GPU"]
23193    fn q_gate_split_refuses_a_separate_gate_wq_instead_of_reading_past_it() {
23194        let e = Engine::new(0).unwrap();
23195        let (head_dim, n_head, t) = (8usize, 4usize, 2usize);
23196        let fused = 2 * head_dim * n_head * t;
23197        let out_n = head_dim * n_head * t;
23198
23199        // half-width `qf` = the separate-gate / ungated layout. MUST be refused.
23200        let narrow = e.htod(&vec![1.0f32; out_n]).unwrap();
23201        let mut q = e.uninit(out_n).unwrap();
23202        let mut gate = e.uninit(out_n).unwrap();
23203        let err = e
23204            .q_gate_split(&narrow, &mut q, &mut gate, head_dim, n_head, t)
23205            .expect_err("half-width wq must be refused, not read past")
23206            .to_string();
23207        assert!(err.contains("NO fused gate"), "{err}");
23208        assert!(err.contains(&format!("{fused}")), "{err}");
23209
23210        // full-width `qf` = a real qwen3.5 fused layout. MUST still run, and split correctly:
23211        // per head hh the block is [q(head_dim) | gate(head_dim)] at stride 2*head_dim.
23212        let host: Vec<f32> = (0..fused).map(|i| i as f32).collect();
23213        let wide = e.htod(&host).unwrap();
23214        e.q_gate_split(&wide, &mut q, &mut gate, head_dim, n_head, t)
23215            .expect("full-width wq splits");
23216        let (qh, gh) = (e.dtoh(&q).unwrap(), e.dtoh(&gate).unwrap());
23217        for tok in 0..t {
23218            for hh in 0..n_head {
23219                for d in 0..head_dim {
23220                    let base = tok * (n_head * 2 * head_dim) + hh * (2 * head_dim);
23221                    let idx = tok * (n_head * head_dim) + hh * head_dim + d;
23222                    assert_eq!(qh[idx], host[base + d], "q t{tok} h{hh} d{d}");
23223                    assert_eq!(gh[idx], host[base + head_dim + d], "gate t{tok} h{hh} d{d}");
23224                }
23225            }
23226        }
23227
23228        // undersized destinations are refused too (the other half of the extent contract)
23229        let mut small = e.uninit(out_n - 1).unwrap();
23230        assert!(
23231            e.q_gate_split(&wide, &mut small, &mut gate, head_dim, n_head, t)
23232                .is_err()
23233        );
23234    }
23235}
23236
23237/// FULL-WIDTH-ROPE CONTRACT on the fused rms_norm+qkv+rope kernels
23238/// (lane/graph-s-key-exactness-20260819, probe O-10 follow-up). CPU-only: the guard runs before
23239/// any launch, so the refusal is testable without a device.
23240#[cfg(test)]
23241mod fused_rope_width_tests {
23242    use super::Engine;
23243
23244    /// gemma-4: rotary width == head width on both classes (GGUF 256/256 and 512/512, and the
23245    /// safetensors route derives the same), which is why the fusion is legal there today.
23246    #[test]
23247    fn full_width_is_accepted() {
23248        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 256, 256).is_ok());
23249        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_cat", 512, 512).is_ok());
23250        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append", 128, 128).is_ok());
23251    }
23252
23253    /// The widths the gemma-4 31B OFFICIAL artifact declares, read from its own GGUF header
23254    /// (`gemma-4-31B-it-official-Q8_0-MTP.gguf`, box3, 2026-08-19):
23255    ///
23256    /// ```text
23257    /// attention.key_length     512   rope.dimension_count     512   (global class)
23258    /// attention.key_length_swa 256   rope.dimension_count_swa 256   (SWA class)
23259    /// ```
23260    ///
23261    /// Both classes satisfy `n_rot == head_dim`, which is why the fusion is legal for gemma and
23262    /// why `HybridModel::gemma4_rope_dims` can feed this guard without refusing what we serve.
23263    /// An artifact that ever declares otherwise gets a loud refusal at the first fused launch
23264    /// instead of a silently over-rotated head.
23265    #[test]
23266    fn gemma4_official_artifact_widths_pass() {
23267        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 512, 512).is_ok());
23268        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append_dc", 256, 256).is_ok());
23269    }
23270
23271    /// The bug this guard exists to make impossible: a partial-rotary arch fused onto a kernel
23272    /// with no `n_dims`, silently rotating the pass-through band.
23273    #[test]
23274    fn partial_rotary_is_refused_with_the_geometry_named() {
23275        // qwen3.5: n_rot 64 of head_dim 256 (the shape probe O-10 pinned in the split path).
23276        let err = Engine::full_width_rope_only("rms_norm_qkv_rope", 64, 256)
23277            .expect_err("partial rotary must refuse");
23278        let msg = err.to_string();
23279        assert!(msg.contains("PARTIAL ROTARY REFUSED"), "{msg}");
23280        assert!(msg.contains("n_rot 64"), "{msg}");
23281        assert!(msg.contains("head_dim 256"), "{msg}");
23282        assert!(
23283            msg.contains("64..256"),
23284            "names the band it would corrupt: {msg}"
23285        );
23286        // step35 full-attn: 64 of 128 (upstream halves n_rot_full).
23287        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append_dc", 64, 128).is_err());
23288        // and the reverse mismatch (a wider rope than the head) is not "close enough" either.
23289        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 256, 128).is_err());
23290    }
23291}