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    /// WEIGHT PREFETCH (SOTA item 3, 2026-07-13, DEFAULT ON): during a bandwidth-idle
1777    /// window (the fa launch reads KV, not weights) prefetch the NEXT matvec's
1778    /// decode-plane bytes into L2 so it reads L2-warm. Value-free scheduling op — same
1779    /// class as prefetch_l2 (numerics untouched by construction). Wired only where it
1780    /// measured positive: the E4B dc attn arm (+0.65%). 26B (flat — MoE ffn dominates),
1781    /// 31B (−0.2% — decode at the DRAM wall) and the ffn gate/up cascade (−1% — 29MB/layer
1782    /// floods the fill path) all probed and NOT wired. MEMRA_WPF=0 rollback seam.
1783    pub fn wpf_level() -> u32 {
1784        static ON: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
1785        *ON.get_or_init(|| {
1786            std::env::var("MEMRA_WPF")
1787                .ok()
1788                .and_then(|v| v.parse().ok())
1789                .unwrap_or(1)
1790        })
1791    }
1792
1793    /// PDL launch arm (SOTA item 2, 2026-07-13, DEFAULT ON): the six MEMRA_PDL_ENTRY glue
1794    /// kernels launch through cuLaunchKernelEx with PROGRAMMATIC_STREAM_SERIALIZATION — the
1795    /// grid launches while the predecessor drains (~120ns/kernel back, pdl_probe), the
1796    /// kernels' entry grid-dep sync restores read order (SASS-audited: ACQBULK precedes
1797    /// every LDG in all six). Valid windows: E4B +1.0-1.2% (128 AND 384-tok gens);
1798    /// 26B/31B/qwen flat no-harm. Battery: kernel-check GREEN, run-gen tokens IDENTICAL x3
1799    /// gemma, spec 64/64 E4B K=1/4/8 + 26B/31B K=4 + qwen PASS. Works eager AND under
1800    /// capture (capture encodes native programmatic edges — the post-capture edge-REWRITE
1801    /// arm died: engine graphs hold cuMemAllocAsync alloc nodes, edge edits on those return
1802    /// CUDA_ERROR_NOT_SUPPORTED). MEMRA_PDL=0 rollback seam.
1803    /// See the `verify_exact` field. Scoped by the dflash round around its t=16 verify.
1804    pub fn set_verify_exact(&self, on: bool) {
1805        self.verify_exact
1806            .store(on, std::sync::atomic::Ordering::Relaxed);
1807    }
1808    pub(crate) fn verify_exact_on(&self) -> bool {
1809        self.verify_exact.load(std::sync::atomic::Ordering::Relaxed)
1810    }
1811
1812    /// m=1 norm+rope+append fold seam (2026-07-23): MEMRA_QKV_APPEND=0 reverts to the
1813    /// fused-norm-rope + standalone-append pair (the exact-oracle bisect arm).
1814    pub fn qkv_append_on() -> bool {
1815        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1816        *ON.get_or_init(|| {
1817            std::env::var("MEMRA_QKV_APPEND")
1818                .map(|v| v != "0")
1819                .unwrap_or(true)
1820        })
1821    }
1822
1823    /// PDL wave-B1a seam: the four dense-glue kernels (rms_norm_f32, add_rms_norm_f32,
1824    /// add_scale_rms_norm_q8_1, quantize_q8_1). MEMRA_PDL_WB=0 reverts alone.
1825    pub fn pdl_wb_on() -> bool {
1826        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1827        *ON.get_or_init(|| {
1828            std::env::var("MEMRA_PDL_WB")
1829                .map(|v| v != "0")
1830                .unwrap_or(true)
1831        })
1832    }
1833
1834    /// Trunk-kernels norm ILP seam (lane/dspark-trunk-kernels-20260820): the T-row verify
1835    /// norms (rms_norm_f32 / add_rms_norm_f32 at grid=T, block=256) are serial-latency
1836    /// chains — 20 strided scalar load->fma rounds measured 11.8-12.2us/inst (nsys-B verify
1837    /// scope: 130 inst/rd = 1.51 ms/rd). The `_v2` twins unroll the element loop 4-deep
1838    /// (independent loads in flight; SAME per-thread element order into ONE accumulator,
1839    /// reduce VERBATIM) — BIT-IDENTICAL per row at every (ncols, blockDim).
1840    /// MEMRA_NORM_ILP=0 reverts to the v1 kernels alone.
1841    pub fn norm_ilp_on() -> bool {
1842        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1843        *ON.get_or_init(|| {
1844            std::env::var("MEMRA_NORM_ILP")
1845                .map(|v| v != "0")
1846                .unwrap_or(true)
1847        })
1848    }
1849
1850    /// Trunk-kernels FFN dual seam (lane/dspark-trunk-kernels-20260820): the qwen35
1851    /// t-parallel verify FFN pair rides the PROVEN dual gate+up doors
1852    /// (`matmul_decode_exact_dual_pre` + `silu_mul_scaled_q8_1`, the q27 verify shape —
1853    /// bit-identical per (tensor,token,row), kernel-check-pinned, MEMRA_SPEC_DUAL_T
1854    /// receipts) instead of two singles + silu_mul + a standalone quantize. The doors
1855    /// existed but the qwen35 body never called them (nsys-B verify scope: gate+up singles
1856    /// = 107 launches/rd at grid 4352). MEMRA_TK_FFN_DUAL=0 reverts to the singles chain.
1857    pub fn tk_ffn_dual_on() -> bool {
1858        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1859        *ON.get_or_init(|| {
1860            std::env::var("MEMRA_TK_FFN_DUAL")
1861                .map(|v| v != "0")
1862                .unwrap_or(true)
1863        })
1864    }
1865
1866    /// PDL wave-A seam: the mmvq matvec PDL launches only (the six glue kernels keep
1867    /// their own MEMRA_PDL master seam). MEMRA_PDL_MMVQ=0 reverts wave-A alone — the
1868    /// per-model no-harm bisect knob.
1869    pub fn pdl_mmvq_on() -> bool {
1870        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1871        *ON.get_or_init(|| {
1872            std::env::var("MEMRA_PDL_MMVQ")
1873                .map(|v| v != "0")
1874                .unwrap_or(true)
1875        })
1876    }
1877
1878    pub fn pdl_on() -> bool {
1879        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1880        *ON.get_or_init(|| std::env::var("MEMRA_PDL").map(|v| v != "0").unwrap_or(true))
1881    }
1882
1883    /// PDL wave-B seam (gap-diagnosis arc, GAP-DIAGNOSIS.md verdict 8): the gemma
1884    /// NVFP4mix decode chain's hot kernels — nvfp4 fused2/mr2 and the q8_0 `_rp`
1885    /// singles — join the wave-A launch class. Scheduling-only (the entry macro waits
1886    /// on the producer before any read), bit-identical by construction.
1887    /// MEMRA_PDL_NVFP4=0 reverts wave-B alone.
1888    pub fn pdl_nvfp4q8_on() -> bool {
1889        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1890        *ON.get_or_init(|| {
1891            std::env::var("MEMRA_PDL_NVFP4")
1892                .map(|v| v != "0")
1893                .unwrap_or(true)
1894        })
1895    }
1896
1897    /// Raw CUfunction for a PDL-attributed launch: the SAME kernels.fatbin loaded once more
1898    /// through the raw driver API (cudarc hides its CUfunction handles; a duplicate module
1899    /// of tiny glue kernels is free). Resolved lazily per name, cached process-wide.
1900    /// Fused t=1 q4_0 mr policy: env MEMRA_Q40_MR wins (1/2); else the per-model
1901    /// FUSED_MR1_DEFAULT (dense gemma = mr1, MoE = mr2 — see the static's doc).
1902    fn q40_mr1_on() -> bool {
1903        static Q40MR: std::sync::OnceLock<Option<u32>> = std::sync::OnceLock::new();
1904        match *Q40MR.get_or_init(|| {
1905            std::env::var("MEMRA_Q40_MR")
1906                .ok()
1907                .and_then(|v| v.parse().ok())
1908        }) {
1909            Some(v) => v == 1,
1910            None => crate::FUSED_MR1_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
1911        }
1912    }
1913
1914    /// PDL wave-B2: flash-module PDL functions. `g` selects the kf8vf8 flavor — the
1915    /// caller MUST pass the SAME flavor its builder launch would resolve (fa_func/func_g
1916    /// mirror); the flavors differ semantically (KV byte formats), a wrong-module launch
1917    /// writes wrong bytes silently.
1918    fn pdl_func_flash(
1919        &self,
1920        g: bool,
1921        name: &'static str,
1922    ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
1923        use cudarc::driver::sys as cu;
1924        // PER-CONTEXT caches (M1-PP2 cross-device fix, 8x box 2026-08-02): CUmodule and
1925        // CUfunction handles are CONTEXT-scoped, and a remote-stage Engine
1926        // (MEMRA_PP_DEVICES=a,b) lives in the other device's primary context. The old
1927        // process-wide OnceLock cache handed stage 1 the dev-a handles, so every stage-1
1928        // launch_pdl* died CUDA_ERROR_INVALID_HANDLE. Key module + function caches by
1929        // this engine's CUcontext; single-context runs behave exactly as before.
1930        static MODS: std::sync::Mutex<Option<std::collections::HashMap<(usize, bool), usize>>> =
1931            std::sync::Mutex::new(None);
1932        static FNS: std::sync::Mutex<
1933            Option<std::collections::HashMap<(usize, bool, &'static str), usize>>,
1934        > = std::sync::Mutex::new(None);
1935        let ctx_key = self.ctx().cu_ctx() as usize;
1936        if let Some(&f) = FNS
1937            .lock()
1938            .unwrap()
1939            .get_or_insert_with(Default::default)
1940            .get(&(ctx_key, g, name))
1941        {
1942            return Ok(f as cu::CUfunction);
1943        }
1944        let module = {
1945            let mut mods = MODS.lock().unwrap();
1946            let map = mods.get_or_insert_with(Default::default);
1947            match map.get(&(ctx_key, g)) {
1948                Some(&m) => m,
1949                None => {
1950                    let m = self.pdl_load_module_in_ctx(if g {
1951                        FLASH_FATBIN_KF8VF8
1952                    } else {
1953                        FLASH_FATBIN
1954                    })?;
1955                    map.insert((ctx_key, g), m);
1956                    m
1957                }
1958            }
1959        };
1960        let cname = std::ffi::CString::new(name)?;
1961        let mut f: cu::CUfunction = std::ptr::null_mut();
1962        let r = unsafe { cu::cuModuleGetFunction(&mut f, module as cu::CUmodule, cname.as_ptr()) };
1963        if r != cu::CUresult::CUDA_SUCCESS {
1964            return Err(format!("pdl_func_flash {name} (g={g}): {r:?}").into());
1965        }
1966        FNS.lock()
1967            .unwrap()
1968            .get_or_insert_with(Default::default)
1969            .insert((ctx_key, g, name), f as usize);
1970        Ok(f)
1971    }
1972
1973    /// Load a fatbin as a raw CUmodule IN THIS ENGINE'S CONTEXT. `cuModuleLoadData` binds
1974    /// the module to the thread's CURRENT context — a remote-stage engine must not
1975    /// inherit the primary's (the INVALID_HANDLE class above). Restores the caller's
1976    /// current context before returning.
1977    fn pdl_load_module_in_ctx(&self, bytes: &[u8]) -> Result<usize, Box<dyn std::error::Error>> {
1978        use cudarc::driver::sys as cu;
1979        let mut prev: cu::CUcontext = std::ptr::null_mut();
1980        unsafe {
1981            cu::cuCtxGetCurrent(&mut prev).result()?;
1982        }
1983        self.ctx().bind_to_thread()?;
1984        let mut m: cu::CUmodule = std::ptr::null_mut();
1985        let r = unsafe { cu::cuModuleLoadData(&mut m, bytes.as_ptr() as *const std::ffi::c_void) };
1986        let restore = if prev.is_null() {
1987            cu::CUresult::CUDA_SUCCESS
1988        } else {
1989            unsafe { cu::cuCtxSetCurrent(prev) }
1990        };
1991        if r != cu::CUresult::CUDA_SUCCESS {
1992            return Err(format!("pdl module load: {r:?}").into());
1993        }
1994        if restore != cu::CUresult::CUDA_SUCCESS {
1995            return Err(format!("pdl module load: ctx restore {restore:?}").into());
1996        }
1997        Ok(m as usize)
1998    }
1999
2000    fn pdl_func(
2001        &self,
2002        name: &'static str,
2003    ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
2004        use cudarc::driver::sys as cu;
2005        // PER-CONTEXT caches — same M1-PP2 cross-device fix as pdl_func_flash (handles
2006        // are context-scoped; key everything by this engine's CUcontext).
2007        static MODULES: std::sync::Mutex<Option<std::collections::HashMap<usize, usize>>> =
2008            std::sync::Mutex::new(None);
2009        // PDL wave-A: the mmvq kernels live in the qmatvec fatbin, not kernels.cu — second
2010        // duplicate module, loaded lazily on the first kernels-module miss.
2011        static QMODULES: std::sync::Mutex<Option<std::collections::HashMap<usize, usize>>> =
2012            std::sync::Mutex::new(None);
2013        static FNS: std::sync::Mutex<
2014            Option<std::collections::HashMap<(usize, &'static str), usize>>,
2015        > = std::sync::Mutex::new(None);
2016        let ctx_key = self.ctx().cu_ctx() as usize;
2017        if let Some(&f) = FNS
2018            .lock()
2019            .unwrap()
2020            .get_or_insert_with(Default::default)
2021            .get(&(ctx_key, name))
2022        {
2023            return Ok(f as cu::CUfunction);
2024        }
2025        let module = {
2026            let mut mods = MODULES.lock().unwrap();
2027            let map = mods.get_or_insert_with(Default::default);
2028            match map.get(&ctx_key) {
2029                Some(&m) => m,
2030                None => {
2031                    let m = self.pdl_load_module_in_ctx(FATBIN)?;
2032                    map.insert(ctx_key, m);
2033                    m
2034                }
2035            }
2036        };
2037        let cname = std::ffi::CString::new(name)?;
2038        let mut f: cu::CUfunction = std::ptr::null_mut();
2039        let mut r =
2040            unsafe { cu::cuModuleGetFunction(&mut f, module as cu::CUmodule, cname.as_ptr()) };
2041        if r == cu::CUresult::CUDA_ERROR_NOT_FOUND {
2042            let qmodule = {
2043                let mut mods = QMODULES.lock().unwrap();
2044                let map = mods.get_or_insert_with(Default::default);
2045                match map.get(&ctx_key) {
2046                    Some(&m) => m,
2047                    None => {
2048                        let m = self.pdl_load_module_in_ctx(QMATVEC_FATBIN)?;
2049                        map.insert(ctx_key, m);
2050                        m
2051                    }
2052                }
2053            };
2054            r = unsafe { cu::cuModuleGetFunction(&mut f, qmodule as cu::CUmodule, cname.as_ptr()) };
2055        }
2056        if r != cu::CUresult::CUDA_SUCCESS {
2057            return Err(format!("pdl_func {name}: {r:?}").into());
2058        }
2059        FNS.lock()
2060            .unwrap()
2061            .get_or_insert_with(Default::default)
2062            .insert((ctx_key, name), f as usize);
2063        Ok(f)
2064    }
2065
2066    /// cuLaunchKernelEx with CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION on the
2067    /// compute stream. ONLY legal for kernels whose entry carries MEMRA_PDL_ENTRY.
2068    ///
2069    /// # Safety
2070    /// `params` must match the kernel's exact parameter list (order, types, count) —
2071    /// a mismatch corrupts the launch silently.
2072    /// Flash-module twin of `launch_pdl` — `g` picks the kf8vf8 flavor (must mirror the
2073    /// builder path's fa_func/func_g choice exactly).
2074    ///
2075    /// # Safety
2076    /// Same contract as `launch_pdl`.
2077    unsafe fn launch_pdl_flash(
2078        &self,
2079        g: bool,
2080        name: &'static str,
2081        grid: (u32, u32, u32),
2082        block: (u32, u32, u32),
2083        smem: u32,
2084        params: &mut [*mut std::ffi::c_void],
2085    ) -> Result<(), Box<dyn std::error::Error>> {
2086        use cudarc::driver::sys as cu;
2087        let f = self.pdl_func_flash(g, name)?;
2088        if smem > 0 {
2089            // mirror the builder path's opt-in ceiling (idempotent host-side set).
2090            let r =
2091                unsafe {
2092                    cu::cuFuncSetAttribute(f,
2093                cu::CUfunction_attribute_enum::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
2094                smem as i32)
2095                };
2096            if r != cu::CUresult::CUDA_SUCCESS {
2097                return Err(format!("pdl smem attr {name}: {r:?}").into());
2098            }
2099        }
2100        let mut attr = cu::CUlaunchAttribute {
2101            id: cu::CUlaunchAttributeID::CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION,
2102            pad: [0; 4],
2103            value: cu::CUlaunchAttributeValue {
2104                programmaticStreamSerializationAllowed: 1,
2105            },
2106        };
2107        let cfg = cu::CUlaunchConfig {
2108            gridDimX: grid.0,
2109            gridDimY: grid.1,
2110            gridDimZ: grid.2,
2111            blockDimX: block.0,
2112            blockDimY: block.1,
2113            blockDimZ: block.2,
2114            sharedMemBytes: smem,
2115            hStream: self.gpu.stream().cu_stream(),
2116            attrs: &mut attr,
2117            numAttrs: 1,
2118        };
2119        let r = unsafe { cu::cuLaunchKernelEx(&cfg, f, params.as_mut_ptr(), std::ptr::null_mut()) };
2120        if r != cu::CUresult::CUDA_SUCCESS {
2121            return Err(format!("launch_pdl_flash {name}: {r:?}").into());
2122        }
2123        Ok(())
2124    }
2125
2126    unsafe fn launch_pdl(
2127        &self,
2128        name: &'static str,
2129        grid: (u32, u32, u32),
2130        block: (u32, u32, u32),
2131        params: &mut [*mut std::ffi::c_void],
2132    ) -> Result<(), Box<dyn std::error::Error>> {
2133        use cudarc::driver::sys as cu;
2134        let f = self.pdl_func(name)?;
2135        let mut attr = cu::CUlaunchAttribute {
2136            id: cu::CUlaunchAttributeID::CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION,
2137            pad: [0; 4],
2138            value: cu::CUlaunchAttributeValue {
2139                programmaticStreamSerializationAllowed: 1,
2140            },
2141        };
2142        let cfg = cu::CUlaunchConfig {
2143            gridDimX: grid.0,
2144            gridDimY: grid.1,
2145            gridDimZ: grid.2,
2146            blockDimX: block.0,
2147            blockDimY: block.1,
2148            blockDimZ: block.2,
2149            sharedMemBytes: 0,
2150            hStream: self.gpu.stream().cu_stream(),
2151            attrs: &mut attr,
2152            numAttrs: 1,
2153        };
2154        let r = unsafe { cu::cuLaunchKernelEx(&cfg, f, params.as_mut_ptr(), std::ptr::null_mut()) };
2155        if r != cu::CUresult::CUDA_SUCCESS {
2156            return Err(format!("launch_pdl {name}: {r:?}").into());
2157        }
2158        Ok(())
2159    }
2160
2161    /// L2-prefetch a quant weight's DECODE plane (the rp4 split-plane mirror when present —
2162    /// that is what the m<=8 dispatch reads — else the raw block bytes). No-op on float arms.
2163    pub fn prefetch_weight_l2(
2164        &self,
2165        w: &crate::model::GpuTensor,
2166    ) -> Result<(), Box<dyn std::error::Error>> {
2167        if let crate::model::GpuTensor::Quant { bytes, rp4, .. } = w {
2168            let p = rp4.as_ref().unwrap_or(bytes);
2169            self.prefetch_l2(p, p.len())?;
2170        }
2171        Ok(())
2172    }
2173
2174    /// DSpark markov chain ops (dflash lane): gather one bf16 row of a [V, rank] table
2175    /// by the DEVICE token id at tok[idx] into f32.
2176    pub fn gather_row_bf16(
2177        &self,
2178        table: &CudaSlice<u8>,
2179        tok: &CudaSlice<u32>,
2180        idx: usize,
2181        dst: &mut CudaSlice<f32>,
2182        ncols: usize,
2183    ) -> Result<(), Box<dyn std::error::Error>> {
2184        let f = self.func("gather_row_bf16_f32");
2185        let cfg = LaunchConfig {
2186            grid_dim: (ncols.div_ceil(256) as u32, 1, 1),
2187            block_dim: (256, 1, 1),
2188            shared_mem_bytes: 0,
2189        };
2190        let (nc, ix) = (ncols as i32, idx as i32);
2191        let __s_b = self.gpu.stream();
2192        let mut b = __s_b.launch_builder(&f);
2193        b.arg(table).arg(tok).arg(&ix).arg(dst).arg(&nc);
2194        unsafe {
2195            b.launch(cfg)?;
2196        }
2197        Ok(())
2198    }
2199
2200    /// DFlash2 grouped dynamic causal conv (dflash lane, DFLASH2-EVAL-20260820.md):
2201    /// out[p,c] = sum_{o<ksize, o<=p} (base[half][o][c] + dyn[p][half][o][group(c)])
2202    /// * x[p-o][c]. `dyn_` is the kernel_projection GEMM output [rows, 2*ksize*groups];
2203    /// `base` is base_kernel [2, ksize, hidden] flattened; `half` picks prepare(0) /
2204    /// finish(1).
2205    #[allow(clippy::too_many_arguments)]
2206    pub fn dflash2_dynconv(
2207        &self,
2208        x: &CudaSlice<f32>,
2209        dyn_: &CudaSlice<f32>,
2210        base: &CudaSlice<f32>,
2211        out: &mut CudaSlice<f32>,
2212        rows: usize,
2213        hidden: usize,
2214        group_size: usize,
2215        ksize: usize,
2216        half: usize,
2217    ) -> Result<(), Box<dyn std::error::Error>> {
2218        assert_eq!(hidden % group_size, 0, "hidden % group_size != 0");
2219        let f = self.func("dflash2_dynconv_f32");
2220        let n = rows * hidden;
2221        let cfg = LaunchConfig {
2222            grid_dim: (n.div_ceil(256) as u32, 1, 1),
2223            block_dim: (256, 1, 1),
2224            shared_mem_bytes: 0,
2225        };
2226        let (ri, hi, gi, ki, hf) = (
2227            rows as i32,
2228            hidden as i32,
2229            group_size as i32,
2230            ksize as i32,
2231            half as i32,
2232        );
2233        let __s_b = self.gpu.stream();
2234        let mut b = __s_b.launch_builder(&f);
2235        b.arg(x)
2236            .arg(dyn_)
2237            .arg(base)
2238            .arg(out)
2239            .arg(&ri)
2240            .arg(&hi)
2241            .arg(&gi)
2242            .arg(&ki)
2243            .arg(&hf);
2244        unsafe {
2245            b.launch(cfg)?;
2246        }
2247        Ok(())
2248    }
2249
2250    /// Per-row top-k (k <= 32) over a [n_rows, n_cols] logits matrix (DFlash2
2251    /// candidate selector). Returns (values [n_rows, k], column indices [n_rows, k]),
2252    /// value-descending, ties to the lower index.
2253    pub fn topk_rows(
2254        &self,
2255        logits: &CudaSlice<f32>,
2256        n_rows: usize,
2257        n_cols: usize,
2258        k: usize,
2259    ) -> Result<(CudaSlice<f32>, CudaSlice<u32>), Box<dyn std::error::Error>> {
2260        assert!(k <= 32 && k >= 1, "topk_rows supports 1..=32, got {k}");
2261        assert!(k <= n_cols, "topk_rows: k {k} > n_cols {n_cols}");
2262        let f = self.func("topk_rows_f32");
2263        let nth = 256usize;
2264        let mut vals = self.uninit(n_rows * k)?;
2265        let mut idxs = self.gpu.stream().alloc_zeros::<u32>(n_rows * k)?;
2266        let cfg = LaunchConfig {
2267            grid_dim: (n_rows as u32, 1, 1),
2268            block_dim: (nth as u32, 1, 1),
2269            shared_mem_bytes: (nth * k * 8) as u32,
2270        };
2271        let (nr, nc, ki) = (n_rows as i32, n_cols as i32, k as i32);
2272        let __s_b = self.gpu.stream();
2273        let mut b = __s_b.launch_builder(&f);
2274        b.arg(logits)
2275            .arg(&nr)
2276            .arg(&nc)
2277            .arg(&ki)
2278            .arg(&mut vals)
2279            .arg(&mut idxs);
2280        unsafe {
2281            b.launch(cfg)?;
2282        }
2283        Ok((vals, idxs))
2284    }
2285
2286    /// logits[row_off .. row_off+n] += bias[0..n] (in place, one row).
2287    pub fn add_row_inplace(
2288        &self,
2289        logits: &mut CudaSlice<f32>,
2290        bias: &CudaSlice<f32>,
2291        n: usize,
2292        row_off: usize,
2293    ) -> Result<(), Box<dyn std::error::Error>> {
2294        let f = self.func("add_row_inplace_f32");
2295        let cfg = LaunchConfig {
2296            grid_dim: (n.div_ceil(256) as u32, 1, 1),
2297            block_dim: (256, 1, 1),
2298            shared_mem_bytes: 0,
2299        };
2300        let (ni, off) = (n as i32, row_off as i64);
2301        let __s_b = self.gpu.stream();
2302        let mut b = __s_b.launch_builder(&f);
2303        b.arg(logits).arg(bias).arg(&ni).arg(&off);
2304        unsafe {
2305            b.launch(cfg)?;
2306        }
2307        Ok(())
2308    }
2309
2310    /// L2 prefetch of a device byte range (latency-hiding arc; value-free scheduling op).
2311    pub fn prefetch_l2(
2312        &self,
2313        p: &CudaSlice<u8>,
2314        n: usize,
2315    ) -> Result<(), Box<dyn std::error::Error>> {
2316        let f = self.func("prefetch_l2_bytes");
2317        let lines = n.div_ceil(128);
2318        let ni = n as i64;
2319        let cfg = LaunchConfig {
2320            grid_dim: (lines.div_ceil(256) as u32, 1, 1),
2321            block_dim: (256, 1, 1),
2322            shared_mem_bytes: 0,
2323        };
2324        let __s_b = self.gpu.stream();
2325        let mut b = __s_b.launch_builder(&f);
2326        b.arg(p).arg(&ni);
2327        unsafe {
2328            b.launch(cfg)?;
2329        }
2330        Ok(())
2331    }
2332
2333    /// MoE router GEMV (MEMRA_ROUTER_KERNEL): deterministic warp-per-(expert,token) f32 dot.
2334    /// Different FP order than the cuBLAS path it replaces — battery-gated numeric config.
2335    pub fn router_gemv(
2336        &self,
2337        w: &CudaSlice<f32>,
2338        x: &CudaSlice<f32>,
2339        n_embd: usize,
2340        n_experts: usize,
2341        t: usize,
2342    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2343        // float4 v2 probed 2026-07-14: +0.25% but flips near-tie routing (new FP order,
2344        // stream differs) — too small to justify a numeric config change; deleted.
2345        // w8 twin (2026-07-31): on the 132-SM H100 the lone-warp form is 14.8% of the q35
2346        // decode step (latency-bound) — the calculus flipped. MEMRA_ROUTER_V2=0 reverts to
2347        // the warp form (rollback seam; new FP order, battery-arbitrated per model).
2348        let w8 = match std::env::var("MEMRA_ROUTER_V2").as_deref() {
2349            Ok("0") => false,
2350            Ok(_) => true,
2351            Err(_) => ROUTER_W8_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
2352        };
2353        // FAST-ROUTER batch twin (lane/fast-router, 2026-08-02): at prefill m the per-(e,tok)
2354        // w8 form re-streams both operand rows per output (GEMV program at GEMM shape — the
2355        // concat-prime exactness fix paid -10% q35 board-2048 prefill through it). The batch
2356        // twin (8x8 expert-x-token register tile) is BIT-IDENTICAL per row (same k order,
2357        // same tree, same fold — kernel-check sweeps m=1..2048 on real router weights), so
2358        // the crossover is pure perf, not a numeric config. MIN_T from the on-box sweep
2359        // (research/fast-router-20260802/crossover-router*.jsonl); decode t=1 and small-t
2360        // spec verify keep the plain w8 form. MEMRA_ROUTER_BATCH=0: rollback seam
2361        // (perf-only, bits equal).
2362        let batch = w8 && t >= ROUTER_BATCH_MIN_T && router_batch_on();
2363        self.router_gemv_form(w, x, n_embd, n_experts, t, w8, batch)
2364    }
2365
2366    /// Form-explicit router GEMV launch (kernel-check bit-identity gate + crossover bench
2367    /// force both forms; `batch` requires `w8`).
2368    pub fn router_gemv_form(
2369        &self,
2370        w: &CudaSlice<f32>,
2371        x: &CudaSlice<f32>,
2372        n_embd: usize,
2373        n_experts: usize,
2374        t: usize,
2375        w8: bool,
2376        batch: bool,
2377    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2378        debug_assert!(!batch || w8, "batch twin exists for the w8 form only");
2379        let mut y = self.alloc_uninit::<f32>(t * n_experts)?;
2380        let f = if batch {
2381            self.func("router_gemv_f32_w8_batch")
2382        } else if w8 {
2383            self.func("router_gemv_f32_w8")
2384        } else {
2385            self.func("router_gemv_f32")
2386        };
2387        let (ne, nx, ti) = (n_embd as i32, n_experts as i32, t as i32);
2388        let cfg = if batch {
2389            LaunchConfig {
2390                grid_dim: (n_experts.div_ceil(8) as u32, t.div_ceil(8) as u32, 1),
2391                block_dim: (32, 8, 1),
2392                shared_mem_bytes: 0,
2393            }
2394        } else {
2395            LaunchConfig {
2396                grid_dim: (n_experts as u32, t as u32, 1),
2397                block_dim: (32, if w8 { 8 } else { 1 }, 1),
2398                shared_mem_bytes: 0,
2399            }
2400        };
2401        let __s_b = self.gpu.stream();
2402        let mut b = __s_b.launch_builder(&f);
2403        b.arg(w).arg(x).arg(&mut y).arg(&ne).arg(&nx).arg(&ti);
2404        unsafe {
2405            b.launch(cfg)?;
2406        }
2407        Ok(y)
2408    }
2409
2410    /// f32 row permute: dst[idx[i], :] = src[i, :] (grouped-GEMM CSR -> pair-id reorder).
2411    pub fn rows_permute(
2412        &self,
2413        src: &CudaSlice<f32>,
2414        idx: &CudaSlice<i32>,
2415        nrows: usize,
2416        ncols: usize,
2417    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2418        let mut dst = self.alloc_uninit::<f32>(nrows * ncols)?;
2419        let f = self.func("rows_permute_f32");
2420        let (nc, nr) = (ncols as i32, nrows as i32);
2421        let cfg = LaunchConfig {
2422            grid_dim: (nrows as u32, 1, 1),
2423            block_dim: (256, 1, 1),
2424            shared_mem_bytes: 0,
2425        };
2426        let __s_b = self.gpu.stream();
2427        let mut b = __s_b.launch_builder(&f);
2428        b.arg(src).arg(idx).arg(&mut dst).arg(&nc).arg(&nr);
2429        unsafe {
2430            b.launch(cfg)?;
2431        }
2432        Ok(dst)
2433    }
2434
2435    /// shexp gate fused dot: g[tok] = sigmoid(dot(x[tok,:], w)) — replaces the per-layer
2436    /// cuBLASLt m=1 GEMM + separate sigmoid launch on the qwen35moe decode path (the
2437    /// splitKreduce x40/step dig, 2026-07-31). One fold order for every t, so the t=1
2438    /// decode chain and the small-t spec-verify chain match per row by construction.
2439    pub fn sigmoid_dot_rows(
2440        &self,
2441        x: &CudaSlice<f32>,
2442        w: &CudaSlice<f32>,
2443        n_embd: usize,
2444        t: usize,
2445    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2446        // MEMRA_SHEXP_DOT=0: rollback seam to the cuBLASLt linear + sigmoid pair (numeric
2447        // config; same class as MEMRA_ROUTER_V2).
2448        static OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2449        if *OFF.get_or_init(|| std::env::var("MEMRA_SHEXP_DOT").as_deref() == Ok("0")) {
2450            let gs = self.linear(x, w, t, n_embd, 1)?;
2451            let mut g = self.uninit(t)?;
2452            self.sigmoid(&gs, &mut g, t)?;
2453            return Ok(g);
2454        }
2455        // FAST-ROUTER lane note (2026-08-02): a register-tiled 8-token batch twin of this
2456        // kernel was built, proven bit-identical, and measured SLOWER at every prefill t on
2457        // the 5090 (0.62-0.89x — launch-latency-bound op, ~7us/layer at m=2048;
2458        // research/fast-router-20260802/crossover-router.jsonl). Dispatch arm killed per
2459        // flags doctrine; this per-token form serves every t.
2460        let mut g = self.alloc_uninit::<f32>(t)?;
2461        let f = self.func("sigmoid_dot_rows_f32");
2462        let (ne, ti) = (n_embd as i32, t as i32);
2463        let cfg = LaunchConfig {
2464            grid_dim: (t as u32, 1, 1),
2465            block_dim: (32, 8, 1),
2466            shared_mem_bytes: 0,
2467        };
2468        let __s_b = self.gpu.stream();
2469        let mut b = __s_b.launch_builder(&f);
2470        b.arg(x).arg(w).arg(&mut g).arg(&ne).arg(&ti);
2471        unsafe {
2472            b.launch(cfg)?;
2473        }
2474        Ok(g)
2475    }
2476
2477    /// ROUND-STREAM stream rollback: all counters <- pos_start + base + n_acc.
2478    pub fn spec_rollback_stream(
2479        &self,
2480        len_ptrs: &CudaSlice<u64>,
2481        pos_start: &CudaSlice<i32>,
2482        acc: &CudaSlice<u32>,
2483        base: usize,
2484        n_rows: usize,
2485    ) -> Result<(), Box<dyn std::error::Error>> {
2486        let f = self.func("spec_rollback_stream");
2487        let (b, nr) = (base as i32, n_rows as i32);
2488        let cfg = LaunchConfig {
2489            grid_dim: (n_rows.div_ceil(64) as u32, 1, 1),
2490            block_dim: (64, 1, 1),
2491            shared_mem_bytes: 0,
2492        };
2493        let __s_bl = self.gpu.stream();
2494        let mut bl = __s_bl.launch_builder(&f);
2495        bl.arg(len_ptrs).arg(pos_start).arg(acc).arg(&b).arg(&nr);
2496        unsafe {
2497            bl.launch(cfg)?;
2498        }
2499        Ok(())
2500    }
2501
2502    /// PLAIN-DECODE GRAPH ring store: ring[(pos_start - base) % cap] = vam[0].
2503    pub fn plain_tok_ring(
2504        &self,
2505        vam: &CudaSlice<u32>,
2506        pos_start: &CudaSlice<i32>,
2507        base: usize,
2508        ring: &mut CudaSlice<u32>,
2509    ) -> Result<(), Box<dyn std::error::Error>> {
2510        let f = self.func("plain_tok_ring");
2511        let (b, cap) = (base as i32, ring.len() as i32);
2512        let cfg = LaunchConfig {
2513            grid_dim: (1, 1, 1),
2514            block_dim: (32, 1, 1),
2515            shared_mem_bytes: 0,
2516        };
2517        let __s_bl = self.gpu.stream();
2518        let mut bl = __s_bl.launch_builder(&f);
2519        bl.arg(vam).arg(pos_start).arg(&b).arg(&mut *ring).arg(&cap);
2520        unsafe {
2521            bl.launch(cfg)?;
2522        }
2523        Ok(())
2524    }
2525
2526    /// ROUND-STREAM stage (c) 4 epilogue: ring commit + tiny counter copies.
2527    pub fn spec_ring_commit(
2528        &self,
2529        vtok: &CudaSlice<u32>,
2530        acc: &CudaSlice<u32>,
2531        brk: &CudaSlice<u32>,
2532        ring: &mut CudaSlice<u32>,
2533        pend: &mut CudaSlice<u32>,
2534    ) -> Result<(), Box<dyn std::error::Error>> {
2535        let f = self.func("spec_ring_commit");
2536        let cfg = LaunchConfig {
2537            grid_dim: (1, 1, 1),
2538            block_dim: (32, 1, 1),
2539            shared_mem_bytes: 0,
2540        };
2541        let __s_b = self.gpu.stream();
2542        let mut b = __s_b.launch_builder(&f);
2543        b.arg(vtok).arg(acc).arg(brk).arg(ring).arg(pend);
2544        unsafe {
2545            b.launch(cfg)?;
2546        }
2547        Ok(())
2548    }
2549    pub fn i32_copy_add(
2550        &self,
2551        src: &CudaSlice<i32>,
2552        dst: &mut CudaSlice<i32>,
2553        delta: i32,
2554    ) -> Result<(), Box<dyn std::error::Error>> {
2555        let f = self.func("i32_copy_add");
2556        let cfg = LaunchConfig {
2557            grid_dim: (1, 1, 1),
2558            block_dim: (32, 1, 1),
2559            shared_mem_bytes: 0,
2560        };
2561        let __s_b = self.gpu.stream();
2562        let mut b = __s_b.launch_builder(&f);
2563        b.arg(src).arg(dst).arg(&delta);
2564        unsafe {
2565            b.launch(cfg)?;
2566        }
2567        Ok(())
2568    }
2569    pub fn u32_copy(
2570        &self,
2571        src: &CudaSlice<u32>,
2572        dst: &mut CudaSlice<u32>,
2573    ) -> Result<(), Box<dyn std::error::Error>> {
2574        let f = self.func("u32_copy");
2575        let cfg = LaunchConfig {
2576            grid_dim: (1, 1, 1),
2577            block_dim: (32, 1, 1),
2578            shared_mem_bytes: 0,
2579        };
2580        let __s_b = self.gpu.stream();
2581        let mut b = __s_b.launch_builder(&f);
2582        b.arg(src).arg(dst);
2583        unsafe {
2584            b.launch(cfg)?;
2585        }
2586        Ok(())
2587    }
2588
2589    /// ROUND-GRAPH adaptive depth: brk[0] <- clamp(acc[0] + 1, floor, cap) — the host
2590    /// adaptive policy as a captured device op (policy-identical: the accept walk depth
2591    /// caps acceptance exactly like drafting fewer tokens).
2592    pub fn spec_adapt_k(
2593        &self,
2594        acc: &CudaSlice<u32>,
2595        brk: &mut CudaSlice<u32>,
2596        floor: usize,
2597        cap: usize,
2598    ) -> Result<(), Box<dyn std::error::Error>> {
2599        let f = self.func("spec_adapt_k");
2600        let (fl, cp) = (floor as i32, cap as i32);
2601        let cfg = LaunchConfig {
2602            grid_dim: (1, 1, 1),
2603            block_dim: (32, 1, 1),
2604            shared_mem_bytes: 0,
2605        };
2606        let __s_b = self.gpu.stream();
2607        let mut b = __s_b.launch_builder(&f);
2608        b.arg(acc).arg(brk).arg(&fl).arg(&cp);
2609        unsafe {
2610            b.launch(cfg)?;
2611        }
2612        Ok(())
2613    }
2614
2615    /// ROUND-STREAM stage (c) 3: accept walk fully device-driven (brk + assembled vtok).
2616    pub fn spec_accept_greedy_dc(
2617        &self,
2618        preds: &CudaSlice<u32>,
2619        vtok: &CudaSlice<u32>,
2620        last_pred: &CudaSlice<u32>,
2621        brk: &CudaSlice<u32>,
2622        out: &mut CudaSlice<u32>,
2623    ) -> Result<(), Box<dyn std::error::Error>> {
2624        let f = self.func("spec_accept_greedy_dc");
2625        let cfg = LaunchConfig {
2626            grid_dim: (1, 1, 1),
2627            block_dim: (32, 1, 1),
2628            shared_mem_bytes: 0,
2629        };
2630        let __s_b = self.gpu.stream();
2631        let mut b = __s_b.launch_builder(&f);
2632        b.arg(preds).arg(vtok).arg(last_pred).arg(brk).arg(out);
2633        unsafe {
2634            b.launch(cfg)?;
2635        }
2636        Ok(())
2637    }
2638
2639    /// ROUND-STREAM stage (c) 2: verify-chain device-pos entries.
2640    pub fn pos_iota(
2641        &self,
2642        pos0: &CudaSlice<i32>,
2643        out: &mut CudaSlice<i32>,
2644        t: usize,
2645    ) -> Result<(), Box<dyn std::error::Error>> {
2646        let f = self.func("pos_iota_i32");
2647        let ti = t as i32;
2648        let cfg = LaunchConfig {
2649            grid_dim: (1, 1, 1),
2650            block_dim: (t.max(1) as u32, 1, 1),
2651            shared_mem_bytes: 0,
2652        };
2653        let __s_b = self.gpu.stream();
2654        let mut b = __s_b.launch_builder(&f);
2655        b.arg(pos0).arg(out).arg(&ti);
2656        unsafe {
2657            b.launch(cfg)?;
2658        }
2659        Ok(())
2660    }
2661    #[allow(clippy::too_many_arguments)]
2662    pub fn append_kv_quantized_rows_dc(
2663        &self,
2664        k_rows: &CudaSlice<f32>,
2665        v_rows: &CudaSlice<f32>,
2666        kc: &mut CudaSlice<u8>,
2667        vc: &mut CudaSlice<u8>,
2668        t0_dev: &CudaSlice<i32>,
2669        t: usize,
2670        kv_dim_k: usize,
2671        kv_dim_v: usize,
2672        k_tok_bytes: usize,
2673        v_tok_bytes: usize,
2674        g: bool,
2675    ) -> Result<(), Box<dyn std::error::Error>> {
2676        let f = if g {
2677            self.func_g("append_quantize_kv_q8_0_q5_1_rows_dc")
2678        } else {
2679            self.func("append_quantize_kv_q8_0_q5_1_rows_dc")
2680        };
2681        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
2682        let cfg = LaunchConfig {
2683            grid_dim: (nblk, t as u32, 1),
2684            block_dim: (32, 1, 1),
2685            shared_mem_bytes: 0,
2686        };
2687        let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
2688        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
2689        let __s_b = self.gpu.stream();
2690        let mut b = __s_b.launch_builder(&f);
2691        b.arg(k_rows)
2692            .arg(v_rows)
2693            .arg(kc)
2694            .arg(vc)
2695            .arg(t0_dev)
2696            .arg(&kdk)
2697            .arg(&kdv)
2698            .arg(&ktb)
2699            .arg(&vtb);
2700        unsafe {
2701            b.launch(cfg)?;
2702        }
2703        Ok(())
2704    }
2705
2706    /// t=1 dc append with a FUSED len_d increment (wave 5c) — one launch replaces
2707    /// append_rows_dc + inc_seqlen. Single block (read-before-inc ordering).
2708    #[allow(clippy::too_many_arguments)]
2709    pub fn append_kv_quantized_row_dc_inc(
2710        &self,
2711        k_row: &CudaSlice<f32>,
2712        v_row: &CudaSlice<f32>,
2713        kc: &mut CudaSlice<u8>,
2714        vc: &mut CudaSlice<u8>,
2715        t0_dev: &mut CudaSlice<i32>,
2716        kv_dim_k: usize,
2717        kv_dim_v: usize,
2718        k_tok_bytes: usize,
2719        v_tok_bytes: usize,
2720        g: bool,
2721    ) -> Result<(), Box<dyn std::error::Error>> {
2722        let f = if g {
2723            self.func_g("append_quantize_kv_q8_0_q5_1_dc_inc")
2724        } else {
2725            self.func("append_quantize_kv_q8_0_q5_1_dc_inc")
2726        };
2727        let nthreads = ((kv_dim_k.max(kv_dim_v) / 32) * 32).min(1024) as u32;
2728        let cfg = LaunchConfig {
2729            grid_dim: (1, 1, 1),
2730            block_dim: (nthreads, 1, 1),
2731            shared_mem_bytes: 0,
2732        };
2733        let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
2734        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
2735        let __s_b = self.gpu.stream();
2736        let mut b = __s_b.launch_builder(&f);
2737        b.arg(k_row)
2738            .arg(v_row)
2739            .arg(kc)
2740            .arg(vc)
2741            .arg(t0_dev)
2742            .arg(&kdk)
2743            .arg(&kdv)
2744            .arg(&ktb)
2745            .arg(&vtb);
2746        unsafe {
2747            b.launch(cfg)?;
2748        }
2749        Ok(())
2750    }
2751
2752    /// ROUND-STREAM: draft-chain pack + in-graph d2t remap (see kernels.cu headers).
2753    pub fn pack_tok_p(
2754        &self,
2755        tok: &CudaSlice<u32>,
2756        p: &CudaSlice<f32>,
2757        out: &mut CudaSlice<u32>,
2758        slot: usize,
2759    ) -> Result<(), Box<dyn std::error::Error>> {
2760        let f = self.func("pack_tok_p");
2761        let sl = slot as i32;
2762        let cfg = LaunchConfig {
2763            grid_dim: (1, 1, 1),
2764            block_dim: (32, 1, 1),
2765            shared_mem_bytes: 0,
2766        };
2767        let __s_b = self.gpu.stream();
2768        let mut b = __s_b.launch_builder(&f);
2769        b.arg(tok).arg(p).arg(out).arg(&sl);
2770        unsafe {
2771            b.launch(cfg)?;
2772        }
2773        Ok(())
2774    }
2775    pub fn tok_map_u32(
2776        &self,
2777        tok: &mut CudaSlice<u32>,
2778        map: &CudaSlice<u32>,
2779    ) -> Result<(), Box<dyn std::error::Error>> {
2780        let f = self.func("tok_map_u32");
2781        let cfg = LaunchConfig {
2782            grid_dim: (1, 1, 1),
2783            block_dim: (32, 1, 1),
2784            shared_mem_bytes: 0,
2785        };
2786        let __s_b = self.gpu.stream();
2787        let mut b = __s_b.launch_builder(&f);
2788        b.arg(tok).arg(map);
2789        unsafe {
2790            b.launch(cfg)?;
2791        }
2792        Ok(())
2793    }
2794
2795    /// ROUND-STREAM stage (c) 1: device verify-token assembly + p-min break derivation.
2796    #[allow(clippy::too_many_arguments)]
2797    pub fn spec_assemble_verify(
2798        &self,
2799        tokp: &CudaSlice<u32>,
2800        pend: &CudaSlice<u32>,
2801        d2t: Option<&CudaSlice<u32>>,
2802        vtok: &mut CudaSlice<u32>,
2803        brk: &mut CudaSlice<u32>,
2804        p_min: f32,
2805        k: usize,
2806        pmin0: bool,
2807    ) -> Result<(), Box<dyn std::error::Error>> {
2808        let f = self.func("spec_assemble_verify");
2809        let (ki, pm) = (k as i32, if pmin0 { 1i32 } else { 0i32 });
2810        let cfg = LaunchConfig {
2811            grid_dim: (1, 1, 1),
2812            block_dim: (32, 1, 1),
2813            shared_mem_bytes: 0,
2814        };
2815        let __s_b = self.gpu.stream();
2816        let mut b = __s_b.launch_builder(&f);
2817        match d2t {
2818            Some(m) => {
2819                b.arg(tokp)
2820                    .arg(pend)
2821                    .arg(m)
2822                    .arg(vtok)
2823                    .arg(brk)
2824                    .arg(&p_min)
2825                    .arg(&ki)
2826                    .arg(&pm);
2827                unsafe {
2828                    b.launch(cfg)?;
2829                }
2830            }
2831            None => {
2832                let null: u64 = 0;
2833                b.arg(tokp)
2834                    .arg(pend)
2835                    .arg(&null)
2836                    .arg(vtok)
2837                    .arg(brk)
2838                    .arg(&p_min)
2839                    .arg(&ki)
2840                    .arg(&pm);
2841                unsafe {
2842                    b.launch(cfg)?;
2843                }
2844            }
2845        }
2846        Ok(())
2847    }
2848
2849    /// ROUND-STREAM stage (b) 3b: recur-restore twins with device-j (see hybrid.cu headers).
2850    #[allow(clippy::too_many_arguments)]
2851    pub fn ssm_conv_ring_rebuild_dc(
2852        &self,
2853        qkv_tm: &CudaSlice<f32>,
2854        ring_old: &CudaSlice<f32>,
2855        conv_state: &mut CudaSlice<f32>,
2856        conv_dim: usize,
2857        acc: &CudaSlice<u32>,
2858        base: usize,
2859        t_v: usize,
2860        d_conv: usize,
2861    ) -> Result<(), Box<dyn std::error::Error>> {
2862        let f = self.func("ssm_conv_ring_rebuild_f32_dc");
2863        let n = conv_dim * (d_conv - 1);
2864        let cfg = LaunchConfig::for_num_elems(n as u32);
2865        let (cd, b0, tv, dc) = (conv_dim as i32, base as i32, t_v as i32, d_conv as i32);
2866        let __s_b = self.gpu.stream();
2867        let mut b = __s_b.launch_builder(&f);
2868        b.arg(qkv_tm)
2869            .arg(ring_old)
2870            .arg(conv_state)
2871            .arg(&cd)
2872            .arg(acc)
2873            .arg(&b0)
2874            .arg(&tv)
2875            .arg(&dc);
2876        unsafe {
2877            b.launch(cfg)?;
2878        }
2879        Ok(())
2880    }
2881    #[allow(clippy::too_many_arguments)]
2882    pub fn gdn_scan_s128_dc(
2883        &self,
2884        q: &CudaSlice<f32>,
2885        k: &CudaSlice<f32>,
2886        v: &CudaSlice<f32>,
2887        g: &CudaSlice<f32>,
2888        beta: &CudaSlice<f32>,
2889        state_in: &CudaSlice<f32>,
2890        state_out: &mut CudaSlice<f32>,
2891        o: &mut CudaSlice<f32>,
2892        n_head: usize,
2893        acc: &CudaSlice<u32>,
2894        base: usize,
2895        t_v: usize,
2896        scale: f32,
2897    ) -> Result<(), Box<dyn std::error::Error>> {
2898        let f = self.func("gdn_scan_s128_dc");
2899        const S_V: u32 = 128;
2900        const WARP: u32 = 32;
2901        const COLS_PER_BLOCK: u32 = 4;
2902        let cfg = LaunchConfig {
2903            grid_dim: (n_head as u32, 1, S_V / COLS_PER_BLOCK),
2904            block_dim: (WARP, COLS_PER_BLOCK, 1),
2905            shared_mem_bytes: 0,
2906        };
2907        let (h, b0, tv) = (n_head as i32, base as i32, t_v as i32);
2908        let __s_b = self.gpu.stream();
2909        let mut b = __s_b.launch_builder(&f);
2910        b.arg(q)
2911            .arg(k)
2912            .arg(v)
2913            .arg(g)
2914            .arg(beta)
2915            .arg(state_in)
2916            .arg(state_out)
2917            .arg(o)
2918            .arg(&h)
2919            .arg(acc)
2920            .arg(&b0)
2921            .arg(&tv)
2922            .arg(&scale);
2923        unsafe {
2924            b.launch(cfg)?;
2925        }
2926        Ok(())
2927    }
2928
2929    /// ROUND-STREAM stage (b) 3a: device per-layer KV-len rollback (see spec_rollback_kv).
2930    pub fn spec_rollback_kv(
2931        &self,
2932        len_ptrs: &CudaSlice<u64>,
2933        saved: &CudaSlice<i32>,
2934        acc: &CudaSlice<u32>,
2935        base: usize,
2936        n_layer: usize,
2937    ) -> Result<(), Box<dyn std::error::Error>> {
2938        let f = self.func("spec_rollback_kv");
2939        let (b, nl) = (base as i32, n_layer as i32);
2940        let cfg = LaunchConfig {
2941            grid_dim: (n_layer.div_ceil(64) as u32, 1, 1),
2942            block_dim: (64, 1, 1),
2943            shared_mem_bytes: 0,
2944        };
2945        let __s_bl = self.gpu.stream();
2946        let mut bl = __s_bl.launch_builder(&f);
2947        bl.arg(len_ptrs).arg(saved).arg(acc).arg(&b).arg(&nl);
2948        unsafe {
2949            bl.launch(cfg)?;
2950        }
2951        Ok(())
2952    }
2953
2954    /// OPTIPIPE increment 1: derive the K=1 successor-valid bit on device.
2955    pub fn spec_fork_valid(
2956        &self,
2957        acc: &CudaSlice<u32>,
2958        optimistic_pending: u32,
2959        valid: &mut CudaSlice<u32>,
2960    ) -> Result<(), Box<dyn std::error::Error>> {
2961        let f = self.func("spec_fork_valid");
2962        let cfg = LaunchConfig {
2963            grid_dim: (1, 1, 1),
2964            block_dim: (1, 1, 1),
2965            shared_mem_bytes: 0,
2966        };
2967        let __s_bl = self.gpu.stream();
2968        let mut bl = __s_bl.launch_builder(&f);
2969        bl.arg(acc).arg(&optimistic_pending).arg(valid);
2970        unsafe {
2971            bl.launch(cfg)?;
2972        }
2973        Ok(())
2974    }
2975
2976    /// OPTIPIPE increment 1: leave stage-local KV lengths on hit, restore them on miss.
2977    pub fn spec_fork_reconcile_kv(
2978        &self,
2979        len_ptrs: &CudaSlice<u64>,
2980        saved: &CudaSlice<i32>,
2981        acc: &CudaSlice<u32>,
2982        valid: &CudaSlice<u32>,
2983        base: usize,
2984        n_layer: usize,
2985    ) -> Result<(), Box<dyn std::error::Error>> {
2986        let f = self.func("spec_fork_reconcile_kv");
2987        let (b, nl) = (base as i32, n_layer as i32);
2988        let cfg = LaunchConfig {
2989            grid_dim: (n_layer.div_ceil(64) as u32, 1, 1),
2990            block_dim: (64, 1, 1),
2991            shared_mem_bytes: 0,
2992        };
2993        let __s_bl = self.gpu.stream();
2994        let mut bl = __s_bl.launch_builder(&f);
2995        bl.arg(len_ptrs)
2996            .arg(saved)
2997            .arg(acc)
2998            .arg(valid)
2999            .arg(&b)
3000            .arg(&nl);
3001        unsafe {
3002            bl.launch(cfg)?;
3003        }
3004        Ok(())
3005    }
3006
3007    /// OPTIPIPE increment 1: conditionally restore one stage-owned recurrent-state buffer.
3008    pub fn spec_fork_restore_f32(
3009        &self,
3010        snapshot: &CudaSlice<f32>,
3011        state: &mut CudaSlice<f32>,
3012        valid: &CudaSlice<u32>,
3013    ) -> Result<(), Box<dyn std::error::Error>> {
3014        assert_eq!(
3015            snapshot.len(),
3016            state.len(),
3017            "fork recurrent snapshot shape mismatch"
3018        );
3019        let f = self.func("spec_fork_restore_f32");
3020        let n = state.len() as i32;
3021        let blocks = state.len().div_ceil(256).min(65535).max(1) as u32;
3022        let cfg = LaunchConfig {
3023            grid_dim: (blocks, 1, 1),
3024            block_dim: (256, 1, 1),
3025            shared_mem_bytes: 0,
3026        };
3027        let __s_bl = self.gpu.stream();
3028        let mut bl = __s_bl.launch_builder(&f);
3029        bl.arg(snapshot).arg(state).arg(valid).arg(&n);
3030        unsafe {
3031            bl.launch(cfg)?;
3032        }
3033        Ok(())
3034    }
3035
3036    /// ROUND-STREAM stage (b): device next-round seed gather (see spec_seed_gather header).
3037    /// Caller D2Ds h_seed into fill_prev after (both slots carry the same value in every arm).
3038    pub fn spec_seed_gather(
3039        &self,
3040        vx: &CudaSlice<f32>,
3041        fill_prev: &CudaSlice<f32>,
3042        acc: &CudaSlice<u32>,
3043        h_seed: &mut CudaSlice<f32>,
3044        base: usize,
3045        n_embd: usize,
3046    ) -> Result<(), Box<dyn std::error::Error>> {
3047        let f = self.func("spec_seed_gather");
3048        let (b, ne) = (base as i32, n_embd as i32);
3049        let cfg = LaunchConfig {
3050            grid_dim: (n_embd.div_ceil(256) as u32, 1, 1),
3051            block_dim: (256, 1, 1),
3052            shared_mem_bytes: 0,
3053        };
3054        let __s_bl = self.gpu.stream();
3055        let mut bl = __s_bl.launch_builder(&f);
3056        bl.arg(vx)
3057            .arg(fill_prev)
3058            .arg(acc)
3059            .arg(h_seed)
3060            .arg(&b)
3061            .arg(&ne);
3062        unsafe {
3063            bl.launch(cfg)?;
3064        }
3065        Ok(())
3066    }
3067
3068    /// ROUND-STREAM stage (a): device greedy accept walk (see spec_accept_greedy header).
3069    pub fn spec_accept_greedy(
3070        &self,
3071        preds: &CudaSlice<u32>,
3072        draft: &CudaSlice<u32>,
3073        last_pred: u32,
3074        base: usize,
3075        k_round: usize,
3076        out: &mut CudaSlice<u32>,
3077    ) -> Result<(), Box<dyn std::error::Error>> {
3078        let f = self.func("spec_accept_greedy");
3079        let (b, k) = (base as i32, k_round as i32);
3080        let cfg = LaunchConfig {
3081            grid_dim: (1, 1, 1),
3082            block_dim: (32, 1, 1),
3083            shared_mem_bytes: 0,
3084        };
3085        let __s_bl = self.gpu.stream();
3086        let mut bl = __s_bl.launch_builder(&f);
3087        bl.arg(preds)
3088            .arg(draft)
3089            .arg(&last_pred)
3090            .arg(&b)
3091            .arg(&k)
3092            .arg(out);
3093        unsafe {
3094            bl.launch(cfg)?;
3095        }
3096        Ok(())
3097    }
3098
3099    // ================= SAMPLED-SPEC PRIMITIVES (spec_sample.cu, piece A) =================
3100    // Counter-based randomness: every call takes (seed, stream_pos) — the caller owns the
3101    // event counter (one per sampled token). temp <= 0 arms are exact greedy limits.
3102
3103    /// y = x/temp + Gumbel(Philox(seed, stream_pos)) over n logits (then run device argmax on y
3104    /// = one categorical sample at temperature `temp`). temp<=0: y = x (pure copy).
3105    pub fn gumbel_perturb(
3106        &self,
3107        x: &CudaSlice<f32>,
3108        y: &mut CudaSlice<f32>,
3109        n: usize,
3110        seed: u64,
3111        stream_pos: u32,
3112        temp: f32,
3113    ) -> Result<(), Box<dyn std::error::Error>> {
3114        let f = self.func("gumbel_perturb_f32");
3115        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3116        let cfg = LaunchConfig {
3117            grid_dim: (n.div_ceil(256) as u32, 1, 1),
3118            block_dim: (256, 1, 1),
3119            shared_mem_bytes: 0,
3120        };
3121        let __s_b = self.gpu.stream();
3122        let mut b = __s_b.launch_builder(&f);
3123        b.arg(x)
3124            .arg(&mut *y)
3125            .arg(&ni)
3126            .arg(&slo)
3127            .arg(&shi)
3128            .arg(&stream_pos)
3129            .arg(&temp);
3130        unsafe {
3131            b.launch(cfg)?;
3132        }
3133        Ok(())
3134    }
3135
3136    /// GRAMMAR TOKEN MASK (constrained decoding, lane/constrained-full): ban every vocab id
3137    /// whose bit is unset in the packed llguidance bitset, IN PLACE on row `col` of a stacked
3138    /// [B, n_vocab] logits buffer. `mask` = the SimpleVob u32 words H2D'd verbatim
3139    /// (~n_vocab/8 bytes/step — trivial on PCIe); ids >= 32*mask_words (padded lm_head tail)
3140    /// are banned too, the device twin of constrained::apply_mask. Banned value -FLT_MAX ==
3141    /// the argmax/gumbel kernels' init sentinel, so a fully-banned tail can never win and
3142    /// ordering matches the host -inf mask bit-for-bit for every finite logit.
3143    pub fn mask_logits_col(
3144        &self,
3145        logits: &mut CudaSlice<f32>,
3146        mask: &CudaSlice<u32>,
3147        col: usize,
3148        n: usize,
3149        mask_words: usize,
3150    ) -> Result<(), Box<dyn std::error::Error>> {
3151        let f = self.func("mask_logits_f32");
3152        let (ci, ni, mw) = (col as i32, n as i32, mask_words as i32);
3153        let cfg = LaunchConfig {
3154            grid_dim: (n.div_ceil(256).min(1024) as u32, 1, 1),
3155            block_dim: (256, 1, 1),
3156            shared_mem_bytes: 0,
3157        };
3158        let __s_b = self.gpu.stream();
3159        let mut b = __s_b.launch_builder(&f);
3160        b.arg(&mut *logits).arg(mask).arg(&ci).arg(&ni).arg(&mw);
3161        unsafe {
3162            b.launch(cfg)?;
3163        }
3164        Ok(())
3165    }
3166
3167    /// Column-`col` twin of `gumbel_perturb` over stacked logits [B, n_vocab] (the batched
3168    /// serving tick's device sampler): y = x[col]/temp + gumbel(seed, stream_pos, lane).
3169    /// SAME kernel/Philox mapping as `gumbel_perturb` — bit-identical perturbation for the
3170    /// same (seed, stream_pos, temp) regardless of which batch column the row sits in
3171    /// (the lane index is the in-row position; `col` only moves the input pointer). That
3172    /// pointer-invariance IS the serving isolation contract for sampled rows.
3173    pub fn gumbel_perturb_col(
3174        &self,
3175        x: &CudaSlice<f32>,
3176        col: usize,
3177        y: &mut CudaSlice<f32>,
3178        n: usize,
3179        seed: u64,
3180        stream_pos: u32,
3181        temp: f32,
3182    ) -> Result<(), Box<dyn std::error::Error>> {
3183        let f = self.func("gumbel_perturb_f32");
3184        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3185        let col_view = x.slice(col * n..(col + 1) * n);
3186        let cfg = LaunchConfig {
3187            grid_dim: (n.div_ceil(256) as u32, 1, 1),
3188            block_dim: (256, 1, 1),
3189            shared_mem_bytes: 0,
3190        };
3191        let __s_b = self.gpu.stream();
3192        let mut b = __s_b.launch_builder(&f);
3193        b.arg(&col_view)
3194            .arg(&mut *y)
3195            .arg(&ni)
3196            .arg(&slo)
3197            .arg(&shi)
3198            .arg(&stream_pos)
3199            .arg(&temp);
3200        unsafe {
3201            b.launch(cfg)?;
3202        }
3203        Ok(())
3204    }
3205
3206    /// Filtered twin of `gumbel_perturb_col`: the per-row (row_max, th) floor comes from
3207    /// DEVICE buffers (`filter_stats` output slots at `stat_idx`) — one filtered draw from
3208    /// the top-k/top-p/min-p-truncated softmax with no stat D2H and no row copy. Same
3209    /// Philox mapping as every gumbel kernel (pointer-invariant across batch columns —
3210    /// the serving isolation contract for sampled rows).
3211    #[allow(clippy::too_many_arguments)]
3212    pub fn gumbel_perturb_filtered_col(
3213        &self,
3214        x: &CudaSlice<f32>,
3215        col: usize,
3216        y: &mut CudaSlice<f32>,
3217        n: usize,
3218        seed: u64,
3219        stream_pos: u32,
3220        temp: f32,
3221        stat_max: &CudaSlice<f32>,
3222        stat_th: &CudaSlice<f32>,
3223        stat_idx: usize,
3224    ) -> Result<(), Box<dyn std::error::Error>> {
3225        let f = self.func("gumbel_perturb_filtered_col_f32");
3226        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3227        let (ci, si) = (col as i32, stat_idx as i32);
3228        let cfg = LaunchConfig {
3229            grid_dim: (n.div_ceil(256) as u32, 1, 1),
3230            block_dim: (256, 1, 1),
3231            shared_mem_bytes: 0,
3232        };
3233        let __s_b = self.gpu.stream();
3234        let mut b = __s_b.launch_builder(&f);
3235        b.arg(x)
3236            .arg(&ci)
3237            .arg(&mut *y)
3238            .arg(&ni)
3239            .arg(&slo)
3240            .arg(&shi)
3241            .arg(&stream_pos)
3242            .arg(&temp)
3243            .arg(stat_max)
3244            .arg(stat_th)
3245            .arg(&si);
3246        unsafe {
3247            b.launch(cfg)?;
3248        }
3249        Ok(())
3250    }
3251
3252    /// In-graph sampling-event counter bump (spec_sample.cu kernel 5): ctr[0] += 1. The sampled
3253    /// graph-draft chain replays with FIXED kernel args, so the Philox event counter must be
3254    /// DEVICE data — the host seeds it once per round; every replay bumps it before the perturb
3255    /// reads it (counter is data, not state — graph-replay-safe).
3256    pub fn sctr_inc(&self, ctr: &mut CudaSlice<u32>) -> Result<(), Box<dyn std::error::Error>> {
3257        let f = self.func("memra_sctr_inc");
3258        let cfg = LaunchConfig {
3259            grid_dim: (1, 1, 1),
3260            block_dim: (1, 1, 1),
3261            shared_mem_bytes: 0,
3262        };
3263        let __s_b = self.gpu.stream();
3264        let mut b = __s_b.launch_builder(&f);
3265        b.arg(&mut *ctr);
3266        unsafe {
3267            b.launch(cfg)?;
3268        }
3269        Ok(())
3270    }
3271
3272    /// Graph-capturable `gumbel_perturb`: the sampling-event counter comes from DEVICE memory
3273    /// (`ctr[0]`) instead of a host scalar. Identical math to `gumbel_perturb` at
3274    /// stream_pos == ctr[0] (same Philox call, same lane mapping) — the eager and graph sampled
3275    /// chains produce bit-identical perturbations for the same (seed, counter, temp).
3276    pub fn gumbel_perturb_ctr(
3277        &self,
3278        x: &CudaSlice<f32>,
3279        y: &mut CudaSlice<f32>,
3280        n: usize,
3281        seed: u64,
3282        ctr: &CudaSlice<u32>,
3283        temp: f32,
3284    ) -> Result<(), Box<dyn std::error::Error>> {
3285        let f = self.func("gumbel_perturb_ctr_f32");
3286        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3287        let cfg = LaunchConfig {
3288            grid_dim: (n.div_ceil(256) as u32, 1, 1),
3289            block_dim: (256, 1, 1),
3290            shared_mem_bytes: 0,
3291        };
3292        let __s_b = self.gpu.stream();
3293        let mut b = __s_b.launch_builder(&f);
3294        b.arg(x)
3295            .arg(&mut *y)
3296            .arg(&ni)
3297            .arg(&slo)
3298            .arg(&shi)
3299            .arg(ctr)
3300            .arg(&temp);
3301        unsafe {
3302            b.launch(cfg)?;
3303        }
3304        Ok(())
3305    }
3306
3307    /// out[pair] = softmax_temp(x[rows[pair]])[ids[pair]] for npair (row, id) pairs; rows index
3308    /// into x with `row_stride` f32s per row. temp<=0: out = 1.0 iff id is the row argmax
3309    /// (smallest-index tie-break — matches the argmax-gate contract).
3310    pub fn softmax_gather(
3311        &self,
3312        x: &CudaSlice<f32>,
3313        row_stride: usize,
3314        ids: &CudaSlice<u32>,
3315        rows: &CudaSlice<i32>,
3316        out: &mut CudaSlice<f32>,
3317        n: usize,
3318        npair: usize,
3319        temp: f32,
3320    ) -> Result<(), Box<dyn std::error::Error>> {
3321        let f = self.func("softmax_gather_f32");
3322        let (ni, rs) = (n as i32, row_stride as i64);
3323        let np = npair as i32;
3324        let cfg = LaunchConfig {
3325            grid_dim: (npair as u32, 1, 1),
3326            block_dim: (256, 1, 1),
3327            shared_mem_bytes: 0,
3328        };
3329        let __s_b = self.gpu.stream();
3330        let mut b = __s_b.launch_builder(&f);
3331        b.arg(x)
3332            .arg(&rs)
3333            .arg(ids)
3334            .arg(rows)
3335            .arg(&mut *out)
3336            .arg(&ni)
3337            .arg(&np)
3338            .arg(&temp);
3339        unsafe {
3340            b.launch(cfg)?;
3341        }
3342        Ok(())
3343    }
3344
3345    /// Sample token from norm(max(0, softmax_temp(p) - softmax_temp(q))) (q = None -> plain
3346    /// categorical from softmax_temp(p)). Row stats (max, sumexp at temp) must be precomputed
3347    /// (softmax_gather's pass-1 values; see spec.rs caller). Deterministic fixed-order CDF walk.
3348    pub fn residual_sample(
3349        &self,
3350        p: &CudaSlice<f32>,
3351        q: Option<&CudaSlice<f32>>,
3352        n: usize,
3353        temp: f32,
3354        seed: u64,
3355        stream_pos: u32,
3356        out_tok: &mut CudaSlice<u32>,
3357    ) -> Result<(), Box<dyn std::error::Error>> {
3358        let f = self.func("residual_sample_f32");
3359        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3360        let nth = 1024u32;
3361        let cfg = LaunchConfig {
3362            grid_dim: (1, 1, 1),
3363            block_dim: (nth, 1, 1),
3364            shared_mem_bytes: 0,
3365        };
3366        let has_q: i32 = q.is_some() as i32;
3367        let qbuf = q.unwrap_or(p); // dummy when absent; kernel gates on has_q
3368        let __s_b = self.gpu.stream();
3369        let mut b = __s_b.launch_builder(&f);
3370        b.arg(p)
3371            .arg(qbuf)
3372            .arg(&has_q)
3373            .arg(&ni)
3374            .arg(&temp)
3375            .arg(&slo)
3376            .arg(&shi)
3377            .arg(&stream_pos)
3378            .arg(&mut *out_tok);
3379        unsafe {
3380            b.launch(cfg)?;
3381        }
3382        Ok(())
3383    }
3384
3385    /// Access the shared MoE residency cache (EDGE-1 §B), building it on first use under
3386    /// MEMRA_MOE_CACHE. The closure runs while the lock is held — keep it to lookup/admit/issue, not
3387    /// the GEMM. `max_block_bytes` sizes the slots (largest of gate/up/down). Returns the closure's
3388    /// result. If MEMRA_MOE_CACHE is unset this is never called (the caller checks the env first).
3389    pub fn with_moe_cache<R>(
3390        &self,
3391        max_block_bytes: usize,
3392        f: impl FnOnce(
3393            &mut crate::moe_cache::MoeSlotCache,
3394            &Engine,
3395        ) -> Result<R, Box<dyn std::error::Error>>,
3396    ) -> Result<R, Box<dyn std::error::Error>> {
3397        let mut guard = self.moe_cache.lock().unwrap();
3398        if guard.is_none() {
3399            *guard = Some(crate::moe_cache::MoeSlotCache::new(self, max_block_bytes)?);
3400        }
3401        let cache = guard.as_mut().unwrap();
3402        f(cache, self)
3403    }
3404
3405    /// Freeze the already-built MoE residency set. This never constructs a cache: callers use it
3406    /// only after a real prefill has populated the machine-specific CPU/GPU working set.
3407    pub fn freeze_moe_cache(&self) {
3408        if let Some(cache) = self.moe_cache.lock().unwrap().as_mut() {
3409            cache.freeze();
3410        }
3411    }
3412
3413    /// The current residency set as (layer, proj, ex) triples, or None if no cache was built.
3414    /// Never constructs a cache.
3415    pub fn export_moe_residency(&self) -> Option<Vec<(u16, u8, u16)>> {
3416        self.moe_cache
3417            .lock()
3418            .unwrap()
3419            .as_ref()
3420            .map(crate::moe_cache::MoeSlotCache::export_residency)
3421    }
3422
3423    pub(crate) fn moe_cache_frozen(&self) -> bool {
3424        self.moe_cache
3425            .lock()
3426            .unwrap()
3427            .as_ref()
3428            .is_some_and(crate::moe_cache::MoeSlotCache::is_frozen)
3429    }
3430
3431    /// A frozen heterogeneous CPU/GPU expert split cannot use Hy3's ordinary batched prefill
3432    /// efficiently: T>=PRIME_MIN_T bypasses the CPU backend and transiently rereads every missing
3433    /// expert through the GPU spill path. Replay the short prompt through decode after freezing,
3434    /// while leaving the profiling warmup's established batched behavior untouched.
3435    /// (`pub`: run-gen's #46 batched-prime gate skips itself when generation will take the
3436    /// tokenwise arm anyway.)
3437    pub fn frozen_cpu_experts_prefer_tokenwise_prime(&self) -> bool {
3438        crate::cpu_experts::configured()
3439            && self.moe_cache_frozen()
3440            && std::env::var("MEMRA_CPU_EXPERT_BATCHED_PRIME").as_deref() != Ok("1")
3441    }
3442
3443    /// Install the loaded model's exact retained expert-block inventory before lazy cache build.
3444    pub(crate) fn configure_moe_cache_layout(&self, block_bytes: Vec<usize>) {
3445        assert!(
3446            self.moe_cache.lock().unwrap().is_none(),
3447            "MoE cache layout configured after cache construction"
3448        );
3449        *self.moe_cache_layout.lock().unwrap() = Some(block_bytes);
3450    }
3451
3452    pub(crate) fn moe_cache_layout(&self) -> Option<Vec<usize>> {
3453        self.moe_cache_layout.lock().unwrap().clone()
3454    }
3455
3456    /// True if the MoE residency cache is enabled (MEMRA_MOE_CACHE set).
3457    pub fn moe_cache_enabled() -> bool {
3458        std::env::var("MEMRA_MOE_CACHE").as_deref() != Ok("0")
3459    }
3460
3461    /// Snapshot the MoE cache counters (hits, misses, staged_bytes, n_slots) for the §D.4 PCIe gate.
3462    /// Returns None if the cache was never built (disabled or no MoE forward ran).
3463    pub fn moe_cache_stats(&self) -> Option<(u64, u64, u64, usize)> {
3464        let guard = self.moe_cache.lock().unwrap();
3465        guard
3466            .as_ref()
3467            .map(|c| (c.hits, c.misses, c.staged_bytes, c.n_slots()))
3468    }
3469
3470    /// Experimental CPU expert backend counters: completed layer calls, experts served, and the
3471    /// sum of backend wall nanoseconds. The timer includes explicit disk->RAM fills on cache misses;
3472    /// callers compare a before/after snapshot around a decode window.
3473    pub fn cpu_expert_stats(
3474        &self,
3475    ) -> Option<(u64, u64, u64, u64, u64, u64, u64, u64, u64, u64, u64)> {
3476        crate::cpu_experts::configured().then(crate::cpu_experts::stats)
3477    }
3478
3479    /// Caller-blocked nanoseconds at CPU expert joins. Compare before/after snapshots to measure
3480    /// the backend tail that resident-GPU expert work did not hide.
3481    pub fn cpu_expert_predictor_stats(&self) -> (u64, u64) {
3482        crate::cpu_experts::predictor_stats()
3483    }
3484
3485    pub fn cpu_expert_exposed_wait_ns(&self) -> Option<u64> {
3486        crate::cpu_experts::configured().then(crate::cpu_experts::exposed_wait_ns)
3487    }
3488
3489    /// CPU-routed expert selections grouped by how many of their three projections were already
3490    /// resident in HBM. This makes otherwise-stranded partial residency visible to tuning runs.
3491    pub fn cpu_expert_gpu_residency_stats(&self) -> Option<(u64, u64, u64)> {
3492        crate::cpu_experts::configured().then(crate::cpu_experts::incomplete_gpu_residency_stats)
3493    }
3494
3495    /// Positioned-read proof-backend counters:
3496    /// `(reads, bytes, read_errors, short_reads, mmap_fallbacks, buffer_waits, ring_full)`.
3497    pub fn moe_pread_stats(&self) -> Option<(u64, u64, u64, u64, u64, u64, u64)> {
3498        let guard = self.moe_cache.lock().unwrap();
3499        guard
3500            .as_ref()
3501            .and_then(|cache| cache.pread_stats())
3502            .map(|stats| {
3503                (
3504                    stats.reads,
3505                    stats.bytes,
3506                    stats.read_errors,
3507                    stats.short_reads,
3508                    stats.fallbacks,
3509                    stats.buffer_waits,
3510                    stats.ring_full,
3511                )
3512            })
3513    }
3514
3515    /// Spill configuration values that warned and substituted their documented defaults.
3516    pub fn spill_config_fallbacks(&self) -> u64 {
3517        crate::spill_pread::config_fallbacks()
3518    }
3519
3520    /// Reset the MoE cache perf counters (to separate warmup from steady-state windows).
3521    pub fn moe_cache_reset_counters(&self) {
3522        if let Some(c) = self.moe_cache.lock().unwrap().as_mut() {
3523            c.reset_counters();
3524        }
3525    }
3526
3527    pub fn htod_bytes(&self, v: &[u8]) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
3528        Ok(self.gpu.stream().clone_htod(v)?)
3529    }
3530
3531    /// `htod_bytes` with a mapped (uninit) tail pad: the wide-load expert dots read up to 6B
3532    /// past the final q4_0 block through their aligned window — the bytes never reach a
3533    /// result (funnelshift discards them) but must be mapped memory.
3534    pub fn htod_bytes_padded(
3535        &self,
3536        v: &[u8],
3537        pad: usize,
3538    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
3539        let mut d = self.alloc_u8_uninit(v.len() + pad)?;
3540        {
3541            let mut view = d.slice_mut(0..v.len());
3542            self.gpu.stream().memcpy_htod(v, &mut view)?;
3543        }
3544        Ok(d)
3545    }
3546
3547    /// Device-to-device copy of `src` into `dst[off..off+len]` (f32). For in-place KV append.
3548    pub fn copy_into(
3549        &self,
3550        dst: &mut CudaSlice<f32>,
3551        off: usize,
3552        src: &CudaSlice<f32>,
3553        len: usize,
3554    ) -> Result<(), Box<dyn std::error::Error>> {
3555        let mut view = dst.slice_mut(off..off + len);
3556        self.gpu
3557            .stream()
3558            .memcpy_dtod(&src.slice(0..len), &mut view)?;
3559        Ok(())
3560    }
3561
3562    /// View a sub-range of a device buffer (for attending over [0..len) of a KV cache).
3563    /// u8 twin of copy_into (D2D byte-range copy at an offset).
3564    pub fn copy_u8_into(
3565        &self,
3566        dst: &mut CudaSlice<u8>,
3567        off: usize,
3568        src: &CudaSlice<u8>,
3569        len: usize,
3570    ) -> Result<(), Box<dyn std::error::Error>> {
3571        let mut view = dst.slice_mut(off..off + len);
3572        self.gpu
3573            .stream()
3574            .memcpy_dtod(&src.slice(0..len), &mut view)?;
3575        Ok(())
3576    }
3577
3578    /// D2D byte-range copy with explicit source and destination offsets.
3579    pub fn copy_u8_range_into(
3580        &self,
3581        dst: &mut CudaSlice<u8>,
3582        dst_off: usize,
3583        src: &CudaSlice<u8>,
3584        src_off: usize,
3585        len: usize,
3586    ) -> Result<(), Box<dyn std::error::Error>> {
3587        let mut dst_view = dst.slice_mut(dst_off..dst_off + len);
3588        self.gpu
3589            .stream()
3590            .memcpy_dtod(&src.slice(src_off..src_off + len), &mut dst_view)?;
3591        Ok(())
3592    }
3593
3594    /// Resolve an absolute append slot to the Step35 SWA layer's physical rows. At wrap, copy
3595    /// only the aligned live prefix through temporary device storage and rebase it at row zero,
3596    /// keeping the audited attention range contiguous without changing its absolute start.
3597    pub fn prepare_kv_append(
3598        &self,
3599        kv: &mut crate::cache::KvLayer,
3600        retain_from: usize,
3601        append_rows: usize,
3602    ) -> Result<usize, Box<dyn std::error::Error>> {
3603        let Some(plan) = kv
3604            .ring
3605            .as_ref()
3606            .map(|ring| ring.append_plan(kv.len, retain_from, append_rows))
3607            .transpose()?
3608        else {
3609            return Ok(kv.len);
3610        };
3611        match plan {
3612            crate::cache::KvRingAppend::Contiguous { write_row } => Ok(write_row),
3613            crate::cache::KvRingAppend::Rebase {
3614                src_row,
3615                keep_rows,
3616                new_base,
3617                write_row,
3618            } => {
3619                if keep_rows > 0 {
3620                    let k_len = keep_rows * kv.k_tok_bytes;
3621                    let v_len = keep_rows * kv.v_tok_bytes;
3622                    let mut k_tmp = self.alloc_u8_uninit(k_len)?;
3623                    let mut v_tmp = self.alloc_u8_uninit(v_len)?;
3624                    self.copy_u8_range_into(&mut k_tmp, 0, &kv.k, src_row * kv.k_tok_bytes, k_len)?;
3625                    self.copy_u8_range_into(&mut v_tmp, 0, &kv.v, src_row * kv.v_tok_bytes, v_len)?;
3626                    self.copy_u8_into(&mut kv.k, 0, &k_tmp, k_len)?;
3627                    self.copy_u8_into(&mut kv.v, 0, &v_tmp, v_len)?;
3628                }
3629                kv.ring.as_mut().unwrap().apply_rebase(new_base);
3630                Ok(write_row)
3631            }
3632        }
3633    }
3634
3635    /// H2D write of `src` into `dst[off..off+src.len()]` (u8). In-place row updates for the
3636    /// adaptive trim head: no realloc, so captured graphs keep their baked addresses.
3637    pub fn htod_u8_into(
3638        &self,
3639        dst: &mut CudaSlice<u8>,
3640        off: usize,
3641        src: &[u8],
3642    ) -> Result<(), Box<dyn std::error::Error>> {
3643        let mut view = dst.slice_mut(off..off + src.len());
3644        self.gpu.stream().memcpy_htod(src, &mut view)?;
3645        Ok(())
3646    }
3647
3648    pub fn view<'a>(&self, b: &'a CudaSlice<f32>, len: usize) -> cudarc::driver::CudaView<'a, f32> {
3649        b.slice(0..len)
3650    }
3651
3652    /// View the first `len` BYTES of a u8 device buffer (quantized KV cache: [0..t_kv*tok_bytes)).
3653    /// Byte-range view (gemma4 R6 window offset into the quantized KV stream).
3654    pub fn view_u8_range<'a>(
3655        &self,
3656        b: &'a CudaSlice<u8>,
3657        start: usize,
3658        end: usize,
3659    ) -> cudarc::driver::CudaView<'a, u8> {
3660        b.slice(start..end)
3661    }
3662    pub fn view_u8<'a>(
3663        &self,
3664        b: &'a CudaSlice<u8>,
3665        len: usize,
3666    ) -> cudarc::driver::CudaView<'a, u8> {
3667        b.slice(0..len)
3668    }
3669
3670    /// Append-quantize ONE token's post-RoPE K (q8_0) and V (q5_1) into the resident byte caches at
3671    /// token index `t` (KVQUANT-PLAN §C). One CTA (one warp) per 32-element block; the kernel writes
3672    /// the f16 scale(s) + packed quants for K and V. k_row/v_row are f32 [kv_dim_k]/[kv_dim_v].
3673    pub fn append_kv_quantized(
3674        &self,
3675        k_row: &CudaSlice<f32>,
3676        v_row: &CudaSlice<f32>,
3677        kc: &mut CudaSlice<u8>,
3678        vc: &mut CudaSlice<u8>,
3679        t: usize,
3680        kv_dim_k: usize,
3681        kv_dim_v: usize,
3682        k_tok_bytes: usize,
3683        v_tok_bytes: usize,
3684        g: bool,
3685    ) -> Result<(), Box<dyn std::error::Error>> {
3686        let f = if g {
3687            self.func_g("append_quantize_kv_q8_0_q5_1")
3688        } else {
3689            self.func("append_quantize_kv_q8_0_q5_1")
3690        };
3691        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
3692        let cfg = LaunchConfig {
3693            grid_dim: (nblk, 1, 1),
3694            block_dim: (32, 1, 1),
3695            shared_mem_bytes: 0,
3696        };
3697        let (ti, kdk, kdv) = (t as i32, kv_dim_k as i32, kv_dim_v as i32);
3698        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3699        let __s_b = self.gpu.stream();
3700        let mut b = __s_b.launch_builder(&f);
3701        b.arg(k_row)
3702            .arg(v_row)
3703            .arg(kc)
3704            .arg(vc)
3705            .arg(&ti)
3706            .arg(&kdk)
3707            .arg(&kdv)
3708            .arg(&ktb)
3709            .arg(&vtb);
3710        unsafe {
3711            b.launch(cfg)?;
3712        }
3713        Ok(())
3714    }
3715
3716    /// Device-counter variant of `append_kv_quantized` (CUDA-GRAPH-PLAN Phase 2): the write slot
3717    /// `t` is read from `t_dev[0]` (a resident device i32[1]) instead of a host int arg, so the
3718    /// launch args are FIXED across decode steps (graph-capturable). Identical quant math.
3719    pub fn append_kv_quantized_dc(
3720        &self,
3721        k_row: &CudaSlice<f32>,
3722        v_row: &CudaSlice<f32>,
3723        kc: &mut CudaSlice<u8>,
3724        vc: &mut CudaSlice<u8>,
3725        t_dev: &CudaSlice<i32>,
3726        kv_dim_k: usize,
3727        kv_dim_v: usize,
3728        k_tok_bytes: usize,
3729        v_tok_bytes: usize,
3730        g: bool,
3731    ) -> Result<(), Box<dyn std::error::Error>> {
3732        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
3733        let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
3734        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3735        // PDL wave-B2: flash-module flavor mirrors the builder path's g flag exactly.
3736        if Self::pdl_on() && Self::pdl_wb_on() {
3737            use cudarc::driver::{DevicePtr, DevicePtrMut};
3738            let s = &self.gpu.stream();
3739            let (pk, _g0) = k_row.device_ptr(s);
3740            let (pv, _g1) = v_row.device_ptr(s);
3741            let (pkc, _g2) = kc.device_ptr_mut(s);
3742            let (pvc, _g3) = vc.device_ptr_mut(s);
3743            let (pt, _g4) = t_dev.device_ptr(s);
3744            let mut ps = [
3745                &pk as *const _ as *mut std::ffi::c_void,
3746                &pv as *const _ as *mut _,
3747                &pkc as *const _ as *mut _,
3748                &pvc as *const _ as *mut _,
3749                &pt as *const _ as *mut _,
3750                &kdk as *const _ as *mut _,
3751                &kdv as *const _ as *mut _,
3752                &ktb as *const _ as *mut _,
3753                &vtb as *const _ as *mut _,
3754            ];
3755            unsafe {
3756                self.launch_pdl_flash(
3757                    g,
3758                    "append_quantize_kv_q8_0_q5_1_dc",
3759                    (nblk, 1, 1),
3760                    (32, 1, 1),
3761                    0,
3762                    &mut ps,
3763                )?;
3764            }
3765            return Ok(());
3766        }
3767        let f = if g {
3768            self.func_g("append_quantize_kv_q8_0_q5_1_dc")
3769        } else {
3770            self.func("append_quantize_kv_q8_0_q5_1_dc")
3771        };
3772        let cfg = LaunchConfig {
3773            grid_dim: (nblk, 1, 1),
3774            block_dim: (32, 1, 1),
3775            shared_mem_bytes: 0,
3776        };
3777        let __s_b = self.gpu.stream();
3778        let mut b = __s_b.launch_builder(&f);
3779        b.arg(k_row)
3780            .arg(v_row)
3781            .arg(kc)
3782            .arg(vc)
3783            .arg(t_dev)
3784            .arg(&kdk)
3785            .arg(&kdv)
3786            .arg(&ktb)
3787            .arg(&vtb);
3788        unsafe {
3789            b.launch(cfg)?;
3790        }
3791        Ok(())
3792    }
3793
3794    /// Append-quantize T token rows in one shot (BATCHED PROMPT PRIME). k_rows/v_rows are
3795    /// token-major [T, kv_dim] post-RoPE f32; rows land at cache slots t0..t0+T. Default = the
3796    /// batched `_rows` kernel: one (nblk, T) launch whose per-(block,token) warp program is the
3797    /// per-token append kernel verbatim -> every written row is BIT-IDENTICAL to T sequential
3798    /// `append_kv_quantized_view` calls (kernel_check pins the bytes). MEMRA_PRIME_APPEND_LOOP=1
3799    /// forces the T-launch per-row loop (the A/B seam that measured the launch overhead).
3800    #[allow(clippy::too_many_arguments)]
3801    pub fn append_kv_quantized_rows(
3802        &self,
3803        k_rows: &CudaSlice<f32>,
3804        v_rows: &CudaSlice<f32>,
3805        kc: &mut CudaSlice<u8>,
3806        vc: &mut CudaSlice<u8>,
3807        t0: usize,
3808        t: usize,
3809        kv_dim_k: usize,
3810        kv_dim_v: usize,
3811        k_tok_bytes: usize,
3812        v_tok_bytes: usize,
3813        g: bool,
3814    ) -> Result<(), Box<dyn std::error::Error>> {
3815        if std::env::var("MEMRA_PRIME_APPEND_LOOP").is_ok() {
3816            for i in 0..t {
3817                let k_row = k_rows.slice(i * kv_dim_k..(i + 1) * kv_dim_k);
3818                let v_row = v_rows.slice(i * kv_dim_v..(i + 1) * kv_dim_v);
3819                self.append_kv_quantized_view(
3820                    &k_row,
3821                    &v_row,
3822                    kc,
3823                    vc,
3824                    t0 + i,
3825                    kv_dim_k,
3826                    kv_dim_v,
3827                    k_tok_bytes,
3828                    v_tok_bytes,
3829                    g,
3830                )?;
3831            }
3832            return Ok(());
3833        }
3834        let f = if g {
3835            self.func_g("append_quantize_kv_q8_0_q5_1_rows")
3836        } else {
3837            self.func("append_quantize_kv_q8_0_q5_1_rows")
3838        };
3839        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
3840        let cfg = LaunchConfig {
3841            grid_dim: (nblk, t as u32, 1),
3842            block_dim: (32, 1, 1),
3843            shared_mem_bytes: 0,
3844        };
3845        let (t0i, kdk, kdv) = (t0 as i32, kv_dim_k as i32, kv_dim_v as i32);
3846        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3847        let __s_b = self.gpu.stream();
3848        let mut b = __s_b.launch_builder(&f);
3849        b.arg(k_rows)
3850            .arg(v_rows)
3851            .arg(kc)
3852            .arg(vc)
3853            .arg(&t0i)
3854            .arg(&kdk)
3855            .arg(&kdv)
3856            .arg(&ktb)
3857            .arg(&vtb);
3858        unsafe {
3859            b.launch(cfg)?;
3860        }
3861        Ok(())
3862    }
3863
3864    /// Increment a device i32[1] counter in place (p[0] += 1) via the resident `inc_i32` kernel.
3865    /// Used to advance the device-resident seqlen/pos counters inside the decode-dc path (and,
3866    /// later, inside a captured graph) without a host round-trip.
3867    pub fn inc_seqlen(&self, p: &mut CudaSlice<i32>) -> Result<(), Box<dyn std::error::Error>> {
3868        let f = self.func("inc_i32");
3869        let cfg = LaunchConfig {
3870            grid_dim: (1, 1, 1),
3871            block_dim: (1, 1, 1),
3872            shared_mem_bytes: 0,
3873        };
3874        let __s_b = self.gpu.stream();
3875        let mut b = __s_b.launch_builder(&f);
3876        b.arg(p);
3877        unsafe {
3878            b.launch(cfg)?;
3879        }
3880        Ok(())
3881    }
3882
3883    /// Like `append_kv_quantized` but k_row/v_row are CudaViews (one token's row sliced out of a
3884    /// token-major [T, kv_dim] activation buffer — the MTP verify path appends T tokens).
3885    pub fn append_kv_quantized_view(
3886        &self,
3887        k_row: &cudarc::driver::CudaView<f32>,
3888        v_row: &cudarc::driver::CudaView<f32>,
3889        kc: &mut CudaSlice<u8>,
3890        vc: &mut CudaSlice<u8>,
3891        t: usize,
3892        kv_dim_k: usize,
3893        kv_dim_v: usize,
3894        k_tok_bytes: usize,
3895        v_tok_bytes: usize,
3896        g: bool,
3897    ) -> Result<(), Box<dyn std::error::Error>> {
3898        let f = if g {
3899            self.func_g("append_quantize_kv_q8_0_q5_1")
3900        } else {
3901            self.func("append_quantize_kv_q8_0_q5_1")
3902        };
3903        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
3904        let cfg = LaunchConfig {
3905            grid_dim: (nblk, 1, 1),
3906            block_dim: (32, 1, 1),
3907            shared_mem_bytes: 0,
3908        };
3909        let (ti, kdk, kdv) = (t as i32, kv_dim_k as i32, kv_dim_v as i32);
3910        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3911        let __s_b = self.gpu.stream();
3912        let mut b = __s_b.launch_builder(&f);
3913        b.arg(k_row)
3914            .arg(v_row)
3915            .arg(kc)
3916            .arg(vc)
3917            .arg(&ti)
3918            .arg(&kdk)
3919            .arg(&kdv)
3920            .arg(&ktb)
3921            .arg(&vtb);
3922        unsafe {
3923            b.launch(cfg)?;
3924        }
3925        Ok(())
3926    }
3927
3928    /// Device-to-device copy of a CudaView `src` into `dst[off..off+len]` (f32). Like `copy_into`
3929    /// but the source is a sub-view (e.g. one column of a token-major activation buffer).
3930    pub fn copy_view_into(
3931        &self,
3932        dst: &mut CudaSlice<f32>,
3933        off: usize,
3934        src: &cudarc::driver::CudaView<f32>,
3935        len: usize,
3936    ) -> Result<(), Box<dyn std::error::Error>> {
3937        let mut view = dst.slice_mut(off..off + len);
3938        self.gpu
3939            .stream()
3940            .memcpy_dtod(&src.slice(0..len), &mut view)?;
3941        Ok(())
3942    }
3943
3944    /// Real device-to-device COPY of `src` into a freshly allocated buffer (NOT an Arc clone).
3945    /// Used for cache snapshots (MTP-PLAN §D.4): `CudaSlice::clone()` only bumps a refcount and
3946    /// would alias the live buffer; this allocs new device memory and memcpy_dtod's the contents.
3947    pub fn clone_dtod(
3948        &self,
3949        src: &CudaSlice<f32>,
3950    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3951        let mut dst = self.gpu.stream().alloc_zeros::<f32>(src.len())?;
3952        self.gpu.stream().memcpy_dtod(src, &mut dst)?;
3953        Ok(dst)
3954    }
3955
3956    /// D2D row extraction: copy a view (e.g. one row of a [B, n] batch buffer) into `dst`.
3957    /// Stream-ordered, async — decode_batch's per-sequence row plumbing.
3958    pub fn dtod_copy_view(
3959        &self,
3960        src: &cudarc::driver::CudaView<f32>,
3961        dst: &mut CudaSlice<f32>,
3962    ) -> Result<(), Box<dyn std::error::Error>> {
3963        self.gpu.stream().memcpy_dtod(src, dst)?;
3964        Ok(())
3965    }
3966
3967    /// D2D i8 twin of `dtod_copy_view` (q8_1 activation rows).
3968    pub fn dtod_copy_view_i8(
3969        &self,
3970        src: &cudarc::driver::CudaView<i8>,
3971        dst: &mut CudaSlice<i8>,
3972    ) -> Result<(), Box<dyn std::error::Error>> {
3973        self.gpu.stream().memcpy_dtod(src, dst)?;
3974        Ok(())
3975    }
3976
3977    /// D2D row placement: copy `src` into `dst[offset .. offset+src.len()]`.
3978    pub fn dtod_copy_into(
3979        &self,
3980        src: &CudaSlice<f32>,
3981        dst: &mut CudaSlice<f32>,
3982        offset: usize,
3983    ) -> Result<(), Box<dyn std::error::Error>> {
3984        let n = src.len();
3985        let mut dv = dst.slice_mut(offset..offset + n);
3986        self.gpu.stream().memcpy_dtod(src, &mut dv)?;
3987        Ok(())
3988    }
3989
3990    /// Batched uniform-size D2D copy (engine-bundle slice 1, DSF-ROUNDCOST-20260820 §1.1):
3991    /// `n` disjoint regions of `words` f32 each; `table` = [src_0..src_{n-1}, dst_0..dst_{n-1}]
3992    /// raw device pointers. ONE kernel launch replaces `n` memcpy_dtod dispatches — the dspark
3993    /// round's snap/commit copy dribble (~0.9 ms/round of dispatch serialization measured).
3994    /// Bytes and stream order are identical to the memcpy sequence it replaces.
3995    pub fn copy_batch_uniform_f32(
3996        &self,
3997        table: &CudaSlice<u64>,
3998        n: usize,
3999        words: usize,
4000    ) -> Result<(), Box<dyn std::error::Error>> {
4001        if n == 0 || words == 0 {
4002            return Ok(());
4003        }
4004        debug_assert!(
4005            table.len() >= 2 * n,
4006            "pointer table must hold n srcs + n dsts"
4007        );
4008        let f = self.func("copy_batch_uniform_f32");
4009        // Enough blocks to stream a multi-MB region, few enough that (chunks x n) stays a
4010        // sane grid: 512K-word ssm regions get 48 grid-striding blocks each.
4011        let chunks = (words / 4).max(1).div_ceil(256).min(48) as u32;
4012        let (ni, wi) = (n as i32, words as i32);
4013        let cfg = LaunchConfig {
4014            grid_dim: (chunks, n as u32, 1),
4015            block_dim: (256, 1, 1),
4016            shared_mem_bytes: 0,
4017        };
4018        let __s = self.gpu.stream();
4019        let mut b = __s.launch_builder(&f);
4020        b.arg(table).arg(&ni).arg(&wi);
4021        unsafe {
4022            b.launch(cfg)?;
4023        }
4024        Ok(())
4025    }
4026
4027    /// H2D refresh of an EXISTING u64 pointer table IN PLACE (stable pointer — the batched
4028    /// state-copy tables are refreshed per round because the GDN ssm handles ping-pong).
4029    pub fn htod_u64_into(
4030        &self,
4031        v: &[u64],
4032        dst: &mut CudaSlice<u64>,
4033    ) -> Result<(), Box<dyn std::error::Error>> {
4034        let mut view = dst.slice_mut(0..v.len());
4035        self.gpu.stream().memcpy_htod(v, &mut view)?;
4036        Ok(())
4037    }
4038
4039    /// Indirect-source copy (engine-bundle slice 3): the src ADDRESS is loaded from a
4040    /// device pointer-table entry at run time, so a captured graph follows the gdn
4041    /// ping-pong through the same table its scan kernels read — a baked memcpy node
4042    /// would keep the capture-time physical buffer. `dst_off`/`words` in f32 elements.
4043    pub fn copy_indirect_src_f32(
4044        &self,
4045        src_entry: &cudarc::driver::CudaView<u64>,
4046        dst: &mut CudaSlice<f32>,
4047        dst_off: usize,
4048        words: usize,
4049    ) -> Result<(), Box<dyn std::error::Error>> {
4050        let f = self.func("copy_indirect_src_f32");
4051        let chunks = (words / 4).max(1).div_ceil(256).min(48) as u32;
4052        let wi = words as i32;
4053        let cfg = LaunchConfig {
4054            grid_dim: (chunks, 1, 1),
4055            block_dim: (256, 1, 1),
4056            shared_mem_bytes: 0,
4057        };
4058        let mut dv = dst.slice_mut(dst_off..dst_off + words);
4059        let __s = self.gpu.stream();
4060        let mut b = __s.launch_builder(&f);
4061        b.arg(src_entry).arg(&mut dv).arg(&wi);
4062        unsafe {
4063            b.launch(cfg)?;
4064        }
4065        Ok(())
4066    }
4067
4068    /// Uninitialized i8 device buffer (decode_batch q8_1 row scratch).
4069    pub fn uninit_i8(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
4070        self.alloc_uninit::<i8>(n)
4071    }
4072
4073    /// Resident-quantized linear (Stage-A: f32 dequant-in-kernel). y[m,out]=x[m,in]@W[out,in]^T.
4074    pub fn qmatvec(
4075        &self,
4076        w: &CudaSlice<u8>,
4077        x: &CudaSlice<f32>,
4078        m: usize,
4079        in_f: usize,
4080        out_f: usize,
4081        qtype: i32,
4082        row_bytes: usize,
4083    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4084        let f = self.func("qmatvec_f32");
4085        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
4086        let cfg = LaunchConfig {
4087            grid_dim: (out_f as u32, m as u32, 1),
4088            block_dim: (256, 1, 1),
4089            shared_mem_bytes: 0,
4090        };
4091        let (inf, outf, mi, qt, rb) =
4092            (in_f as i32, out_f as i32, m as i32, qtype, row_bytes as i64);
4093        let __s_b = self.gpu.stream();
4094        let mut b = __s_b.launch_builder(&f);
4095        b.arg(w)
4096            .arg(x)
4097            .arg(&mut y)
4098            .arg(&inf)
4099            .arg(&outf)
4100            .arg(&mi)
4101            .arg(&qt)
4102            .arg(&rb);
4103        unsafe {
4104            b.launch(cfg)?;
4105        }
4106        Ok(y)
4107    }
4108
4109    /// Allocate a reusable u8 GPU scratch buffer (for staged expert weights).
4110    pub fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
4111        let s = self.gpu.stream().alloc_zeros::<u8>(n)?;
4112        self.keep_if_capturing(&s);
4113        Ok(s)
4114    }
4115
4116    /// Uninitialized u8 scratch — skips alloc_zeros' memset. ONLY for staging buffers whose read
4117    /// range is fully overwritten by a stage_expert H2D before any kernel reads it (LAUNCH-STRUCTURE
4118    /// STAGE 2: the per-layer MoE scratch trio was 3 dead ~1MB memsets per layer per decode token).
4119    pub fn alloc_u8_uninit(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
4120        let s = unsafe { self.gpu.stream().alloc::<u8>(n)? };
4121        self.keep_if_capturing(&s);
4122        Ok(s)
4123    }
4124
4125    /// Zero a SUB-RANGE of an f32 buffer (CudaViewMut) — the row-sized memset the moe_out
4126    /// memset-elision uses for tokens that fall off the gdec fast path (LAUNCH-STRUCTURE STAGE 2).
4127    pub fn memset_zeros_view(
4128        &self,
4129        dst: &mut cudarc::driver::CudaViewMut<f32>,
4130    ) -> Result<(), Box<dyn std::error::Error>> {
4131        self.gpu.stream().memset_zeros(dst)?;
4132        Ok(())
4133    }
4134
4135    /// EDGE-1 staging: copy `host_bytes` (a sub-slice of a HostExps buffer) into `scratch`
4136    /// at byte offset `off` (async H2D on the default stream). Length is host_bytes.len().
4137    /// The qmatvec_view that reads `scratch[off..]` is enqueued on the SAME stream after this,
4138    /// so ordering is guaranteed without an explicit sync (Stage-1; Stage-2 prefetch on a 2nd
4139    /// stream would require an event).
4140    pub fn stage_expert(
4141        &self,
4142        host_bytes: &[u8],
4143        scratch: &mut CudaSlice<u8>,
4144        off: usize,
4145    ) -> Result<(), Box<dyn std::error::Error>> {
4146        let mut dst = scratch.slice_mut(off..off + host_bytes.len()); // CudaViewMut<u8>
4147        self.gpu.stream().memcpy_htod(host_bytes, &mut dst)?; // accepts &[u8] HostSlice src
4148        Ok(())
4149    }
4150
4151    /// EDGE-1 §A: fused MoE router. `logits` is the router output [t, n_expert] (device, f32, the
4152    /// `gate_inp @ z` result). Returns (sel_idx [t, n_used] i32, sel_w [t, n_used] f32): the top-k
4153    /// expert ids (DESC by prob, ascending-index tiebreak) and renormalized weights. Replaces the
4154    /// host dtoh + softmax-256 + stable DESC top-8 sort + renorm (hybrid_forward.rs ~281-298).
4155    /// One CTA per token row, 256 threads (one per expert).
4156    pub fn moe_router_topk(
4157        &self,
4158        logits: &CudaSlice<f32>,
4159        t: usize,
4160        n_expert: usize,
4161        n_used: usize,
4162    ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4163        let f = self.func("moe_router_topk_f32");
4164        let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?; // kernel fully overwrites
4165        let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?; // kernel fully overwrites
4166        let cfg = LaunchConfig {
4167            grid_dim: (t as u32, 1, 1),
4168            block_dim: (n_expert as u32, 1, 1),
4169            shared_mem_bytes: 0,
4170        };
4171        let (ne, nu) = (n_expert as i32, n_used as i32);
4172        let __s_b = self.gpu.stream();
4173        let mut b = __s_b.launch_builder(&f);
4174        b.arg(logits)
4175            .arg(&mut sel_idx)
4176            .arg(&mut sel_w)
4177            .arg(&ne)
4178            .arg(&nu);
4179        unsafe {
4180            b.launch(cfg)?;
4181        }
4182        Ok((sel_idx, sel_w))
4183    }
4184
4185    /// gemma4 twin: per-expert output scale folded into the topk renorm write (replaces the
4186    /// separate moe_w_exscale launch; value chain identical: (w/ws) * s[sel]).
4187    pub fn moe_router_topk_scaled(
4188        &self,
4189        logits: &CudaSlice<f32>,
4190        t: usize,
4191        n_expert: usize,
4192        n_used: usize,
4193        ex_scale: &CudaSlice<f32>,
4194    ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4195        // barrier-lean v2 twin (per-warp top-k + one-warp merge) FALSIFIED 2026-07-14:
4196        // bit-identical streams but −1.4% (26B plain N=3 interleaved) — at t=1 the grid is
4197        // ONE block, so the 6.6us is launch/dependency overhead, not the barrier chain;
4198        // fewer barriers bought nothing and the merge structure cost. jsonl is the record.
4199        let f = self.func("moe_router_topk_scaled_f32");
4200        let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?;
4201        let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?;
4202        let cfg = LaunchConfig {
4203            grid_dim: (t as u32, 1, 1),
4204            block_dim: (n_expert as u32, 1, 1),
4205            shared_mem_bytes: 0,
4206        };
4207        let (ne, nu) = (n_expert as i32, n_used as i32);
4208        let __s_b = self.gpu.stream();
4209        let mut b = __s_b.launch_builder(&f);
4210        b.arg(logits)
4211            .arg(&mut sel_idx)
4212            .arg(&mut sel_w)
4213            .arg(&ne)
4214            .arg(&nu)
4215            .arg(ex_scale);
4216        unsafe {
4217            b.launch(cfg)?;
4218        }
4219        Ok((sel_idx, sel_w))
4220    }
4221
4222    /// LAUNCH-STRUCTURE STAGE 1 (2026-07-05): fused router + SINGLE-SYNC host readback. The old
4223    /// MEMRA_FUSED_ROUTER path lost 2% at t=1 because it paid TWO full stream syncs (dtoh_i32 then
4224    /// dtoh, each = clone_dtoh + synchronize) + two alloc_zeros memsets per MoE layer, where the
4225    /// host route pays ONE sync on the 1KB logits dtoh. This variant: uninit outputs (kernel fully
4226    /// overwrites), both DtoH copies issued ASYNC into a persistent PINNED host staging buffer
4227    /// (flags=0 — cacheable, NOT cudarc's WRITECOMBINED default, so the host-side reads of sel/w
4228    /// stay cached), then ONE synchronize. Numerics identical to `moe_router_topk` (same kernel).
4229    pub fn moe_router_topk_host(
4230        &self,
4231        logits: &CudaSlice<f32>,
4232        t: usize,
4233        n_expert: usize,
4234        n_used: usize,
4235    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
4236        let f = self.func("moe_router_topk_f32");
4237        let n = t * n_used;
4238        let mut sel_idx = self.alloc_uninit::<i32>(n)?;
4239        let mut sel_w = self.alloc_uninit::<f32>(n)?;
4240        let cfg = LaunchConfig {
4241            grid_dim: (t as u32, 1, 1),
4242            block_dim: (n_expert as u32, 1, 1),
4243            shared_mem_bytes: 0,
4244        };
4245        let (ne, nu) = (n_expert as i32, n_used as i32);
4246        let __s_b = self.gpu.stream();
4247        let mut b = __s_b.launch_builder(&f);
4248        b.arg(logits)
4249            .arg(&mut sel_idx)
4250            .arg(&mut sel_w)
4251            .arg(&ne)
4252            .arg(&nu);
4253        unsafe {
4254            b.launch(cfg)?;
4255        }
4256        // single-sync readback: sel (i32) at offset 0, w (f32) at offset n*4 of the pinned stage.
4257        let bytes = n * 8;
4258        let mut guard = self.router_stage.lock().unwrap();
4259        if guard.as_ref().map(|p| p.cap < bytes).unwrap_or(true) {
4260            *guard = Some(PinnedStage::new(bytes.max(4096))?);
4261        }
4262        let stage = guard.as_mut().unwrap();
4263        let (si, sw) = unsafe {
4264            (
4265                std::slice::from_raw_parts_mut(stage.ptr as *mut i32, n),
4266                std::slice::from_raw_parts_mut(stage.ptr.add(n * 4) as *mut f32, n),
4267            )
4268        };
4269        self.gpu.stream().memcpy_dtoh(&sel_idx, si)?; // async (pinned dst)
4270        self.gpu.stream().memcpy_dtoh(&sel_w, sw)?; // async (pinned dst)
4271        self.gpu.stream().synchronize()?; // ONE sync for both
4272        Ok((si.iter().map(|&i| i as u32).collect(), sw.to_vec()))
4273    }
4274
4275    /// Device sigmoid router for Step-3.7 / DeepSeek-V3-class MoEs. `correction_bias` is added
4276    /// only to the top-k key; returned weights use the un-biased sigmoid score. `active` masks
4277    /// original expert ids before top-k. Exact key ties choose the smaller original id.
4278    #[allow(clippy::too_many_arguments)]
4279    pub fn moe_router_sigmoid_topk(
4280        &self,
4281        logits: &CudaSlice<f32>,
4282        t: usize,
4283        n_expert: usize,
4284        n_used: usize,
4285        active_count: usize,
4286        correction_bias: &CudaSlice<f32>,
4287        active: &CudaSlice<u8>,
4288        scaling_factor: f32,
4289        route_norm: bool,
4290    ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4291        crate::sigrouter_contract::validate_active_count(n_used, active_count)?;
4292        if n_expert == 0 || n_expert > 1024 || n_used == 0 || n_used > n_expert {
4293            return Err(format!(
4294                "sigmoid router shape unsupported: n_expert={n_expert}, n_used={n_used}",
4295            )
4296            .into());
4297        }
4298        if logits.len() < t * n_expert
4299            || correction_bias.len() != n_expert
4300            || active.len() != n_expert
4301        {
4302            return Err(format!(
4303                "sigmoid router buffer mismatch: logits={} bias={} active={} expected logits>={} row={}",
4304                logits.len(), correction_bias.len(), active.len(), t * n_expert, n_expert,
4305            ).into());
4306        }
4307        let f = self.func("moe_router_sigmoid_topk_f32");
4308        let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?;
4309        let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?;
4310        let threads = n_expert.div_ceil(32) * 32;
4311        let cfg = LaunchConfig {
4312            grid_dim: (t as u32, 1, 1),
4313            block_dim: (threads as u32, 1, 1),
4314            shared_mem_bytes: 0,
4315        };
4316        let (ne, nu, rn) = (n_expert as i32, n_used as i32, i32::from(route_norm));
4317        let __s_b = self.gpu.stream();
4318        let mut b = __s_b.launch_builder(&f);
4319        b.arg(logits)
4320            .arg(correction_bias)
4321            .arg(active)
4322            .arg(&mut sel_idx)
4323            .arg(&mut sel_w)
4324            .arg(&ne)
4325            .arg(&nu)
4326            .arg(&scaling_factor)
4327            .arg(&rn);
4328        unsafe {
4329            b.launch(cfg)?;
4330        }
4331        Ok((sel_idx, sel_w))
4332    }
4333
4334    /// Single-sync pinned readback twin of `moe_router_sigmoid_topk`. This preserves the existing
4335    /// grouped/staged dispatch contract while replacing the full-logit DtoH plus host sigmoid/sort.
4336    #[allow(clippy::too_many_arguments)]
4337    pub fn moe_router_sigmoid_topk_host(
4338        &self,
4339        logits: &CudaSlice<f32>,
4340        t: usize,
4341        n_expert: usize,
4342        n_used: usize,
4343        active_count: usize,
4344        correction_bias: &CudaSlice<f32>,
4345        active: &CudaSlice<u8>,
4346        scaling_factor: f32,
4347        route_norm: bool,
4348    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
4349        let (sel_idx, sel_w) = self.moe_router_sigmoid_topk(
4350            logits,
4351            t,
4352            n_expert,
4353            n_used,
4354            active_count,
4355            correction_bias,
4356            active,
4357            scaling_factor,
4358            route_norm,
4359        )?;
4360        let n = t * n_used;
4361        let bytes = n * 8;
4362        let mut guard = self.router_stage.lock().unwrap();
4363        if guard.as_ref().map(|p| p.cap < bytes).unwrap_or(true) {
4364            *guard = Some(PinnedStage::new(bytes.max(4096))?);
4365        }
4366        let stage = guard.as_mut().unwrap();
4367        let (si, sw) = unsafe {
4368            (
4369                std::slice::from_raw_parts_mut(stage.ptr as *mut i32, n),
4370                std::slice::from_raw_parts_mut(stage.ptr.add(n * 4) as *mut f32, n),
4371            )
4372        };
4373        self.gpu.stream().memcpy_dtoh(&sel_idx, si)?;
4374        self.gpu.stream().memcpy_dtoh(&sel_w, sw)?;
4375        self.gpu.stream().synchronize()?;
4376        Ok((si.iter().map(|&i| i as u32).collect(), sw.to_vec()))
4377    }
4378
4379    /// EDGE-1 §C.2: async H2D of `host_bytes` into `scratch[off..]` on the COPY stream, returning a
4380    /// recorded event the compute stream can `wait` on before the dependent GEMM. Used for in-token
4381    /// expert prefetch (pipeline by one). `host_bytes` should be pinned for a true DMA (§C.1).
4382    pub fn stage_expert_async(
4383        &self,
4384        host_bytes: &[u8],
4385        scratch: &mut CudaSlice<u8>,
4386        off: usize,
4387    ) -> Result<cudarc::driver::CudaEvent, Box<dyn std::error::Error>> {
4388        let mut dst = scratch.slice_mut(off..off + host_bytes.len());
4389        self.copy_stream.memcpy_htod(host_bytes, &mut dst)?;
4390        Ok(self.copy_stream.record_event(None)?)
4391    }
4392
4393    /// Make the compute stream wait for an async copy event (the consumer side of `stage_expert_async`).
4394    pub fn compute_wait(
4395        &self,
4396        ev: &cudarc::driver::CudaEvent,
4397    ) -> Result<(), Box<dyn std::error::Error>> {
4398        self.gpu.stream().wait(ev)?;
4399        Ok(())
4400    }
4401
4402    /// qmatvec over a byte sub-range of a (resident/scratch) CudaSlice<u8> holding ONE expert
4403    /// matrix. x is a CudaView<f32> (a sliced row of z, or a sliced activation). Reuses the
4404    /// validated qmatvec_f32 dequant path (NOT a fast path — the correctness gate). The
4405    /// CudaView base+offset pointer is honored by the launch arg.
4406    pub fn qmatvec_view(
4407        &self,
4408        w: &CudaSlice<u8>,
4409        range: std::ops::Range<usize>,
4410        x: &cudarc::driver::CudaView<f32>,
4411        m: usize,
4412        in_f: usize,
4413        out_f: usize,
4414        qtype: i32,
4415        row_bytes: usize,
4416    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4417        let f = self.func("qmatvec_f32");
4418        let wv = w.slice(range); // CudaView<u8>, offset honored
4419        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
4420        let cfg = LaunchConfig {
4421            grid_dim: (out_f as u32, m as u32, 1),
4422            block_dim: (256, 1, 1),
4423            shared_mem_bytes: 0,
4424        };
4425        let (inf, outf, mi, qt, rb) =
4426            (in_f as i32, out_f as i32, m as i32, qtype, row_bytes as i64);
4427        let __s_b = self.gpu.stream();
4428        let mut b = __s_b.launch_builder(&f);
4429        b.arg(&wv)
4430            .arg(x)
4431            .arg(&mut y)
4432            .arg(&inf)
4433            .arg(&outf)
4434            .arg(&mi)
4435            .arg(&qt)
4436            .arg(&rb);
4437        unsafe {
4438            b.launch(cfg)?;
4439        }
4440        Ok(y)
4441    }
4442
4443    /// STAGE-2 GROUPED DECODE (2026-07-04): one MoE layer's gate+up+SiLU for all `n_used` routed
4444    /// experts of ONE token in ONE launch (replaces 8x qmatvec(gate) + 8x qmatvec(up) + 8x
4445    /// silu_mul = 24 launches). `gp`/`up` are the 8 expert weight-block device pointers (SLRU
4446    /// cache slots — fixed-address, stable for the launch). Returns act [n_used, n_ff].
4447    /// BIT-IDENTICAL to the sequential chain: each dot reproduces qmatvec_f32's exact 256-thread
4448    /// reduction; the SiLU epilogue is silu_mul_f32's exact expression (see kernel header).
4449    #[allow(clippy::too_many_arguments)]
4450    /// dp4a q8 twins (MoE expert dp4a arc, 2026-07-06): same contract as the _f32 versions but
4451    /// consume a PRE-QUANTIZED q8_1 activation. FP-order differs from _f32 (int dot + warp tree)
4452    /// — the argmax/stream-identity battery arbitrates; MEMRA_MOE_Q8=0 restores f32.
4453    pub fn moe_gate_up_silu8_q8(
4454        &self,
4455        gp: WPtr8,
4456        up: WPtr8,
4457        aq: &CudaSlice<i8>,
4458        ad: &CudaSlice<f32>,
4459        in_f: usize,
4460        n_ff: usize,
4461        n_used: usize,
4462        qt_g: i32,
4463        qt_u: i32,
4464        rb_g: usize,
4465        rb_u: usize,
4466    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4467        let f = self.func("moe_gate_up_silu8_q8");
4468        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
4469        let cfg = LaunchConfig {
4470            grid_dim: (n_ff as u32, n_used as u32, 1),
4471            block_dim: (32, 1, 1),
4472            shared_mem_bytes: 0,
4473        };
4474        let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
4475        let __s_b = self.gpu.stream();
4476        let mut b = __s_b.launch_builder(&f);
4477        b.arg(&gp)
4478            .arg(&up)
4479            .arg(aq)
4480            .arg(ad)
4481            .arg(&mut act)
4482            .arg(&inf)
4483            .arg(&nff)
4484            .arg(&qt_g)
4485            .arg(&qt_u)
4486            .arg(&rbg)
4487            .arg(&rbu);
4488        unsafe {
4489            b.launch(cfg)?;
4490        }
4491        Ok(act)
4492    }
4493
4494    #[allow(clippy::too_many_arguments)]
4495    pub fn moe_down8_fma_q8(
4496        &self,
4497        dp: WPtr8,
4498        w: F32x8,
4499        aq2: &CudaSlice<i8>,
4500        ad2: &CudaSlice<f32>,
4501        dst: &mut cudarc::driver::CudaViewMut<f32>,
4502        in_f: usize,
4503        out_f: usize,
4504        n_used: usize,
4505        qt: i32,
4506        rb: usize,
4507    ) -> Result<(), Box<dyn std::error::Error>> {
4508        let f = self.func("moe_down8_fma_q8");
4509        let cfg = LaunchConfig {
4510            grid_dim: (out_f as u32, 1, 1),
4511            block_dim: (32, 1, 1),
4512            shared_mem_bytes: 0,
4513        };
4514        let (inf, outf, nu, rbi) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
4515        let __s_b = self.gpu.stream();
4516        let mut b = __s_b.launch_builder(&f);
4517        b.arg(&dp)
4518            .arg(&w)
4519            .arg(aq2)
4520            .arg(ad2)
4521            .arg(dst)
4522            .arg(&inf)
4523            .arg(&outf)
4524            .arg(&nu)
4525            .arg(&qt)
4526            .arg(&rbi);
4527        unsafe {
4528            b.launch(cfg)?;
4529        }
4530        Ok(())
4531    }
4532
4533    /// q8 sequential expert matvec (staged path twin of qmatvec_view for IQ3_S/IQ4_XS).
4534    pub fn qmatvec_expert_q8(
4535        &self,
4536        w: &CudaSlice<u8>,
4537        range: std::ops::Range<usize>,
4538        aq: &CudaSlice<i8>,
4539        ad: &CudaSlice<f32>,
4540        m: usize,
4541        in_f: usize,
4542        out_f: usize,
4543        qtype: i32,
4544        row_bytes: usize,
4545    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4546        let f = self.func("qmatvec_expert_q8");
4547        let wv = w.slice(range);
4548        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
4549        const ROWS: u32 = 4; // MEMRA_MMVQ_ROWS
4550        let cfg = LaunchConfig {
4551            grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, m as u32, 1),
4552            block_dim: (32, ROWS, 1),
4553            shared_mem_bytes: 0,
4554        };
4555        let (inf, outf, mi, rbi) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
4556        let __s_b = self.gpu.stream();
4557        let mut b = __s_b.launch_builder(&f);
4558        b.arg(&wv)
4559            .arg(aq)
4560            .arg(ad)
4561            .arg(&mut y)
4562            .arg(&inf)
4563            .arg(&outf)
4564            .arg(&mi)
4565            .arg(&qtype)
4566            .arg(&rbi);
4567        unsafe {
4568            b.launch(cfg)?;
4569        }
4570        Ok(y)
4571    }
4572
4573    pub fn moe_gate_up_silu8(
4574        &self,
4575        gp: WPtr8,
4576        up: WPtr8,
4577        x: &cudarc::driver::CudaView<f32>,
4578        in_f: usize,
4579        n_ff: usize,
4580        n_used: usize,
4581        qt_g: i32,
4582        qt_u: i32,
4583        rb_g: usize,
4584        rb_u: usize,
4585    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4586        let f = self.func("moe_gate_up_silu8_f32");
4587        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?; // fully overwritten
4588        let cfg = LaunchConfig {
4589            grid_dim: (n_ff as u32, n_used as u32, 1),
4590            block_dim: (256, 1, 1),
4591            shared_mem_bytes: 0,
4592        };
4593        let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
4594        let __s_b = self.gpu.stream();
4595        let mut b = __s_b.launch_builder(&f);
4596        b.arg(&gp)
4597            .arg(&up)
4598            .arg(x)
4599            .arg(&mut act)
4600            .arg(&inf)
4601            .arg(&nff)
4602            .arg(&qt_g)
4603            .arg(&qt_u)
4604            .arg(&rbg)
4605            .arg(&rbu);
4606        unsafe {
4607            b.launch(cfg)?;
4608        }
4609        Ok(act)
4610    }
4611
4612    /// STAGE-2 GROUPED DECODE: one MoE layer's down-proj + weighted accumulation for all `n_used`
4613    /// routed experts in ONE launch (replaces 8x qmatvec(down) + 8x axpy = 16 launches), writing
4614    /// the token's moe_out row DIRECTLY (`dst` is the zeroed row; the in-kernel slot-ordered
4615    /// __fmaf_rn chain starting at 0.0f reproduces the sequential axpy_f32 accumulation into the
4616    /// zeroed row bit-for-bit — the A2 byte-identity scheme at m=1).
4617    #[allow(clippy::too_many_arguments)]
4618    pub fn moe_down8_fma_into(
4619        &self,
4620        dp: WPtr8,
4621        w: F32x8,
4622        act: &CudaSlice<f32>,
4623        dst: &mut cudarc::driver::CudaViewMut<f32>,
4624        in_f: usize,
4625        out_f: usize,
4626        n_used: usize,
4627        qt: i32,
4628        rb: usize,
4629    ) -> Result<(), Box<dyn std::error::Error>> {
4630        let f = self.func("moe_down8_fma_f32");
4631        let cfg = LaunchConfig {
4632            grid_dim: (out_f as u32, 1, 1),
4633            block_dim: (256, 1, 1),
4634            shared_mem_bytes: 0,
4635        };
4636        let (inf, outf, nu, rbv) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
4637        let __s_b = self.gpu.stream();
4638        let mut b = __s_b.launch_builder(&f);
4639        b.arg(&dp)
4640            .arg(&w)
4641            .arg(act)
4642            .arg(dst)
4643            .arg(&inf)
4644            .arg(&outf)
4645            .arg(&nu)
4646            .arg(&qt)
4647            .arg(&rbv);
4648        unsafe {
4649            b.launch(cfg)?;
4650        }
4651        Ok(())
4652    }
4653
4654    /// LAUNCH-STRUCTURE STAGE 3: device-dispatch twin of `moe_gate_up_silu8` for FULLY-RESIDENT
4655    /// layers. The expert ids come from the router kernel's DEVICE `sel` output (no DtoH) and the
4656    /// weight pointers from the per-layer device table `[3, n_expert]` of slot base addresses.
4657    /// BIT-IDENTICAL math (same grid/block/reduction; only the pointer/id source differs).
4658    #[allow(clippy::too_many_arguments)]
4659    /// dp4a q8 twin of the _dev pair (resident-experts arc).
4660    ///
4661    /// GEOMETRY VARIANTS (multirow/occupancy arc 2026-07-05): all outputs are BIT-IDENTICAL to
4662    /// the base one-warp-per-(row,slot) kernel (same expert_dot_g g-order + warp tree per row;
4663    /// down's FMA chain stays slot-ordered serial). Seams:
4664    ///   MEMRA_MOE_DEVQ8_GU   = 0(base) | 1 | 2 | 4 -> _r{1,2,4} multirow twin (RPW rows/warp)
4665    ///                       | s2 (gate/up warp split) | s2z (s2 + WPB rows packed per block)
4666    ///                       | gs4 (gate/up x low/high-group 4-warp split, nsb==64 only)
4667    ///                       | u64 (nsb==64 unrolled ILP twin, geometry unchanged)
4668    ///   MEMRA_MOE_DEVQ8_WPB  = warps per block for _r twins / z-rows for s2z (default 4)
4669    ///   MEMRA_MOE_DEVQ8_DOWN = auto(default: w8h2 when in_f==512 & n_used<=8 — measured +3.8%
4670    ///                       decode on 35B/G7e) | 0 (base one-warp serial-slot) | 1 | 2 | 4 ->
4671    ///                       _w8r{1,2,4} slot-parallel twin | h2 (half-warp dual-row, nsb==16
4672    ///                       only) | w8h2 (h2 x slot-parallel)
4673    #[allow(clippy::too_many_arguments)]
4674    /// MoE PREFILL pair-batch matvec: one launch covers all (token,expert) pairs for one proj.
4675    #[allow(clippy::too_many_arguments)]
4676    pub fn moe_pairs_matvec_q8(
4677        &self,
4678        table: &CudaSlice<u64>,
4679        proj: i32,
4680        pair_tok: &CudaSlice<i32>,
4681        pair_ex: &CudaSlice<i32>,
4682        aq: &CudaSlice<i8>,
4683        ad: &CudaSlice<f32>,
4684        in_f: usize,
4685        out_f: usize,
4686        n_expert: usize,
4687        n_pairs: usize,
4688        qtype: i32,
4689        row_bytes: usize,
4690    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4691        let f = self.func("moe_pairs_matvec_q8");
4692        let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
4693        const ROWS: u32 = 4;
4694        let cfg = LaunchConfig {
4695            grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_pairs as u32, 1),
4696            block_dim: (32, ROWS, 1),
4697            shared_mem_bytes: 0,
4698        };
4699        let (inf, outf, ne, np, rbi) = (
4700            in_f as i32,
4701            out_f as i32,
4702            n_expert as i32,
4703            n_pairs as i32,
4704            row_bytes as i64,
4705        );
4706        let __s_b = self.gpu.stream();
4707        let mut b = __s_b.launch_builder(&f);
4708        b.arg(table)
4709            .arg(&proj)
4710            .arg(pair_tok)
4711            .arg(pair_ex)
4712            .arg(aq)
4713            .arg(ad)
4714            .arg(&mut y)
4715            .arg(&inf)
4716            .arg(&outf)
4717            .arg(&ne)
4718            .arg(&np)
4719            .arg(&qtype)
4720            .arg(&rbi);
4721        unsafe {
4722            b.launch(cfg)?;
4723        }
4724        Ok(y)
4725    }
4726
4727    /// Expert-major pair matvec (weight-reuse across each expert's token group).
4728    #[allow(clippy::too_many_arguments)]
4729    pub fn moe_pairs_matvec_q8_em(
4730        &self,
4731        table: &CudaSlice<u64>,
4732        proj: i32,
4733        ex_ids: &CudaSlice<i32>,
4734        ex_off: &CudaSlice<i32>,
4735        ex_pairs: &CudaSlice<i32>,
4736        pair_tok: &CudaSlice<i32>,
4737        aq: &CudaSlice<i8>,
4738        ad: &CudaSlice<f32>,
4739        in_f: usize,
4740        out_f: usize,
4741        n_expert: usize,
4742        n_active: usize,
4743        n_pairs: usize,
4744        qtype: i32,
4745        row_bytes: usize,
4746    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4747        let f = self.func("moe_pairs_matvec_q8_em");
4748        let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
4749        const ROWS: u32 = 4;
4750        let cfg = LaunchConfig {
4751            grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_active as u32, 1),
4752            block_dim: (32, ROWS, 1),
4753            shared_mem_bytes: 0,
4754        };
4755        let (inf, outf, ne, na, rbi) = (
4756            in_f as i32,
4757            out_f as i32,
4758            n_expert as i32,
4759            n_active as i32,
4760            row_bytes as i64,
4761        );
4762        let __s_b = self.gpu.stream();
4763        let mut b = __s_b.launch_builder(&f);
4764        b.arg(table)
4765            .arg(&proj)
4766            .arg(ex_ids)
4767            .arg(ex_off)
4768            .arg(ex_pairs)
4769            .arg(pair_tok)
4770            .arg(aq)
4771            .arg(ad)
4772            .arg(&mut y)
4773            .arg(&inf)
4774            .arg(&outf)
4775            .arg(&ne)
4776            .arg(&na)
4777            .arg(&qtype)
4778            .arg(&rbi);
4779        unsafe {
4780            b.launch(cfg)?;
4781        }
4782        Ok(y)
4783    }
4784
4785    // Decode-once expert-major MMQ (rung 3). Same CSR inputs/geometry as _em; kernel dequants each
4786    // weight group once per (row,group) then dp4a's across the expert's token group.
4787    #[allow(clippy::too_many_arguments)]
4788    pub fn moe_pairs_matvec_q8_dec(
4789        &self,
4790        table: &CudaSlice<u64>,
4791        proj: i32,
4792        ex_ids: &CudaSlice<i32>,
4793        ex_off: &CudaSlice<i32>,
4794        ex_pairs: &CudaSlice<i32>,
4795        pair_tok: &CudaSlice<i32>,
4796        aq: &CudaSlice<i8>,
4797        ad: &CudaSlice<f32>,
4798        in_f: usize,
4799        out_f: usize,
4800        n_expert: usize,
4801        n_active: usize,
4802        n_pairs: usize,
4803        qtype: i32,
4804        row_bytes: usize,
4805    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4806        let f = self.func("moe_pairs_matvec_q8_dec");
4807        let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
4808        const ROWS: u32 = 4;
4809        let cfg = LaunchConfig {
4810            grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_active as u32, 1),
4811            block_dim: (32, ROWS, 1),
4812            shared_mem_bytes: 0,
4813        };
4814        let (inf, outf, ne, na, rbi) = (
4815            in_f as i32,
4816            out_f as i32,
4817            n_expert as i32,
4818            n_active as i32,
4819            row_bytes as i64,
4820        );
4821        let __s_b = self.gpu.stream();
4822        let mut b = __s_b.launch_builder(&f);
4823        b.arg(table)
4824            .arg(&proj)
4825            .arg(ex_ids)
4826            .arg(ex_off)
4827            .arg(ex_pairs)
4828            .arg(pair_tok)
4829            .arg(aq)
4830            .arg(ad)
4831            .arg(&mut y)
4832            .arg(&inf)
4833            .arg(&outf)
4834            .arg(&ne)
4835            .arg(&na)
4836            .arg(&qtype)
4837            .arg(&rbi);
4838        unsafe {
4839            b.launch(cfg)?;
4840        }
4841        Ok(y)
4842    }
4843
4844    pub fn moe_pairs_gelu_mul(
4845        &self,
4846        gate: &CudaSlice<f32>,
4847        up: &CudaSlice<f32>,
4848        n: usize,
4849    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4850        let f = self.func("moe_pairs_gelu_mul");
4851        let mut act = self.alloc_uninit::<f32>(n)?;
4852        let cfg = LaunchConfig::for_num_elems(n as u32);
4853        let nl = n as i64;
4854        let __s_b = self.gpu.stream();
4855        let mut b = __s_b.launch_builder(&f);
4856        b.arg(gate).arg(up).arg(&mut act).arg(&nl);
4857        unsafe {
4858            b.launch(cfg)?;
4859        }
4860        Ok(act)
4861    }
4862
4863    pub fn moe_pairs_silu_mul(
4864        &self,
4865        gate: &CudaSlice<f32>,
4866        up: &CudaSlice<f32>,
4867        n: usize,
4868    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4869        let f = self.func("moe_pairs_silu_mul");
4870        let mut act = self.alloc_uninit::<f32>(n)?;
4871        let cfg = LaunchConfig::for_num_elems(n as u32);
4872        let nl = n as i64;
4873        let __s_b = self.gpu.stream();
4874        let mut b = __s_b.launch_builder(&f);
4875        b.arg(gate).arg(up).arg(&mut act).arg(&nl);
4876        unsafe {
4877            b.launch(cfg)?;
4878        }
4879        Ok(act)
4880    }
4881
4882    #[allow(clippy::too_many_arguments)]
4883    pub fn moe_pairs_scatter(
4884        &self,
4885        y_down: &CudaSlice<f32>,
4886        pair_w: &CudaSlice<f32>,
4887        tok_pair_off: &CudaSlice<i32>,
4888        tok_pair_ids: &CudaSlice<i32>,
4889        moe_out: &mut CudaSlice<f32>,
4890        t: usize,
4891        n_embd: usize,
4892    ) -> Result<(), Box<dyn std::error::Error>> {
4893        let f = self.func("moe_pairs_scatter");
4894        let cfg = LaunchConfig {
4895            grid_dim: (((n_embd + 255) / 256) as u32, t as u32, 1),
4896            block_dim: (256, 1, 1),
4897            shared_mem_bytes: 0,
4898        };
4899        let ne = n_embd as i32;
4900        let __s_b = self.gpu.stream();
4901        let mut b = __s_b.launch_builder(&f);
4902        b.arg(y_down)
4903            .arg(pair_w)
4904            .arg(tok_pair_off)
4905            .arg(tok_pair_ids)
4906            .arg(moe_out)
4907            .arg(&ne);
4908        unsafe {
4909            b.launch(cfg)?;
4910        }
4911        Ok(())
4912    }
4913
4914    /// gemma4 GELU twin of moe_gate_up_silu8_dev_q8 (base geometry — slot-packed j8/j8r2
4915    /// twins probed 2026-08-01 g26 decode dig: bit-identical rows, -2.5%/-2.9% whole-model
4916    /// decode x3 interleaved -> refuted and killed; research/g26-decode-20260801/receipts.md).
4917    #[allow(clippy::too_many_arguments)]
4918    pub fn moe_gate_up_gelu8_dev_q8(
4919        &self,
4920        table: &CudaSlice<u64>,
4921        sel: &cudarc::driver::CudaView<i32>,
4922        aq: &CudaSlice<i8>,
4923        ad: &CudaSlice<f32>,
4924        in_f: usize,
4925        n_ff: usize,
4926        n_used: usize,
4927        n_expert: usize,
4928        qt_g: i32,
4929        qt_u: i32,
4930        rb_g: usize,
4931        rb_u: usize,
4932    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4933        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
4934        let (inf, nff, ne, rbg, rbu) = (
4935            in_f as i32,
4936            n_ff as i32,
4937            n_expert as i32,
4938            rb_g as i64,
4939            rb_u as i64,
4940        );
4941        let f = self.func("moe_gate_up_gelu8_dev_q8");
4942        let cfg = LaunchConfig {
4943            grid_dim: (n_ff as u32, n_used as u32, 1),
4944            block_dim: (32, 1, 1),
4945            shared_mem_bytes: 0,
4946        };
4947        let __s_b = self.gpu.stream();
4948        let mut b = __s_b.launch_builder(&f);
4949        b.arg(table)
4950            .arg(sel)
4951            .arg(aq)
4952            .arg(ad)
4953            .arg(&mut act)
4954            .arg(&inf)
4955            .arg(&nff)
4956            .arg(&ne)
4957            .arg(&qt_g)
4958            .arg(&qt_u)
4959            .arg(&rbg)
4960            .arg(&rbu);
4961        unsafe {
4962            b.launch(cfg)?;
4963        }
4964        Ok(act)
4965    }
4966
4967    /// gemma4 GELU rows twin (verify): one launch over (n_ff, n_used, t).
4968    #[allow(clippy::too_many_arguments)]
4969    pub fn moe_gate_up_gelu8_dev_q8_rows(
4970        &self,
4971        table: &CudaSlice<u64>,
4972        sel: &CudaSlice<i32>,
4973        aq: &CudaSlice<i8>,
4974        ad: &CudaSlice<f32>,
4975        t: usize,
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>(t * n_used * n_ff)?;
4986        let (inf, nff, ne, rbg, rbu, nu) = (
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            n_used as i32,
4993        );
4994        let f = self.func("moe_gate_up_gelu8_dev_q8_rows");
4995        let cfg = LaunchConfig {
4996            grid_dim: (n_ff as u32, n_used as u32, t as u32),
4997            block_dim: (32, 1, 1),
4998            shared_mem_bytes: 0,
4999        };
5000        let __s_b = self.gpu.stream();
5001        let mut b = __s_b.launch_builder(&f);
5002        b.arg(table)
5003            .arg(sel)
5004            .arg(aq)
5005            .arg(ad)
5006            .arg(&mut act)
5007            .arg(&inf)
5008            .arg(&nff)
5009            .arg(&ne)
5010            .arg(&qt_g)
5011            .arg(&qt_u)
5012            .arg(&rbg)
5013            .arg(&rbu)
5014            .arg(&nu);
5015        unsafe {
5016            b.launch(cfg)?;
5017        }
5018        Ok(act)
5019    }
5020
5021    /// gemma4 GELU CSR twin (verify dedup: owner block serves every pair of its expert).
5022    #[allow(clippy::too_many_arguments)]
5023    pub fn moe_gate_up_gelu8_dev_q8_csr(
5024        &self,
5025        table: &CudaSlice<u64>,
5026        sel: &CudaSlice<i32>,
5027        aq: &CudaSlice<i8>,
5028        ad: &CudaSlice<f32>,
5029        n_pairs: usize,
5030        in_f: usize,
5031        n_ff: usize,
5032        n_used: usize,
5033        n_expert: usize,
5034        qt_g: i32,
5035        qt_u: i32,
5036        rb_g: usize,
5037        rb_u: usize,
5038    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5039        let mut act = self.alloc_uninit::<f32>(n_pairs * n_ff)?;
5040        let (inf, nff, ne, rbg, rbu, nu, npi) = (
5041            in_f as i32,
5042            n_ff as i32,
5043            n_expert as i32,
5044            rb_g as i64,
5045            rb_u as i64,
5046            n_used as i32,
5047            n_pairs as i32,
5048        );
5049        let f = self.func("moe_gate_up_gelu8_dev_q8_csr");
5050        let cfg = LaunchConfig {
5051            grid_dim: (n_ff as u32, n_pairs as u32, 1),
5052            block_dim: (32, 1, 1),
5053            shared_mem_bytes: 0,
5054        };
5055        let __s_b = self.gpu.stream();
5056        let mut b = __s_b.launch_builder(&f);
5057        b.arg(table)
5058            .arg(sel)
5059            .arg(aq)
5060            .arg(ad)
5061            .arg(&mut act)
5062            .arg(&inf)
5063            .arg(&nff)
5064            .arg(&ne)
5065            .arg(&qt_g)
5066            .arg(&qt_u)
5067            .arg(&rbg)
5068            .arg(&rbu)
5069            .arg(&nu)
5070            .arg(&npi);
5071        unsafe {
5072            b.launch(cfg)?;
5073        }
5074        Ok(act)
5075    }
5076
5077    /// gemma4 generic down rows twin (verify): one launch over (out_f, 1, t).
5078    #[allow(clippy::too_many_arguments)]
5079    pub fn moe_down8_fma_dev_q8_rows_g(
5080        &self,
5081        table: &CudaSlice<u64>,
5082        sel: &CudaSlice<i32>,
5083        w: &CudaSlice<f32>,
5084        aq2: &CudaSlice<i8>,
5085        ad2: &CudaSlice<f32>,
5086        dst: &mut CudaSlice<f32>,
5087        t: usize,
5088        in_f: usize,
5089        out_f: usize,
5090        n_used: usize,
5091        n_expert: usize,
5092        qt: i32,
5093        rb: usize,
5094    ) -> Result<(), Box<dyn std::error::Error>> {
5095        let (inf, outf, nu, ne, rbi) = (
5096            in_f as i32,
5097            out_f as i32,
5098            n_used as i32,
5099            n_expert as i32,
5100            rb as i64,
5101        );
5102        // Exact Step-3.7 B=1 shape: expose the eight independent slot dots as
5103        // eight warps, then replay the original slot-ordered FMA chain. Every
5104        // other shape retains the generic one-warp rows kernel.
5105        let step_b1_w8 = t == 1 && in_f == 1280 && out_f == 4096 && n_used == 8 && qt == QT_IQ4_XS;
5106        let f = self.func(if step_b1_w8 {
5107            "moe_down8_fma_dev_q8_rows_w8"
5108        } else {
5109            "moe_down8_fma_dev_q8_rows_g"
5110        });
5111        let cfg = LaunchConfig {
5112            grid_dim: (out_f as u32, 1, t as u32),
5113            block_dim: (32, if step_b1_w8 { 8 } else { 1 }, 1),
5114            shared_mem_bytes: 0,
5115        };
5116        let __s_b = self.gpu.stream();
5117        let mut b = __s_b.launch_builder(&f);
5118        b.arg(table)
5119            .arg(sel)
5120            .arg(w)
5121            .arg(aq2)
5122            .arg(ad2)
5123            .arg(dst)
5124            .arg(&inf)
5125            .arg(&outf)
5126            .arg(&nu)
5127            .arg(&ne)
5128            .arg(&qt)
5129            .arg(&rbi);
5130        unsafe {
5131            b.launch(cfg)?;
5132        }
5133        Ok(())
5134    }
5135
5136    /// rp_q4 microprobe (2026-07-10 verify-trunk lever): b4 GGUF-block layout vs the Q4_0
5137    /// split-plane twin on the wq-class shape. Returns (blk_us, rp_us) after asserting bitwise
5138    /// identity. Bench-only surface (rp_q4_probe bin); no production dispatch reads this.
5139    pub fn rp_probe_q4(&self, m: usize) -> Result<(f64, f64), Box<dyn std::error::Error>> {
5140        let (out_f, in_f) = (2048usize, 2816usize);
5141        let nblk = in_f / 32;
5142        let mut seed = 0x9E3779B97F4A7C15u64;
5143        let mut rng = move || {
5144            seed = seed
5145                .wrapping_mul(6364136223846793005)
5146                .wrapping_add(1442695040888963407);
5147            (seed >> 33) as u8
5148        };
5149        let mut w = vec![0u8; out_f * nblk * 18];
5150        for b in w.iter_mut() {
5151            *b = rng();
5152        }
5153        for r in 0..out_f {
5154            for g in 0..nblk {
5155                let off = (r * nblk + g) * 18;
5156                w[off] = 0x00;
5157                w[off + 1] = 0x2C; // sane half d
5158            }
5159        }
5160        let qplane = out_f * nblk * 16;
5161        let mut wrp = vec![0u8; w.len()];
5162        for r in 0..out_f {
5163            for g in 0..nblk {
5164                let src = &w[(r * nblk + g) * 18..(r * nblk + g) * 18 + 18];
5165                wrp[qplane + (r * nblk + g) * 2..qplane + (r * nblk + g) * 2 + 2]
5166                    .copy_from_slice(&src[0..2]);
5167                wrp[(r * nblk + g) * 16..(r * nblk + g) * 16 + 16].copy_from_slice(&src[2..18]);
5168            }
5169        }
5170        let w_d = self.htod_bytes(&w)?;
5171        let wrp_d = self.htod_bytes(&wrp)?;
5172        let mut aq = vec![0i8; m * in_f];
5173        for v in aq.iter_mut() {
5174            *v = rng() as i8;
5175        }
5176        let aq_d = self.htod_i8(&aq)?;
5177        let ad_d = self.htod(&vec![0.03125f32; m * nblk])?;
5178        let mut y0 = self.alloc_uninit::<f32>(m * out_f)?;
5179        let mut y1 = self.alloc_uninit::<f32>(m * out_f)?;
5180        const RPB: u32 = 4;
5181        let cfg = LaunchConfig {
5182            grid_dim: ((out_f as u32).div_ceil(RPB), 1, 1),
5183            block_dim: (32, RPB, 1),
5184            shared_mem_bytes: 0,
5185        };
5186        let (inf, outf, mi) = (in_f as i32, out_f as i32, m as i32);
5187        let (rb, qp) = ((nblk * 18) as i64, qplane as i64);
5188        let fb = self.func("qmatvec_q4_0_mmvq_b4");
5189        let fr = self.func("qmatvec_q4_0_mmvq_b4_rp");
5190        {
5191            let __s_b = self.gpu.stream();
5192            let mut b = __s_b.launch_builder(&fb);
5193            b.arg(&w_d)
5194                .arg(&aq_d)
5195                .arg(&ad_d)
5196                .arg(&mut y0)
5197                .arg(&inf)
5198                .arg(&outf)
5199                .arg(&mi)
5200                .arg(&rb);
5201            unsafe {
5202                b.launch(cfg)?;
5203            }
5204            let __s_b = self.gpu.stream();
5205            let mut b = __s_b.launch_builder(&fr);
5206            b.arg(&wrp_d)
5207                .arg(&aq_d)
5208                .arg(&ad_d)
5209                .arg(&mut y1)
5210                .arg(&inf)
5211                .arg(&outf)
5212                .arg(&mi)
5213                .arg(&qp);
5214            unsafe {
5215                b.launch(cfg)?;
5216            }
5217        }
5218        self.gpu.stream().synchronize()?;
5219        let (h0, h1) = (self.dtoh(&y0)?, self.dtoh(&y1)?);
5220        let nd = h0
5221            .iter()
5222            .zip(&h1)
5223            .filter(|(a, b)| a.to_bits() != b.to_bits())
5224            .count();
5225        if nd != 0 {
5226            return Err(format!("rp twin not bitwise: {nd}/{} diffs", h0.len()).into());
5227        }
5228        let mut time = |rp: bool| -> Result<f64, Box<dyn std::error::Error>> {
5229            self.gpu.stream().synchronize()?;
5230            let t0 = std::time::Instant::now();
5231            for _ in 0..500 {
5232                if rp {
5233                    let __s_b = self.gpu.stream();
5234                    let mut b = __s_b.launch_builder(&fr);
5235                    b.arg(&wrp_d)
5236                        .arg(&aq_d)
5237                        .arg(&ad_d)
5238                        .arg(&mut y1)
5239                        .arg(&inf)
5240                        .arg(&outf)
5241                        .arg(&mi)
5242                        .arg(&qp);
5243                    unsafe {
5244                        b.launch(cfg)?;
5245                    }
5246                } else {
5247                    let __s_b = self.gpu.stream();
5248                    let mut b = __s_b.launch_builder(&fb);
5249                    b.arg(&w_d)
5250                        .arg(&aq_d)
5251                        .arg(&ad_d)
5252                        .arg(&mut y0)
5253                        .arg(&inf)
5254                        .arg(&outf)
5255                        .arg(&mi)
5256                        .arg(&rb);
5257                    unsafe {
5258                        b.launch(cfg)?;
5259                    }
5260                }
5261            }
5262            self.gpu.stream().synchronize()?;
5263            Ok(t0.elapsed().as_secs_f64() * 1e6 / 500.0)
5264        };
5265        let _ = time(false)?;
5266        let _ = time(true)?; // warm
5267        Ok((time(false)?, time(true)?))
5268    }
5269
5270    /// Build the Q4_0 split-plane decode mirror for a 2D Quant tensor (device-side permutation,
5271    /// q4_0_split_rp_build). Raw bytes stay resident (prefill/gemm/Stage-A); the m<=8 decode
5272    /// dispatch prefers the mirror (_rp twins). No-op unless (Q4_0, 2D, mirror absent).
5273    /// VRAM cost == the tensor's weight size. MEMRA_Q4RP=0 disables at the call sites.
5274    pub fn build_q4_rp4(
5275        &self,
5276        t: &mut crate::model::GpuTensor,
5277    ) -> Result<(), Box<dyn std::error::Error>> {
5278        use crate::model::GpuTensor;
5279        let GpuTensor::Quant {
5280            bytes,
5281            qtype,
5282            row_bytes,
5283            ne,
5284            rp4,
5285            ..
5286        } = t
5287        else {
5288            return Ok(());
5289        };
5290        if *qtype != QT_Q4_0 || rp4.is_some() || ne.len() != 2 {
5291            return Ok(());
5292        }
5293        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5294        if in_f % 32 != 0 || *row_bytes != (in_f / 32) * 18 {
5295            return Ok(());
5296        }
5297        let nblk = in_f / 32;
5298        let mut dst = self.alloc_uninit::<u8>(out_f * nblk * 18)?;
5299        let f = self.func("q4_0_split_rp_build");
5300        let n = (out_f * nblk) as i32;
5301        let cfg = LaunchConfig {
5302            grid_dim: (((out_f * nblk) as u32).div_ceil(256), 1, 1),
5303            block_dim: (256, 1, 1),
5304            shared_mem_bytes: 0,
5305        };
5306        let (of, nb) = (out_f as i32, nblk as i32);
5307        let _ = n;
5308        let __s_b = self.gpu.stream();
5309        let mut b = __s_b.launch_builder(&f);
5310        b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
5311        unsafe {
5312            b.launch(cfg)?;
5313        }
5314        *rp4 = Some(dst);
5315        Ok(())
5316    }
5317
5318    /// Q8_0 twin of `build_q4_rp4` (H100 coalescing fix, 2026-07-26 ncu: GGUF 34B-stride
5319    /// weight loads hold Max Bandwidth at 41-46%; the split mirror makes them aligned 16B
5320    /// ldcs). Raw bytes stay resident (prefill GEMM/MMQ/fused m=1 launches read GGUF layout);
5321    /// the mmvq/batched decode arms prefer the mirror via `rp4`. Bit-identical outputs.
5322    pub fn build_q8_rp4(
5323        &self,
5324        t: &mut crate::model::GpuTensor,
5325    ) -> Result<(), Box<dyn std::error::Error>> {
5326        use crate::model::GpuTensor;
5327        let GpuTensor::Quant {
5328            bytes,
5329            qtype,
5330            row_bytes,
5331            ne,
5332            rp4,
5333            ..
5334        } = t
5335        else {
5336            return Ok(());
5337        };
5338        if *qtype != QT_Q8_0 || rp4.is_some() || ne.len() != 2 {
5339            return Ok(());
5340        }
5341        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5342        if in_f % 32 != 0 || *row_bytes != (in_f / 32) * 34 {
5343            return Ok(());
5344        }
5345        *rp4 = Some(self.build_q8_rp4_raw(bytes, in_f, out_f)?);
5346        Ok(())
5347    }
5348
5349    /// Raw rp-mirror build for gates/benches: split GGUF Q8_0 bytes into the qplane+dplane
5350    /// mirror without a GpuTensor (same kernel the loader path above uses).
5351    pub fn build_q8_rp4_raw(
5352        &self,
5353        bytes: &CudaSlice<u8>,
5354        in_f: usize,
5355        out_f: usize,
5356    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
5357        assert!(in_f % 32 == 0);
5358        let nblk = in_f / 32;
5359        let mut dst = self.alloc_uninit::<u8>(out_f * nblk * 34)?;
5360        let f = self.func("q8_0_split_rp_build");
5361        let cfg = LaunchConfig {
5362            grid_dim: (((out_f * nblk) as u32).div_ceil(256), 1, 1),
5363            block_dim: (256, 1, 1),
5364            shared_mem_bytes: 0,
5365        };
5366        let (of, nb) = (out_f as i32, nblk as i32);
5367        let __s_b = self.gpu.stream();
5368        let mut b = __s_b.launch_builder(&f);
5369        b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
5370        unsafe {
5371            b.launch(cfg)?;
5372        }
5373        Ok(dst)
5374    }
5375
5376    /// K-quant twins of `build_q8_rp4` (H100 K-quant coalescing fix, 2026-08-01 ncu on the
5377    /// q27 Q4_K_M decode: q4_K mmvq DRAM 41-54% with 65% excessive sectors, q6_K 40% with
5378    /// 78% — the 144B/210B superblock strides land every 4B weight load off-sector). The
5379    /// mirror re-packs each tensor into planes (q4_K: qs ++ 16B meta; q6_K: ql ++ qh ++
5380    /// scales ++ d — same total bytes) so every quant fetch is an aligned 16B ldcs. Raw
5381    /// bytes stay resident (prefill GEMM/dequant/Stage-A read GGUF layout); the mmvq/batched
5382    /// decode arms prefer the mirror via `rp4`. Bit-identical outputs.
5383    pub fn build_q4k_rp4(
5384        &self,
5385        t: &mut crate::model::GpuTensor,
5386    ) -> Result<(), Box<dyn std::error::Error>> {
5387        use crate::model::GpuTensor;
5388        let GpuTensor::Quant {
5389            bytes,
5390            qtype,
5391            row_bytes,
5392            ne,
5393            rp4,
5394            ..
5395        } = t
5396        else {
5397            return Ok(());
5398        };
5399        if *qtype != QT_Q4_K || rp4.is_some() || ne.len() != 2 {
5400            return Ok(());
5401        }
5402        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5403        if in_f % 256 != 0 || *row_bytes != (in_f / 256) * 144 {
5404            return Ok(());
5405        }
5406        *rp4 = Some(self.build_kq_rp4_raw(bytes, in_f, out_f, QT_Q4_K)?);
5407        Ok(())
5408    }
5409
5410    pub fn build_q6k_rp4(
5411        &self,
5412        t: &mut crate::model::GpuTensor,
5413    ) -> Result<(), Box<dyn std::error::Error>> {
5414        use crate::model::GpuTensor;
5415        let GpuTensor::Quant {
5416            bytes,
5417            qtype,
5418            row_bytes,
5419            ne,
5420            rp4,
5421            ..
5422        } = t
5423        else {
5424            return Ok(());
5425        };
5426        if *qtype != QT_Q6_K || rp4.is_some() || ne.len() != 2 {
5427            return Ok(());
5428        }
5429        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5430        if in_f % 256 != 0 || *row_bytes != (in_f / 256) * 210 {
5431            return Ok(());
5432        }
5433        *rp4 = Some(self.build_kq_rp4_raw(bytes, in_f, out_f, QT_Q6_K)?);
5434        Ok(())
5435    }
5436
5437    /// Raw K-quant rp-mirror build for gates/benches (same kernels the loader path uses).
5438    pub fn build_kq_rp4_raw(
5439        &self,
5440        bytes: &CudaSlice<u8>,
5441        in_f: usize,
5442        out_f: usize,
5443        qtype: i32,
5444    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
5445        assert!(in_f % 256 == 0);
5446        let nsbk = in_f / 256;
5447        let (sb_bytes, kname) = match qtype {
5448            QT_Q4_K => (144usize, "q4_K_split_rp_build"),
5449            QT_Q6_K => (210usize, "q6_K_split_rp_build"),
5450            _ => return Err(format!("build_kq_rp4_raw: qtype {qtype} has no rp mirror").into()),
5451        };
5452        let mut dst = self.alloc_uninit::<u8>(out_f * nsbk * sb_bytes)?;
5453        let f = self.func(kname);
5454        let cfg = LaunchConfig {
5455            grid_dim: (((out_f * nsbk) as u32).div_ceil(256), 1, 1),
5456            block_dim: (256, 1, 1),
5457            shared_mem_bytes: 0,
5458        };
5459        let (of, nb) = (out_f as i32, nsbk as i32);
5460        let __s_b = self.gpu.stream();
5461        let mut b = __s_b.launch_builder(&f);
5462        b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
5463        unsafe {
5464            b.launch(cfg)?;
5465        }
5466        Ok(dst)
5467    }
5468
5469    /// MEMRA_KQRP seam: the K-quant (q4_K/q6_K) split-plane decode mirrors at model load.
5470    /// Default follows the Q8RP convention — ON on the Hopper lane (80GB pays the mirror
5471    /// VRAM), OFF elsewhere (a 24GB card cannot hold model + mirror + KV for the big trunks).
5472    pub fn kqrp_enabled() -> bool {
5473        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5474        *ON.get_or_init(|| match std::env::var("MEMRA_KQRP").as_deref() {
5475            Ok("0") => false,
5476            Ok(_) => true,
5477            Err(_) => cfg!(memra_hopper_mma),
5478        })
5479    }
5480
5481    /// IN-PLACE split-plane swap (the 31B dense arc): build the split layout and REPLACE the
5482    /// GGUF bytes (zero extra steady-state VRAM — the transient peak is one tensor's size).
5483    /// The tensor's `rp` flag then routes every consumer (mmvq/batched `_rp` twins, the
5484    /// `qmatvec_gemm_q4_0_rp` prefill kernel). Callers gate on the fast path being active —
5485    /// the Stage-A f32 oracle (`MEMRA_FAST=0`) reads GGUF layout and must never see a swap.
5486    pub fn build_q4_rp_swap(
5487        &self,
5488        t: &mut crate::model::GpuTensor,
5489    ) -> Result<bool, Box<dyn std::error::Error>> {
5490        use crate::model::GpuTensor;
5491        // QTYPE GUARD IN THE SWAP ITSELF (gemma4 NVFP4mix prefill-NaN, 2026-08-17):
5492        // `rp4` is a SHARED field — the Q8RP walk parks Q8_0 MIRRORS there, and this
5493        // fn used to `take()` whatever it found. On a Q8_0-carrying gemma4-dense trunk
5494        // the swap hijacked those mirrors: `bytes` became split-plane in place, the
5495        // m<=16 `_rp` dispatch read them correctly (masking the corruption from every
5496        // decode pin), and every GGUF-layout prefill consumer (MMQ + GEMM) read the
5497        // fp16 d-plane as weights -> layer-0 NaN, <pad>-spam serving. Only a tensor
5498        // this fn's OWN builder serves may ever be swapped; everything else refuses
5499        // here, regardless of walk ordering.
5500        if !matches!(t, GpuTensor::Quant { qtype, .. } if *qtype == QT_Q4_0) {
5501            return Ok(false);
5502        }
5503        self.build_q4_rp4(t)?;
5504        self.gpu.stream().synchronize()?; // build kernel reads the GGUF bytes — drain BEFORE dropping them
5505        let GpuTensor::Quant { bytes, rp4, rp, .. } = t else {
5506            return Ok(false);
5507        };
5508        match rp4.take() {
5509            Some(split) => {
5510                *bytes = split; // the GGUF-layout buffer drops here
5511                *rp = true;
5512                Ok(true)
5513            }
5514            None => Ok(false),
5515        }
5516    }
5517
5518    /// MEMRA_Q4RP seam (default ON): the Q4_0 split-plane decode mirror at model load.
5519    pub fn q4rp_enabled() -> bool {
5520        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5521        *ON.get_or_init(|| {
5522            std::env::var("MEMRA_Q4RP")
5523                .map(|v| v != "0")
5524                .unwrap_or(true)
5525        })
5526    }
5527
5528    /// gemma4-E4B: dense [t][row_elems] gather of layer il's rows from the strided prologue
5529    /// buffer ([t][n_layer][n_epl]; off = il*n_epl, stride = n_layer*n_epl).
5530    pub fn copy_rows_strided(
5531        &self,
5532        src: &CudaSlice<f32>,
5533        dst: &mut CudaSlice<f32>,
5534        row_elems: usize,
5535        n_rows: usize,
5536        src_stride: usize,
5537        src_off: usize,
5538    ) -> Result<(), Box<dyn std::error::Error>> {
5539        let f = self.func("copy_rows_strided_f32");
5540        let cfg = LaunchConfig {
5541            grid_dim: (((row_elems as u32 + 255) / 256).max(1), n_rows as u32, 1),
5542            block_dim: (256, 1, 1),
5543            shared_mem_bytes: 0,
5544        };
5545        let (re, nr) = (row_elems as i32, n_rows as i32);
5546        let (st, off) = (src_stride as i64, src_off as i64);
5547        let __s_b = self.gpu.stream();
5548        let mut b = __s_b.launch_builder(&f);
5549        b.arg(src)
5550            .arg(&mut *dst)
5551            .arg(&re)
5552            .arg(&nr)
5553            .arg(&st)
5554            .arg(&off);
5555        unsafe {
5556            b.launch(cfg)?;
5557        }
5558        Ok(())
5559    }
5560
5561    /// Async device u32 store (value rides the kernel ARG — no host-memory transfer/sync).
5562    pub fn u32_set_k(
5563        &self,
5564        dst: &mut CudaSlice<u32>,
5565        v: u32,
5566        idx: usize,
5567    ) -> Result<(), Box<dyn std::error::Error>> {
5568        let f = self.func("u32_set_k");
5569        let cfg = LaunchConfig {
5570            grid_dim: (1, 1, 1),
5571            block_dim: (1, 1, 1),
5572            shared_mem_bytes: 0,
5573        };
5574        let ii = idx as i32;
5575        let __s_b = self.gpu.stream();
5576        let mut b = __s_b.launch_builder(&f);
5577        b.arg(dst).arg(&v).arg(&ii);
5578        unsafe {
5579            b.launch(cfg)?;
5580        }
5581        Ok(())
5582    }
5583
5584    /// counter += v (device-slot append advance; the +1 twin is `inc_seqlen`).
5585    pub fn i32_add_k(
5586        &self,
5587        d: &mut CudaSlice<i32>,
5588        v: i32,
5589    ) -> Result<(), Box<dyn std::error::Error>> {
5590        let f = self.func("i32_add_k");
5591        let cfg = LaunchConfig {
5592            grid_dim: (1, 1, 1),
5593            block_dim: (32, 1, 1),
5594            shared_mem_bytes: 0,
5595        };
5596        let __s_b = self.gpu.stream();
5597        let mut b = __s_b.launch_builder(&f);
5598        b.arg(d).arg(&v);
5599        unsafe {
5600            b.launch(cfg)?;
5601        }
5602        Ok(())
5603    }
5604
5605    /// pos rows from a device counter: dst[i] = ctr[0] + i (verify-stream rope positions).
5606    pub fn i32_iota_from(
5607        &self,
5608        ctr: &CudaSlice<i32>,
5609        dst: &mut CudaSlice<i32>,
5610        n: usize,
5611    ) -> Result<(), Box<dyn std::error::Error>> {
5612        let f = self.func("i32_iota_from");
5613        let cfg = LaunchConfig::for_num_elems(n as u32);
5614        let ni = n as i32;
5615        let __s_b = self.gpu.stream();
5616        let mut b = __s_b.launch_builder(&f);
5617        b.arg(ctr).arg(dst).arg(&ni);
5618        unsafe {
5619            b.launch(cfg)?;
5620        }
5621        Ok(())
5622    }
5623
5624    /// In-place trim-id translate: buf[idx] = map[buf[idx]] (FR-Spec d2t, async single-slot).
5625    pub fn u32_map_k(
5626        &self,
5627        buf: &mut CudaSlice<u32>,
5628        map: &CudaSlice<u32>,
5629        idx: usize,
5630    ) -> Result<(), Box<dyn std::error::Error>> {
5631        let f = self.func("u32_map_k");
5632        let cfg = LaunchConfig {
5633            grid_dim: (1, 1, 1),
5634            block_dim: (1, 1, 1),
5635            shared_mem_bytes: 0,
5636        };
5637        let ii = idx as i32;
5638        let __s_b = self.gpu.stream();
5639        let mut b = __s_b.launch_builder(&f);
5640        b.arg(buf).arg(map).arg(&ii);
5641        unsafe {
5642            b.launch(cfg)?;
5643        }
5644        Ok(())
5645    }
5646
5647    /// Pack a[off..off+n1] ++ b[0..n2] into one buffer (single dtoh follows).
5648    #[allow(clippy::too_many_arguments)]
5649    pub fn u32_pack2(
5650        &self,
5651        a: &CudaSlice<u32>,
5652        off_a: usize,
5653        n1: usize,
5654        b_in: &CudaSlice<u32>,
5655        n2: usize,
5656        out: &mut CudaSlice<u32>,
5657    ) -> Result<(), Box<dyn std::error::Error>> {
5658        let f = self.func("u32_pack2");
5659        let cfg = LaunchConfig::for_num_elems((n1 + n2) as u32);
5660        let (oa, i1, i2) = (off_a as i32, n1 as i32, n2 as i32);
5661        let __s_b = self.gpu.stream();
5662        let mut b = __s_b.launch_builder(&f);
5663        b.arg(a).arg(&oa).arg(&i1).arg(b_in).arg(&i2).arg(out);
5664        unsafe {
5665            b.launch(cfg)?;
5666        }
5667        Ok(())
5668    }
5669
5670    /// gemma4 R3 device fold: w[i] *= s[sel[i]] over the router's [n] (sel, w) pair.
5671    pub fn moe_w_exscale(
5672        &self,
5673        w: &mut CudaSlice<f32>,
5674        sel: &CudaSlice<i32>,
5675        s: &CudaSlice<f32>,
5676        n: usize,
5677    ) -> Result<(), Box<dyn std::error::Error>> {
5678        let f = self.func("moe_w_exscale");
5679        let cfg = LaunchConfig::for_num_elems(n as u32);
5680        let ni = n as i32;
5681        let __s_b = self.gpu.stream();
5682        let mut b = __s_b.launch_builder(&f);
5683        b.arg(w).arg(sel).arg(s).arg(&ni);
5684        unsafe {
5685            b.launch(cfg)?;
5686        }
5687        Ok(())
5688    }
5689
5690    /// Down-projection macro fold: w[i] *= macros[2*n_expert + sel[i]] on the device router
5691    /// weights (one launch per MoE layer, only for macro-carrying artifacts — see MoeWeights).
5692    pub fn moe_w_scale_by_expert(
5693        &self,
5694        w: &mut CudaSlice<f32>,
5695        sel: &CudaSlice<i32>,
5696        macros: &CudaSlice<f32>,
5697        n_expert: usize,
5698        n: usize,
5699    ) -> Result<(), Box<dyn std::error::Error>> {
5700        let f = self.func("moe_w_scale_by_expert");
5701        let cfg = LaunchConfig {
5702            grid_dim: (n.div_ceil(64) as u32, 1, 1),
5703            block_dim: (64, 1, 1),
5704            shared_mem_bytes: 0,
5705        };
5706        let (ne, nn) = (n_expert as i32, n as i32);
5707        let __s_b = self.gpu.stream();
5708        let mut b = __s_b.launch_builder(&f);
5709        b.arg(w).arg(sel).arg(macros).arg(&ne).arg(&nn);
5710        unsafe {
5711            b.launch(cfg)?;
5712        }
5713        Ok(())
5714    }
5715
5716    pub fn moe_gate_up_silu8_dev_q8(
5717        &self,
5718        table: &CudaSlice<u64>,
5719        sel: &cudarc::driver::CudaView<i32>,
5720        aq: &CudaSlice<i8>,
5721        ad: &CudaSlice<f32>,
5722        in_f: usize,
5723        n_ff: usize,
5724        n_used: usize,
5725        n_expert: usize,
5726        qt_g: i32,
5727        qt_u: i32,
5728        rb_g: usize,
5729        rb_u: usize,
5730        macros: &CudaSlice<f32>,
5731    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5732        static GU: std::sync::OnceLock<(String, u32)> = std::sync::OnceLock::new();
5733        let (mode, wpb) = GU.get_or_init(|| {
5734            let mode = std::env::var("MEMRA_MOE_DEVQ8_GU").unwrap_or_default();
5735            let wpb = std::env::var("MEMRA_MOE_DEVQ8_WPB")
5736                .ok()
5737                .and_then(|v| v.parse().ok())
5738                .unwrap_or(4u32)
5739                .clamp(1, 16);
5740            (mode, wpb)
5741        });
5742        let (mode, wpb) = (mode.as_str(), *wpb);
5743        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
5744        let (inf, nff, ne, rbg, rbu) = (
5745            in_f as i32,
5746            n_ff as i32,
5747            n_expert as i32,
5748            rb_g as i64,
5749            rb_u as i64,
5750        );
5751        let (f, cfg) = match mode {
5752            "1" | "2" | "4" => {
5753                let rpw: u32 = mode.parse().unwrap();
5754                let f = self.func(match rpw {
5755                    1 => "moe_gate_up_silu8_dev_q8_r1",
5756                    2 => "moe_gate_up_silu8_dev_q8_r2",
5757                    _ => "moe_gate_up_silu8_dev_q8_r4",
5758                });
5759                let rows_per_block = (rpw * wpb) as usize;
5760                let gx = n_ff.div_ceil(rows_per_block) as u32;
5761                (
5762                    f,
5763                    LaunchConfig {
5764                        grid_dim: (gx, n_used as u32, 1),
5765                        block_dim: (32, wpb, 1),
5766                        shared_mem_bytes: 0,
5767                    },
5768                )
5769            }
5770            "j8" if n_used <= 32 => (
5771                self.func("moe_gate_up_silu8_dev_q8_j8"),
5772                LaunchConfig {
5773                    grid_dim: (n_ff as u32, 1, 1),
5774                    block_dim: (32, n_used as u32, 1),
5775                    shared_mem_bytes: 0,
5776                },
5777            ),
5778            // SMEM-GRID twins (IQ3_S 2KB grid copied to shared, static smem — bit-identical dots)
5779            "vsm2" => {
5780                let f = self.func("moe_gate_up_silu8_dev_q8_vsm2");
5781                let sh = (rb_g + rb_u) as u32;
5782                use cudarc::driver::sys::CUfunction_attribute_enum as A;
5783                f.set_attribute(
5784                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
5785                    sh as i32,
5786                )?;
5787                (
5788                    f,
5789                    LaunchConfig {
5790                        grid_dim: (n_ff as u32, n_used as u32, 1),
5791                        block_dim: (32, 1, 1),
5792                        shared_mem_bytes: sh,
5793                    },
5794                )
5795            }
5796            "vsm" => {
5797                let f = self.func("moe_gate_up_silu8_dev_q8_vsm");
5798                let sh = (rb_g + rb_u) as u32;
5799                use cudarc::driver::sys::CUfunction_attribute_enum as A;
5800                f.set_attribute(
5801                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
5802                    sh as i32,
5803                )?;
5804                (
5805                    f,
5806                    LaunchConfig {
5807                        grid_dim: (n_ff as u32, n_used as u32, 1),
5808                        block_dim: (32, 1, 1),
5809                        shared_mem_bytes: sh,
5810                    },
5811                )
5812            }
5813            "sg" => (
5814                self.func("moe_gate_up_silu8_dev_q8_sg"),
5815                LaunchConfig {
5816                    grid_dim: (n_ff as u32, n_used as u32, 1),
5817                    block_dim: (32, 1, 1),
5818                    shared_mem_bytes: 0,
5819                },
5820            ),
5821            "j8sg" if n_used <= 32 => (
5822                self.func("moe_gate_up_silu8_dev_q8_j8sg"),
5823                LaunchConfig {
5824                    grid_dim: (n_ff as u32, 1, 1),
5825                    block_dim: (32, n_used as u32, 1),
5826                    shared_mem_bytes: 0,
5827                },
5828            ),
5829            "u64" if in_f == 2048 => (
5830                self.func("moe_gate_up_silu8_dev_q8_u64"),
5831                LaunchConfig {
5832                    grid_dim: (n_ff as u32, n_used as u32, 1),
5833                    block_dim: (32, 1, 1),
5834                    shared_mem_bytes: 0,
5835                },
5836            ),
5837            "gs4" if in_f == 2048 => (
5838                self.func("moe_gate_up_silu8_dev_q8_gs4"),
5839                LaunchConfig {
5840                    grid_dim: (n_ff as u32, n_used as u32, 1),
5841                    block_dim: (32, 4, 1),
5842                    shared_mem_bytes: 0,
5843                },
5844            ),
5845            // _v twin (down8 lane 2026-07-08): wide-load IQ4_XS dot, base geometry, bit-identical.
5846            "v" | "" => (
5847                self.func("moe_gate_up_silu8_dev_q8_v"),
5848                LaunchConfig {
5849                    grid_dim: (n_ff as u32, n_used as u32, 1),
5850                    block_dim: (32, 1, 1),
5851                    shared_mem_bytes: 0,
5852                },
5853            ),
5854            "s2" => (
5855                self.func("moe_gate_up_silu8_dev_q8_s2"),
5856                LaunchConfig {
5857                    grid_dim: (n_ff as u32, n_used as u32, 1),
5858                    block_dim: (32, 2, 1),
5859                    shared_mem_bytes: 0,
5860                },
5861            ),
5862            "s2z" => {
5863                let rz = wpb.min(16); // s2z smem tile is [16][2]
5864                (
5865                    self.func("moe_gate_up_silu8_dev_q8_s2z"),
5866                    LaunchConfig {
5867                        grid_dim: (n_ff.div_ceil(rz as usize) as u32, n_used as u32, 1),
5868                        block_dim: (32, 2, rz),
5869                        shared_mem_bytes: 0,
5870                    },
5871                )
5872            }
5873            _ => (
5874                self.func("moe_gate_up_silu8_dev_q8"),
5875                LaunchConfig {
5876                    grid_dim: (n_ff as u32, n_used as u32, 1),
5877                    block_dim: (32, 1, 1),
5878                    shared_mem_bytes: 0,
5879                },
5880            ),
5881        };
5882        let __s_b = self.gpu.stream();
5883        let mut b = __s_b.launch_builder(&f);
5884        b.arg(table)
5885            .arg(sel)
5886            .arg(aq)
5887            .arg(ad)
5888            .arg(&mut act)
5889            .arg(&inf)
5890            .arg(&nff)
5891            .arg(&ne)
5892            .arg(&qt_g)
5893            .arg(&qt_u)
5894            .arg(&rbg)
5895            .arg(&rbu)
5896            .arg(macros);
5897        unsafe {
5898            b.launch(cfg)?;
5899        }
5900        Ok(act)
5901    }
5902
5903    #[allow(clippy::too_many_arguments)]
5904    pub fn moe_down8_fma_dev_q8(
5905        &self,
5906        table: &CudaSlice<u64>,
5907        sel: &cudarc::driver::CudaView<i32>,
5908        w: &cudarc::driver::CudaView<f32>,
5909        aq2: &CudaSlice<i8>,
5910        ad2: &CudaSlice<f32>,
5911        dst: &mut cudarc::driver::CudaViewMut<f32>,
5912        in_f: usize,
5913        out_f: usize,
5914        n_used: usize,
5915        n_expert: usize,
5916        qt: i32,
5917        rb: usize,
5918    ) -> Result<(), Box<dyn std::error::Error>> {
5919        static DOWN: std::sync::OnceLock<String> = std::sync::OnceLock::new();
5920        let mode = DOWN.get_or_init(|| std::env::var("MEMRA_MOE_DEVQ8_DOWN").unwrap_or_default());
5921        let (inf, outf, nu, ne, rbi) = (
5922            in_f as i32,
5923            out_f as i32,
5924            n_used as i32,
5925            n_expert as i32,
5926            rb as i64,
5927        );
5928        // the w8 twins' smem tile is [RPW][8] — n_used must fit the 8-slot tile;
5929        // the h2 twins are nsb==16 (in_f==512) shape-gated.
5930        let (f, cfg) = match mode.as_str() {
5931            m @ ("1" | "2" | "4") if n_used <= 8 => {
5932                let rpw: usize = m.parse().unwrap();
5933                let f = self.func(match rpw {
5934                    1 => "moe_down8_fma_dev_q8_w8r1",
5935                    2 => "moe_down8_fma_dev_q8_w8r2",
5936                    _ => "moe_down8_fma_dev_q8_w8r4",
5937                });
5938                (
5939                    f,
5940                    LaunchConfig {
5941                        grid_dim: (out_f.div_ceil(rpw) as u32, 1, 1),
5942                        block_dim: (32, n_used as u32, 1),
5943                        shared_mem_bytes: 0,
5944                    },
5945                )
5946            }
5947            "h2" if in_f == 512 => (
5948                self.func("moe_down8_fma_dev_q8_h2"),
5949                LaunchConfig {
5950                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
5951                    block_dim: (32, 1, 1),
5952                    shared_mem_bytes: 0,
5953                },
5954            ),
5955            // "" = AUTO gemma shape (in_f==704): w8r2 measured +1 tok/s vs base (sweep
5956            // 1/2/4 -> 133.6/134.2/133.6, 2026-07-10); slot-ordered chain preserved.
5957            "" if in_f == 704 && n_used <= 8 => (
5958                self.func("moe_down8_fma_dev_q8_w8r2"),
5959                LaunchConfig {
5960                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
5961                    block_dim: (32, n_used as u32, 1),
5962                    shared_mem_bytes: 0,
5963                },
5964            ),
5965            // "" = AUTO: the measured winner for the 35B expert shape (arc 2026-07-05, +3.8%);
5966            // any shape the h2 kernels can't take (nsb!=16 / n_used>8) falls to base via `_`.
5967            // _v twins (down8 lane 2026-07-08): wide-load IQ4_XS dot, bit-identical outputs.
5968            "w8h2v" | "" if in_f == 512 && n_used <= 8 => (
5969                self.func("moe_down8_fma_dev_q8_w8h2v"),
5970                LaunchConfig {
5971                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
5972                    block_dim: (32, n_used as u32, 1),
5973                    shared_mem_bytes: 0,
5974                },
5975            ),
5976            "w8h2r2v" if in_f == 512 && n_used <= 8 => (
5977                self.func("moe_down8_fma_dev_q8_w8h2r2v"),
5978                LaunchConfig {
5979                    grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
5980                    block_dim: (32, n_used as u32, 1),
5981                    shared_mem_bytes: 0,
5982                },
5983            ),
5984            "w8h2r2" if in_f == 512 && n_used <= 8 => (
5985                self.func("moe_down8_fma_dev_q8_w8h2r2"),
5986                LaunchConfig {
5987                    grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
5988                    block_dim: (32, n_used as u32, 1),
5989                    shared_mem_bytes: 0,
5990                },
5991            ),
5992            "w8h2" if in_f == 512 && n_used <= 8 => (
5993                self.func("moe_down8_fma_dev_q8_w8h2"),
5994                LaunchConfig {
5995                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
5996                    block_dim: (32, n_used as u32, 1),
5997                    shared_mem_bytes: 0,
5998                },
5999            ),
6000            _ => (
6001                self.func("moe_down8_fma_dev_q8"),
6002                LaunchConfig {
6003                    grid_dim: (out_f as u32, 1, 1),
6004                    block_dim: (32, 1, 1),
6005                    shared_mem_bytes: 0,
6006                },
6007            ),
6008        };
6009        let __s_b = self.gpu.stream();
6010        let mut b = __s_b.launch_builder(&f);
6011        b.arg(table)
6012            .arg(sel)
6013            .arg(w)
6014            .arg(aq2)
6015            .arg(ad2)
6016            .arg(dst)
6017            .arg(&inf)
6018            .arg(&outf)
6019            .arg(&nu)
6020            .arg(&ne)
6021            .arg(&qt)
6022            .arg(&rbi);
6023        unsafe {
6024            b.launch(cfg)?;
6025        }
6026        Ok(())
6027    }
6028
6029    /// SMALL-M VERIFY rows twin (MEMRA_SPEC_M2, lane/spec-m2): ONE launch covers all `t` tokens
6030    /// of the spec verify's MoE dev gate/up (grid.z = token) — the _v geometry per token, with
6031    /// tok-offset sel/aq/ad/act pointers matching the serial loop's slices. BIT-IDENTICAL per
6032    /// token (see the kernel header). aq/ad are the BATCHED z-quantize ([t, in_f] rows —
6033    /// quantize_q8_1's per-32-block program is row-independent, so batched rows == the serial
6034    /// loop's per-token quantize_q8_1_view bytes). Returns act [t, n_used, n_ff].
6035    #[allow(clippy::too_many_arguments)]
6036    pub fn moe_gate_up_silu8_dev_q8_rows(
6037        &self,
6038        table: &CudaSlice<u64>,
6039        sel: &CudaSlice<i32>,
6040        aq: &CudaSlice<i8>,
6041        ad: &CudaSlice<f32>,
6042        t: usize,
6043        in_f: usize,
6044        n_ff: usize,
6045        n_used: usize,
6046        n_expert: usize,
6047        qt_g: i32,
6048        qt_u: i32,
6049        rb_g: usize,
6050        rb_u: usize,
6051        macros: &CudaSlice<f32>,
6052    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6053        let f = self.func("moe_gate_up_silu8_dev_q8_v_rows");
6054        let mut act = self.alloc_uninit::<f32>(t * n_used * n_ff)?;
6055        let cfg = LaunchConfig {
6056            grid_dim: (n_ff as u32, n_used as u32, t as u32),
6057            block_dim: (32, 1, 1),
6058            shared_mem_bytes: 0,
6059        };
6060        let (inf, nff, ne, nu, rbg, rbu) = (
6061            in_f as i32,
6062            n_ff as i32,
6063            n_expert as i32,
6064            n_used as i32,
6065            rb_g as i64,
6066            rb_u as i64,
6067        );
6068        let __s_b = self.gpu.stream();
6069        let mut b = __s_b.launch_builder(&f);
6070        b.arg(table)
6071            .arg(sel)
6072            .arg(aq)
6073            .arg(ad)
6074            .arg(&mut act)
6075            .arg(&inf)
6076            .arg(&nff)
6077            .arg(&ne)
6078            .arg(&qt_g)
6079            .arg(&qt_u)
6080            .arg(&rbg)
6081            .arg(&rbu)
6082            .arg(&nu)
6083            .arg(macros);
6084        unsafe {
6085            b.launch(cfg)?;
6086        }
6087        Ok(act)
6088    }
6089
6090    /// SMALL-M VERIFY rows twin of the down proj: w8h2v geometry per token on a grid.z token
6091    /// axis. Caller gates the w8h2v shape contract (in_f == 512, n_used <= 8) — same gate as
6092    /// the AUTO dispatch in `moe_down8_fma_dev_q8`. aq2/ad2 = batched act quantize
6093    /// ([t*n_used, in_f] rows). dst rows are FULLY overwritten per token.
6094    #[allow(clippy::too_many_arguments)]
6095    pub fn moe_down8_fma_dev_q8_rows(
6096        &self,
6097        table: &CudaSlice<u64>,
6098        sel: &CudaSlice<i32>,
6099        w: &CudaSlice<f32>,
6100        aq2: &CudaSlice<i8>,
6101        ad2: &CudaSlice<f32>,
6102        dst: &mut CudaSlice<f32>,
6103        t: usize,
6104        in_f: usize,
6105        out_f: usize,
6106        n_used: usize,
6107        n_expert: usize,
6108        qt: i32,
6109        rb: usize,
6110    ) -> Result<(), Box<dyn std::error::Error>> {
6111        assert!(
6112            in_f == 512 && n_used <= 8,
6113            "down rows twin is w8h2v shape-gated"
6114        );
6115        let f = self.func("moe_down8_fma_dev_q8_w8h2v_rows");
6116        let cfg = LaunchConfig {
6117            grid_dim: (out_f.div_ceil(2) as u32, 1, t as u32),
6118            block_dim: (32, n_used as u32, 1),
6119            shared_mem_bytes: 0,
6120        };
6121        let (inf, outf, nu, ne, rbi) = (
6122            in_f as i32,
6123            out_f as i32,
6124            n_used as i32,
6125            n_expert as i32,
6126            rb as i64,
6127        );
6128        let __s_b = self.gpu.stream();
6129        let mut b = __s_b.launch_builder(&f);
6130        b.arg(table)
6131            .arg(sel)
6132            .arg(w)
6133            .arg(aq2)
6134            .arg(ad2)
6135            .arg(dst)
6136            .arg(&inf)
6137            .arg(&outf)
6138            .arg(&nu)
6139            .arg(&ne)
6140            .arg(&qt)
6141            .arg(&rbi);
6142        unsafe {
6143            b.launch(cfg)?;
6144        }
6145        Ok(())
6146    }
6147
6148    /// CSR gate/up v3 (owner-scan dedup, no build kernel): qtypes {IQ4_XS, IQ3_S} (caller
6149    /// gates), grid.y = pair index; the first pair of each expert serves all its pairs.
6150    /// Bit-identical to moe_gate_up_silu8_dev_q8_v_rows (explicit-intrinsic accumulate).
6151    #[allow(clippy::too_many_arguments)]
6152    pub fn moe_gate_up_silu8_dev_q8_csr(
6153        &self,
6154        table: &CudaSlice<u64>,
6155        sel: &CudaSlice<i32>,
6156        aq: &CudaSlice<i8>,
6157        ad: &CudaSlice<f32>,
6158        n_pairs: usize,
6159        in_f: usize,
6160        n_ff: usize,
6161        n_used: usize,
6162        n_expert: usize,
6163        qt_g: i32,
6164        qt_u: i32,
6165        rb_g: usize,
6166        rb_u: usize,
6167    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6168        // NVFP4 experts take the NVFP4-specialized owner-scan twin (lane/moebatch-q35moe);
6169        // host gate guarantees qt_g == qt_u within a supported class.
6170        let f = if qt_g == crate::QT_NVFP4 {
6171            self.func("moe_gate_up_silu8_dev_q8_csr_nvfp4")
6172        } else {
6173            self.func("moe_gate_up_silu8_dev_q8_csr_iq4")
6174        };
6175        let mut act = self.alloc_uninit::<f32>(n_pairs * n_ff)?;
6176        let cfg = LaunchConfig {
6177            grid_dim: (n_ff as u32, n_pairs as u32, 1),
6178            block_dim: (32, 1, 1),
6179            shared_mem_bytes: 0,
6180        };
6181        let (inf, nff, ne, nu, npi, rbg, rbu) = (
6182            in_f as i32,
6183            n_ff as i32,
6184            n_expert as i32,
6185            n_used as i32,
6186            n_pairs as i32,
6187            rb_g as i64,
6188            rb_u as i64,
6189        );
6190        let __s_b = self.gpu.stream();
6191        let mut b = __s_b.launch_builder(&f);
6192        b.arg(table)
6193            .arg(sel)
6194            .arg(aq)
6195            .arg(ad)
6196            .arg(&mut act)
6197            .arg(&inf)
6198            .arg(&nff)
6199            .arg(&ne)
6200            .arg(&qt_g)
6201            .arg(&qt_u)
6202            .arg(&rbg)
6203            .arg(&rbu)
6204            .arg(&nu)
6205            .arg(&npi);
6206        unsafe {
6207            b.launch(cfg)?;
6208        }
6209        Ok(act)
6210    }
6211
6212    /// TEST SEAM (down8 lane 2026-07-08): launch a down dev_q8 variant BY NAME with its
6213    /// canonical geometry, bypassing the env-cached dispatch so moe-devq8-check can byte-
6214    /// compare variants in one process. Variants: "base", "w8h2", "w8h2r2", "w8h2v", "w8h2r2v".
6215    #[allow(clippy::too_many_arguments)]
6216    pub fn moe_down8_fma_dev_q8_variant(
6217        &self,
6218        variant: &str,
6219        table: &CudaSlice<u64>,
6220        sel: &cudarc::driver::CudaView<i32>,
6221        w: &cudarc::driver::CudaView<f32>,
6222        aq2: &CudaSlice<i8>,
6223        ad2: &CudaSlice<f32>,
6224        dst: &mut cudarc::driver::CudaViewMut<f32>,
6225        in_f: usize,
6226        out_f: usize,
6227        n_used: usize,
6228        n_expert: usize,
6229        qt: i32,
6230        rb: usize,
6231    ) -> Result<(), Box<dyn std::error::Error>> {
6232        let (inf, outf, nu, ne, rbi) = (
6233            in_f as i32,
6234            out_f as i32,
6235            n_used as i32,
6236            n_expert as i32,
6237            rb as i64,
6238        );
6239        let (f, cfg) = match variant {
6240            "w8h2" | "w8h2v" => (
6241                self.func(if variant == "w8h2" {
6242                    "moe_down8_fma_dev_q8_w8h2"
6243                } else {
6244                    "moe_down8_fma_dev_q8_w8h2v"
6245                }),
6246                LaunchConfig {
6247                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6248                    block_dim: (32, n_used as u32, 1),
6249                    shared_mem_bytes: 0,
6250                },
6251            ),
6252            "w8h2r2" | "w8h2r2v" => (
6253                self.func(if variant == "w8h2r2" {
6254                    "moe_down8_fma_dev_q8_w8h2r2"
6255                } else {
6256                    "moe_down8_fma_dev_q8_w8h2r2v"
6257                }),
6258                LaunchConfig {
6259                    grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
6260                    block_dim: (32, n_used as u32, 1),
6261                    shared_mem_bytes: 0,
6262                },
6263            ),
6264            _ => (
6265                self.func("moe_down8_fma_dev_q8"),
6266                LaunchConfig {
6267                    grid_dim: (out_f as u32, 1, 1),
6268                    block_dim: (32, 1, 1),
6269                    shared_mem_bytes: 0,
6270                },
6271            ),
6272        };
6273        let __s_b = self.gpu.stream();
6274        let mut b = __s_b.launch_builder(&f);
6275        b.arg(table)
6276            .arg(sel)
6277            .arg(w)
6278            .arg(aq2)
6279            .arg(ad2)
6280            .arg(dst)
6281            .arg(&inf)
6282            .arg(&outf)
6283            .arg(&nu)
6284            .arg(&ne)
6285            .arg(&qt)
6286            .arg(&rbi);
6287        unsafe {
6288            b.launch(cfg)?;
6289        }
6290        Ok(())
6291    }
6292
6293    /// TEST SEAM (down8 lane): gate_up twin of the above. Variants: "base", "v".
6294    #[allow(clippy::too_many_arguments)]
6295    pub fn moe_gate_up_silu8_dev_q8_variant(
6296        &self,
6297        variant: &str,
6298        table: &CudaSlice<u64>,
6299        sel: &cudarc::driver::CudaView<i32>,
6300        aq: &CudaSlice<i8>,
6301        ad: &CudaSlice<f32>,
6302        in_f: usize,
6303        n_ff: usize,
6304        n_used: usize,
6305        n_expert: usize,
6306        qt_g: i32,
6307        qt_u: i32,
6308        rb_g: usize,
6309        rb_u: usize,
6310    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6311        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
6312        let (inf, nff, ne, rbg, rbu) = (
6313            in_f as i32,
6314            n_ff as i32,
6315            n_expert as i32,
6316            rb_g as i64,
6317            rb_u as i64,
6318        );
6319        let f = self.func(if variant == "v" {
6320            "moe_gate_up_silu8_dev_q8_v"
6321        } else {
6322            "moe_gate_up_silu8_dev_q8"
6323        });
6324        let cfg = LaunchConfig {
6325            grid_dim: (n_ff as u32, n_used as u32, 1),
6326            block_dim: (32, 1, 1),
6327            shared_mem_bytes: 0,
6328        };
6329        let __s_b = self.gpu.stream();
6330        let mut b = __s_b.launch_builder(&f);
6331        b.arg(table)
6332            .arg(sel)
6333            .arg(aq)
6334            .arg(ad)
6335            .arg(&mut act)
6336            .arg(&inf)
6337            .arg(&nff)
6338            .arg(&ne)
6339            .arg(&qt_g)
6340            .arg(&qt_u)
6341            .arg(&rbg)
6342            .arg(&rbu);
6343        unsafe {
6344            b.launch(cfg)?;
6345        }
6346        Ok(act)
6347    }
6348
6349    pub fn moe_gate_up_silu8_dev(
6350        &self,
6351        table: &CudaSlice<u64>,
6352        sel: &cudarc::driver::CudaView<i32>,
6353        x: &cudarc::driver::CudaView<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        macros: &CudaSlice<f32>,
6363    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6364        let f = self.func("moe_gate_up_silu8_dev");
6365        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?; // fully overwritten
6366        let cfg = LaunchConfig {
6367            grid_dim: (n_ff as u32, n_used as u32, 1),
6368            block_dim: (256, 1, 1),
6369            shared_mem_bytes: 0,
6370        };
6371        let (inf, nff, ne, rbg, rbu) = (
6372            in_f as i32,
6373            n_ff as i32,
6374            n_expert as i32,
6375            rb_g as i64,
6376            rb_u as i64,
6377        );
6378        let __s_b = self.gpu.stream();
6379        let mut b = __s_b.launch_builder(&f);
6380        b.arg(table)
6381            .arg(sel)
6382            .arg(x)
6383            .arg(&mut act)
6384            .arg(&inf)
6385            .arg(&nff)
6386            .arg(&ne)
6387            .arg(&qt_g)
6388            .arg(&qt_u)
6389            .arg(&rbg)
6390            .arg(&rbu)
6391            .arg(macros);
6392        unsafe {
6393            b.launch(cfg)?;
6394        }
6395        Ok(act)
6396    }
6397
6398    /// LAUNCH-STRUCTURE STAGE 3: device-dispatch twin of `moe_down8_fma_into` — expert ids AND
6399    /// renormalized weights read from the router kernel's device output. BIT-IDENTICAL chain.
6400    #[allow(clippy::too_many_arguments)]
6401    pub fn moe_down8_fma_dev(
6402        &self,
6403        table: &CudaSlice<u64>,
6404        sel: &cudarc::driver::CudaView<i32>,
6405        w: &cudarc::driver::CudaView<f32>,
6406        act: &CudaSlice<f32>,
6407        dst: &mut cudarc::driver::CudaViewMut<f32>,
6408        in_f: usize,
6409        out_f: usize,
6410        n_used: usize,
6411        n_expert: usize,
6412        qt: i32,
6413        rb: usize,
6414    ) -> Result<(), Box<dyn std::error::Error>> {
6415        let f = self.func("moe_down8_fma_dev");
6416        let cfg = LaunchConfig {
6417            grid_dim: (out_f as u32, 1, 1),
6418            block_dim: (256, 1, 1),
6419            shared_mem_bytes: 0,
6420        };
6421        let (inf, outf, nu, ne, rbv) = (
6422            in_f as i32,
6423            out_f as i32,
6424            n_used as i32,
6425            n_expert as i32,
6426            rb as i64,
6427        );
6428        let __s_b = self.gpu.stream();
6429        let mut b = __s_b.launch_builder(&f);
6430        b.arg(table)
6431            .arg(sel)
6432            .arg(w)
6433            .arg(act)
6434            .arg(dst)
6435            .arg(&inf)
6436            .arg(&outf)
6437            .arg(&nu)
6438            .arg(&ne)
6439            .arg(&qt)
6440            .arg(&rbv);
6441        unsafe {
6442            b.launch(cfg)?;
6443        }
6444        Ok(())
6445    }
6446
6447    /// dst[i] += alpha * src[i], i in 0..n. dst is a CudaViewMut (a row of moe_out).
6448    pub fn axpy_into(
6449        &self,
6450        src: &CudaSlice<f32>,
6451        alpha: f32,
6452        dst: &mut cudarc::driver::CudaViewMut<f32>,
6453        n: usize,
6454    ) -> Result<(), Box<dyn std::error::Error>> {
6455        let f = self.func("axpy_f32");
6456        let cfg = LaunchConfig::for_num_elems(n as u32);
6457        let (a, ni) = (alpha, n as i32);
6458        let __s_b = self.gpu.stream();
6459        let mut b = __s_b.launch_builder(&f);
6460        b.arg(src).arg(dst).arg(&a).arg(&ni);
6461        unsafe {
6462            b.launch(cfg)?;
6463        }
6464        Ok(())
6465    }
6466
6467    /// dst[r*ncols + c] += src[r*ncols + c] * scale[r]. Per-row scalar accumulate (shared expert).
6468    pub fn add_scaled_rows(
6469        &self,
6470        src: &CudaSlice<f32>,
6471        scale: &CudaSlice<f32>,
6472        dst: &mut CudaSlice<f32>,
6473        ncols: usize,
6474        nrows: usize,
6475    ) -> Result<(), Box<dyn std::error::Error>> {
6476        let f = self.func("add_scaled_rows_f32");
6477        let cfg = LaunchConfig::for_num_elems((ncols * nrows) as u32);
6478        let (nc, nr) = (ncols as i32, nrows as i32);
6479        let __s_b = self.gpu.stream();
6480        let mut b = __s_b.launch_builder(&f);
6481        b.arg(src).arg(scale).arg(dst).arg(&nc).arg(&nr);
6482        unsafe {
6483            b.launch(cfg)?;
6484        }
6485        Ok(())
6486    }
6487
6488    // ======== A2 GROUPED MoE PREFILL KERNELS ========
6489
6490    /// Gather m_e rows from src[T, ncols] into dst[m_e, ncols] using index array idx[m_e].
6491    pub fn gather_rows(
6492        &self,
6493        src: &CudaSlice<f32>,
6494        idx: &CudaSlice<i32>,
6495        dst: &mut CudaSlice<f32>,
6496        ncols: usize,
6497        m_e: usize,
6498    ) -> Result<(), Box<dyn std::error::Error>> {
6499        let f = self.func("gather_rows_f32");
6500        let cfg = LaunchConfig::for_num_elems((m_e * ncols) as u32);
6501        let (nc, me) = (ncols as i32, m_e as i32);
6502        let __s_b = self.gpu.stream();
6503        let mut b = __s_b.launch_builder(&f);
6504        b.arg(src).arg(idx).arg(dst).arg(&nc).arg(&me);
6505        unsafe {
6506            b.launch(cfg)?;
6507        }
6508        Ok(())
6509    }
6510
6511    /// Scatter expert outputs into per-token slots: dst[tok_idx[r], slot_idx[r], :] = src[r, :] * weight[r].
6512    /// dst is [T, n_used, ncols], zero-initialized. Each (expert, token) pair maps to a unique slot.
6513    /// Scatter expert outputs into per-token slots (raw copy, no weight multiply).
6514    /// Weight stored into wbuf[tok*n_used + slot] for FMA in reduce step.
6515    pub fn scatter_slot(
6516        &self,
6517        src: &CudaSlice<f32>,
6518        tok_idx: &CudaSlice<i32>,
6519        slot_idx: &CudaSlice<i32>,
6520        weight: &CudaSlice<f32>,
6521        dst: &mut CudaSlice<f32>,
6522        wbuf: &mut CudaSlice<f32>,
6523        ncols: usize,
6524        n_used: usize,
6525        m_e: usize,
6526    ) -> Result<(), Box<dyn std::error::Error>> {
6527        let f = self.func("scatter_add_slot_f32");
6528        let cfg = LaunchConfig::for_num_elems((m_e * ncols) as u32);
6529        let (nc, nu, me) = (ncols as i32, n_used as i32, m_e as i32);
6530        let __s_b = self.gpu.stream();
6531        let mut b = __s_b.launch_builder(&f);
6532        b.arg(src)
6533            .arg(tok_idx)
6534            .arg(slot_idx)
6535            .arg(weight)
6536            .arg(dst)
6537            .arg(wbuf)
6538            .arg(&nc)
6539            .arg(&nu)
6540            .arg(&me);
6541        unsafe {
6542            b.launch(cfg)?;
6543        }
6544        Ok(())
6545    }
6546
6547    /// Reduce n_used slots per token: dst[t, col] = sum_s slots[t, s, col].
6548    /// Reduce n_used slots per token: dst[t, col] = sum_s FMA(wbuf[t,s], slots[t,s,col], acc).
6549    /// Uses FMA for bit-identity with the sequential axpy path.
6550    pub fn reduce_slots(
6551        &self,
6552        slots: &CudaSlice<f32>,
6553        wbuf: &CudaSlice<f32>,
6554        dst: &mut CudaSlice<f32>,
6555        ncols: usize,
6556        n_used: usize,
6557        t: usize,
6558    ) -> Result<(), Box<dyn std::error::Error>> {
6559        let f = self.func("reduce_slots_f32");
6560        let cfg = LaunchConfig::for_num_elems((t * ncols) as u32);
6561        let (nc, nu, ti) = (ncols as i32, n_used as i32, t as i32);
6562        let __s_b = self.gpu.stream();
6563        let mut b = __s_b.launch_builder(&f);
6564        b.arg(slots).arg(wbuf).arg(dst).arg(&nc).arg(&nu).arg(&ti);
6565        unsafe {
6566            b.launch(cfg)?;
6567        }
6568        Ok(())
6569    }
6570
6571    /// Stage-B: quantize activation [m,in] f32 -> q8_1 (int8 qs + per-block f32 scale).
6572    /// Quantize an activation [m, in_f] to q8_1 (int8 qs + per-32 f32 scale). Public so the
6573    /// forward can quantize a SHARED activation ONCE and feed it to several matmuls (gate+up
6574    /// share `z`; q/k/v and wqkv/gate/beta/alpha share `h`) — quantize_q8_1 was 13.5% of decode
6575    /// GPU time, ~half of it redundant re-quantization of the same row.
6576    /// quantize_q8_1 over a CudaView (a sliced z-row) — same kernel, offset-honoring arg.
6577    pub fn quantize_q8_1_view(
6578        &self,
6579        x: &cudarc::driver::CudaView<f32>,
6580        m: usize,
6581        in_f: usize,
6582    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6583        let f = self.func("quantize_q8_1");
6584        let nblk = in_f / 32;
6585        let mut q = self.alloc_uninit::<i8>(m * in_f)?;
6586        let mut d = self.alloc_uninit::<f32>(m * nblk)?;
6587        let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
6588        let (inf, mi) = (in_f as i32, m as i32);
6589        let __s_b = self.gpu.stream();
6590        let mut b = __s_b.launch_builder(&f);
6591        b.arg(x).arg(&mut q).arg(&mut d).arg(&inf).arg(&mi);
6592        unsafe {
6593            b.launch(cfg)?;
6594        }
6595        Ok((q, d))
6596    }
6597
6598    pub fn quantize_q8_1(
6599        &self,
6600        x: &CudaSlice<f32>,
6601        m: usize,
6602        in_f: usize,
6603    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6604        let nblk = in_f / 32;
6605        let mut q = self.alloc_uninit::<i8>(m * in_f)?; // full-overwrite output: skip memset
6606        let mut d = self.alloc_uninit::<f32>(m * nblk)?; // full-overwrite output: skip memset
6607        // WARP-PER-BLOCK kernel: one warp per 32-block -> m*in_f threads total.
6608        let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
6609        let (inf, mi) = (in_f as i32, m as i32);
6610        if Self::pdl_on() && Self::pdl_wb_on() {
6611            {
6612                use cudarc::driver::{DevicePtr, DevicePtrMut};
6613                let s = &self.gpu.stream();
6614                let (px, _g0) = x.device_ptr(s);
6615                let (pq, _g1) = q.device_ptr_mut(s);
6616                let (pd, _g2) = d.device_ptr_mut(s);
6617                let mut ps = [
6618                    &px as *const _ as *mut std::ffi::c_void,
6619                    &pq as *const _ as *mut _,
6620                    &pd as *const _ as *mut _,
6621                    &inf as *const _ as *mut _,
6622                    &mi as *const _ as *mut _,
6623                ];
6624                unsafe {
6625                    self.launch_pdl("quantize_q8_1", cfg.grid_dim, cfg.block_dim, &mut ps)?;
6626                }
6627            }
6628            return Ok((q, d));
6629        }
6630        let f = self.func("quantize_q8_1");
6631        let __s_b = self.gpu.stream();
6632        let mut b = __s_b.launch_builder(&f);
6633        b.arg(x).arg(&mut q).arg(&mut d).arg(&inf).arg(&mi);
6634        unsafe {
6635            b.launch(cfg)?;
6636        }
6637        Ok((q, d))
6638    }
6639
6640    /// Stage-C FP4: quantize activation [m,in] f32 -> e2m1 nibbles (aq4: u32 [m, in/8]) + per-16
6641    /// UE4M3 scale (ad4: u8 [m, in/16]), the layout the mxf4nvf4 block-scale GEMM B-operand wants.
6642    /// in_f must be a multiple of 64 (one NVFP4 K-block). One thread per (token, 16-block).
6643    pub fn quantize_fp4_act(
6644        &self,
6645        x: &CudaSlice<f32>,
6646        m: usize,
6647        in_f: usize,
6648    ) -> Result<(CudaSlice<u32>, CudaSlice<u8>), Box<dyn std::error::Error>> {
6649        let f = self.func("quantize_fp4_act");
6650        let nb16 = in_f / 16;
6651        let mut aq4 = self.alloc_uninit::<u32>(m * (in_f / 8))?; // full-overwrite output: skip memset
6652        let mut ad4 = self.alloc_uninit::<u8>(m * nb16)?; // full-overwrite output: skip memset
6653        let cfg = LaunchConfig::for_num_elems((m * nb16) as u32);
6654        let (inf, mi) = (in_f as i32, m as i32);
6655        let __s_b = self.gpu.stream();
6656        let mut b = __s_b.launch_builder(&f);
6657        b.arg(x).arg(&mut aq4).arg(&mut ad4).arg(&inf).arg(&mi);
6658        unsafe {
6659            b.launch(cfg)?;
6660        }
6661        Ok((aq4, ad4))
6662    }
6663
6664    /// Stage-C FP4 GEMM (NVFP4 weights): native mxf4nvf4 block-scale tensor-core matmul. Feeds raw
6665    /// e2m1 weight nibbles + raw UE4M3 micro-scales directly to mma.sync.m16n8k64 (762 TFLOP/s peak,
6666    /// 3.5x int8). Activation `x` is quantized to FP4 e2m1 here. NVFP4 per-tensor macro-scale applied
6667    /// post (scale==1.0 -> no-op). `bytes` = raw NVFP4 weight rows. Used by the MEMRA_FP4 prefill path.
6668    pub fn qmatvec_gemm_nvfp4_fp4(
6669        &self,
6670        bytes: &CudaSlice<u8>,
6671        x: &CudaSlice<f32>,
6672        m: usize,
6673        in_f: usize,
6674        out_f: usize,
6675        row_bytes: usize,
6676        scale: f32,
6677    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6678        assert!(
6679            in_f % 64 == 0,
6680            "FP4 GEMM requires in_f % 64 == 0, got {in_f}"
6681        );
6682        let (aq4, ad4) = self.quantize_fp4_act(x, m, in_f)?;
6683        let mut y = self.fp4_gemm_launch(bytes, &aq4, &ad4, m, in_f, out_f, row_bytes)?;
6684        if scale != 1.0 {
6685            self.scale_inplace(&mut y, scale, m * out_f)?;
6686        }
6687        Ok(y)
6688    }
6689
6690    /// Shared mxf4 GEMM launch (pre-quantized FP4 activation aq4/ad4). Same CTA tile as the int8 GEMM
6691    /// (BM=64 rows x BN=128 tokens, 4 warps). No macro-scale applied here.
6692    fn fp4_gemm_launch(
6693        &self,
6694        bytes: &CudaSlice<u8>,
6695        aq4: &CudaSlice<u32>,
6696        ad4: &CudaSlice<u8>,
6697        m: usize,
6698        in_f: usize,
6699        out_f: usize,
6700        row_bytes: usize,
6701    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6702        let f = self.func("qmatvec_gemm_nvfp4_fp4");
6703        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
6704        const BM: u32 = 64;
6705        const BN: u32 = 256;
6706        let cfg = LaunchConfig {
6707            grid_dim: ((out_f as u32 + BM - 1) / BM, (m as u32 + BN - 1) / BN, 1),
6708            block_dim: (32, 4, 1),
6709            shared_mem_bytes: 0,
6710        };
6711        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
6712        let __s_b = self.gpu.stream();
6713        let mut b = __s_b.launch_builder(&f);
6714        b.arg(bytes)
6715            .arg(aq4)
6716            .arg(ad4)
6717            .arg(&mut y)
6718            .arg(&inf)
6719            .arg(&outf)
6720            .arg(&mi)
6721            .arg(&rb);
6722        unsafe {
6723            b.launch(cfg)?;
6724        }
6725        Ok(y)
6726    }
6727
6728    /// Test entry (kernel_check): run the FP4 GEMM from raw bytes; NO macro-scale (caller compares bare).
6729    pub fn qmatvec_gemm_nvfp4_fp4_raw(
6730        &self,
6731        bytes: &CudaSlice<u8>,
6732        x: &CudaSlice<f32>,
6733        m: usize,
6734        in_f: usize,
6735        out_f: usize,
6736        row_bytes: usize,
6737    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6738        assert!(
6739            in_f % 64 == 0,
6740            "FP4 GEMM requires in_f % 64 == 0, got {in_f}"
6741        );
6742        let (aq4, ad4) = self.quantize_fp4_act(x, m, in_f)?;
6743        self.fp4_gemm_launch(bytes, &aq4, &ad4, m, in_f, out_f, row_bytes)
6744    }
6745
6746    /// Stage-B: Q8_0 weight x q8_1 activation int8 dp4a matmul. y[m,out]=x@W^T.
6747    pub fn qmatvec_q8_0_fast(
6748        &self,
6749        w: &CudaSlice<u8>,
6750        x: &CudaSlice<f32>,
6751        m: usize,
6752        in_f: usize,
6753        out_f: usize,
6754        row_bytes: usize,
6755    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6756        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
6757        let f = self.func("qmatvec_q8_0_dp4a");
6758        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
6759        let cfg = LaunchConfig {
6760            grid_dim: (out_f as u32, m as u32, 1),
6761            block_dim: (128, 1, 1),
6762            shared_mem_bytes: 0,
6763        };
6764        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
6765        let __s_b = self.gpu.stream();
6766        let mut b = __s_b.launch_builder(&f);
6767        b.arg(w)
6768            .arg(&aq)
6769            .arg(&ad)
6770            .arg(&mut y)
6771            .arg(&inf)
6772            .arg(&outf)
6773            .arg(&mi)
6774            .arg(&rb);
6775        unsafe {
6776            b.launch(cfg)?;
6777        }
6778        Ok(y)
6779    }
6780
6781    /// Stage-B: Q4_K weight x q8_1 activation int8 dp4a (decode). Min-offset via q8_1 sum term.
6782    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
6783    pub fn qmatvec_q4_K_fast(
6784        &self,
6785        w: &CudaSlice<u8>,
6786        x: &CudaSlice<f32>,
6787        m: usize,
6788        in_f: usize,
6789        out_f: usize,
6790        row_bytes: usize,
6791    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6792        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
6793        let f = self.func("qmatvec_q4_K_dp4a");
6794        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
6795        let cfg = LaunchConfig {
6796            grid_dim: (out_f as u32, m as u32, 1),
6797            block_dim: (128, 1, 1),
6798            shared_mem_bytes: 0,
6799        };
6800        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
6801        let __s_b = self.gpu.stream();
6802        let mut b = __s_b.launch_builder(&f);
6803        b.arg(w)
6804            .arg(&aq)
6805            .arg(&ad)
6806            .arg(&mut y)
6807            .arg(&inf)
6808            .arg(&outf)
6809            .arg(&mi)
6810            .arg(&rb);
6811        unsafe {
6812            b.launch(cfg)?;
6813        }
6814        Ok(y)
6815    }
6816
6817    /// Stage-B: Q6_K weight x q8_1 activation int8 dp4a (decode, symmetric).
6818    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
6819    pub fn qmatvec_q6_K_fast(
6820        &self,
6821        w: &CudaSlice<u8>,
6822        x: &CudaSlice<f32>,
6823        m: usize,
6824        in_f: usize,
6825        out_f: usize,
6826        row_bytes: usize,
6827    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6828        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
6829        let f = self.func("qmatvec_q6_K_dp4a");
6830        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
6831        let cfg = LaunchConfig {
6832            grid_dim: (out_f as u32, m as u32, 1),
6833            block_dim: (128, 1, 1),
6834            shared_mem_bytes: 0,
6835        };
6836        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
6837        let __s_b = self.gpu.stream();
6838        let mut b = __s_b.launch_builder(&f);
6839        b.arg(w)
6840            .arg(&aq)
6841            .arg(&ad)
6842            .arg(&mut y)
6843            .arg(&inf)
6844            .arg(&outf)
6845            .arg(&mi)
6846            .arg(&rb);
6847        unsafe {
6848            b.launch(cfg)?;
6849        }
6850        Ok(y)
6851    }
6852
6853    /// Stage-B: Q5_K weight x q8_1 activation int8 dp4a (decode). Min-offset via q8_1 sum term.
6854    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
6855    pub fn qmatvec_q5_K_fast(
6856        &self,
6857        w: &CudaSlice<u8>,
6858        x: &CudaSlice<f32>,
6859        m: usize,
6860        in_f: usize,
6861        out_f: usize,
6862        row_bytes: usize,
6863    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6864        self.qmatvec_dp4a_named("qmatvec_q5_K_dp4a", w, x, m, in_f, out_f, row_bytes)
6865    }
6866    /// Stage-B: Q3_K weight x q8_1 activation int8 dp4a (decode, symmetric).
6867    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
6868    pub fn qmatvec_q3_K_fast(
6869        &self,
6870        w: &CudaSlice<u8>,
6871        x: &CudaSlice<f32>,
6872        m: usize,
6873        in_f: usize,
6874        out_f: usize,
6875        row_bytes: usize,
6876    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6877        self.qmatvec_dp4a_named("qmatvec_q3_K_dp4a", w, x, m, in_f, out_f, row_bytes)
6878    }
6879    /// A6 split-plane twin of `qmatvec_nvfp4_fast` (weights repacked; used by the rp gates).
6880    pub fn qmatvec_nvfp4_fast_rp(
6881        &self,
6882        w: &CudaSlice<u8>,
6883        x: &CudaSlice<f32>,
6884        m: usize,
6885        in_f: usize,
6886        out_f: usize,
6887        row_bytes: usize,
6888    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6889        assert!(
6890            in_f % 64 == 0,
6891            "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
6892        );
6893        self.qmatvec_dp4a_named("qmatvec_nvfp4_dp4a_rp", w, x, m, in_f, out_f, row_bytes)
6894    }
6895    /// Stage-B: NVFP4 weight x q8_1 activation int8 dp4a (decode, symmetric, codebook lookup).
6896    pub fn qmatvec_nvfp4_fast(
6897        &self,
6898        w: &CudaSlice<u8>,
6899        x: &CudaSlice<f32>,
6900        m: usize,
6901        in_f: usize,
6902        out_f: usize,
6903        row_bytes: usize,
6904    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6905        // B1: the NVFP4 dp4a kernel maps two 32-elem q8_1 blocks onto one 64-elem block_nvfp4
6906        // (sblk = g >> 1). in_f must be a multiple of 64 or the last block reads a partial superblock.
6907        assert!(
6908            in_f % 64 == 0,
6909            "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
6910        );
6911        self.qmatvec_dp4a_named("qmatvec_nvfp4_dp4a", w, x, m, in_f, out_f, row_bytes)
6912    }
6913    /// Stage-B (optional perf): IQ4_XS codebook int8 dp4a.
6914    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
6915    pub fn qmatvec_iq4_XS_fast(
6916        &self,
6917        w: &CudaSlice<u8>,
6918        x: &CudaSlice<f32>,
6919        m: usize,
6920        in_f: usize,
6921        out_f: usize,
6922        row_bytes: usize,
6923    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6924        self.qmatvec_dp4a_named("qmatvec_iq4_XS_dp4a", w, x, m, in_f, out_f, row_bytes)
6925    }
6926
6927    /// Shared dp4a launcher: quantize_q8_1 then call the named kernel (grid (out,m), block 64).
6928    fn qmatvec_dp4a_named(
6929        &self,
6930        name: &str,
6931        w: &CudaSlice<u8>,
6932        x: &CudaSlice<f32>,
6933        m: usize,
6934        in_f: usize,
6935        out_f: usize,
6936        row_bytes: usize,
6937    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6938        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
6939        let f = self.func(name);
6940        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
6941        let cfg = LaunchConfig {
6942            grid_dim: (out_f as u32, m as u32, 1),
6943            block_dim: (128, 1, 1),
6944            shared_mem_bytes: 0,
6945        };
6946        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
6947        let __s_b = self.gpu.stream();
6948        let mut b = __s_b.launch_builder(&f);
6949        b.arg(w)
6950            .arg(&aq)
6951            .arg(&ad)
6952            .arg(&mut y)
6953            .arg(&inf)
6954            .arg(&outf)
6955            .arg(&mi)
6956            .arg(&rb);
6957        unsafe {
6958            b.launch(cfg)?;
6959        }
6960        Ok(y)
6961    }
6962
6963    pub fn htod(&self, v: &[f32]) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6964        Ok(self.gpu.stream().clone_htod(v)?)
6965    }
6966    pub fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
6967        Ok(self.gpu.stream().clone_htod(v)?)
6968    }
6969    /// i8 upload (moe-devq8-check: synthetic q8_1 activation bytes).
6970    pub fn htod_i8(&self, v: &[i8]) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
6971        Ok(self.gpu.stream().clone_htod(v)?)
6972    }
6973    pub fn htod_u64(&self, v: &[u64]) -> Result<CudaSlice<u64>, Box<dyn std::error::Error>> {
6974        Ok(self.gpu.stream().clone_htod(v)?)
6975    }
6976    /// View twin of `dtoh` (lean-logits component 3: D2H one row of a [B, n_vocab] stack).
6977    pub fn dtoh_view(
6978        &self,
6979        d: &cudarc::driver::CudaView<f32>,
6980    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
6981        let v = self.gpu.stream().clone_dtoh(d)?;
6982        self.gpu.stream().synchronize()?;
6983        Ok(v)
6984    }
6985    pub fn dtoh(&self, d: &CudaSlice<f32>) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
6986        let v = self.gpu.stream().clone_dtoh(d)?;
6987        self.gpu.stream().synchronize()?;
6988        Ok(v)
6989    }
6990    /// Queue two f32 device-to-host copies on the compute stream, then establish one host
6991    /// boundary for both. Hy3's CPU/GPU expert split needs the router logits and the MoE input;
6992    /// issuing them together avoids a second stream synchronization in every trunk layer.
6993    pub fn dtoh_pair(
6994        &self,
6995        a: &CudaSlice<f32>,
6996        b: &CudaSlice<f32>,
6997    ) -> Result<(Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
6998        let av = self.gpu.stream().clone_dtoh(a)?;
6999        let bv = self.gpu.stream().clone_dtoh(b)?;
7000        self.gpu.stream().synchronize()?;
7001        Ok((av, bv))
7002    }
7003    /// Device-to-host copy of an i32 buffer (fused-router sel_idx readback).
7004    pub fn dtoh_i32(&self, d: &CudaSlice<i32>) -> Result<Vec<i32>, Box<dyn std::error::Error>> {
7005        let v = self.gpu.stream().clone_dtoh(d)?;
7006        self.gpu.stream().synchronize()?;
7007        Ok(v)
7008    }
7009    /// Device-to-host copy of a u8 buffer (used to read back the quantized KV cache for validation).
7010    pub fn dtoh_u8(&self, d: &CudaSlice<u8>) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
7011        let v = self.gpu.stream().clone_dtoh(d)?;
7012        self.gpu.stream().synchronize()?;
7013        Ok(v)
7014    }
7015    pub fn dtoh_u8_view(
7016        &self,
7017        d: &cudarc::driver::CudaView<u8>,
7018    ) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
7019        let v = self.gpu.stream().clone_dtoh(d)?;
7020        self.gpu.stream().synchronize()?;
7021        Ok(v)
7022    }
7023    pub fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7024        let s = self.gpu.stream().alloc_zeros::<f32>(n)?;
7025        self.keep_if_capturing(&s);
7026        Ok(s)
7027    }
7028
7029    /// GPU-resident greedy argmax (CUDA-GRAPH-PLAN Phase 1): logits[n_vocab] -> token id in a
7030    /// resident device u32 [1]. PARALLEL 2-pass (RANK1 LEVER): the old single-CTA scan (one 256-thread
7031    /// block on one SM over 248K logits) was memory-starved at ~426us/token. Now pass 1 fans NB=256
7032    /// blocks across the SMs to saturate HBM, pass 2 reduces the NB partials. Bit-identical to host
7033    /// `argmax` (smallest index on tie). The whole point is NOT to dtoh logits — only a [1] u32 is read
7034    /// back (or kept resident for graph replay). Returns the device token buffer.
7035    /// Softmax probability of the (already-argmaxed) token `tok` under `logits` — the spec-decode
7036    /// p-min confidence signal. 2-pass like the parallel argmax; returns a device [1] f32.
7037    pub fn prob_of_token_device(
7038        &self,
7039        logits: &CudaSlice<f32>,
7040        tok: &CudaSlice<u32>,
7041        n_vocab: usize,
7042    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7043        let nb = ARGMAX_NB;
7044        let mut part = self.alloc_uninit::<f32>(nb)?;
7045        let mut p = self.alloc_uninit::<f32>(1)?;
7046        let f1 = self.func("prob_of_token_partial_f32");
7047        let cfg1 = LaunchConfig {
7048            grid_dim: (nb as u32, 1, 1),
7049            block_dim: (256, 1, 1),
7050            shared_mem_bytes: 0,
7051        };
7052        let nv = n_vocab as i32;
7053        let __s_b1 = self.gpu.stream();
7054        let mut b1 = __s_b1.launch_builder(&f1);
7055        b1.arg(logits).arg(tok).arg(&mut part).arg(&nv);
7056        unsafe {
7057            b1.launch(cfg1)?;
7058        }
7059        let f2 = self.func("prob_of_token_final_f32");
7060        let cfg2 = LaunchConfig {
7061            grid_dim: (1, 1, 1),
7062            block_dim: (256, 1, 1),
7063            shared_mem_bytes: 0,
7064        };
7065        let nbi = nb as i32;
7066        let __s_b2 = self.gpu.stream();
7067        let mut b2 = __s_b2.launch_builder(&f2);
7068        b2.arg(&part).arg(&mut p).arg(&nbi);
7069        unsafe {
7070            b2.launch(cfg2)?;
7071        }
7072        Ok(p)
7073    }
7074
7075    /// Like `prob_of_token_device` but writes into a PERSISTENT `p_out` buffer (stable pointer).
7076    /// Required for CUDA-graph capture of the draft chain: the captured prob kernels must write
7077    /// where the host reads the p-min confidence between replays. Same kernels, same math.
7078    /// Slot-addressed twin of `prob_of_token_device_into`: token read from `tok_all[tok_idx]`
7079    /// (a view at the slot), probability written to `p_out[p_idx]` — same two kernels, the
7080    /// pointers just land mid-buffer. Zero-sync (gemma confidence-adaptive draft depth).
7081    pub fn prob_of_token_device_col(
7082        &self,
7083        logits: &CudaSlice<f32>,
7084        tok_all: &CudaSlice<u32>,
7085        tok_idx: usize,
7086        p_out: &mut CudaSlice<f32>,
7087        p_idx: usize,
7088        n_vocab: usize,
7089    ) -> Result<(), Box<dyn std::error::Error>> {
7090        let tok_v = tok_all.slice(tok_idx..tok_idx + 1);
7091        let mut p_v = p_out.slice_mut(p_idx..p_idx + 1);
7092        let nb = ARGMAX_NB;
7093        let mut part = self.alloc_uninit::<f32>(nb)?;
7094        let f1 = self.func("prob_of_token_partial_f32");
7095        let cfg1 = LaunchConfig {
7096            grid_dim: (nb as u32, 1, 1),
7097            block_dim: (256, 1, 1),
7098            shared_mem_bytes: 0,
7099        };
7100        let nv = n_vocab as i32;
7101        let __s_b1 = self.gpu.stream();
7102        let mut b1 = __s_b1.launch_builder(&f1);
7103        b1.arg(logits).arg(&tok_v).arg(&mut part).arg(&nv);
7104        unsafe {
7105            b1.launch(cfg1)?;
7106        }
7107        let f2 = self.func("prob_of_token_final_f32");
7108        let cfg2 = LaunchConfig {
7109            grid_dim: (1, 1, 1),
7110            block_dim: (256, 1, 1),
7111            shared_mem_bytes: 0,
7112        };
7113        let nbi = nb as i32;
7114        let __s_b2 = self.gpu.stream();
7115        let mut b2 = __s_b2.launch_builder(&f2);
7116        b2.arg(&part).arg(&mut p_v).arg(&nbi);
7117        unsafe {
7118            b2.launch(cfg2)?;
7119        }
7120        Ok(())
7121    }
7122
7123    pub fn prob_of_token_device_into(
7124        &self,
7125        logits: &CudaSlice<f32>,
7126        tok: &CudaSlice<u32>,
7127        p_out: &mut CudaSlice<f32>,
7128        n_vocab: usize,
7129    ) -> Result<(), Box<dyn std::error::Error>> {
7130        let nb = ARGMAX_NB;
7131        let mut part = self.alloc_uninit::<f32>(nb)?;
7132        let f1 = self.func("prob_of_token_partial_f32");
7133        let cfg1 = LaunchConfig {
7134            grid_dim: (nb as u32, 1, 1),
7135            block_dim: (256, 1, 1),
7136            shared_mem_bytes: 0,
7137        };
7138        let nv = n_vocab as i32;
7139        let __s_b1 = self.gpu.stream();
7140        let mut b1 = __s_b1.launch_builder(&f1);
7141        b1.arg(logits).arg(tok).arg(&mut part).arg(&nv);
7142        unsafe {
7143            b1.launch(cfg1)?;
7144        }
7145        let f2 = self.func("prob_of_token_final_f32");
7146        let cfg2 = LaunchConfig {
7147            grid_dim: (1, 1, 1),
7148            block_dim: (256, 1, 1),
7149            shared_mem_bytes: 0,
7150        };
7151        let nbi = nb as i32;
7152        let __s_b2 = self.gpu.stream();
7153        let mut b2 = __s_b2.launch_builder(&f2);
7154        b2.arg(&part).arg(p_out).arg(&nbi);
7155        unsafe {
7156            b2.launch(cfg2)?;
7157        }
7158        Ok(())
7159    }
7160
7161    pub fn argmax_token_device(
7162        &self,
7163        logits: &CudaSlice<f32>,
7164        n_vocab: usize,
7165    ) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
7166        let mut tok = unsafe { self.gpu.stream().alloc::<u32>(1)? };
7167        self.argmax_token_device_into(logits, &mut tok, n_vocab)?;
7168        Ok(tok)
7169    }
7170    /// Like `argmax_token_device` but writes into a PERSISTENT `tok` buffer (stable pointer) instead
7171    /// of allocating a fresh one. Required for CUDA-graph capture: the captured argmax must write the
7172    /// next token into the SAME device buffer the next replay's embed_gather reads, so the buffer
7173    /// pointer is baked once and the token id never round-trips to host inside steady state. The
7174    /// pass-1 partials scratch (`argmax_partials`) is also a resident stable-pointer buffer so both
7175    /// captured passes bake fixed addresses.
7176    pub fn argmax_token_device_into(
7177        &self,
7178        logits: &CudaSlice<f32>,
7179        tok: &mut CudaSlice<u32>,
7180        n_vocab: usize,
7181    ) -> Result<(), Box<dyn std::error::Error>> {
7182        let nb = ARGMAX_NB;
7183        let f1 = self.func("argmax_partial_f32");
7184        let f2 = self.func("argmax_final_f32");
7185        let mut guard = self.argmax_partials.lock().unwrap();
7186        if guard.is_none() {
7187            // allocate ONCE; under generate_graph this runs in the tracking-off prime window so the
7188            // buffers carry no cudarc events (illegal inside capture).
7189            let pv = self.gpu.stream().alloc_zeros::<f32>(nb)?;
7190            let pi = self.gpu.stream().alloc_zeros::<i32>(nb)?;
7191            *guard = Some((pv, pi));
7192        }
7193        let (part_v, part_i) = guard.as_mut().unwrap();
7194        let nv = n_vocab as i32;
7195        let nbi = nb as i32;
7196        // pass 1: NB blocks x 256 threads grid-stride scan -> per-block (val, idx) partials.
7197        let cfg1 = LaunchConfig {
7198            grid_dim: (nb as u32, 1, 1),
7199            block_dim: (256, 1, 1),
7200            shared_mem_bytes: 0,
7201        };
7202        let __s_b1 = self.gpu.stream();
7203        let mut b1 = __s_b1.launch_builder(&f1);
7204        b1.arg(logits).arg(&mut *part_v).arg(&mut *part_i).arg(&nv);
7205        unsafe {
7206            b1.launch(cfg1)?;
7207        }
7208        // pass 2: one block reduces NB partials -> token_out[0].
7209        let cfg2 = LaunchConfig {
7210            grid_dim: (1, 1, 1),
7211            block_dim: (256, 1, 1),
7212            shared_mem_bytes: 0,
7213        };
7214        let __s_b2 = self.gpu.stream();
7215        let mut b2 = __s_b2.launch_builder(&f2);
7216        b2.arg(&*part_v).arg(&*part_i).arg(tok).arg(&nbi);
7217        unsafe {
7218            b2.launch(cfg2)?;
7219        }
7220        Ok(())
7221    }
7222    /// Column-`col` device argmax over a stacked verify-logits buffer [t, n_vocab] (spec accept
7223    /// walk): toks[out_idx] = argmax(logits[col*n_vocab .. (col+1)*n_vocab]). SAME 2-pass kernels
7224    /// and tie-break contract as `argmax_token_device_into` (bit-identical to host argmax,
7225    /// argmax_gate-validated) — only the input pointer (a column view) and the output slot differ.
7226    /// Lets the accept walk read ONE [t] u32 instead of dtoh'ing the full [t, n_vocab] logits.
7227    pub fn argmax_token_device_col(
7228        &self,
7229        logits: &CudaSlice<f32>,
7230        col: usize,
7231        n_vocab: usize,
7232        toks: &mut CudaSlice<u32>,
7233        out_idx: usize,
7234    ) -> Result<(), Box<dyn std::error::Error>> {
7235        let nb = ARGMAX_NB;
7236        let f1 = self.func("argmax_partial_f32");
7237        let f2 = self.func("argmax_final_f32");
7238        let mut guard = self.argmax_partials.lock().unwrap();
7239        if guard.is_none() {
7240            let pv = self.gpu.stream().alloc_zeros::<f32>(nb)?;
7241            let pi = self.gpu.stream().alloc_zeros::<i32>(nb)?;
7242            *guard = Some((pv, pi));
7243        }
7244        let (part_v, part_i) = guard.as_mut().unwrap();
7245        let col_view = logits.slice(col * n_vocab..(col + 1) * n_vocab);
7246        let nv = n_vocab as i32;
7247        let nbi = nb as i32;
7248        let cfg1 = LaunchConfig {
7249            grid_dim: (nb as u32, 1, 1),
7250            block_dim: (256, 1, 1),
7251            shared_mem_bytes: 0,
7252        };
7253        let __s_b1 = self.gpu.stream();
7254        let mut b1 = __s_b1.launch_builder(&f1);
7255        b1.arg(&col_view)
7256            .arg(&mut *part_v)
7257            .arg(&mut *part_i)
7258            .arg(&nv);
7259        unsafe {
7260            b1.launch(cfg1)?;
7261        }
7262        let mut tok_view = toks.slice_mut(out_idx..out_idx + 1);
7263        let cfg2 = LaunchConfig {
7264            grid_dim: (1, 1, 1),
7265            block_dim: (256, 1, 1),
7266            shared_mem_bytes: 0,
7267        };
7268        let __s_b2 = self.gpu.stream();
7269        let mut b2 = __s_b2.launch_builder(&f2);
7270        b2.arg(&*part_v).arg(&*part_i).arg(&mut tok_view).arg(&nbi);
7271        unsafe {
7272            b2.launch(cfg2)?;
7273        }
7274        Ok(())
7275    }
7276    /// Read back a device u32 buffer (the spec accept walk's [t] per-column argmax tokens).
7277    pub fn htod_u32_v(&self, v: &[u32]) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
7278        Ok(self.gpu.stream().clone_htod(v)?)
7279    }
7280    pub fn dtoh_u32(&self, d: &CudaSlice<u32>) -> Result<Vec<u32>, Box<dyn std::error::Error>> {
7281        let v = self.gpu.stream().clone_dtoh(d)?;
7282        self.gpu.stream().synchronize()?;
7283        Ok(v)
7284    }
7285    /// Allocate a zeroed device u32 buffer (persistent spec-loop prediction slots).
7286    /// H2D into an EXISTING u32 buffer (stable pointer — the per-step grammar-mask upload:
7287    /// contents change every step, the address must not, so a captured graph can read it).
7288    pub fn htod_u32_into(
7289        &self,
7290        dst: &mut CudaSlice<u32>,
7291        src: &[u32],
7292    ) -> Result<(), Box<dyn std::error::Error>> {
7293        let mut view = dst.slice_mut(0..src.len());
7294        self.gpu.stream().memcpy_htod(src, &mut view)?;
7295        Ok(())
7296    }
7297
7298    /// H2D into an existing i32 buffer. OPTIPIPE uses this to refresh a stage-local saved-len
7299    /// table without changing the device address its reconcile kernel consumes.
7300    pub fn htod_i32_into(
7301        &self,
7302        dst: &mut CudaSlice<i32>,
7303        src: &[i32],
7304    ) -> Result<(), Box<dyn std::error::Error>> {
7305        let mut view = dst.slice_mut(0..src.len());
7306        self.gpu.stream().memcpy_htod(src, &mut view)?;
7307        Ok(())
7308    }
7309
7310    pub fn alloc_u32_zeroed(&self, n: usize) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
7311        let s = self.gpu.stream().alloc_zeros::<u32>(n)?;
7312        self.keep_if_capturing(&s);
7313        Ok(s)
7314    }
7315    /// embed_gather into a PERSISTENT `x_out` buffer (stable pointer) for CUDA-graph capture (the
7316    /// embed output starts the per-step kernel chain and must be at a fixed address across replays).
7317    pub fn embed_gather_device_into(
7318        &self,
7319        embd: &CudaSlice<u8>,
7320        token_d: &CudaSlice<u32>,
7321        x_out: &mut CudaSlice<f32>,
7322        n_embd: usize,
7323        qtype: i32,
7324        row_bytes: usize,
7325    ) -> Result<(), Box<dyn std::error::Error>> {
7326        let f = self.func("embed_gather_u32");
7327        let cfg = LaunchConfig {
7328            grid_dim: (((n_embd as u32 + 255) / 256).max(1), 1, 1),
7329            block_dim: (256, 1, 1),
7330            shared_mem_bytes: 0,
7331        };
7332        let (ne, qt, rb) = (n_embd as i32, qtype, row_bytes as i64);
7333        let __s_b = self.gpu.stream();
7334        let mut b = __s_b.launch_builder(&f);
7335        b.arg(embd)
7336            .arg(token_d)
7337            .arg(x_out)
7338            .arg(&ne)
7339            .arg(&qt)
7340            .arg(&rb);
7341        unsafe {
7342            b.launch(cfg)?;
7343        }
7344        Ok(())
7345    }
7346    /// Read a [1] i32 device counter (pos / seqlen) back to host. Tiny D2H + sync.
7347    pub fn dtoh_i32_one(&self, d: &CudaSlice<i32>) -> Result<i32, Box<dyn std::error::Error>> {
7348        let v = self.gpu.stream().clone_dtoh(d)?;
7349        self.gpu.stream().synchronize()?;
7350        Ok(v[0])
7351    }
7352    /// Set a [1] i32 device counter IN PLACE (keeps the buffer pointer stable — required for the
7353    /// graph-resident pos/seqlen counters whose addresses are baked into captured graphs). Restores
7354    /// the counter value after the throwaway capture warmups corrupt it.
7355    /// ASYNC i32 single-slot store (value rides the kernel arg — no host-memory transfer/sync).
7356    /// The graph-arc device-len counters use this; set_i32_one below is the SYNCING pageable
7357    /// copy (fine at stream-idle boundaries, poison mid-round).
7358    pub fn i32_set_k(
7359        &self,
7360        dst: &mut CudaSlice<i32>,
7361        v: i32,
7362    ) -> Result<(), Box<dyn std::error::Error>> {
7363        let f = self.func("i32_set_k");
7364        let cfg = LaunchConfig {
7365            grid_dim: (1, 1, 1),
7366            block_dim: (1, 1, 1),
7367            shared_mem_bytes: 0,
7368        };
7369        let idx = 0i32;
7370        let __s_b = self.gpu.stream();
7371        let mut b = __s_b.launch_builder(&f);
7372        b.arg(dst).arg(&v).arg(&idx);
7373        unsafe {
7374            b.launch(cfg)?;
7375        }
7376        Ok(())
7377    }
7378
7379    pub fn set_i32_one(
7380        &self,
7381        d: &mut CudaSlice<i32>,
7382        v: i32,
7383    ) -> Result<(), Box<dyn std::error::Error>> {
7384        self.gpu.stream().memcpy_htod(&[v], d)?;
7385        Ok(())
7386    }
7387    /// Set a [1] u32 device buffer IN PLACE (stable pointer) — for the resident `token_d` counter
7388    /// during priming / capture-state restore.
7389    pub fn set_u32_one(
7390        &self,
7391        d: &mut CudaSlice<u32>,
7392        v: u32,
7393    ) -> Result<(), Box<dyn std::error::Error>> {
7394        self.gpu.stream().memcpy_htod(&[v], d)?;
7395        Ok(())
7396    }
7397    /// Read back a [1] u32 device buffer (the argmax token). One tiny D2H + sync.
7398    pub fn dtoh_u32_one(&self, d: &CudaSlice<u32>) -> Result<u32, Box<dyn std::error::Error>> {
7399        let v = self.gpu.stream().clone_dtoh(d)?;
7400        self.gpu.stream().synchronize()?;
7401        Ok(v[0])
7402    }
7403    /// Upload raw bytes to a resident device u8 buffer (e.g. the embed table for device gather).
7404    pub fn upload_u8(&self, bytes: &[u8]) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
7405        Ok(self.gpu.stream().clone_htod(bytes)?)
7406    }
7407    /// Embed-from-device (CUDA-GRAPH-PLAN Phase 1): gather+dequant the row for the token id in
7408    /// `token_d[0]` from the resident embed table `embd` -> x_out[n_embd]. Bit-identical to host
7409    /// EmbedHost::gather (same per-dtype `deq`). No host round-trip of the token id.
7410    pub fn embed_gather_device(
7411        &self,
7412        embd: &CudaSlice<u8>,
7413        token_d: &CudaSlice<u32>,
7414        n_embd: usize,
7415        qtype: i32,
7416        row_bytes: usize,
7417    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7418        let f = self.func("embed_gather_u32");
7419        let mut x = self.alloc_uninit::<f32>(n_embd)?;
7420        let cfg = LaunchConfig {
7421            grid_dim: (((n_embd as u32 + 255) / 256).max(1), 1, 1),
7422            block_dim: (256, 1, 1),
7423            shared_mem_bytes: 0,
7424        };
7425        let (ne, qt, rb) = (n_embd as i32, qtype, row_bytes as i64);
7426        let __s_b = self.gpu.stream();
7427        let mut b = __s_b.launch_builder(&f);
7428        b.arg(embd)
7429            .arg(token_d)
7430            .arg(&mut x)
7431            .arg(&ne)
7432            .arg(&qt)
7433            .arg(&rb);
7434        unsafe {
7435            b.launch(cfg)?;
7436        }
7437        Ok(x)
7438    }
7439
7440    /// T-token device embed gather (spec verify/replay): tokens uploaded as a tiny [T] u32 htod,
7441    /// rows dequanted on-device -> x[T, n_embd]. Replaces host per-row dequant + T*n_embd*4B htod
7442    /// (nsys: 84% of spec API time was HtoD). Bit-identical rows (same per-dtype deq).
7443    pub fn embed_gather_device_t(
7444        &self,
7445        embd: &CudaSlice<u8>,
7446        tokens: &[u32],
7447        n_embd: usize,
7448        qtype: i32,
7449        row_bytes: usize,
7450    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7451        let t = tokens.len();
7452        let tok_d = self.gpu.stream().clone_htod(tokens)?;
7453        let f = self.func("embed_gather_u32_t");
7454        let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
7455        let cfg = LaunchConfig {
7456            grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
7457            block_dim: (256, 1, 1),
7458            shared_mem_bytes: 0,
7459        };
7460        let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
7461        let __s_b = self.gpu.stream();
7462        let mut b = __s_b.launch_builder(&f);
7463        b.arg(embd)
7464            .arg(&tok_d)
7465            .arg(&mut x)
7466            .arg(&ne)
7467            .arg(&qt)
7468            .arg(&rb)
7469            .arg(&ti);
7470        unsafe {
7471            b.launch(cfg)?;
7472        }
7473        Ok(x)
7474    }
7475
7476    /// T-token embed gather from a DEVICE token buffer (round-stream stage c: the verify tokens
7477    /// are assembled on-device from the draft-chain pack slots; no host round trip). Same kernel
7478    /// as embed_gather_device_t — bit-identical rows.
7479    /// embed_gather over a token VIEW (spec round: tokens live in the round's batch buffer).
7480    pub fn embed_gather_device_tv(
7481        &self,
7482        embd: &CudaSlice<u8>,
7483        tok_v: &cudarc::driver::CudaView<u32>,
7484        t: usize,
7485        n_embd: usize,
7486        qtype: i32,
7487        row_bytes: usize,
7488    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7489        let f = self.func("embed_gather_u32_t");
7490        let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
7491        let cfg = LaunchConfig {
7492            grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
7493            block_dim: (256, 1, 1),
7494            shared_mem_bytes: 0,
7495        };
7496        let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
7497        let __s_b = self.gpu.stream();
7498        let mut b = __s_b.launch_builder(&f);
7499        b.arg(embd)
7500            .arg(tok_v)
7501            .arg(&mut x)
7502            .arg(&ne)
7503            .arg(&qt)
7504            .arg(&rb)
7505            .arg(&ti);
7506        unsafe {
7507            b.launch(cfg)?;
7508        }
7509        Ok(x)
7510    }
7511
7512    pub fn embed_gather_device_td(
7513        &self,
7514        embd: &CudaSlice<u8>,
7515        tok_d: &CudaSlice<u32>,
7516        t: usize,
7517        n_embd: usize,
7518        qtype: i32,
7519        row_bytes: usize,
7520    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7521        let f = self.func("embed_gather_u32_t");
7522        let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
7523        let cfg = LaunchConfig {
7524            grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
7525            block_dim: (256, 1, 1),
7526            shared_mem_bytes: 0,
7527        };
7528        let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
7529        let __s_b = self.gpu.stream();
7530        let mut b = __s_b.launch_builder(&f);
7531        b.arg(embd)
7532            .arg(tok_d)
7533            .arg(&mut x)
7534            .arg(&ne)
7535            .arg(&qt)
7536            .arg(&rb)
7537            .arg(&ti);
7538        unsafe {
7539            b.launch(cfg)?;
7540        }
7541        Ok(x)
7542    }
7543
7544    /// Uninitialized device buffer — SKIPS the memset that `alloc_zeros` always issues. Decode
7545    /// profile (nsys): ~1050 memsets/token = 6.5% of decode GPU time + ~half the launch count, the
7546    /// dominant contributor to the 19% inter-kernel idle gap and a blocker for clean CUDA-graph
7547    /// capture. Use ONLY for buffers a kernel FULLY overwrites (every element written, no `+=`).
7548    /// SAFETY: caller guarantees the producing kernel writes every element before any read.
7549    #[inline]
7550    /// Keep an allocation alive for the current capture (no-op when retain mode is off).
7551    fn keep_if_capturing<T: cudarc::driver::DeviceRepr + Send + 'static>(&self, s: &CudaSlice<T>) {
7552        if self
7553            .capture_keep_on
7554            .load(std::sync::atomic::Ordering::Relaxed)
7555        {
7556            self.capture_keep.lock().unwrap().push(Box::new(s.clone()));
7557        }
7558    }
7559
7560    fn alloc_uninit<T: cudarc::driver::DeviceRepr + Send + 'static>(
7561        &self,
7562        n: usize,
7563    ) -> Result<CudaSlice<T>, Box<dyn std::error::Error>> {
7564        let mut s = unsafe { self.gpu.stream().alloc::<T>(n)? };
7565        // MEMRA_DEBUG_ZERO_ALLOCS=1 (task #14 defect hunt): memset EVERY engine allocation —
7566        // the global uninit-read discriminator (the prime-fn-scoped zeroing experiment could
7567        // not cover engine-internal buffers). Debug-only: massive launch overhead.
7568        {
7569            static Z: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7570            if *Z.get_or_init(|| std::env::var("MEMRA_DEBUG_ZERO_ALLOCS").as_deref() == Ok("1")) {
7571                // raw D8 memset (T lacks ValidAsZeroBits in the generic bound)
7572                use cudarc::driver::DevicePtrMut;
7573                let n_bytes = s.len() * std::mem::size_of::<T>();
7574                let stream = self.gpu.stream();
7575                let (p_, _g) = s.device_ptr_mut(&stream);
7576                unsafe {
7577                    cudarc::driver::sys::cuMemsetD8Async(p_, 0, n_bytes, stream.cu_stream())
7578                        .result()?;
7579                }
7580            }
7581        }
7582        self.keep_if_capturing(&s);
7583        Ok(s)
7584    }
7585
7586    /// Public f32 uninitialized scratch (see `alloc_uninit`). For decode/forward scratch a kernel
7587    /// fully overwrites. SAFETY: producing kernel must write every element before any read.
7588    /// Uninitialized q8_1 activation pair (int8 + per-32 scales) — the fa combine q8-emit
7589    /// consumers alloc through this (m=1 decode arms).
7590    pub fn uninit_q8_pair(
7591        &self,
7592        n: usize,
7593    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7594        Ok((
7595            self.alloc_uninit::<i8>(n)?,
7596            self.alloc_uninit::<f32>(n / 32)?,
7597        ))
7598    }
7599
7600    pub fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7601        self.alloc_uninit::<f32>(n)
7602    }
7603
7604    /// i8 uninitialized scratch (same contract as `uninit`).
7605    pub fn alloc_i8_uninit(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
7606        self.alloc_uninit::<i8>(n)
7607    }
7608
7609    /// RMSNorm: x[ncols,nrows] row-major, weight[ncols] -> dst. One block/row, 256 threads.
7610    /// gemma4: 3 rms_norms of the SAME input in one launch (one reduction, three weights).
7611    /// Per-output bit-identical to three rms_norm calls (verbatim reduction/scale chain).
7612    #[allow(clippy::too_many_arguments)]
7613    pub fn rms_norm3(
7614        &self,
7615        x: &CudaSlice<f32>,
7616        w0: &CudaSlice<f32>,
7617        w1: &CudaSlice<f32>,
7618        w2: &CudaSlice<f32>,
7619        d0: &mut CudaSlice<f32>,
7620        d1: &mut CudaSlice<f32>,
7621        d2: &mut CudaSlice<f32>,
7622        ncols: usize,
7623        nrows: usize,
7624        eps: f32,
7625    ) -> Result<(), Box<dyn std::error::Error>> {
7626        let f = self.func("rms_norm3_f32");
7627        let cfg = LaunchConfig {
7628            grid_dim: (nrows as u32, 1, 1),
7629            block_dim: (rms_block(), 1, 1),
7630            shared_mem_bytes: 0,
7631        };
7632        let (nc, e) = (ncols as i32, eps);
7633        let __s_b = self.gpu.stream();
7634        let mut b = __s_b.launch_builder(&f);
7635        b.arg(x)
7636            .arg(w0)
7637            .arg(w1)
7638            .arg(w2)
7639            .arg(d0)
7640            .arg(d1)
7641            .arg(d2)
7642            .arg(&nc)
7643            .arg(&e);
7644        unsafe {
7645            b.launch(cfg)?;
7646        }
7647        Ok(())
7648    }
7649
7650    /// gemma4 fused q/k/v head norms (one launch, per-row rms_norm_f32-verbatim).
7651    #[allow(clippy::too_many_arguments)]
7652    /// True when the warp-per-row qkv norm would engage for (rows, ncols) — the emit lane
7653    /// piggybacks on the same conditions.
7654    pub fn qkvnorm_w_on_prefill(rows: usize, ncols: usize) -> bool {
7655        static WARP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7656        *WARP_ON.get_or_init(|| {
7657            std::env::var("MEMRA_QKVNORM_W")
7658                .map(|v| v != "0")
7659                .unwrap_or(true)
7660        }) && ncols % 4 == 0
7661            && rows >= 64
7662    }
7663
7664    /// w4 norm with bf16 V EMIT (31B glue lane): the v segment also writes its normed rows as
7665    /// bf16 (the FA V operand — bit-identical to a post-hoc f32_to_bf16). Prefill-depth only.
7666    #[allow(clippy::too_many_arguments)]
7667    pub fn rms_norm_qkv_w4b(
7668        &self,
7669        q: &CudaSlice<f32>,
7670        k: &CudaSlice<f32>,
7671        v: &CudaSlice<f32>,
7672        wq: &CudaSlice<f32>,
7673        wk: &CudaSlice<f32>,
7674        wv: &CudaSlice<f32>,
7675        dq: &mut CudaSlice<f32>,
7676        dk: &mut CudaSlice<f32>,
7677        dv: &mut CudaSlice<f32>,
7678        dvb: &mut CudaSlice<u8>,
7679        ncols: usize,
7680        rq: usize,
7681        rk: usize,
7682        eps: f32,
7683        vf16: bool,
7684    ) -> Result<(), Box<dyn std::error::Error>> {
7685        assert!(ncols % 4 == 0 && rq + 2 * rk >= 64);
7686        let f = self.func("rms_norm_qkv_w4b_f32");
7687        let rows = (rq + 2 * rk) as u32;
7688        let cfg = LaunchConfig {
7689            grid_dim: (rows.div_ceil(8), 1, 1),
7690            block_dim: (256, 1, 1),
7691            shared_mem_bytes: 0,
7692        };
7693        let (nc, rqi, rki, rvi, e) = (ncols as i32, rq as i32, rk as i32, rk as i32, eps);
7694        let vf = vf16 as i32;
7695        let __s_b = self.gpu.stream();
7696        let mut b = __s_b.launch_builder(&f);
7697        b.arg(q)
7698            .arg(k)
7699            .arg(v)
7700            .arg(wq)
7701            .arg(wk)
7702            .arg(wv)
7703            .arg(dq)
7704            .arg(dk)
7705            .arg(dv)
7706            .arg(&mut *dvb)
7707            .arg(&nc)
7708            .arg(&rqi)
7709            .arg(&rki)
7710            .arg(&rvi)
7711            .arg(&e)
7712            .arg(&vf);
7713        unsafe {
7714            b.launch(cfg)?;
7715        }
7716        Ok(())
7717    }
7718
7719    pub fn rms_norm_qkv(
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        ncols: usize,
7731        rq: usize,
7732        rk: usize,
7733        eps: f32,
7734    ) -> Result<(), Box<dyn std::error::Error>> {
7735        // Warp-per-row float4 twin (default; MEMRA_QKVNORM_W=0 reverts): the block-per-row form
7736        // spends 767us/launch on 17k+ 2KB rows at prefill depth (launch/reduce latency-bound,
7737        // ~92GB/s). Own numeric config (reduce order differs) — battery-gated.
7738        static WARP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7739        let warp_on = *WARP_ON.get_or_init(|| {
7740            std::env::var("MEMRA_QKVNORM_W")
7741                .map(|v| v != "0")
7742                .unwrap_or(true)
7743        });
7744        // rows >= 64 keeps decode (nh + 2*nkv rows) on the block-tree kernel — decode/verify/
7745        // replay numerics are untouched on every model; only prefill depth takes the new config.
7746        if warp_on && ncols % 4 == 0 && rq + 2 * rk >= 64 {
7747            let f = self.func("rms_norm_qkv_w4_f32");
7748            let rows = (rq + 2 * rk) as u32;
7749            let cfg = LaunchConfig {
7750                grid_dim: (rows.div_ceil(8), 1, 1),
7751                block_dim: (256, 1, 1),
7752                shared_mem_bytes: 0,
7753            };
7754            let (nc, rqi, rki, rvi, e) = (ncols as i32, rq as i32, rk as i32, rk as i32, eps);
7755            let __s_b = self.gpu.stream();
7756            let mut b = __s_b.launch_builder(&f);
7757            b.arg(q)
7758                .arg(k)
7759                .arg(v)
7760                .arg(wq)
7761                .arg(wk)
7762                .arg(wv)
7763                .arg(dq)
7764                .arg(dk)
7765                .arg(dv)
7766                .arg(&nc)
7767                .arg(&rqi)
7768                .arg(&rki)
7769                .arg(&rvi)
7770                .arg(&e);
7771            unsafe {
7772                b.launch(cfg)?;
7773            }
7774            return Ok(());
7775        }
7776        let f = self.func("rms_norm_qkv_f32");
7777        let grid = (rq + 2 * rk) as u32;
7778        let cfg = LaunchConfig {
7779            grid_dim: (grid, 1, 1),
7780            block_dim: (rms_block(), 1, 1),
7781            shared_mem_bytes: 0,
7782        };
7783        let (nc, rqi, rki, e) = (ncols as i32, rq as i32, rk as i32, eps);
7784        let __s_b = self.gpu.stream();
7785        let mut b = __s_b.launch_builder(&f);
7786        b.arg(q)
7787            .arg(k)
7788            .arg(v)
7789            .arg(wq)
7790            .arg(wk)
7791            .arg(wv)
7792            .arg(dq)
7793            .arg(dk)
7794            .arg(dv)
7795            .arg(&nc)
7796            .arg(&rqi)
7797            .arg(&rki)
7798            .arg(&e);
7799        unsafe {
7800            b.launch(cfg)?;
7801        }
7802        Ok(())
7803    }
7804
7805    /// gemma4 fused pair of rms_norms over two different inputs (same width).
7806    #[allow(clippy::too_many_arguments)]
7807    pub fn rms_norm2x(
7808        &self,
7809        a: &CudaSlice<f32>,
7810        bb: &CudaSlice<f32>,
7811        wa: &CudaSlice<f32>,
7812        wb: &CudaSlice<f32>,
7813        da: &mut CudaSlice<f32>,
7814        db: &mut CudaSlice<f32>,
7815        ncols: usize,
7816        nrows: usize,
7817        eps: f32,
7818    ) -> Result<(), Box<dyn std::error::Error>> {
7819        let f = self.func("rms_norm2x_f32");
7820        let cfg = LaunchConfig {
7821            grid_dim: (2 * nrows as u32, 1, 1),
7822            block_dim: (rms_block(), 1, 1),
7823            shared_mem_bytes: 0,
7824        };
7825        let (nc, nr, e) = (ncols as i32, nrows as i32, eps);
7826        let __s_b = self.gpu.stream();
7827        let mut b = __s_b.launch_builder(&f);
7828        b.arg(a)
7829            .arg(bb)
7830            .arg(wa)
7831            .arg(wb)
7832            .arg(da)
7833            .arg(db)
7834            .arg(&nc)
7835            .arg(&nr)
7836            .arg(&e);
7837        unsafe {
7838            b.launch(cfg)?;
7839        }
7840        Ok(())
7841    }
7842
7843    /// gemma4 R4: in-place final-logit softcap y = cap*tanh(y/cap).
7844    pub fn softcap(
7845        &self,
7846        y: &mut CudaSlice<f32>,
7847        cap: f32,
7848        n: usize,
7849    ) -> Result<(), Box<dyn std::error::Error>> {
7850        let f = self.func("softcap_f32");
7851        let cfg = LaunchConfig::for_num_elems(n as u32);
7852        let ni = n as i32;
7853        let __s_b = self.gpu.stream();
7854        let mut b = __s_b.launch_builder(&f);
7855        b.arg(y).arg(&cap).arg(&ni);
7856        unsafe {
7857            b.launch(cfg)?;
7858        }
7859        Ok(())
7860    }
7861
7862    /// gemma4 suppress-token mask: y[row][ids[j]] = -inf over t logits rows (fixed-arg launch —
7863    /// graph-capture safe; NOT monotonic like softcap, so it must run before any argmax).
7864    pub fn mask_ids_rows(
7865        &self,
7866        y: &mut CudaSlice<f32>,
7867        ids: &CudaSlice<i32>,
7868        n_ids: usize,
7869        n_vocab: usize,
7870        t: usize,
7871    ) -> Result<(), Box<dyn std::error::Error>> {
7872        let f = self.func("mask_ids_rows_f32");
7873        let cfg = LaunchConfig::for_num_elems((n_ids * t) as u32);
7874        let (ni, nv, ti) = (n_ids as i32, n_vocab as i32, t as i32);
7875        let __s_b = self.gpu.stream();
7876        let mut b = __s_b.launch_builder(&f);
7877        b.arg(y).arg(ids).arg(&ni).arg(&nv).arg(&ti);
7878        unsafe {
7879            b.launch(cfg)?;
7880        }
7881        Ok(())
7882    }
7883
7884    /// gemma4: res = (a+b)*c AND dst = rms_norm(res, w) in one launch.
7885    #[allow(clippy::too_many_arguments)]
7886    pub fn add_scale_rms_norm(
7887        &self,
7888        a: &CudaSlice<f32>,
7889        b_in: &CudaSlice<f32>,
7890        c: f32,
7891        w: &CudaSlice<f32>,
7892        res: &mut CudaSlice<f32>,
7893        dst: &mut CudaSlice<f32>,
7894        ncols: usize,
7895        nrows: usize,
7896        eps: f32,
7897    ) -> Result<(), Box<dyn std::error::Error>> {
7898        let f = self.func("add_scale_rms_norm_f32");
7899        let cfg = LaunchConfig {
7900            grid_dim: (nrows as u32, 1, 1),
7901            block_dim: (rms_block(), 1, 1),
7902            shared_mem_bytes: 0,
7903        };
7904        let (nc, e2) = (ncols as i32, eps);
7905        let __s_b = self.gpu.stream();
7906        let mut b = __s_b.launch_builder(&f);
7907        b.arg(a)
7908            .arg(b_in)
7909            .arg(&c)
7910            .arg(w)
7911            .arg(res)
7912            .arg(dst)
7913            .arg(&nc)
7914            .arg(&e2);
7915        unsafe {
7916            b.launch(cfg)?;
7917        }
7918        Ok(())
7919    }
7920
7921    /// gemma4: res = (a+b)*c AND the next layer's attn_norm EMITTED q8_1 in one launch.
7922    /// Quantize epilogue bit-identical to quantize_q8_1 (the rms_norm_q8_1 form).
7923    #[allow(clippy::too_many_arguments)]
7924    pub fn add_scale_rms_norm_q8_1(
7925        &self,
7926        a: &CudaSlice<f32>,
7927        b_in: &CudaSlice<f32>,
7928        c: f32,
7929        w: &CudaSlice<f32>,
7930        res: &mut CudaSlice<f32>,
7931        ncols: usize,
7932        nrows: usize,
7933        eps: f32,
7934    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7935        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
7936        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
7937        let (nc, e2) = (ncols as i32, eps);
7938        if Self::pdl_on() && Self::pdl_wb_on() {
7939            {
7940                use cudarc::driver::{DevicePtr, DevicePtrMut};
7941                let s = &self.gpu.stream();
7942                let (pa, _g0) = a.device_ptr(s);
7943                let (pb, _g1) = b_in.device_ptr(s);
7944                let (pw, _g2) = w.device_ptr(s);
7945                let (pr, _g3) = res.device_ptr_mut(s);
7946                let (pq, _g4) = out_q.device_ptr_mut(s);
7947                let (pd, _g5) = out_d.device_ptr_mut(s);
7948                let mut ps = [
7949                    &pa as *const _ as *mut std::ffi::c_void,
7950                    &pb as *const _ as *mut _,
7951                    &c as *const _ as *mut _,
7952                    &pw as *const _ as *mut _,
7953                    &pr as *const _ as *mut _,
7954                    &pq as *const _ as *mut _,
7955                    &pd as *const _ as *mut _,
7956                    &nc as *const _ as *mut _,
7957                    &e2 as *const _ as *mut _,
7958                ];
7959                unsafe {
7960                    self.launch_pdl(
7961                        "add_scale_rms_norm_q8_1",
7962                        (nrows as u32, 1, 1),
7963                        (rms_block(), 1, 1),
7964                        &mut ps,
7965                    )?;
7966                }
7967            }
7968            return Ok((out_q, out_d));
7969        }
7970        let f = self.func("add_scale_rms_norm_q8_1");
7971        let cfg = LaunchConfig {
7972            grid_dim: (nrows as u32, 1, 1),
7973            block_dim: (rms_block(), 1, 1),
7974            shared_mem_bytes: 0,
7975        };
7976        let __s_b = self.gpu.stream();
7977        let mut b = __s_b.launch_builder(&f);
7978        b.arg(a)
7979            .arg(b_in)
7980            .arg(&c)
7981            .arg(w)
7982            .arg(res)
7983            .arg(&mut out_q)
7984            .arg(&mut out_d)
7985            .arg(&nc)
7986            .arg(&e2);
7987        unsafe {
7988            b.launch(cfg)?;
7989        }
7990        Ok((out_q, out_d))
7991    }
7992
7993    /// Slot-fed add_scale_rms_norm_q8_1 twin (alloc-free capture lane).
7994    #[allow(clippy::too_many_arguments)]
7995    pub fn add_scale_rms_norm_q8_1_into(
7996        &self,
7997        a: &CudaSlice<f32>,
7998        b_in: &CudaSlice<f32>,
7999        c: f32,
8000        w: &CudaSlice<f32>,
8001        res: &mut CudaSlice<f32>,
8002        ncols: usize,
8003        nrows: usize,
8004        eps: f32,
8005        out_q: &mut CudaSlice<i8>,
8006        out_d: &mut CudaSlice<f32>,
8007    ) -> Result<(), Box<dyn std::error::Error>> {
8008        debug_assert!(out_q.len() >= nrows * ncols && out_d.len() >= nrows * (ncols / 32));
8009        let (nc, e2) = (ncols as i32, eps);
8010        if Self::pdl_on() && Self::pdl_wb_on() {
8011            use cudarc::driver::{DevicePtr, DevicePtrMut};
8012            let s = &self.gpu.stream();
8013            let (pa, _g0) = a.device_ptr(s);
8014            let (pb, _g1) = b_in.device_ptr(s);
8015            let (pw, _g2) = w.device_ptr(s);
8016            let (pr, _g3) = res.device_ptr_mut(s);
8017            let (pq, _g4) = out_q.device_ptr_mut(s);
8018            let (pd, _g5) = out_d.device_ptr_mut(s);
8019            let mut ps = [
8020                &pa as *const _ as *mut std::ffi::c_void,
8021                &pb as *const _ as *mut _,
8022                &c as *const _ as *mut _,
8023                &pw as *const _ as *mut _,
8024                &pr as *const _ as *mut _,
8025                &pq as *const _ as *mut _,
8026                &pd as *const _ as *mut _,
8027                &nc as *const _ as *mut _,
8028                &e2 as *const _ as *mut _,
8029            ];
8030            unsafe {
8031                self.launch_pdl(
8032                    "add_scale_rms_norm_q8_1",
8033                    (nrows as u32, 1, 1),
8034                    (rms_block(), 1, 1),
8035                    &mut ps,
8036                )?;
8037            }
8038            return Ok(());
8039        }
8040        let f = self.func("add_scale_rms_norm_q8_1");
8041        let cfg = LaunchConfig {
8042            grid_dim: (nrows as u32, 1, 1),
8043            block_dim: (rms_block(), 1, 1),
8044            shared_mem_bytes: 0,
8045        };
8046        let __s_b = self.gpu.stream();
8047        let mut b = __s_b.launch_builder(&f);
8048        b.arg(a)
8049            .arg(b_in)
8050            .arg(&c)
8051            .arg(w)
8052            .arg(res)
8053            .arg(&mut *out_q)
8054            .arg(&mut *out_d)
8055            .arg(&nc)
8056            .arg(&e2);
8057        unsafe {
8058            b.launch(cfg)?;
8059        }
8060        Ok(())
8061    }
8062
8063    /// E4B glue fusion: rms(a, wa) prologue + the add_scale_rms_norm_q8_1 program — one launch
8064    /// replaces the per-layer rms_norm_f32(y) + emit pair in the PLE tail.
8065    #[allow(clippy::too_many_arguments)]
8066    pub fn rms_pre_add_scale_rms_norm_q8_1(
8067        &self,
8068        a: &CudaSlice<f32>,
8069        wa: &CudaSlice<f32>,
8070        b_in: &CudaSlice<f32>,
8071        c: f32,
8072        w: &CudaSlice<f32>,
8073        res: &mut CudaSlice<f32>,
8074        ncols: usize,
8075        nrows: usize,
8076        eps: f32,
8077    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8078        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
8079        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
8080        let (nc, e2) = (ncols as i32, eps);
8081        if Self::pdl_on() {
8082            {
8083                use cudarc::driver::{DevicePtr, DevicePtrMut};
8084                let s = &self.gpu.stream();
8085                let (pa, _g0) = a.device_ptr(s);
8086                let (pwa, _g1) = wa.device_ptr(s);
8087                let (pb, _g2) = b_in.device_ptr(s);
8088                let (pw, _g3) = w.device_ptr(s);
8089                let (pr, _g4) = res.device_ptr_mut(s);
8090                let (pq, _g5) = out_q.device_ptr_mut(s);
8091                let (pd, _g6) = out_d.device_ptr_mut(s);
8092                let mut ps = [
8093                    &pa as *const _ as *mut std::ffi::c_void,
8094                    &pwa as *const _ as *mut _,
8095                    &pb as *const _ as *mut _,
8096                    &c as *const _ as *mut _,
8097                    &pw as *const _ as *mut _,
8098                    &pr as *const _ as *mut _,
8099                    &pq as *const _ as *mut _,
8100                    &pd as *const _ as *mut _,
8101                    &nc as *const _ as *mut _,
8102                    &e2 as *const _ as *mut _,
8103                ];
8104                unsafe {
8105                    self.launch_pdl(
8106                        "rms_pre_add_scale_rms_norm_q8_1",
8107                        (nrows as u32, 1, 1),
8108                        (rms_block(), 1, 1),
8109                        &mut ps,
8110                    )?;
8111                }
8112            }
8113            return Ok((out_q, out_d));
8114        }
8115        let f = self.func("rms_pre_add_scale_rms_norm_q8_1");
8116        let cfg = LaunchConfig {
8117            grid_dim: (nrows as u32, 1, 1),
8118            block_dim: (rms_block(), 1, 1),
8119            shared_mem_bytes: 0,
8120        };
8121        let __s_b = self.gpu.stream();
8122        let mut b = __s_b.launch_builder(&f);
8123        b.arg(a)
8124            .arg(wa)
8125            .arg(b_in)
8126            .arg(&c)
8127            .arg(w)
8128            .arg(res)
8129            .arg(&mut out_q)
8130            .arg(&mut out_d)
8131            .arg(&nc)
8132            .arg(&e2);
8133        unsafe {
8134            b.launch(cfg)?;
8135        }
8136        Ok((out_q, out_d))
8137    }
8138
8139    /// GELU(tanh)*up with the activation emitted q8_1 alongside f32 (glue-fusion lane): the
8140    /// consumer matmul rides matmul_pre, killing its standalone quantize_q8_1 launch.
8141    pub fn gelu_tanh_mul_q8_1(
8142        &self,
8143        gate: &CudaSlice<f32>,
8144        up: &cudarc::driver::CudaView<f32>,
8145        act: &mut CudaSlice<f32>,
8146        ncols: usize,
8147        nrows: usize,
8148    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8149        debug_assert!(ncols % 128 == 0);
8150        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
8151        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
8152        let nc = ncols as i32;
8153        if Self::pdl_on() {
8154            {
8155                use cudarc::driver::{DevicePtr, DevicePtrMut};
8156                let s = &self.gpu.stream();
8157                let (pg, _g0) = gate.device_ptr(s);
8158                let (pu, _g1) = up.device_ptr(s);
8159                let (pact, _g2) = act.device_ptr_mut(s);
8160                let (pq, _g3) = out_q.device_ptr_mut(s);
8161                let (pd, _g4) = out_d.device_ptr_mut(s);
8162                let mut ps = [
8163                    &pg as *const _ as *mut std::ffi::c_void,
8164                    &pu as *const _ as *mut _,
8165                    &pact as *const _ as *mut _,
8166                    &pq as *const _ as *mut _,
8167                    &pd as *const _ as *mut _,
8168                    &nc as *const _ as *mut _,
8169                ];
8170                unsafe {
8171                    self.launch_pdl(
8172                        "gelu_tanh_mul_q8_1",
8173                        (nrows as u32, 1, 1),
8174                        (rms_block(), 1, 1),
8175                        &mut ps,
8176                    )?;
8177                }
8178            }
8179            return Ok((out_q, out_d));
8180        }
8181        let f = self.func("gelu_tanh_mul_q8_1");
8182        let cfg = LaunchConfig {
8183            grid_dim: (nrows as u32, 1, 1),
8184            block_dim: (rms_block(), 1, 1),
8185            shared_mem_bytes: 0,
8186        };
8187        let __s_b = self.gpu.stream();
8188        let mut b = __s_b.launch_builder(&f);
8189        b.arg(gate)
8190            .arg(up)
8191            .arg(act)
8192            .arg(&mut out_q)
8193            .arg(&mut out_d)
8194            .arg(&nc);
8195        unsafe {
8196            b.launch(cfg)?;
8197        }
8198        Ok((out_q, out_d))
8199    }
8200
8201    /// Slot-fed gelu_tanh_mul_q8_1 twin (alloc-free capture lane; incl. the PDL arm).
8202    #[allow(clippy::too_many_arguments)]
8203    pub fn gelu_tanh_mul_q8_1_into(
8204        &self,
8205        gate: &CudaSlice<f32>,
8206        up: &cudarc::driver::CudaView<f32>,
8207        act: &mut CudaSlice<f32>,
8208        ncols: usize,
8209        nrows: usize,
8210        out_q: &mut CudaSlice<i8>,
8211        out_d: &mut CudaSlice<f32>,
8212    ) -> Result<(), Box<dyn std::error::Error>> {
8213        debug_assert!(ncols % 128 == 0);
8214        debug_assert!(out_q.len() >= nrows * ncols && out_d.len() >= nrows * (ncols / 32));
8215        let nc = ncols as i32;
8216        if Self::pdl_on() {
8217            use cudarc::driver::{DevicePtr, DevicePtrMut};
8218            let s = &self.gpu.stream();
8219            let (pg, _g0) = gate.device_ptr(s);
8220            let (pu, _g1) = up.device_ptr(s);
8221            let (pact, _g2) = act.device_ptr_mut(s);
8222            let (pq, _g3) = out_q.device_ptr_mut(s);
8223            let (pd, _g4) = out_d.device_ptr_mut(s);
8224            let mut ps = [
8225                &pg as *const _ as *mut std::ffi::c_void,
8226                &pu as *const _ as *mut _,
8227                &pact as *const _ as *mut _,
8228                &pq as *const _ as *mut _,
8229                &pd as *const _ as *mut _,
8230                &nc as *const _ as *mut _,
8231            ];
8232            unsafe {
8233                self.launch_pdl(
8234                    "gelu_tanh_mul_q8_1",
8235                    (nrows as u32, 1, 1),
8236                    (rms_block(), 1, 1),
8237                    &mut ps,
8238                )?;
8239            }
8240            return Ok(());
8241        }
8242        let f = self.func("gelu_tanh_mul_q8_1");
8243        let cfg = LaunchConfig {
8244            grid_dim: (nrows as u32, 1, 1),
8245            block_dim: (rms_block(), 1, 1),
8246            shared_mem_bytes: 0,
8247        };
8248        let __s_b = self.gpu.stream();
8249        let mut b = __s_b.launch_builder(&f);
8250        b.arg(gate)
8251            .arg(up)
8252            .arg(&mut *act)
8253            .arg(&mut *out_q)
8254            .arg(&mut *out_d)
8255            .arg(&nc);
8256        unsafe {
8257            b.launch(cfg)?;
8258        }
8259        Ok(())
8260    }
8261
8262    /// gemma4: add + rms_norm3 with outputs 0/2 emitted q8_1 (zsh + moe_in) and 1 f32 (router).
8263    #[allow(clippy::too_many_arguments)]
8264    pub fn add_rms_norm3_q8z(
8265        &self,
8266        a: &CudaSlice<f32>,
8267        b_in: &CudaSlice<f32>,
8268        w0: &CudaSlice<f32>,
8269        w1: &CudaSlice<f32>,
8270        w2: &CudaSlice<f32>,
8271        res: &mut CudaSlice<f32>,
8272        out1: &mut CudaSlice<f32>,
8273        ncols: usize,
8274        nrows: usize,
8275        eps: f32,
8276    ) -> Result<
8277        (
8278            (CudaSlice<i8>, CudaSlice<f32>),
8279            (CudaSlice<i8>, CudaSlice<f32>),
8280        ),
8281        Box<dyn std::error::Error>,
8282    > {
8283        let mut q0 = self.alloc_uninit::<i8>(nrows * ncols)?;
8284        let mut d0 = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
8285        let mut q2 = self.alloc_uninit::<i8>(nrows * ncols)?;
8286        let mut d2 = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
8287        let f = self.func("add_rms_norm3_q8z_f32");
8288        let cfg = LaunchConfig {
8289            grid_dim: (nrows as u32, 1, 1),
8290            block_dim: (rms_block(), 1, 1),
8291            shared_mem_bytes: 0,
8292        };
8293        let (nc, e2) = (ncols as i32, eps);
8294        let __s_b = self.gpu.stream();
8295        let mut b = __s_b.launch_builder(&f);
8296        b.arg(a)
8297            .arg(b_in)
8298            .arg(w0)
8299            .arg(w1)
8300            .arg(w2)
8301            .arg(res)
8302            .arg(&mut q0)
8303            .arg(&mut d0)
8304            .arg(out1)
8305            .arg(&mut q2)
8306            .arg(&mut d2)
8307            .arg(&nc)
8308            .arg(&e2);
8309        unsafe {
8310            b.launch(cfg)?;
8311        }
8312        Ok(((q0, d0), (q2, d2)))
8313    }
8314
8315    /// gemma4: res = a+b AND the three rms_norms of res in one launch.
8316    #[allow(clippy::too_many_arguments)]
8317    pub fn add_rms_norm3(
8318        &self,
8319        a: &CudaSlice<f32>,
8320        b_in: &CudaSlice<f32>,
8321        w0: &CudaSlice<f32>,
8322        w1: &CudaSlice<f32>,
8323        w2: &CudaSlice<f32>,
8324        res: &mut CudaSlice<f32>,
8325        d0: &mut CudaSlice<f32>,
8326        d1: &mut CudaSlice<f32>,
8327        d2: &mut CudaSlice<f32>,
8328        ncols: usize,
8329        nrows: usize,
8330        eps: f32,
8331    ) -> Result<(), Box<dyn std::error::Error>> {
8332        let f = self.func("add_rms_norm3_f32");
8333        let cfg = LaunchConfig {
8334            grid_dim: (nrows as u32, 1, 1),
8335            block_dim: (rms_block(), 1, 1),
8336            shared_mem_bytes: 0,
8337        };
8338        let (nc, e2) = (ncols as i32, eps);
8339        let __s_b = self.gpu.stream();
8340        let mut b = __s_b.launch_builder(&f);
8341        b.arg(a)
8342            .arg(b_in)
8343            .arg(w0)
8344            .arg(w1)
8345            .arg(w2)
8346            .arg(res)
8347            .arg(d0)
8348            .arg(d1)
8349            .arg(d2)
8350            .arg(&nc)
8351            .arg(&e2);
8352        unsafe {
8353            b.launch(cfg)?;
8354        }
8355        Ok(())
8356    }
8357
8358    /// dst = (a + b) * c (residual add + layer scale, one launch).
8359    pub fn add_scale(
8360        &self,
8361        a: &CudaSlice<f32>,
8362        b_in: &CudaSlice<f32>,
8363        c: f32,
8364        dst: &mut CudaSlice<f32>,
8365        n: usize,
8366    ) -> Result<(), Box<dyn std::error::Error>> {
8367        let f = self.func("add_scale_f32");
8368        let cfg = LaunchConfig::for_num_elems(n as u32);
8369        let ni = n as i32;
8370        let __s_b = self.gpu.stream();
8371        let mut b = __s_b.launch_builder(&f);
8372        b.arg(a).arg(b_in).arg(&c).arg(dst).arg(&ni);
8373        unsafe {
8374            b.launch(cfg)?;
8375        }
8376        Ok(())
8377    }
8378
8379    /// Vision-tower LayerNorm (with bias) over [nrows, ncols] — lane/vision.
8380    pub fn layer_norm_bias(
8381        &self,
8382        x: &CudaSlice<f32>,
8383        w: &CudaSlice<f32>,
8384        b: &CudaSlice<f32>,
8385        dst: &mut CudaSlice<f32>,
8386        ncols: usize,
8387        nrows: usize,
8388        eps: f32,
8389    ) -> Result<(), Box<dyn std::error::Error>> {
8390        let f = self.func("layer_norm_bias_f32");
8391        let (nc, e) = (ncols as i32, eps);
8392        let cfg = LaunchConfig {
8393            grid_dim: (nrows as u32, 1, 1),
8394            block_dim: (256, 1, 1),
8395            shared_mem_bytes: 0,
8396        };
8397        let __s_b = self.gpu.stream();
8398        let mut lb = __s_b.launch_builder(&f);
8399        lb.arg(x).arg(w).arg(b).arg(&mut *dst).arg(&nc).arg(&e);
8400        unsafe {
8401            lb.launch(cfg)?;
8402        }
8403        Ok(())
8404    }
8405
8406    /// gelu_pytorch_tanh elementwise (vision tower MLP activation).
8407    pub fn gelu_tanh(
8408        &self,
8409        x: &CudaSlice<f32>,
8410        dst: &mut CudaSlice<f32>,
8411        n: usize,
8412    ) -> Result<(), Box<dyn std::error::Error>> {
8413        let f = self.func("gelu_tanh_f32");
8414        let ni = n as i64;
8415        let cfg = LaunchConfig {
8416            grid_dim: (n.div_ceil(256) as u32, 1, 1),
8417            block_dim: (256, 1, 1),
8418            shared_mem_bytes: 0,
8419        };
8420        let __s_b = self.gpu.stream();
8421        let mut lb = __s_b.launch_builder(&f);
8422        lb.arg(x).arg(&mut *dst).arg(&ni);
8423        unsafe {
8424            lb.launch(cfg)?;
8425        }
8426        Ok(())
8427    }
8428
8429    /// In-place row softmax over [nrows, ncols] (bidirectional vision attention).
8430    pub fn row_softmax(
8431        &self,
8432        x: &mut CudaSlice<f32>,
8433        ncols: usize,
8434        nrows: usize,
8435    ) -> Result<(), Box<dyn std::error::Error>> {
8436        let f = self.func("row_softmax_f32");
8437        let nc = ncols as i32;
8438        let cfg = LaunchConfig {
8439            grid_dim: (nrows as u32, 1, 1),
8440            block_dim: (256, 1, 1),
8441            shared_mem_bytes: 0,
8442        };
8443        let __s_b = self.gpu.stream();
8444        let mut lb = __s_b.launch_builder(&f);
8445        lb.arg(&mut *x).arg(&nc);
8446        unsafe {
8447            lb.launch(cfg)?;
8448        }
8449        Ok(())
8450    }
8451
8452    pub fn rms_norm(
8453        &self,
8454        x: &CudaSlice<f32>,
8455        w: &CudaSlice<f32>,
8456        dst: &mut CudaSlice<f32>,
8457        ncols: usize,
8458        nrows: usize,
8459        eps: f32,
8460    ) -> Result<(), Box<dyn std::error::Error>> {
8461        let (nc, e) = (ncols as i32, eps);
8462        let kname = if Self::norm_ilp_on() {
8463            "rms_norm_f32_v2"
8464        } else {
8465            "rms_norm_f32"
8466        };
8467        if Self::pdl_on() && Self::pdl_wb_on() {
8468            use cudarc::driver::{DevicePtr, DevicePtrMut};
8469            let s = &self.gpu.stream();
8470            let (px, _g0) = x.device_ptr(s);
8471            let (pw, _g1) = w.device_ptr(s);
8472            let (pd, _g2) = dst.device_ptr_mut(s);
8473            let mut ps = [
8474                &px as *const _ as *mut std::ffi::c_void,
8475                &pw as *const _ as *mut _,
8476                &pd as *const _ as *mut _,
8477                &nc as *const _ as *mut _,
8478                &e as *const _ as *mut _,
8479            ];
8480            unsafe {
8481                self.launch_pdl(kname, (nrows as u32, 1, 1), (rms_block(), 1, 1), &mut ps)?;
8482            }
8483            return Ok(());
8484        }
8485        let f = self.func(kname);
8486        let cfg = LaunchConfig {
8487            grid_dim: (nrows as u32, 1, 1),
8488            block_dim: (rms_block(), 1, 1),
8489            shared_mem_bytes: 0,
8490        };
8491        let __s_b = self.gpu.stream();
8492        let mut b = __s_b.launch_builder(&f);
8493        b.arg(x).arg(w).arg(dst).arg(&nc).arg(&e);
8494        unsafe {
8495            b.launch(cfg)?;
8496        }
8497        Ok(())
8498    }
8499
8500    /// RMS-norm with blockDim=1024 — BIT-IDENTICAL to the fused `rms_norm_q8_1` and
8501    /// `add_rms_norm_q8_1` kernels' sum-of-squares reduction. The spec verify path MUST use this
8502    /// to match decode's FP accumulation order: the standard `rms_norm` at blockDim=256 has a
8503    /// different per-thread stride (ncols/256 partials vs ncols/1024 partials) and therefore a
8504    /// different shfl-tree reduction that can shift `scale = rsqrt(sum/n + eps)` by ULPs, causing
8505    /// divergence through the GDN scan and argmax flips on the 9B text prompt. The underlying
8506    /// `rms_norm_f32` kernel supports any blockDim (generic reduce with shared[32]).
8507    pub fn rms_norm_decode(
8508        &self,
8509        x: &CudaSlice<f32>,
8510        w: &CudaSlice<f32>,
8511        dst: &mut CudaSlice<f32>,
8512        ncols: usize,
8513        nrows: usize,
8514        eps: f32,
8515    ) -> Result<(), Box<dyn std::error::Error>> {
8516        let f = self.func(if Self::norm_ilp_on() {
8517            "rms_norm_f32_v2"
8518        } else {
8519            "rms_norm_f32"
8520        });
8521        let cfg = LaunchConfig {
8522            grid_dim: (nrows as u32, 1, 1),
8523            block_dim: (1024, 1, 1),
8524            shared_mem_bytes: 0,
8525        };
8526        let (nc, e) = (ncols as i32, eps);
8527        let __s_b = self.gpu.stream();
8528        let mut b = __s_b.launch_builder(&f);
8529        b.arg(x).arg(w).arg(dst).arg(&nc).arg(&e);
8530        unsafe {
8531            b.launch(cfg)?;
8532        }
8533        Ok(())
8534    }
8535
8536    /// DECODE GLUE-FUSION LEVER: `z = rms_norm(x)*w` emitted DIRECTLY as q8_1 (no f32 `z` materialized,
8537    /// no standalone quantize_q8_1 launch). Returns (out_q [nrows*ncols i8], out_d [nrows*nblk f32])
8538    /// ready to feed matmul_pre. BIT-IDENTICAL to rms_norm + quantize_q8_1. ncols % 32 == 0.
8539    pub fn rms_norm_q8_1(
8540        &self,
8541        x: &CudaSlice<f32>,
8542        w: &CudaSlice<f32>,
8543        ncols: usize,
8544        nrows: usize,
8545        eps: f32,
8546    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8547        let nblk = ncols / 32;
8548        let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
8549        let mut d = self.alloc_uninit::<f32>(nrows * nblk)?;
8550        let (nc, e) = (ncols as i32, eps);
8551        if Self::pdl_on() {
8552            {
8553                use cudarc::driver::{DevicePtr, DevicePtrMut};
8554                let s = &self.gpu.stream();
8555                let (px, _g0) = x.device_ptr(s);
8556                let (pw, _g1) = w.device_ptr(s);
8557                let (pq, _g2) = q.device_ptr_mut(s);
8558                let (pd, _g3) = d.device_ptr_mut(s);
8559                let mut ps = [
8560                    &px as *const _ as *mut std::ffi::c_void,
8561                    &pw as *const _ as *mut _,
8562                    &pq as *const _ as *mut _,
8563                    &pd as *const _ as *mut _,
8564                    &nc as *const _ as *mut _,
8565                    &e as *const _ as *mut _,
8566                ];
8567                unsafe {
8568                    self.launch_pdl("rms_norm_q8_1", (nrows as u32, 1, 1), (1024, 1, 1), &mut ps)?;
8569                }
8570            }
8571            return Ok((q, d));
8572        }
8573        let f = self.func("rms_norm_q8_1");
8574        // 1024 threads: decode is nrows=1 -> ONE CTA; 32 warps hide the pass1->pass2 latency
8575        // (s[32] reduce already sized for 32 warps). Same shape math at any blockDim.
8576        let cfg = LaunchConfig {
8577            grid_dim: (nrows as u32, 1, 1),
8578            block_dim: (1024, 1, 1),
8579            shared_mem_bytes: 0,
8580        };
8581        let __s_b = self.gpu.stream();
8582        let mut b = __s_b.launch_builder(&f);
8583        b.arg(x).arg(w).arg(&mut q).arg(&mut d).arg(&nc).arg(&e);
8584        unsafe {
8585            b.launch(cfg)?;
8586        }
8587        Ok((q, d))
8588    }
8589
8590    /// Slot-fed rms_norm_q8_1 twin (alloc-free capture lane): identical launch (incl. the
8591    /// PDL arm), caller-owned outputs.
8592    pub fn rms_norm_q8_1_into(
8593        &self,
8594        x: &CudaSlice<f32>,
8595        w: &CudaSlice<f32>,
8596        ncols: usize,
8597        nrows: usize,
8598        eps: f32,
8599        q: &mut CudaSlice<i8>,
8600        d: &mut CudaSlice<f32>,
8601    ) -> Result<(), Box<dyn std::error::Error>> {
8602        let nblk = ncols / 32;
8603        debug_assert!(q.len() >= nrows * ncols && d.len() >= nrows * nblk);
8604        let (nc, e) = (ncols as i32, eps);
8605        if Self::pdl_on() {
8606            use cudarc::driver::{DevicePtr, DevicePtrMut};
8607            let s = &self.gpu.stream();
8608            let (px, _g0) = x.device_ptr(s);
8609            let (pw, _g1) = w.device_ptr(s);
8610            let (pq, _g2) = q.device_ptr_mut(s);
8611            let (pd, _g3) = d.device_ptr_mut(s);
8612            let mut ps = [
8613                &px as *const _ as *mut std::ffi::c_void,
8614                &pw as *const _ as *mut _,
8615                &pq as *const _ as *mut _,
8616                &pd as *const _ as *mut _,
8617                &nc as *const _ as *mut _,
8618                &e as *const _ as *mut _,
8619            ];
8620            unsafe {
8621                self.launch_pdl("rms_norm_q8_1", (nrows as u32, 1, 1), (1024, 1, 1), &mut ps)?;
8622            }
8623            return Ok(());
8624        }
8625        let f = self.func("rms_norm_q8_1");
8626        let cfg = LaunchConfig {
8627            grid_dim: (nrows as u32, 1, 1),
8628            block_dim: (1024, 1, 1),
8629            shared_mem_bytes: 0,
8630        };
8631        let __s_b = self.gpu.stream();
8632        let mut b = __s_b.launch_builder(&f);
8633        b.arg(x).arg(w).arg(&mut *q).arg(&mut *d).arg(&nc).arg(&e);
8634        unsafe {
8635            b.launch(cfg)?;
8636        }
8637        Ok(())
8638    }
8639
8640    /// Slot-fed quantize_q8_1 twin (alloc-free capture lane).
8641    pub fn quantize_q8_1_into(
8642        &self,
8643        x: &CudaSlice<f32>,
8644        m: usize,
8645        in_f: usize,
8646        q: &mut CudaSlice<i8>,
8647        d: &mut CudaSlice<f32>,
8648    ) -> Result<(), Box<dyn std::error::Error>> {
8649        let nblk = in_f / 32;
8650        debug_assert!(q.len() >= m * in_f && d.len() >= m * nblk);
8651        let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
8652        let (inf, mi) = (in_f as i32, m as i32);
8653        if Self::pdl_on() && Self::pdl_wb_on() {
8654            use cudarc::driver::{DevicePtr, DevicePtrMut};
8655            let s = &self.gpu.stream();
8656            let (px, _g0) = x.device_ptr(s);
8657            let (pq, _g1) = q.device_ptr_mut(s);
8658            let (pd, _g2) = d.device_ptr_mut(s);
8659            let mut ps = [
8660                &px as *const _ as *mut std::ffi::c_void,
8661                &pq as *const _ as *mut _,
8662                &pd as *const _ as *mut _,
8663                &inf as *const _ as *mut _,
8664                &mi as *const _ as *mut _,
8665            ];
8666            unsafe {
8667                self.launch_pdl("quantize_q8_1", cfg.grid_dim, cfg.block_dim, &mut ps)?;
8668            }
8669            return Ok(());
8670        }
8671        let f = self.func("quantize_q8_1");
8672        let __s_b = self.gpu.stream();
8673        let mut b = __s_b.launch_builder(&f);
8674        b.arg(x).arg(&mut *q).arg(&mut *d).arg(&inf).arg(&mi);
8675        unsafe {
8676            b.launch(cfg)?;
8677        }
8678        Ok(())
8679    }
8680
8681    /// DECODE GLUE-FUSION LEVER: `res = a+b; z = rms_norm(res)*w` with z emitted as q8_1. `res` is
8682    /// still written (the post-ffn residual add reads it). Fuses add_rms_norm + quantize_q8_1.
8683    /// Returns (out_q, out_d) for matmul_pre. BIT-IDENTICAL. ncols % 32 == 0.
8684    pub fn add_rms_norm_q8_1(
8685        &self,
8686        a: &CudaSlice<f32>,
8687        b_in: &CudaSlice<f32>,
8688        w: &CudaSlice<f32>,
8689        res: &mut CudaSlice<f32>,
8690        ncols: usize,
8691        nrows: usize,
8692        eps: f32,
8693    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8694        let nblk = ncols / 32;
8695        let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
8696        let mut d = self.alloc_uninit::<f32>(nrows * nblk)?;
8697        let f = self.func("add_rms_norm_q8_1");
8698        // 1024 threads: same single-CTA-at-decode reasoning as rms_norm_q8_1.
8699        let cfg = LaunchConfig {
8700            grid_dim: (nrows as u32, 1, 1),
8701            block_dim: (1024, 1, 1),
8702            shared_mem_bytes: 0,
8703        };
8704        let (nc, e) = (ncols as i32, eps);
8705        let __s_bld = self.gpu.stream();
8706        let mut bld = __s_bld.launch_builder(&f);
8707        bld.arg(a)
8708            .arg(b_in)
8709            .arg(w)
8710            .arg(res)
8711            .arg(&mut q)
8712            .arg(&mut d)
8713            .arg(&nc)
8714            .arg(&e);
8715        unsafe {
8716            bld.launch(cfg)?;
8717        }
8718        Ok((q, d))
8719    }
8720
8721    /// RANK3 LEVER (add+rmsnorm fuse): `res = a + b; dst = rms_norm(res) * w` in ONE launch. Fuses
8722    /// e.add(a,b,res) + e.rms_norm(res,w,dst), removing one launch + one HBM read of the residual per
8723    /// residual+norm pair. BIT-IDENTICAL to the two-kernel sequence (same IEEE add, same reduction).
8724    pub fn add_rms_norm(
8725        &self,
8726        a: &CudaSlice<f32>,
8727        b: &CudaSlice<f32>,
8728        w: &CudaSlice<f32>,
8729        res: &mut CudaSlice<f32>,
8730        dst: &mut CudaSlice<f32>,
8731        ncols: usize,
8732        nrows: usize,
8733        eps: f32,
8734    ) -> Result<(), Box<dyn std::error::Error>> {
8735        let (nc, e) = (ncols as i32, eps);
8736        let kname = if Self::norm_ilp_on() {
8737            "add_rms_norm_f32_v2"
8738        } else {
8739            "add_rms_norm_f32"
8740        };
8741        if Self::pdl_on() && Self::pdl_wb_on() {
8742            use cudarc::driver::{DevicePtr, DevicePtrMut};
8743            let s = &self.gpu.stream();
8744            let (pa, _g0) = a.device_ptr(s);
8745            let (pb, _g1) = b.device_ptr(s);
8746            let (pw, _g2) = w.device_ptr(s);
8747            let (pr, _g3) = res.device_ptr_mut(s);
8748            let (pd, _g4) = dst.device_ptr_mut(s);
8749            let mut ps = [
8750                &pa as *const _ as *mut std::ffi::c_void,
8751                &pb as *const _ as *mut _,
8752                &pw as *const _ as *mut _,
8753                &pr as *const _ as *mut _,
8754                &pd as *const _ as *mut _,
8755                &nc as *const _ as *mut _,
8756                &e as *const _ as *mut _,
8757            ];
8758            unsafe {
8759                self.launch_pdl(kname, (nrows as u32, 1, 1), (rms_block(), 1, 1), &mut ps)?;
8760            }
8761            return Ok(());
8762        }
8763        let f = self.func(kname);
8764        let cfg = LaunchConfig {
8765            grid_dim: (nrows as u32, 1, 1),
8766            block_dim: (rms_block(), 1, 1),
8767            shared_mem_bytes: 0,
8768        };
8769        let __s_b2 = self.gpu.stream();
8770        let mut b2 = __s_b2.launch_builder(&f);
8771        b2.arg(a)
8772            .arg(b)
8773            .arg(w)
8774            .arg(&mut *res)
8775            .arg(&mut *dst)
8776            .arg(&nc)
8777            .arg(&e);
8778        unsafe {
8779            b2.launch(cfg)?;
8780        }
8781        Ok(())
8782    }
8783
8784    /// E4B glue fusion: rms(a, wa) prologue + add_rms_norm — folds the post-attn norm into
8785    /// the tail entry (res = rms(a)*wa + b; dst = rms(res)*w).
8786    #[allow(clippy::too_many_arguments)]
8787    pub fn rms_pre_add_rms_norm(
8788        &self,
8789        a: &CudaSlice<f32>,
8790        wa: &CudaSlice<f32>,
8791        b: &CudaSlice<f32>,
8792        w: &CudaSlice<f32>,
8793        res: &mut CudaSlice<f32>,
8794        dst: &mut CudaSlice<f32>,
8795        ncols: usize,
8796        nrows: usize,
8797        eps: f32,
8798    ) -> Result<(), Box<dyn std::error::Error>> {
8799        let f = self.func("rms_pre_add_rms_norm_f32");
8800        let cfg = LaunchConfig {
8801            grid_dim: (nrows as u32, 1, 1),
8802            block_dim: (rms_block(), 1, 1),
8803            shared_mem_bytes: 0,
8804        };
8805        let (nc, e) = (ncols as i32, eps);
8806        let __s_b2 = self.gpu.stream();
8807        let mut b2 = __s_b2.launch_builder(&f);
8808        b2.arg(a)
8809            .arg(wa)
8810            .arg(b)
8811            .arg(w)
8812            .arg(&mut *res)
8813            .arg(&mut *dst)
8814            .arg(&nc)
8815            .arg(&e);
8816        unsafe {
8817            b2.launch(cfg)?;
8818        }
8819        Ok(())
8820    }
8821
8822    /// wave-2 fold: rms(a,wa) + add + ffn-norm with zsh EMITTED q8_1 (fused2 consumes it).
8823    #[allow(clippy::too_many_arguments)]
8824    pub fn rms_pre_add_rms_norm_q8z(
8825        &self,
8826        a: &CudaSlice<f32>,
8827        wa: &CudaSlice<f32>,
8828        b: &CudaSlice<f32>,
8829        w: &CudaSlice<f32>,
8830        res: &mut CudaSlice<f32>,
8831        dst: &mut CudaSlice<f32>,
8832        ncols: usize,
8833        nrows: usize,
8834        eps: f32,
8835    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8836        debug_assert!(ncols % 128 == 0);
8837        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
8838        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
8839        let (nc, e) = (ncols as i32, eps);
8840        if Self::pdl_on() {
8841            {
8842                use cudarc::driver::{DevicePtr, DevicePtrMut};
8843                let s = &self.gpu.stream();
8844                let (pa, _g0) = a.device_ptr(s);
8845                let (pwa, _g1) = wa.device_ptr(s);
8846                let (pb, _g2) = b.device_ptr(s);
8847                let (pw, _g3) = w.device_ptr(s);
8848                let (pr, _g4) = res.device_ptr_mut(s);
8849                let (pdst, _g5) = dst.device_ptr_mut(s);
8850                let (pq, _g6) = out_q.device_ptr_mut(s);
8851                let (pd, _g7) = out_d.device_ptr_mut(s);
8852                let mut ps = [
8853                    &pa as *const _ as *mut std::ffi::c_void,
8854                    &pwa as *const _ as *mut _,
8855                    &pb as *const _ as *mut _,
8856                    &pw as *const _ as *mut _,
8857                    &pr as *const _ as *mut _,
8858                    &pdst as *const _ as *mut _,
8859                    &pq as *const _ as *mut _,
8860                    &pd as *const _ as *mut _,
8861                    &nc as *const _ as *mut _,
8862                    &e as *const _ as *mut _,
8863                ];
8864                unsafe {
8865                    self.launch_pdl(
8866                        "rms_pre_add_rms_norm_q8z_f32",
8867                        (nrows as u32, 1, 1),
8868                        (rms_block(), 1, 1),
8869                        &mut ps,
8870                    )?;
8871                }
8872            }
8873            return Ok((out_q, out_d));
8874        }
8875        let f = self.func("rms_pre_add_rms_norm_q8z_f32");
8876        let cfg = LaunchConfig {
8877            grid_dim: (nrows as u32, 1, 1),
8878            block_dim: (rms_block(), 1, 1),
8879            shared_mem_bytes: 0,
8880        };
8881        let __s_b2 = self.gpu.stream();
8882        let mut b2 = __s_b2.launch_builder(&f);
8883        b2.arg(a)
8884            .arg(wa)
8885            .arg(b)
8886            .arg(w)
8887            .arg(&mut *res)
8888            .arg(&mut *dst)
8889            .arg(&mut out_q)
8890            .arg(&mut out_d)
8891            .arg(&nc)
8892            .arg(&e);
8893        unsafe {
8894            b2.launch(cfg)?;
8895        }
8896        Ok((out_q, out_d))
8897    }
8898
8899    /// Slot-fed twin of `rms_pre_add_rms_norm_q8z` (gemma4 pn-fold, slotted/graph arm):
8900    /// identical kernel, caller-owned outputs, PLAIN launch only — the dc_slotted capture
8901    /// body must stay attribute-free (the fused2_into precedent).
8902    #[allow(clippy::too_many_arguments)]
8903    pub fn rms_pre_add_rms_norm_q8z_into(
8904        &self,
8905        a: &CudaSlice<f32>,
8906        wa: &CudaSlice<f32>,
8907        b: &CudaSlice<f32>,
8908        w: &CudaSlice<f32>,
8909        res: &mut CudaSlice<f32>,
8910        dst: &mut CudaSlice<f32>,
8911        ncols: usize,
8912        nrows: usize,
8913        eps: f32,
8914        out_q: &mut CudaSlice<i8>,
8915        out_d: &mut CudaSlice<f32>,
8916    ) -> Result<(), Box<dyn std::error::Error>> {
8917        debug_assert!(ncols % 128 == 0);
8918        let (nc, e) = (ncols as i32, eps);
8919        let f = self.func("rms_pre_add_rms_norm_q8z_f32");
8920        let cfg = LaunchConfig {
8921            grid_dim: (nrows as u32, 1, 1),
8922            block_dim: (rms_block(), 1, 1),
8923            shared_mem_bytes: 0,
8924        };
8925        let __s_b = self.gpu.stream();
8926        let mut b2 = __s_b.launch_builder(&f);
8927        b2.arg(a)
8928            .arg(wa)
8929            .arg(b)
8930            .arg(w)
8931            .arg(&mut *res)
8932            .arg(&mut *dst)
8933            .arg(&mut *out_q)
8934            .arg(&mut *out_d)
8935            .arg(&nc)
8936            .arg(&e);
8937        unsafe {
8938            b2.launch(cfg)?;
8939        }
8940        Ok(())
8941    }
8942
8943    /// Slot-fed twin of `rms_pre_add_scale_rms_norm_q8_1` (gemma4 pn-fold exit, slotted
8944    /// arm): identical kernel, caller-owned outputs, PLAIN launch (capture-safe).
8945    #[allow(clippy::too_many_arguments)]
8946    pub fn rms_pre_add_scale_rms_norm_q8_1_into(
8947        &self,
8948        a: &CudaSlice<f32>,
8949        wa: &CudaSlice<f32>,
8950        b_in: &CudaSlice<f32>,
8951        c: f32,
8952        w: &CudaSlice<f32>,
8953        res: &mut CudaSlice<f32>,
8954        ncols: usize,
8955        nrows: usize,
8956        eps: f32,
8957        out_q: &mut CudaSlice<i8>,
8958        out_d: &mut CudaSlice<f32>,
8959    ) -> Result<(), Box<dyn std::error::Error>> {
8960        debug_assert!(ncols % 128 == 0);
8961        let (nc, e2) = (ncols as i32, eps);
8962        let f = self.func("rms_pre_add_scale_rms_norm_q8_1");
8963        let cfg = LaunchConfig {
8964            grid_dim: (nrows as u32, 1, 1),
8965            block_dim: (rms_block(), 1, 1),
8966            shared_mem_bytes: 0,
8967        };
8968        let __s_b = self.gpu.stream();
8969        let mut b2 = __s_b.launch_builder(&f);
8970        b2.arg(a)
8971            .arg(wa)
8972            .arg(b_in)
8973            .arg(&c)
8974            .arg(w)
8975            .arg(&mut *res)
8976            .arg(&mut *out_q)
8977            .arg(&mut *out_d)
8978            .arg(&nc)
8979            .arg(&e2);
8980        unsafe {
8981            b2.launch(cfg)?;
8982        }
8983        Ok(())
8984    }
8985
8986    /// gemma4 pn-fold seam (GAP-DIAGNOSIS verdict 7, the E4B glue backport): the dense
8987    /// decode/verify/slotted trio folds post_attn_norm into the tail entry
8988    /// (rms_pre_add_rms_norm[_q8z]) and post_ffw_norm into the residual exit
8989    /// (rms_pre_add_scale_rms_norm_q8_1). BITS-CHANGING vs the two-launch chain (the
8990    /// single-phase reduction's expansion rounding — E4B receipts); every arm moves
8991    /// together so decode == verify == graph parity holds BY CONSTRUCTION within either
8992    /// seam value. MEMRA_G4_PNFOLD=0 restores the unfused chain everywhere.
8993    pub fn g4_pnfold_on() -> bool {
8994        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8995        *ON.get_or_init(|| {
8996            std::env::var("MEMRA_G4_PNFOLD")
8997                .map(|v| v != "0")
8998                .unwrap_or(true)
8999        })
9000    }
9001
9002    /// wave-4b: OUT-dim concat of three Q4_0 tensors (same in_features; rows are independent
9003    /// blocks, so the concat is a D2D byte concat of the GGUF-layout planes). Returns None
9004    /// off-class (non-Q4_0, mismatched widths, or any tensor already rp-swapped in place).
9005    pub fn build_q4_out_concat3(
9006        &self,
9007        w0: &crate::model::GpuTensor,
9008        w1: &crate::model::GpuTensor,
9009        w2: &crate::model::GpuTensor,
9010    ) -> Result<Option<crate::model::GpuTensor>, Box<dyn std::error::Error>> {
9011        use crate::model::GpuTensor;
9012        let part = |w: &GpuTensor| -> Option<(usize, usize)> {
9013            match w {
9014                GpuTensor::Quant {
9015                    qtype,
9016                    row_bytes,
9017                    rp,
9018                    ..
9019                } if *qtype == QT_Q4_0 && !*rp => Some((*row_bytes, w.out_features())),
9020                _ => None,
9021            }
9022        };
9023        let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (part(w0), part(w1), part(w2))
9024        else {
9025            return Ok(None);
9026        };
9027        if rb0 != rb1
9028            || rb0 != rb2
9029            || w0.in_features() != w1.in_features()
9030            || w0.in_features() != w2.in_features()
9031        {
9032            return Ok(None);
9033        }
9034        fn bytes_of(w: &crate::model::GpuTensor) -> &CudaSlice<u8> {
9035            match w {
9036                crate::model::GpuTensor::Quant { bytes, .. } => bytes,
9037                _ => unreachable!(),
9038            }
9039        }
9040        let (b0, b1, b2) = (bytes_of(w0), bytes_of(w1), bytes_of(w2));
9041        let total = rb0 * (o0 + o1 + o2);
9042        let mut cat = self.alloc_u8(total)?;
9043        self.copy_u8_into(&mut cat, 0, b0, rb0 * o0)?;
9044        self.copy_u8_into(&mut cat, rb0 * o0, b1, rb1 * o1)?;
9045        self.copy_u8_into(&mut cat, rb0 * (o0 + o1), b2, rb2 * o2)?;
9046        Ok(Some(GpuTensor::Quant {
9047            bytes: cat,
9048            qtype: QT_Q4_0,
9049            row_bytes: rb0,
9050            ne: vec![w0.in_features() as u64, (o0 + o1 + o2) as u64],
9051            scale: 1.0,
9052            rp: false,
9053            #[cfg(memra_cutlass)]
9054            cutlass: None,
9055            fp8: None,
9056            blk: None,
9057            rp4: None,
9058            f16: None,
9059        }))
9060    }
9061
9062    /// FULL-WIDTH-ROPE CONTRACT for the fused rms_norm+qkv+rope kernels
9063    /// (lane/graph-s-key-exactness-20260819, probe O-10 follow-up).
9064    ///
9065    /// `rms_norm_qkv_rope_f32` / `_cat_f32` (`cu/kernels.cu`) and
9066    /// `rms_norm_qkv_rope_append_body` (`cu/flash_attn.cu`) take NO `n_dims`/`n_rot` argument.
9067    /// They compute `int half = ncols / 2` and rotate the FULL head width by construction — the
9068    /// standalone `rope_neox*` kernels take `n_dims` and early-return above it, these do not.
9069    ///
9070    /// Every call site today is a gemma-4 arm, where `n_rot == head_dim`, so nothing is wrong
9071    /// now. But this is the n_rot bug ONE FUSION away: fuse a partial-rotary arch onto these
9072    /// kernels (qwen3.5 = 64 rotary dims of a 256-wide head; step35 full-attn = 64 of 128) and
9073    /// 192 dims that must pass through unrotated get rotated silently — no error, no NaN, just a
9074    /// wrong model. The n_rot lane already paid for that class once, in the config derivation.
9075    ///
9076    /// So the fusions now take the layer's DERIVED rope width and refuse anything but full
9077    /// width. A future partial-rotary caller fails at its first launch with the geometry named
9078    /// instead of serving quietly wrong logits.
9079    fn full_width_rope_only(
9080        kernel: &str,
9081        n_rot: usize,
9082        head_dim: usize,
9083    ) -> Result<(), Box<dyn std::error::Error>> {
9084        if n_rot == head_dim {
9085            return Ok(());
9086        }
9087        Err(format!(
9088            "{kernel}: PARTIAL ROTARY REFUSED — n_rot {n_rot} != head_dim {head_dim}. This fused \
9089             rms_norm+qkv+rope kernel carries no n_dims parameter and rotates the full head \
9090             width (half = ncols/2), so it would rotate dims {n_rot}..{head_dim} that must pass \
9091             through unrotated. Use the split path (rms_norm_qkv + rope_neox/rope_neox2 with \
9092             n_dims={n_rot}), or add an n_dims early-return to the kernel and widen this guard."
9093        )
9094        .into())
9095    }
9096
9097    /// wave-4b: the qkv-cat twin — one contiguous [rq+2*rk, hd] input from the concat matvec.
9098    /// `n_rot` is the layer's derived rotary width and MUST equal `head_dim`
9099    /// ([`Engine::full_width_rope_only`]).
9100    #[allow(clippy::too_many_arguments)]
9101    pub fn rms_norm_qkv_rope_cat(
9102        &self,
9103        qkv: &CudaSlice<f32>,
9104        wq: &CudaSlice<f32>,
9105        wk: &CudaSlice<f32>,
9106        wv: &CudaSlice<f32>,
9107        q: &mut CudaSlice<f32>,
9108        k: &mut CudaSlice<f32>,
9109        v: &mut CudaSlice<f32>,
9110        head_dim: usize,
9111        n_rot: usize,
9112        rq: usize,
9113        rk: usize,
9114        pos: &CudaSlice<i32>,
9115        nh_q: usize,
9116        nh_k: usize,
9117        base: f32,
9118        freq_scale: f32,
9119        ff: Option<&CudaSlice<f32>>,
9120        eps: f32,
9121    ) -> Result<(), Box<dyn std::error::Error>> {
9122        Self::full_width_rope_only("rms_norm_qkv_rope_cat", n_rot, head_dim)?;
9123        let rows = rq + rk + rk;
9124        let theta_scale = base.powf(-2.0 / head_dim as f32);
9125        let (nc, rqi, rki, nhq, nhk) = (
9126            head_dim as i32,
9127            rq as i32,
9128            rk as i32,
9129            nh_q as i32,
9130            nh_k as i32,
9131        );
9132        if Self::pdl_on() {
9133            use cudarc::driver::{DevicePtr, DevicePtrMut};
9134            let s = &self.gpu.stream();
9135            let (pqkv, _g0) = qkv.device_ptr(s);
9136            let (pwq, _g1) = wq.device_ptr(s);
9137            let (pwk, _g2) = wk.device_ptr(s);
9138            let (pwv, _g3) = wv.device_ptr(s);
9139            let (pq, _g4) = q.device_ptr_mut(s);
9140            let (pk, _g5) = k.device_ptr_mut(s);
9141            let (pv, _g6) = v.device_ptr_mut(s);
9142            let (ppos, _g7) = pos.device_ptr(s);
9143            let (pff, _g8) = match ff {
9144                Some(t) => {
9145                    let (p, g) = t.device_ptr(s);
9146                    (p, Some(g))
9147                }
9148                None => (0, None),
9149            };
9150            let mut ps = [
9151                &pqkv as *const _ as *mut std::ffi::c_void,
9152                &pwq as *const _ as *mut _,
9153                &pwk as *const _ as *mut _,
9154                &pwv as *const _ as *mut _,
9155                &pq as *const _ as *mut _,
9156                &pk as *const _ as *mut _,
9157                &pv as *const _ as *mut _,
9158                &nc as *const _ as *mut _,
9159                &rqi as *const _ as *mut _,
9160                &rki as *const _ as *mut _,
9161                &ppos as *const _ as *mut _,
9162                &nhq as *const _ as *mut _,
9163                &nhk as *const _ as *mut _,
9164                &theta_scale as *const _ as *mut _,
9165                &freq_scale as *const _ as *mut _,
9166                &pff as *const _ as *mut _,
9167                &eps as *const _ as *mut _,
9168            ];
9169            unsafe {
9170                self.launch_pdl(
9171                    "rms_norm_qkv_rope_cat_f32",
9172                    (rows as u32, 1, 1),
9173                    (rms_block(), 1, 1),
9174                    &mut ps,
9175                )?;
9176            }
9177            return Ok(());
9178        }
9179        let f = self.func("rms_norm_qkv_rope_cat_f32");
9180        let cfg = LaunchConfig {
9181            grid_dim: (rows as u32, 1, 1),
9182            block_dim: (rms_block(), 1, 1),
9183            shared_mem_bytes: 0,
9184        };
9185        let __s_b = self.gpu.stream();
9186        let mut b = __s_b.launch_builder(&f);
9187        match ff {
9188            Some(t) => {
9189                b.arg(qkv)
9190                    .arg(wq)
9191                    .arg(wk)
9192                    .arg(wv)
9193                    .arg(&mut *q)
9194                    .arg(&mut *k)
9195                    .arg(&mut *v)
9196                    .arg(&nc)
9197                    .arg(&rqi)
9198                    .arg(&rki)
9199                    .arg(pos)
9200                    .arg(&nhq)
9201                    .arg(&nhk)
9202                    .arg(&theta_scale)
9203                    .arg(&freq_scale)
9204                    .arg(t)
9205                    .arg(&eps);
9206                unsafe {
9207                    b.launch(cfg)?;
9208                }
9209            }
9210            None => {
9211                let null: u64 = 0;
9212                b.arg(qkv)
9213                    .arg(wq)
9214                    .arg(wk)
9215                    .arg(wv)
9216                    .arg(&mut *q)
9217                    .arg(&mut *k)
9218                    .arg(&mut *v)
9219                    .arg(&nc)
9220                    .arg(&rqi)
9221                    .arg(&rki)
9222                    .arg(pos)
9223                    .arg(&nhq)
9224                    .arg(&nhk)
9225                    .arg(&theta_scale)
9226                    .arg(&freq_scale)
9227                    .arg(&null)
9228                    .arg(&eps);
9229                unsafe {
9230                    b.launch(cfg)?;
9231                }
9232            }
9233        }
9234        Ok(())
9235    }
9236
9237    /// wave-3 fold: rms_norm_qkv + rope_neox2 in ONE launch (n_dims == head_dim; ff nullable).
9238    /// `n_rot` is the layer's derived rotary width and MUST equal `head_dim`
9239    /// ([`Engine::full_width_rope_only`]).
9240    #[allow(clippy::too_many_arguments)]
9241    pub fn rms_norm_qkv_rope(
9242        &self,
9243        q0: &CudaSlice<f32>,
9244        k0: &CudaSlice<f32>,
9245        v0: &CudaSlice<f32>,
9246        wq: &CudaSlice<f32>,
9247        wk: &CudaSlice<f32>,
9248        wv: &CudaSlice<f32>,
9249        q: &mut CudaSlice<f32>,
9250        k: &mut CudaSlice<f32>,
9251        v: &mut CudaSlice<f32>,
9252        head_dim: usize,
9253        n_rot: usize,
9254        rq: usize,
9255        rk: usize,
9256        pos: &CudaSlice<i32>,
9257        nh_q: usize,
9258        nh_k: usize,
9259        base: f32,
9260        freq_scale: f32,
9261        ff: Option<&CudaSlice<f32>>,
9262        eps: f32,
9263    ) -> Result<(), Box<dyn std::error::Error>> {
9264        Self::full_width_rope_only("rms_norm_qkv_rope", n_rot, head_dim)?;
9265        let f = self.func("rms_norm_qkv_rope_f32");
9266        let rows = rq + rk + rk; // q rows + k rows + v rows (rk == rv)
9267        let cfg = LaunchConfig {
9268            grid_dim: (rows as u32, 1, 1),
9269            block_dim: (rms_block(), 1, 1),
9270            shared_mem_bytes: 0,
9271        };
9272        let theta_scale = base.powf(-2.0 / head_dim as f32);
9273        let (nc, rqi, rki, nhq, nhk) = (
9274            head_dim as i32,
9275            rq as i32,
9276            rk as i32,
9277            nh_q as i32,
9278            nh_k as i32,
9279        );
9280        let __s_b = self.gpu.stream();
9281        let mut b = __s_b.launch_builder(&f);
9282        match ff {
9283            Some(t) => {
9284                b.arg(q0)
9285                    .arg(k0)
9286                    .arg(v0)
9287                    .arg(wq)
9288                    .arg(wk)
9289                    .arg(wv)
9290                    .arg(&mut *q)
9291                    .arg(&mut *k)
9292                    .arg(&mut *v)
9293                    .arg(&nc)
9294                    .arg(&rqi)
9295                    .arg(&rki)
9296                    .arg(pos)
9297                    .arg(&nhq)
9298                    .arg(&nhk)
9299                    .arg(&theta_scale)
9300                    .arg(&freq_scale)
9301                    .arg(t)
9302                    .arg(&eps);
9303                unsafe {
9304                    b.launch(cfg)?;
9305                }
9306            }
9307            None => {
9308                let null: u64 = 0;
9309                b.arg(q0)
9310                    .arg(k0)
9311                    .arg(v0)
9312                    .arg(wq)
9313                    .arg(wk)
9314                    .arg(wv)
9315                    .arg(&mut *q)
9316                    .arg(&mut *k)
9317                    .arg(&mut *v)
9318                    .arg(&nc)
9319                    .arg(&rqi)
9320                    .arg(&rki)
9321                    .arg(pos)
9322                    .arg(&nhq)
9323                    .arg(&nhk)
9324                    .arg(&theta_scale)
9325                    .arg(&freq_scale)
9326                    .arg(&null)
9327                    .arg(&eps);
9328                unsafe {
9329                    b.launch(cfg)?;
9330                }
9331            }
9332        }
9333        Ok(())
9334    }
9335
9336    /// FUSED norm+rope+APPEND (m=1 decode, 2026-07-23): one launch replaces the
9337    /// rms_norm_qkv_rope + append_kv_quantized_dc pair. Kernel lives in the flash fatbins
9338    /// (format-flavored quant tail) — `g` must mirror the append path's flavor exactly.
9339    /// `n_rot` is the layer's derived rotary width and MUST equal `head_dim`
9340    /// ([`Engine::full_width_rope_only`]).
9341    #[allow(clippy::too_many_arguments)]
9342    pub fn rms_norm_qkv_rope_append_dc(
9343        &self,
9344        q0: &CudaSlice<f32>,
9345        k0: &CudaSlice<f32>,
9346        v0: &CudaSlice<f32>,
9347        wq: &CudaSlice<f32>,
9348        wk: &CudaSlice<f32>,
9349        wv: &CudaSlice<f32>,
9350        q: &mut CudaSlice<f32>,
9351        k: &mut CudaSlice<f32>,
9352        v: &mut CudaSlice<f32>,
9353        head_dim: usize,
9354        n_rot: usize,
9355        rq: usize,
9356        rk: usize,
9357        pos: &CudaSlice<i32>,
9358        nh_q: usize,
9359        nh_k: usize,
9360        base: f32,
9361        freq_scale: f32,
9362        ff: Option<&CudaSlice<f32>>,
9363        eps: f32,
9364        kc: &mut CudaSlice<u8>,
9365        vc: &mut CudaSlice<u8>,
9366        t_dev: &CudaSlice<i32>,
9367        k_tok_bytes: usize,
9368        v_tok_bytes: usize,
9369        g: bool,
9370    ) -> Result<(), Box<dyn std::error::Error>> {
9371        Self::full_width_rope_only("rms_norm_qkv_rope_append_dc", n_rot, head_dim)?;
9372        let rows = rq + rk + rk;
9373        let theta_scale = base.powf(-2.0 / head_dim as f32);
9374        let (nc, rqi, rki, nhq, nhk) = (
9375            head_dim as i32,
9376            rq as i32,
9377            rk as i32,
9378            nh_q as i32,
9379            nh_k as i32,
9380        );
9381        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
9382        if Self::pdl_on() && Self::pdl_wb_on() {
9383            use cudarc::driver::{DevicePtr, DevicePtrMut};
9384            let s = &self.gpu.stream();
9385            let (p0, _a0) = q0.device_ptr(s);
9386            let (p1, _a1) = k0.device_ptr(s);
9387            let (p2, _a2) = v0.device_ptr(s);
9388            let (pwq, _a3) = wq.device_ptr(s);
9389            let (pwk, _a4) = wk.device_ptr(s);
9390            let (pwv, _a5) = wv.device_ptr(s);
9391            let (pq, _a6) = q.device_ptr_mut(s);
9392            let (pk, _a7) = k.device_ptr_mut(s);
9393            let (pv, _a8) = v.device_ptr_mut(s);
9394            let (pp, _a9) = pos.device_ptr(s);
9395            let pff: u64 = match ff {
9396                Some(t) => {
9397                    let (p, _gg) = t.device_ptr(s);
9398                    p as u64
9399                }
9400                None => 0,
9401            };
9402            let (pkc, _a10) = kc.device_ptr_mut(s);
9403            let (pvc, _a11) = vc.device_ptr_mut(s);
9404            let (pt, _a12) = t_dev.device_ptr(s);
9405            let mut ps = [
9406                &p0 as *const _ as *mut std::ffi::c_void,
9407                &p1 as *const _ as *mut _,
9408                &p2 as *const _ as *mut _,
9409                &pwq as *const _ as *mut _,
9410                &pwk as *const _ as *mut _,
9411                &pwv as *const _ as *mut _,
9412                &pq as *const _ as *mut _,
9413                &pk as *const _ as *mut _,
9414                &pv as *const _ as *mut _,
9415                &nc as *const _ as *mut _,
9416                &rqi as *const _ as *mut _,
9417                &rki as *const _ as *mut _,
9418                &pp as *const _ as *mut _,
9419                &nhq as *const _ as *mut _,
9420                &nhk as *const _ as *mut _,
9421                &theta_scale as *const _ as *mut _,
9422                &freq_scale as *const _ as *mut _,
9423                &pff as *const _ as *mut _,
9424                &eps as *const _ as *mut _,
9425                &pkc as *const _ as *mut _,
9426                &pvc as *const _ as *mut _,
9427                &pt as *const _ as *mut _,
9428                &ktb as *const _ as *mut _,
9429                &vtb as *const _ as *mut _,
9430            ];
9431            unsafe {
9432                self.launch_pdl_flash(
9433                    g,
9434                    "rms_norm_qkv_rope_append_dc_f32",
9435                    (rows as u32, 1, 1),
9436                    (rms_block(), 1, 1),
9437                    0,
9438                    &mut ps,
9439                )?;
9440            }
9441            return Ok(());
9442        }
9443        let f = if g {
9444            self.func_g("rms_norm_qkv_rope_append_dc_f32")
9445        } else {
9446            self.func("rms_norm_qkv_rope_append_dc_f32")
9447        };
9448        let cfg = LaunchConfig {
9449            grid_dim: (rows as u32, 1, 1),
9450            block_dim: (rms_block(), 1, 1),
9451            shared_mem_bytes: 0,
9452        };
9453        let __s_b = self.gpu.stream();
9454        let mut b = __s_b.launch_builder(&f);
9455        match ff {
9456            Some(t) => {
9457                b.arg(q0)
9458                    .arg(k0)
9459                    .arg(v0)
9460                    .arg(wq)
9461                    .arg(wk)
9462                    .arg(wv)
9463                    .arg(&mut *q)
9464                    .arg(&mut *k)
9465                    .arg(&mut *v)
9466                    .arg(&nc)
9467                    .arg(&rqi)
9468                    .arg(&rki)
9469                    .arg(pos)
9470                    .arg(&nhq)
9471                    .arg(&nhk)
9472                    .arg(&theta_scale)
9473                    .arg(&freq_scale)
9474                    .arg(t)
9475                    .arg(&eps)
9476                    .arg(&mut *kc)
9477                    .arg(&mut *vc)
9478                    .arg(t_dev)
9479                    .arg(&ktb)
9480                    .arg(&vtb);
9481                unsafe {
9482                    b.launch(cfg)?;
9483                }
9484            }
9485            None => {
9486                let null: u64 = 0;
9487                b.arg(q0)
9488                    .arg(k0)
9489                    .arg(v0)
9490                    .arg(wq)
9491                    .arg(wk)
9492                    .arg(wv)
9493                    .arg(&mut *q)
9494                    .arg(&mut *k)
9495                    .arg(&mut *v)
9496                    .arg(&nc)
9497                    .arg(&rqi)
9498                    .arg(&rki)
9499                    .arg(pos)
9500                    .arg(&nhq)
9501                    .arg(&nhk)
9502                    .arg(&theta_scale)
9503                    .arg(&freq_scale)
9504                    .arg(&null)
9505                    .arg(&eps)
9506                    .arg(&mut *kc)
9507                    .arg(&mut *vc)
9508                    .arg(t_dev)
9509                    .arg(&ktb)
9510                    .arg(&vtb);
9511                unsafe {
9512                    b.launch(cfg)?;
9513                }
9514            }
9515        }
9516        Ok(())
9517    }
9518
9519    /// wave-2 fold: a + b with the sum emitted q8_1 alongside f32.
9520    /// Host-len twin of `rms_norm_qkv_rope_append_dc` (zoo-fusion arc): the eager decode
9521    /// arm tracks kv length on host (`kvl.len`), so the append slot rides the launch arg
9522    /// instead of the device counter. Kernel body is SHARED with the _dc entry (one
9523    /// inlined body — bit-identical to the rms_norm_qkv_rope + append pair it replaces,
9524    /// same law as the dc fold). `n_rot` is the layer's derived rotary width and MUST equal
9525    /// `head_dim` ([`Engine::full_width_rope_only`]).
9526    #[allow(clippy::too_many_arguments)]
9527    pub fn rms_norm_qkv_rope_append(
9528        &self,
9529        q0: &CudaSlice<f32>,
9530        k0: &CudaSlice<f32>,
9531        v0: &CudaSlice<f32>,
9532        wq: &CudaSlice<f32>,
9533        wk: &CudaSlice<f32>,
9534        wv: &CudaSlice<f32>,
9535        q: &mut CudaSlice<f32>,
9536        k: &mut CudaSlice<f32>,
9537        v: &mut CudaSlice<f32>,
9538        head_dim: usize,
9539        n_rot: usize,
9540        rq: usize,
9541        rk: usize,
9542        pos: &CudaSlice<i32>,
9543        nh_q: usize,
9544        nh_k: usize,
9545        base: f32,
9546        freq_scale: f32,
9547        ff: Option<&CudaSlice<f32>>,
9548        eps: f32,
9549        kc: &mut CudaSlice<u8>,
9550        vc: &mut CudaSlice<u8>,
9551        t: usize,
9552        k_tok_bytes: usize,
9553        v_tok_bytes: usize,
9554        g: bool,
9555    ) -> Result<(), Box<dyn std::error::Error>> {
9556        Self::full_width_rope_only("rms_norm_qkv_rope_append", n_rot, head_dim)?;
9557        let rows = rq + rk + rk;
9558        let theta_scale = base.powf(-2.0 / head_dim as f32);
9559        let (nc, rqi, rki, nhq, nhk) = (
9560            head_dim as i32,
9561            rq as i32,
9562            rk as i32,
9563            nh_q as i32,
9564            nh_k as i32,
9565        );
9566        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
9567        let ti = t as i32;
9568        if Self::pdl_on() && Self::pdl_wb_on() {
9569            use cudarc::driver::{DevicePtr, DevicePtrMut};
9570            let s = &self.gpu.stream();
9571            let (p0, _a0) = q0.device_ptr(s);
9572            let (p1, _a1) = k0.device_ptr(s);
9573            let (p2, _a2) = v0.device_ptr(s);
9574            let (pwq, _a3) = wq.device_ptr(s);
9575            let (pwk, _a4) = wk.device_ptr(s);
9576            let (pwv, _a5) = wv.device_ptr(s);
9577            let (pq, _a6) = q.device_ptr_mut(s);
9578            let (pk, _a7) = k.device_ptr_mut(s);
9579            let (pv, _a8) = v.device_ptr_mut(s);
9580            let (pp, _a9) = pos.device_ptr(s);
9581            let pff: u64 = match ff {
9582                Some(t) => {
9583                    let (p, _gg) = t.device_ptr(s);
9584                    p as u64
9585                }
9586                None => 0,
9587            };
9588            let (pkc, _a10) = kc.device_ptr_mut(s);
9589            let (pvc, _a11) = vc.device_ptr_mut(s);
9590            let mut ps = [
9591                &p0 as *const _ as *mut std::ffi::c_void,
9592                &p1 as *const _ as *mut _,
9593                &p2 as *const _ as *mut _,
9594                &pwq as *const _ as *mut _,
9595                &pwk as *const _ as *mut _,
9596                &pwv as *const _ as *mut _,
9597                &pq as *const _ as *mut _,
9598                &pk as *const _ as *mut _,
9599                &pv as *const _ as *mut _,
9600                &nc as *const _ as *mut _,
9601                &rqi as *const _ as *mut _,
9602                &rki as *const _ as *mut _,
9603                &pp as *const _ as *mut _,
9604                &nhq as *const _ as *mut _,
9605                &nhk as *const _ as *mut _,
9606                &theta_scale as *const _ as *mut _,
9607                &freq_scale as *const _ as *mut _,
9608                &pff as *const _ as *mut _,
9609                &eps as *const _ as *mut _,
9610                &pkc as *const _ as *mut _,
9611                &pvc as *const _ as *mut _,
9612                &ti as *const _ as *mut _,
9613                &ktb as *const _ as *mut _,
9614                &vtb as *const _ as *mut _,
9615            ];
9616            unsafe {
9617                self.launch_pdl_flash(
9618                    g,
9619                    "rms_norm_qkv_rope_append_f32",
9620                    (rows as u32, 1, 1),
9621                    (rms_block(), 1, 1),
9622                    0,
9623                    &mut ps,
9624                )?;
9625            }
9626            return Ok(());
9627        }
9628        let f = if g {
9629            self.func_g("rms_norm_qkv_rope_append_f32")
9630        } else {
9631            self.func("rms_norm_qkv_rope_append_f32")
9632        };
9633        let cfg = LaunchConfig {
9634            grid_dim: (rows as u32, 1, 1),
9635            block_dim: (rms_block(), 1, 1),
9636            shared_mem_bytes: 0,
9637        };
9638        let __s_b = self.gpu.stream();
9639        let mut b = __s_b.launch_builder(&f);
9640        let null: u64 = 0;
9641        b.arg(q0)
9642            .arg(k0)
9643            .arg(v0)
9644            .arg(wq)
9645            .arg(wk)
9646            .arg(wv)
9647            .arg(&mut *q)
9648            .arg(&mut *k)
9649            .arg(&mut *v)
9650            .arg(&nc)
9651            .arg(&rqi)
9652            .arg(&rki)
9653            .arg(pos)
9654            .arg(&nhq)
9655            .arg(&nhk)
9656            .arg(&theta_scale)
9657            .arg(&freq_scale);
9658        match ff {
9659            Some(t) => {
9660                b.arg(t);
9661            }
9662            None => {
9663                b.arg(&null);
9664            }
9665        }
9666        b.arg(&eps)
9667            .arg(&mut *kc)
9668            .arg(&mut *vc)
9669            .arg(&ti)
9670            .arg(&ktb)
9671            .arg(&vtb);
9672        unsafe {
9673            b.launch(cfg)?;
9674        }
9675        Ok(())
9676    }
9677
9678    pub fn add_q8_1(
9679        &self,
9680        a: &CudaSlice<f32>,
9681        b: &CudaSlice<f32>,
9682        res: &mut CudaSlice<f32>,
9683        ncols: usize,
9684        nrows: usize,
9685    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9686        debug_assert!(ncols % 128 == 0);
9687        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
9688        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9689        let f = self.func("add_q8_1_f32");
9690        let cfg = LaunchConfig {
9691            grid_dim: (nrows as u32, 1, 1),
9692            block_dim: (rms_block(), 1, 1),
9693            shared_mem_bytes: 0,
9694        };
9695        let nc = ncols as i32;
9696        let __s_b2 = self.gpu.stream();
9697        let mut b2 = __s_b2.launch_builder(&f);
9698        b2.arg(a)
9699            .arg(b)
9700            .arg(&mut *res)
9701            .arg(&mut out_q)
9702            .arg(&mut out_d)
9703            .arg(&nc);
9704        unsafe {
9705            b2.launch(cfg)?;
9706        }
9707        Ok((out_q, out_d))
9708    }
9709
9710    /// E4B FFN-tail exit fusion (glue wave 5): resid = b + rms(a, wa) emitted f32 + q8_1 pair
9711    /// in ONE launch — replaces rms_norm(a,wa->sn) + add_q8_1(sn,b). Same rms_block() config
9712    /// as both parents (bit-identity: identical reduction + quad-walk quantize).
9713    pub fn rms_pre_add_q8_1(
9714        &self,
9715        a: &CudaSlice<f32>,
9716        wa: &CudaSlice<f32>,
9717        b: &CudaSlice<f32>,
9718        res: &mut CudaSlice<f32>,
9719        ncols: usize,
9720        nrows: usize,
9721        eps: f32,
9722    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9723        debug_assert!(ncols % 128 == 0);
9724        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
9725        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9726        let f = self.func("rms_pre_add_q8_1_f32");
9727        let cfg = LaunchConfig {
9728            grid_dim: (nrows as u32, 1, 1),
9729            block_dim: (rms_block(), 1, 1),
9730            shared_mem_bytes: 0,
9731        };
9732        let (nc, ep) = (ncols as i32, eps);
9733        let __s_b2 = self.gpu.stream();
9734        let mut b2 = __s_b2.launch_builder(&f);
9735        b2.arg(a)
9736            .arg(wa)
9737            .arg(b)
9738            .arg(&mut *res)
9739            .arg(&mut out_q)
9740            .arg(&mut out_d)
9741            .arg(&nc)
9742            .arg(&ep);
9743        unsafe {
9744            b2.launch(cfg)?;
9745        }
9746        Ok((out_q, out_d))
9747    }
9748
9749    /// L2 norm per row (head_dim), no weight.
9750    /// PREFILL l2 dispatch (round 27): the warp-per-row float4 v2 when the numeric-config
9751    /// seam allows (MEMRA_L2_V2, default ON, d_state==128 only); else the strided kernel.
9752    pub fn l2_v2_on(ncols: usize) -> bool {
9753        ncols == 128 && std::env::var("MEMRA_L2_V2").as_deref() != Ok("0")
9754    }
9755
9756    pub fn l2_norm_pp(
9757        &self,
9758        x: &CudaSlice<f32>,
9759        dst: &mut CudaSlice<f32>,
9760        dst16: Option<&mut CudaSlice<u8>>,
9761        ncols: usize,
9762        nrows: usize,
9763        eps: f32,
9764    ) -> Result<(), Box<dyn std::error::Error>> {
9765        if Self::l2_v2_on(ncols) {
9766            let f = self.func("l2_norm_pp_v2_f32");
9767            let rows_per_block = 8u32; // 256 threads = 8 warps = 8 rows
9768            let cfg = LaunchConfig {
9769                grid_dim: ((nrows as u32).div_ceil(rows_per_block), 1, 1),
9770                block_dim: (256, 1, 1),
9771                shared_mem_bytes: 0,
9772            };
9773            let (nc, nr, e) = (ncols as i32, nrows as i32, eps);
9774            // mirror-fold: bf16 twin address by value (0 = skip; matches the nullable param)
9775            let d16: u64 = match dst16 {
9776                Some(d) => self.addr_u8(d),
9777                None => 0,
9778            };
9779            let __s_b = self.gpu.stream();
9780            let mut b = __s_b.launch_builder(&f);
9781            b.arg(x).arg(dst).arg(&d16).arg(&nc).arg(&nr).arg(&e);
9782            unsafe {
9783                b.launch(cfg)?;
9784            }
9785            return Ok(());
9786        }
9787        self.l2_norm(x, dst, ncols, nrows, eps)
9788    }
9789
9790    pub fn l2_norm(
9791        &self,
9792        x: &CudaSlice<f32>,
9793        dst: &mut CudaSlice<f32>,
9794        ncols: usize,
9795        nrows: usize,
9796        eps: f32,
9797    ) -> Result<(), Box<dyn std::error::Error>> {
9798        let f = self.func("l2_norm_f32");
9799        let cfg = LaunchConfig {
9800            grid_dim: (nrows as u32, 1, 1),
9801            block_dim: (256, 1, 1),
9802            shared_mem_bytes: 0,
9803        };
9804        let (nc, e) = (ncols as i32, eps);
9805        let __s_b = self.gpu.stream();
9806        let mut b = __s_b.launch_builder(&f);
9807        b.arg(x).arg(dst).arg(&nc).arg(&e);
9808        unsafe {
9809            b.launch(cfg)?;
9810        }
9811        Ok(())
9812    }
9813
9814    /// L2-norm with blockDim=32 (warp-tree reduction) — BIT-IDENTICAL to gdn_prep_decode_f32's
9815    /// per-warp L2 norm. The verify path MUST use this to match decode's FP accumulation order:
9816    /// l2_norm at blockDim=256 produces a different shfl-tree reduction of the 128-element
9817    /// squared-sum (pairwise tree vs serial-4-then-warp-tree), causing ULP differences that
9818    /// propagate through gdn_scan and flip argmax on marginal logits.
9819    pub fn l2_norm_decode(
9820        &self,
9821        x: &CudaSlice<f32>,
9822        dst: &mut CudaSlice<f32>,
9823        ncols: usize,
9824        nrows: usize,
9825        eps: f32,
9826    ) -> Result<(), Box<dyn std::error::Error>> {
9827        let f = self.func("l2_norm_f32");
9828        let cfg = LaunchConfig {
9829            grid_dim: (nrows as u32, 1, 1),
9830            block_dim: (32, 1, 1),
9831            shared_mem_bytes: 0,
9832        };
9833        let (nc, e) = (ncols as i32, eps);
9834        let __s_b = self.gpu.stream();
9835        let mut b = __s_b.launch_builder(&f);
9836        b.arg(x).arg(dst).arg(&nc).arg(&e);
9837        unsafe {
9838            b.launch(cfg)?;
9839        }
9840        Ok(())
9841    }
9842
9843    /// RoPE NEOX in-place. x:[head_dim, n_heads, n_tokens], pos:[n_tokens].
9844    pub fn rope_neox(
9845        &self,
9846        x: &mut CudaSlice<f32>,
9847        pos: &CudaSlice<i32>,
9848        head_dim: usize,
9849        n_dims: usize,
9850        n_heads: usize,
9851        n_tokens: usize,
9852        freq_base: f32,
9853        freq_scale: f32,
9854    ) -> Result<(), Box<dyn std::error::Error>> {
9855        let f = self.func("rope_neox_f32");
9856        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
9857        let grid = (n_heads * n_tokens) as u32;
9858        let cfg = LaunchConfig {
9859            grid_dim: (grid, 1, 1),
9860            block_dim: ((head_dim / 2) as u32, 1, 1),
9861            shared_mem_bytes: 0,
9862        };
9863        let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
9864        let __s_b = self.gpu.stream();
9865        let mut b = __s_b.launch_builder(&f);
9866        b.arg(x)
9867            .arg(pos)
9868            .arg(&hd)
9869            .arg(&nd)
9870            .arg(&nh)
9871            .arg(&theta_scale)
9872            .arg(&freq_scale);
9873        unsafe {
9874            b.launch(cfg)?;
9875        }
9876        Ok(())
9877    }
9878
9879    /// RoPE NEOX with per-dim freq factors (gemma4 global layers, rope_freqs.weight [n_dims/2]).
9880    pub fn rope_neox_ff(
9881        &self,
9882        x: &mut CudaSlice<f32>,
9883        pos: &CudaSlice<i32>,
9884        head_dim: usize,
9885        n_dims: usize,
9886        n_heads: usize,
9887        n_tokens: usize,
9888        freq_base: f32,
9889        freq_scale: f32,
9890        ff: &CudaSlice<f32>,
9891    ) -> Result<(), Box<dyn std::error::Error>> {
9892        let f = self.func("rope_neox_ff_f32");
9893        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
9894        let grid = (n_heads * n_tokens) as u32;
9895        let cfg = LaunchConfig {
9896            grid_dim: (grid, 1, 1),
9897            block_dim: ((head_dim / 2) as u32, 1, 1),
9898            shared_mem_bytes: 0,
9899        };
9900        let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
9901        let __s_b = self.gpu.stream();
9902        let mut b = __s_b.launch_builder(&f);
9903        b.arg(x)
9904            .arg(pos)
9905            .arg(&hd)
9906            .arg(&nd)
9907            .arg(&nh)
9908            .arg(&theta_scale)
9909            .arg(&freq_scale)
9910            .arg(ff);
9911        unsafe {
9912            b.launch(cfg)?;
9913        }
9914        Ok(())
9915    }
9916
9917    /// gemma4: rope q and k in one launch (per-row chain = rope_neox / rope_neox_ff verbatim).
9918    #[allow(clippy::too_many_arguments)]
9919    pub fn rope_neox2(
9920        &self,
9921        q: &mut CudaSlice<f32>,
9922        k: &mut CudaSlice<f32>,
9923        pos: &CudaSlice<i32>,
9924        head_dim: usize,
9925        n_dims: usize,
9926        nh_q: usize,
9927        nh_k: usize,
9928        n_tokens: usize,
9929        freq_base: f32,
9930        freq_scale: f32,
9931        ff: Option<&CudaSlice<f32>>,
9932    ) -> Result<(), Box<dyn std::error::Error>> {
9933        let f = self.func("rope_neox2_f32");
9934        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
9935        let grid = ((nh_q + nh_k) * n_tokens) as u32;
9936        let cfg = LaunchConfig {
9937            grid_dim: (grid, 1, 1),
9938            block_dim: ((head_dim / 2) as u32, 1, 1),
9939            shared_mem_bytes: 0,
9940        };
9941        let (hd, nd, nq, nk, nt) = (
9942            head_dim as i32,
9943            n_dims as i32,
9944            nh_q as i32,
9945            nh_k as i32,
9946            n_tokens as i32,
9947        );
9948        let __s_b = self.gpu.stream();
9949        let mut b = __s_b.launch_builder(&f);
9950        b.arg(q)
9951            .arg(k)
9952            .arg(pos)
9953            .arg(&hd)
9954            .arg(&nd)
9955            .arg(&nq)
9956            .arg(&nk)
9957            .arg(&nt)
9958            .arg(&theta_scale)
9959            .arg(&freq_scale);
9960        match ff {
9961            Some(ffv) => {
9962                b.arg(ffv);
9963                unsafe {
9964                    b.launch(cfg)?;
9965                }
9966            }
9967            None => {
9968                let null: u64 = 0;
9969                b.arg(&null);
9970                unsafe {
9971                    b.launch(cfg)?;
9972                }
9973            }
9974        }
9975        Ok(())
9976    }
9977
9978    /// gemma4 R1: dst = GELU_tanh(gate) * up.
9979    pub fn gelu_tanh_mul(
9980        &self,
9981        gate: &CudaSlice<f32>,
9982        up: &CudaSlice<f32>,
9983        dst: &mut CudaSlice<f32>,
9984        n: usize,
9985    ) -> Result<(), Box<dyn std::error::Error>> {
9986        let f = self.func("gelu_tanh_mul_f32");
9987        let cfg = LaunchConfig::for_num_elems(n as u32);
9988        let ni = n as i32;
9989        let __s_b = self.gpu.stream();
9990        let mut b = __s_b.launch_builder(&f);
9991        b.arg(gate).arg(up).arg(dst).arg(&ni);
9992        unsafe {
9993            b.launch(cfg)?;
9994        }
9995        Ok(())
9996    }
9997
9998    pub fn silu_mul(
9999        &self,
10000        gate: &CudaSlice<f32>,
10001        up: &CudaSlice<f32>,
10002        dst: &mut CudaSlice<f32>,
10003        n: usize,
10004    ) -> Result<(), Box<dyn std::error::Error>> {
10005        let f = self.func("silu_mul_f32");
10006        // float4 kernel: one thread per 4 elements (tail handled in-kernel)
10007        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
10008        let ni = n as i32;
10009        let __s_b = self.gpu.stream();
10010        let mut b = __s_b.launch_builder(&f);
10011        b.arg(gate).arg(up).arg(dst).arg(&ni);
10012        unsafe {
10013            b.launch(cfg)?;
10014        }
10015        Ok(())
10016    }
10017
10018    /// f16out twin of `silu_mul` (task #17): the epilogue also emits the fp16 GEMM operand
10019    /// for the down projection — kills the standalone convert pass. Bit-identical class.
10020    pub fn silu_mul_f16out(
10021        &self,
10022        gate: &CudaSlice<f32>,
10023        up: &CudaSlice<f32>,
10024        dst: &mut CudaSlice<f32>,
10025        dst16: &mut CudaSlice<u8>,
10026        n: usize,
10027    ) -> Result<(), Box<dyn std::error::Error>> {
10028        let f = self.func("silu_mul_f16out_f32");
10029        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
10030        let ni = n as i32;
10031        let __s_b = self.gpu.stream();
10032        let mut b = __s_b.launch_builder(&f);
10033        b.arg(gate).arg(up).arg(dst).arg(dst16).arg(&ni);
10034        unsafe {
10035            b.launch(cfg)?;
10036        }
10037        Ok(())
10038    }
10039
10040    /// FFN SwiGLU epilogue fusion (RANK3 LEVER 2): `dst = silu(gate*gs) * (up*us)` in ONE launch,
10041    /// folding the per-tensor NVFP4 macro-scale (`gs`,`us`) that would otherwise be two separate
10042    /// `scale_inplace` launches on the gate/up matmul outputs. BIT-IDENTICAL to
10043    /// scale_inplace(gate,gs); scale_inplace(up,us); silu_mul(gate,up,dst) — identical float ops in
10044    /// identical order. For non-NVFP4 weights gs==us==1.0 -> identical to `silu_mul`. Net: -2
10045    /// launches per dense FFN layer (the gate+up post-matmul scales).
10046    pub fn silu_mul_scaled(
10047        &self,
10048        gate: &CudaSlice<f32>,
10049        up: &CudaSlice<f32>,
10050        gs: f32,
10051        us: f32,
10052        dst: &mut CudaSlice<f32>,
10053        n: usize,
10054    ) -> Result<(), Box<dyn std::error::Error>> {
10055        let f = self.func("silu_mul_scaled_f32");
10056        let cfg = LaunchConfig::for_num_elems(n as u32);
10057        let ni = n as i32;
10058        let (gsf, usf) = (gs, us);
10059        let __s_b = self.gpu.stream();
10060        let mut b = __s_b.launch_builder(&f);
10061        b.arg(gate).arg(up).arg(&gsf).arg(&usf).arg(dst).arg(&ni);
10062        unsafe {
10063            b.launch(cfg)?;
10064        }
10065        Ok(())
10066    }
10067
10068    /// swigluoai (MiniMax-M3 / GPT-OSS): clamped SwiGLU epilogue, math 1:1 vs llama.cpp
10069    /// ggml_cuda_op_swiglu_oai_single. `dst = swish_alpha(clamp(gate*gs)) * (1 + clamp(up*us))`.
10070    /// gs/us fold the NVFP4 macro-scales exactly like `silu_mul_scaled`.
10071    #[allow(clippy::too_many_arguments)]
10072    pub fn swigluoai_mul_scaled(
10073        &self,
10074        gate: &CudaSlice<f32>,
10075        up: &CudaSlice<f32>,
10076        gs: f32,
10077        us: f32,
10078        alpha: f32,
10079        limit: f32,
10080        dst: &mut CudaSlice<f32>,
10081        n: usize,
10082    ) -> Result<(), Box<dyn std::error::Error>> {
10083        let f = self.func("swigluoai_mul_scaled_f32");
10084        let cfg = LaunchConfig::for_num_elems(n as u32);
10085        let ni = n as i32;
10086        let __s_b = self.gpu.stream();
10087        let mut b = __s_b.launch_builder(&f);
10088        b.arg(gate)
10089            .arg(up)
10090            .arg(&gs)
10091            .arg(&us)
10092            .arg(&alpha)
10093            .arg(&limit)
10094            .arg(dst)
10095            .arg(&ni);
10096        unsafe {
10097            b.launch(cfg)?;
10098        }
10099        Ok(())
10100    }
10101
10102    /// RANK2 LEVER (q8_1 quant-fold): SwiGLU epilogue that EMITS the q8_1 quantization of `act`
10103    /// directly (aq int8 [n] + ad f32 [n/32]), so ffn_down's standalone `quantize_q8_1` launch is
10104    /// removed — the down-proj activation has one consumer, so the quant folds into the producer for
10105    /// free (no extra HBM read; no f32 `act` write). gs/us fold the gate/up NVFP4 macro-scales like
10106    /// `silu_mul_scaled`. BIT-IDENTICAL q8_1 to silu_mul_scaled(...) then quantize_q8_1(...). Only
10107    /// valid when ffn_down uses the q8_1 dp4a/mmvq path; the caller checks `uses_q8_1_fast(ffn_down)`.
10108    /// n must be a multiple of 32 (n_ff always is).
10109    pub fn silu_mul_scaled_q8_1(
10110        &self,
10111        gate: &CudaSlice<f32>,
10112        up: &CudaSlice<f32>,
10113        gs: f32,
10114        us: f32,
10115        n: usize,
10116    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10117        let f = self.func("silu_mul_scaled_q8_1");
10118        let nblk = n / 32;
10119        let mut aq = self.alloc_uninit::<i8>(n)?; // full-overwrite output
10120        let mut ad = self.alloc_uninit::<f32>(nblk)?; // full-overwrite output
10121        // WARP-PER-BLOCK kernel: one warp (32 lanes) per 32-block -> n threads total.
10122        let cfg = LaunchConfig::for_num_elems(n as u32);
10123        let (gsf, usf, ni) = (gs, us, n as i32);
10124        let __s_b = self.gpu.stream();
10125        let mut b = __s_b.launch_builder(&f);
10126        b.arg(gate)
10127            .arg(up)
10128            .arg(&gsf)
10129            .arg(&usf)
10130            .arg(&mut aq)
10131            .arg(&mut ad)
10132            .arg(&ni);
10133        unsafe {
10134            b.launch(cfg)?;
10135        }
10136        Ok((aq, ad))
10137    }
10138
10139    pub fn add(
10140        &self,
10141        a: &CudaSlice<f32>,
10142        b_in: &CudaSlice<f32>,
10143        dst: &mut CudaSlice<f32>,
10144        n: usize,
10145    ) -> Result<(), Box<dyn std::error::Error>> {
10146        let f = self.func("add_f32");
10147        // float4 kernel: one thread per 4 elements (tail handled in-kernel)
10148        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
10149        let ni = n as i32;
10150        let __s_bld = self.gpu.stream();
10151        let mut bld = __s_bld.launch_builder(&f);
10152        bld.arg(a).arg(b_in).arg(dst).arg(&ni);
10153        unsafe {
10154            bld.launch(cfg)?;
10155        }
10156        Ok(())
10157    }
10158
10159    pub fn mul(
10160        &self,
10161        a: &CudaSlice<f32>,
10162        b_in: &CudaSlice<f32>,
10163        dst: &mut CudaSlice<f32>,
10164        n: usize,
10165    ) -> Result<(), Box<dyn std::error::Error>> {
10166        let f = self.func("mul_f32");
10167        let cfg = LaunchConfig::for_num_elems(n as u32);
10168        let ni = n as i32;
10169        let __s_bld = self.gpu.stream();
10170        let mut bld = __s_bld.launch_builder(&f);
10171        bld.arg(a).arg(b_in).arg(dst).arg(&ni);
10172        unsafe {
10173            bld.launch(cfg)?;
10174        }
10175        Ok(())
10176    }
10177
10178    /// Unified weight-tensor matmul: dispatches quant tensors to qmatvec (weights packed) and
10179    /// float tensors to cuBLASLt. y[m,out] = x[m,in] @ W[out,in]^T.
10180    pub fn matmul(
10181        &self,
10182        w: &crate::model::GpuTensor,
10183        x: &CudaSlice<f32>,
10184        m: usize,
10185    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10186        use crate::model::GpuTensor;
10187        let in_f = w.in_features();
10188        let out_f = w.out_features();
10189        // PREFILL (T>1) ROOT FIX: batched tensor-core int8 GEMM. Decodes each weight tile to int8
10190        // in smem ONCE and reuses across all tokens via mma — vs the dp4a matvec's per-token weight
10191        // re-read. Only the 4 daily-hot dtypes; m=1 decode keeps dp4a (it's bandwidth-bound, mma
10192        // gives nothing). Quantize the activation once here then call the GEMM.
10193        // m cutoff FIXED at 16: the m=4 MMA-verify A/B (2026-07-06, was MEMRA_GEMM_M) measured
10194        // NEGATIVE — the MMA tile grid starves at m=4 (BN=256 -> grid.y=1) and its FP order
10195        // shifted verify argmax at tight margins. Do not lower without re-running that battery.
10196        #[allow(non_snake_case)]
10197        // VERIFY-EXACT scope pushes the GEMM crossover out of reach (usize::MAX) — the
10198        // t>=16 dflash verify must ride the decode-exact batched class (parity law).
10199        let GEMM_M_THRESHOLD = if self.verify_exact_on() {
10200            usize::MAX
10201        } else {
10202            16usize
10203        };
10204
10205        // PREFILL GEMM (m>=16). ACCURACY-FIRST dispatch (2026-06-28, prefill-gemm-beat-research wf
10206        // wllbyo6vc step 1): the int8 W4A8 GEMM (qmatvec_gemm, q8_1 activation, s32 accumulate) is
10207        // ACCURATE (prefill logit maxdiff 0.159, < dp4a 0.55) and the default. The FP4 W4A4 mxf4 path
10208        // (try_fp4_gemm) quantizes the ACTIVATION to e2m1 4-bit (8 magnitude levels) -> maxdiff 1.0
10209        // when combined — a real accuracy loss, NOT a math bug. So FP4-W4A4 is taken ONLY under the
10210        // explicit MEMRA_FP4 opt-in AND it must come SECOND (int8 W4A8 is the correct default for NVFP4).
10211        // The workflow plan rebuilds the FP4 path (kill per-K repack, widen K, deepen pipeline, TMA) to
10212        // be both fast AND accurate; until then NVFP4 prefill defaults to the accurate int8 GEMM.
10213        // TINY-OUT_F GUARD (2026-06-28, ncu trace): the tiling GEMM's grid is (ceil(out_f/BM=64),
10214        // ceil(m/BN=256)). For tiny out_f (ssm_beta/ssm_alpha out_f=num_v_heads~32), grid.x=1 -> only
10215        // ceil(m/256) CTAs (e.g. 2 for m=512) on 82 SMs = 0.39% SM throughput, 852us EACH (measured
10216        // worst offender). The dp4a path grids (out_f, m) = far more CTAs, filling the GPU. So route
10217        // out_f < 2*BM to dp4a (skip the tiling GEMM which structurally can't fill the SMs here).
10218        const GEMM_MIN_OUT_F: usize = 128; // 2*BM; below this the GEMM grid.x starves the 82 SMs
10219        // VENDORED llama MMQ prefill GEMMs. NVFP4 W4A8 is DEFAULT-ON (2026-07-05 flip: same int8
10220        // accuracy class as the int8 GEMM below at ~1.9x pp512, rp-loader coexists with the A6
10221        // repack; MEMRA_MMQ_W4A8=0 = escape hatch). W4A4 mxf4nvf4 + Q4_K/Q5_K stay behind MEMRA_MMQ=1.
10222        // The env policy lives in mmq_supports/qmatvec_mmq. Feeds raw f32 activation `x` (the
10223        // launcher quantizes internally). out_f>=MMQ_Y/2 keeps the tile grid from starving the SMs.
10224        // FP8-ACT PREFILL (MEMRA_PP_FP8=1, probe verdict 2026-07-08): F8-E4M3-origin projections
10225        // carry their raw e4m3 device bytes (the `fp8` operand stashed at load next to the Q8_0
10226        // re-encode) — cuBLASLt FP8 TN at 620-795 TF vs 47-72 TF for this class's int8 GEMM.
10227        // Weight side EXACT (checkpoint bytes); activation rides ONE per-batch e4m3 scale
10228        // (amax/448) folded with weight_scale in-GEMM. Prefill only; decode keeps Q8_0 untouched.
10229        if m >= GEMM_M_THRESHOLD {
10230            if let Some(y) = self.try_fp8_gemm(w, x, m)? {
10231                return Ok(y);
10232            }
10233            // PER-BLOCK FP8 MMQ (lane/fp8-mmq): the block-128 class try_fp8_gemm skips (cuBLASLt
10234            // takes no block grid on sm_120). Exact per block — the checkpoint's e4m3 bytes and its
10235            // f32 grid go into the tile unchanged. TWO SOURCES, TWO DEFAULTS: the load-time stash is
10236            // opt-in (MEMRA_FP8_MMQ=1), the native-resident QT_F8_E4M3_BLK grid is DEFAULT ON
10237            // (MEMRA_FP8_MMQ=0 reverts it to dequant-per-call) — see fp8_ffi.rs for why the same
10238            // tile defaults differently by operand source.
10239            if let Some(y) = self.try_fp8_blk_mmq(w, x, m)? {
10240                return Ok(y);
10241            }
10242            // FP16-mirror prefill (MEMRA_PP_F16=1, probe 2026-07-26: 3.2-3.7x the MMQ class).
10243            // Mirror presence IS the gate (only built under the env). Decode never reaches here.
10244            if let Some(y) = self.try_f16_gemm(w, x, m)? {
10245                return Ok(y);
10246            }
10247        }
10248        // F8-E4M3 BLOCK-128 (QT_F8_E4M3_BLK, lane/fp8-blk128-decode). TWO arms, split at the SAME
10249        // m threshold the rest of this method uses:
10250        //   * m >= threshold (prefill): dequant-per-call to the ARM B' Q8_0 slab and recurse, so
10251        //     prefill keeps the floor's kernels AND the floor's bits (try_e4m3_blk_prefill).
10252        //   * m <  threshold: the native per-block GEMV — m=1 decode and the m=2..15 verify tiers.
10253        //     grid.y=m runs the exact m=1 program per (token,row), so the decode-parity law holds
10254        //     across every tier by construction with no batched twin needed.
10255        //
10256        // NOT gated on `fast`: this dtype has no dp4a twin and no Stage-A f32-dequant oracle (the
10257        // generic `deq()` switch has no block-scale input), exactly as QT_F8_E4M3 has none, so
10258        // MEMRA_FAST=0 cannot route it anywhere else. Placed before every GEMM/MMQ arm below
10259        // because gemm_supports/mmq_supports/mmvq_supports all deliberately REFUSE this qtype —
10260        // reaching the generic tail would panic rather than produce wrong numbers, and this pair of
10261        // arms is what makes sure it never gets there.
10262        if let GpuTensor::Quant { qtype, .. } = w {
10263            if *qtype == QT_F8_E4M3_BLK {
10264                if m >= GEMM_M_THRESHOLD {
10265                    if let Some(y) = self.try_e4m3_blk_prefill(w, x, m)? {
10266                        return Ok(y);
10267                    }
10268                }
10269                let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10270                if let Some(y) = self.try_e4m3_blk_pre(w, &aq, &ad, m)? {
10271                    return Ok(y);
10272                }
10273            }
10274        }
10275        if m >= GEMM_M_THRESHOLD && out_f >= GEMM_MIN_OUT_F && self.mmq_supports(w) {
10276            return self.qmatvec_mmq(w, x, m);
10277        }
10278        if m >= GEMM_M_THRESHOLD && out_f >= GEMM_MIN_OUT_F && self.gemm_supports(w) {
10279            let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10280            return self.qmatvec_gemm(w, &aq, &ad, m);
10281        }
10282        // FP4 W4A4 only as an explicit speed/accuracy tradeoff opt-in, and only if the int8 GEMM
10283        // above didn't already handle this weight (e.g. NVFP4 with in_f%64!=0, or MEMRA_NO_GEMM set).
10284        if m >= GEMM_M_THRESHOLD {
10285            if let Some(y) = self.try_fp4_gemm(w, x, m, in_f, out_f)? {
10286                return Ok(y);
10287            }
10288        }
10289        // Stage-B fast int8 dp4a is the DEFAULT since 2026-07-08 (it has been the daily path
10290        // for weeks; the old opt-in flag was a silent-slow-path landmine). MEMRA_FAST=0 reverts
10291        // to Stage-A f32-dequant (the correctness oracle path).
10292        let fast = std::env::var("MEMRA_FAST").as_deref() != Ok("0");
10293        // PERF-3 decode-GEMV: m=1 warp-per-row MMVQ (MEMRA_MMVQ). The big decode matvecs reach
10294        // `matmul` directly (ffn_down, lm_head output, wo), so route them here too — not only the
10295        // matmul_pre siblings. qmatvec_mmvq_raw quantizes the activation internally (q8_1) like the
10296        // _fast paths; the NVFP4 macro-scale is applied by the `scale != 1.0` block below.
10297        if m == 1 && fast {
10298            if let GpuTensor::Quant {
10299                bytes,
10300                qtype,
10301                row_bytes,
10302                rp,
10303                rp4,
10304                scale,
10305                ..
10306            } = w
10307            {
10308                if self.mmvq_supports(*qtype) {
10309                    // NVFP4 macro-scale rides the kernel's fused epilogue arg (one launch total);
10310                    // non-NVFP4 has scale==1.0 so qmatvec_mmvq skips scale_inplace either way.
10311                    // Q4_0 split-plane mirror (rp4): the decode arm reads it via the _rp twins.
10312                    let (bytes, rp) = match rp4 {
10313                        Some(m4) => (m4, true),
10314                        None => (bytes, *rp),
10315                    };
10316                    let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10317                    return self.qmatvec_mmvq(
10318                        bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, *scale, rp,
10319                    );
10320                }
10321            }
10322        }
10323        // BATCHED weight-resident matvec for the m=2-4 band (the MTP/verify forward's ffn_down, wo, and
10324        // lm_head `output` reach `matmul` directly at m=T=2..4). Walks the weight ONCE, dp4a vs all m
10325        // activation columns -> 1 weight read for m tokens (vs grid.y=m re-reading m times below). Quant
10326        // the activation once here (q8_1) like the _fast paths; macro-scale applied via the scale!=1.0
10327        // block below. MEMRA_NO_BATCHED -> per-m path.
10328        //
10329        // DECODE-PARITY GATE (2026-07-07, the 9B synth K=3/4/6 spec FAIL root cause): the batched
10330        // kernels are bit-identical per (token,row) to MMVQ's 32-thread warp reduce, NOT to the
10331        // dp4a kernels' 128-thread two-level reduce. Without MEMRA_MMVQ the m=1 decode chain rides
10332        // dp4a, so a verify riding batched here has a DIFFERENT FP order than the decode it must
10333        // match bit-for-bit — greedy spec flips at tight-margin tokens (the old HANDOVER "ENV LAW:
10334        // FAST+MMVQ both required" footgun, closed here). Parity law: the m>1 kernel CLASS must be
10335        // a pure function of (dtype, env) equal to the m=1 class — batched iff MMVQ. Without MMVQ
10336        // the verify falls to the per-m grid.y=m dp4a path below (each column = the exact m=1
10337        // dp4a program). MEMRA_MMVQ=1 (the daily config) is dispatch-unchanged.
10338        if (2..=16).contains(&m)
10339            && fast
10340            && std::env::var("MEMRA_NO_BATCHED").is_err()
10341            && (m <= 4 || Self::b8_enabled())
10342        {
10343            // b16 tier (2026-07-11, spec K>7): Q4_0/Q6_K have base+_rp b16 kernels; Q8_0's
10344            // b16 exists only as the split-plane _rp twin, so it joins iff the q8rp mirror
10345            // is present (rp4) — the mirror pick below then routes to the _rp family.
10346            // QT_F8_E4M3 joins unconditionally (lane/rp-on-st): its b16 IS the base kernel,
10347            // because the native e4m3 row layout is already aligned and needs no mirror.
10348            // NVFP4/Q4_K/Q8_0 all join unconditionally now (lane/rp-on-st): each has base + _rp
10349            // b16 twins, so either residency layout has its aligned form at this width. Q8_0's
10350            // old `rp4.is_some()` precondition is GONE — the mirror is a bandwidth lever, not the
10351            // exact tier's admission ticket (it was refusing FP8-ST over 23.9 MiB of ssm_beta).
10352            let m_ok = m <= 8
10353                || matches!(w, GpuTensor::Quant { qtype, .. }
10354                if *qtype == QT_Q4_0 || *qtype == QT_Q6_K || *qtype == QT_F8_E4M3
10355                    || *qtype == QT_NVFP4 || *qtype == QT_Q4_K || *qtype == QT_Q5_K || *qtype == QT_Q8_0);
10356            if m_ok {
10357                if let GpuTensor::Quant {
10358                    bytes,
10359                    qtype,
10360                    row_bytes,
10361                    rp,
10362                    rp4,
10363                    ..
10364                } = w
10365                {
10366                    if self.batched_supports(*qtype) && self.mmvq_supports(*qtype) {
10367                        let (bytes, rp) = match rp4 {
10368                            Some(m4) => (m4, true),
10369                            None => (bytes, *rp),
10370                        };
10371                        let mcols = Self::batched_mcols(m);
10372                        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10373                        let mut y = self.qmatvec_mmvq_batched(
10374                            bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, mcols, 1.0, rp,
10375                        )?;
10376                        if let GpuTensor::Quant { scale, .. } = w {
10377                            if *scale != 1.0 {
10378                                self.scale_inplace(&mut y, *scale, m * out_f)?;
10379                            }
10380                        }
10381                        return Ok(y);
10382                    }
10383                }
10384            }
10385        }
10386        // F8-E4M3 (MEMRA_ST_E4M3) catch-all for the m<16 band the arms above didn't take (m=9..15,
10387        // the K=8 verify tier; or m=2..8 under MEMRA_NO_BATCHED/MEMRA_B8=0): grid.y=m e4m3 mmvq —
10388        // the SAME per-(token,row) program as the m=1 decode launch (bit-identical by construction),
10389        // weight re-read m times (rare tier; exactness over bandwidth here). There is no _dp4a twin
10390        // for this dtype, so the generic match below must never see it under `fast`.
10391        if fast {
10392            if let GpuTensor::Quant {
10393                bytes,
10394                qtype,
10395                row_bytes,
10396                scale,
10397                ..
10398            } = w
10399            {
10400                if *qtype == QT_F8_E4M3 {
10401                    let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10402                    return self.qmatvec_mmvq(
10403                        bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, *scale, false,
10404                    );
10405                }
10406            }
10407        }
10408        let mut y = match w {
10409            GpuTensor::Quant {
10410                bytes,
10411                qtype,
10412                row_bytes,
10413                ..
10414            } if fast && *qtype == QT_Q8_0 => {
10415                self.qmatvec_q8_0_fast(bytes, x, m, in_f, out_f, *row_bytes)?
10416            }
10417            GpuTensor::Quant {
10418                bytes,
10419                qtype,
10420                row_bytes,
10421                ..
10422            } if fast && *qtype == QT_Q4_K => {
10423                self.qmatvec_q4_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
10424            }
10425            GpuTensor::Quant {
10426                bytes,
10427                qtype,
10428                row_bytes,
10429                ..
10430            } if fast && *qtype == QT_Q6_K => {
10431                self.qmatvec_q6_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
10432            }
10433            GpuTensor::Quant {
10434                bytes,
10435                qtype,
10436                row_bytes,
10437                ..
10438            } if fast && *qtype == QT_Q5_K => {
10439                self.qmatvec_q5_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
10440            }
10441            GpuTensor::Quant {
10442                bytes,
10443                qtype,
10444                row_bytes,
10445                ..
10446            } if fast && *qtype == QT_Q3_K => {
10447                self.qmatvec_q3_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
10448            }
10449            GpuTensor::Quant {
10450                bytes,
10451                qtype,
10452                row_bytes,
10453                rp,
10454                ..
10455            } if fast && *qtype == QT_NVFP4 => self.qmatvec_dp4a_named(
10456                if *rp {
10457                    "qmatvec_nvfp4_dp4a_rp"
10458                } else {
10459                    "qmatvec_nvfp4_dp4a"
10460                },
10461                bytes,
10462                x,
10463                m,
10464                in_f,
10465                out_f,
10466                *row_bytes,
10467            )?,
10468            // IQ4_XS trunk fast path — DEFAULT ON since 2026-08-02 (MEMRA_IQ_FAST=0 reverts to
10469            // Stage-A; see iq_fast_enabled). The old opt-in default was the KAT-Coder decode
10470            // anomaly (research/kat-anomaly-20260802/).
10471            GpuTensor::Quant {
10472                bytes,
10473                qtype,
10474                row_bytes,
10475                ..
10476            } if fast && *qtype == QT_IQ4_XS && Self::iq_fast_enabled() => {
10477                self.qmatvec_iq4_XS_fast(bytes, x, m, in_f, out_f, *row_bytes)?
10478            }
10479            // B3: IQ3_S uses the Stage-A f32 dequant-in-kernel path. There is NO
10480            // qmatvec_iq3_s_dp4a kernel — do NOT add a `*qtype == QT_IQ3_S` fast guard here
10481            // without first writing the matching kernel, or func() will panic
10482            // "kernel ... not in any fatbin".
10483            GpuTensor::Quant {
10484                bytes,
10485                qtype,
10486                row_bytes,
10487                rp,
10488                ..
10489            } =>
10490            // Stage-A generic: repacked NVFP4 uses the device-side split-plane tag (the
10491            // deq(row,j) form cannot address the planes; same value/product order).
10492            {
10493                self.qmatvec(
10494                    bytes,
10495                    x,
10496                    m,
10497                    in_f,
10498                    out_f,
10499                    if *rp && *qtype == QT_NVFP4 {
10500                        QT_NVFP4_RP
10501                    } else {
10502                        *qtype
10503                    },
10504                    *row_bytes,
10505                )?
10506            }
10507            GpuTensor::Float { data, .. } => self.linear(x, data, m, in_f, out_f)?,
10508            // MEMRA_FULL_PREC bf16-resident weight: dequant-on-use to f32 scratch, then the same
10509            // cuBLASLt f32 GEMV as the Float arm.
10510            GpuTensor::FloatBf16 { data, .. } => {
10511                self.linear_bf16_chunked(x, data, m, in_f, out_f, false)?
10512            }
10513        };
10514        // NVFP4 per-tensor macro-scale (post-matmul). scale==1.0 for all other quants/float -> no-op.
10515        if let GpuTensor::Quant { scale, .. } = w {
10516            if *scale != 1.0 {
10517                self.scale_inplace(&mut y, *scale, m * out_f)?;
10518            }
10519        }
10520        Ok(y)
10521    }
10522
10523    /// True when `MEMRA_FAST=0`, i.e. the Stage-A f32 oracle is the requested arithmetic and every
10524    /// `matmul_pre` call will take the raw-f32 escape rather than the q8_1 pair.
10525    ///
10526    /// WHY THIS EXISTS AS ITS OWN PREDICATE: `uses_q8_1_fast` needs a weight, but the callers that
10527    /// have to DECIDE WHETHER TO MATERIALIZE an f32 activation sit one level above any weight — the
10528    /// gemma-4 decode trunk emits a q8_1 pair per layer and hands it down, so the f32 has to be
10529    /// produced (or not) by the loop that owns the residual. Cached in a OnceLock like every other
10530    /// `*_on()` flag here: this is read once per layer per token on the decode path, and the daily
10531    /// path must not pay an env lookup for a flag that is off.
10532    pub fn stage_a_raw_needed() -> bool {
10533        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
10534        *ON.get_or_init(|| std::env::var("MEMRA_FAST").as_deref() == Ok("0"))
10535    }
10536
10537    /// True if `w` would take the int8-dp4a fast path under MEMRA_FAST (so its activation can be
10538    /// pre-quantized once and shared across sibling matmuls via `matmul_pre`).
10539    pub fn uses_q8_1_fast(&self, w: &crate::model::GpuTensor) -> bool {
10540        use crate::model::GpuTensor;
10541        if std::env::var("MEMRA_FAST").as_deref() == Ok("0") {
10542            return false;
10543        }
10544        match w {
10545            // QT_F8_E4M3_BLK is admitted for the same reason QT_F8_E4M3 is: its ONLY kernel class
10546            // takes the shared q8_1 activation, so callers may pre-quantize once and share it
10547            // across siblings. It is NOT admitted to any of the fused/dual epilogue doors those
10548            // siblings can then open (`q8_fused_params`, `e4m3_fused_params` and
10549            // `matmul_pre_dual_noscale` all match on their own qtype and refuse this one) — the
10550            // block class has no fused twin yet, so each of its projections takes its own launch.
10551            GpuTensor::Quant { qtype, .. } => {
10552                matches!(
10553                    *qtype,
10554                    QT_Q8_0
10555                        | QT_Q4_K
10556                        | QT_Q6_K
10557                        | QT_Q5_K
10558                        | QT_Q3_K
10559                        | QT_NVFP4
10560                        | QT_F8_E4M3
10561                        | QT_F8_E4M3_BLK
10562                        | QT_Q4_0
10563                ) || (*qtype == QT_IQ4_XS && Self::iq_fast_enabled())
10564            }
10565            GpuTensor::Float { .. } | GpuTensor::FloatBf16 { .. } => false,
10566        }
10567    }
10568
10569    /// matmul with a PRE-QUANTIZED q8_1 activation (aq,ad from `quantize_q8_1`). Skips the
10570    /// per-matmul re-quantize so sibling matmuls that share an input (gate+up share `z`;
10571    /// q/k/v + wqkv/gate/beta/alpha share `h`) quantize ONCE. Caller MUST have checked
10572    /// `uses_q8_1_fast(w)`; falls back to plain `matmul` otherwise (Stage-A / Float / non-fast).
10573    pub fn matmul_pre(
10574        &self,
10575        w: &crate::model::GpuTensor,
10576        aq: &CudaSlice<i8>,
10577        ad: &CudaSlice<f32>,
10578        x_fallback: &CudaSlice<f32>,
10579        m: usize,
10580    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10581        use crate::model::GpuTensor;
10582        // Every raw-f32 arm below (fp8/f16/MMQ/fp4) reads m*in_f from x_fallback. Callers that
10583        // pre-quantized and dropped the f32 input pass an EMPTY x_fallback (E4B's fusion port:
10584        // h = zeros(0)) — the length guard keeps those on the aq/ad GEMM instead of feeding a
10585        // 0-byte buffer to a convert kernel (illegal address -> cublasLt status 13; the E4B
10586        // rc=30013 dig, 2026-07-31).
10587        let x_raw_ok = x_fallback.len() >= m * w.in_features();
10588        // FP8-ACT PREFILL (MEMRA_PP_FP8=1): same arm as `matmul` — the fp8 operand needs the RAW
10589        // f32 activation (per-batch e4m3 quant differs from q8_1), so x_fallback not aq/ad.
10590        if m >= 16 && x_raw_ok && !self.verify_exact_on() {
10591            if let Some(y) = self.try_fp8_gemm(w, x_fallback, m)? {
10592                return Ok(y);
10593            }
10594            // PER-BLOCK FP8 MMQ — same arm as `matmul` (stash opt-in, native-resident default ON);
10595            // its own quantizer wants the RAW f32 activation, so x_fallback not aq/ad.
10596            if let Some(y) = self.try_fp8_blk_mmq(w, x_fallback, m)? {
10597                return Ok(y);
10598            }
10599            // FP16-mirror prefill (same arm as `matmul` — fp16 wants the RAW f32 activation).
10600            if let Some(y) = self.try_f16_gemm(w, x_fallback, m)? {
10601                return Ok(y);
10602            }
10603        }
10604        // BLOCK-128 e4m3 (QT_F8_E4M3_BLK) — the same two arms as `matmul`, split at the same m, and
10605        // placed at the same point in the order (after the prefill GEMM hooks, before every arm
10606        // that refuses this qtype). The prefill arm needs the RAW f32 activation for the Q8_0
10607        // dispatch it recurses into, so it takes x_fallback and is skipped when that is empty
10608        // (a pre-quantized caller that dropped its f32 input never runs at prefill m anyway).
10609        if m >= 16 && x_raw_ok && !self.verify_exact_on() {
10610            if let Some(y) = self.try_e4m3_blk_prefill(w, x_fallback, m)? {
10611                return Ok(y);
10612            }
10613        }
10614        if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
10615            return Ok(y);
10616        }
10617        // VENDORED llama MMQ prefill GEMMs (NVFP4 W4A8 default-on; W4A4/k-quant behind MEMRA_MMQ=1
10618        // — policy in mmq_supports) — use the RAW f32 activation (their own internal quant:
10619        // q8_1 D4 for NVFP4 W4A8, FP8/UE4M3 for W4A4, q8_1 DS4 for Q4_K/Q5_K), so x_fallback not
10620        // aq/ad.
10621        if m >= 16
10622            && w.out_features() >= 128
10623            && self.mmq_supports(w)
10624            && !self.verify_exact_on()
10625            && x_raw_ok
10626        {
10627            return self.qmatvec_mmq(w, x_fallback, m);
10628        }
10629        // Stage-C FP4 prefill (MEMRA_FP4): native mxf4 GEMM needs the f32 activation (FP4-quant differs
10630        // from q8_1), so re-quantize from x_fallback rather than reuse aq/ad. NVFP4 only, m>=16.
10631        if m >= 16 && x_raw_ok && !self.verify_exact_on() {
10632            if let Some(y) =
10633                self.try_fp4_gemm(w, x_fallback, m, w.in_features(), w.out_features())?
10634            {
10635                return Ok(y);
10636            }
10637        }
10638        // Prefill GEMM root fix: if T>1 and the dtype has a GEMM kernel, batch via tensor cores
10639        // (reuses the already-quantized aq/ad — no extra quantize). m=1 falls through to dp4a.
10640        if m >= 16 && self.gemm_supports(w) && !self.verify_exact_on() {
10641            return self.qmatvec_gemm(w, aq, ad, m);
10642        }
10643        // THE SEVENTH ARM. Every raw-f32 arm above is guarded by `x_raw_ok`; this one — the
10644        // Stage-A / Float escape — was not, and it is the ONLY one `MEMRA_FAST=0` opens. So the
10645        // 2026-07-31 E4B rc=30013 fix (add the length guard, keep empty-fallback callers off the
10646        // raw-f32 arms) protected the six arms the FAST path can reach and missed the one the
10647        // ORACLE path reaches. Consequence before this guard: the gemma-4 decode arms, which pass
10648        // `e.zeros(0)` because their trunk emits only a q8_1 pair and never materializes an f32
10649        // attn-normed activation, fell through to Stage-A `qmatvec_f32` (cu/qmatvec.cu:5442),
10650        // which reads `m * in_f` floats out of a 0-byte allocation ->
10651        // CUDA_ERROR_ILLEGAL_ADDRESS at layer 0 of the first decode token. The fault is STICKY:
10652        // it poisons the context, so every LATER request in that process fails with an unrelated
10653        // message ("cache alloc failed: ...") and the true cause appears exactly once, in the
10654        // first failure. That is what made `MEMRA_FAST=0` — the reference named in
10655        // `tools/argmax-margin-gate.sh`'s own header — return an opaque HTTP 500 on the gemma
10656        // dense artifact and left the arm with no working truth instrument.
10657        //
10658        // Refuse loudly instead of reading out of bounds. A named error at the true call site is
10659        // strictly better than an illegal address surfacing later at an unrelated sync point, and
10660        // an oracle that cannot run must say so rather than corrupt the context it runs in.
10661        if !self.uses_q8_1_fast(w) {
10662            if !x_raw_ok {
10663                return Err(format!(
10664                    "matmul_pre: q8_1-fast is off for this weight but x_fallback holds {} f32 \
10665                     (need m*in_f = {}*{} = {}). This call site pre-quantized its activation and \
10666                     dropped the f32, so there is nothing to fall back to — pass the real f32 \
10667                     activation (see Engine::rms_norm_decode, which is bit-identical to \
10668                     rms_norm_q8_1's reduction) or keep the weight on the q8_1 path.",
10669                    x_fallback.len(),
10670                    m,
10671                    w.in_features(),
10672                    m * w.in_features()
10673                )
10674                .into());
10675            }
10676            return self.matmul(w, x_fallback, m);
10677        }
10678        let in_f = w.in_features();
10679        let out_f = w.out_features();
10680        let (bytes, qtype, row_bytes, scale, rp) = match w {
10681            GpuTensor::Quant {
10682                bytes,
10683                qtype,
10684                row_bytes,
10685                scale,
10686                rp,
10687                ..
10688            } => (bytes, *qtype, *row_bytes, *scale, *rp),
10689            _ => unreachable!("uses_q8_1_fast guaranteed Quant"),
10690        };
10691        // Q4_0 split-plane mirror: only the mmvq/batched decode arms read it (the _rp twins);
10692        // the dp4a/oracle tails below keep the raw GGUF bytes.
10693        let (mbytes, mrp) = match w {
10694            GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
10695            _ => (bytes, rp),
10696        };
10697        // PERF-3 decode-GEMV: warp-per-row MMVQ for the m=1 decode arm, gated behind MEMRA_MMVQ.
10698        // Only the 4 daily-hot dtypes have an _mmvq kernel (Q8_0/Q4_K/Q6_K/NVFP4); Q5_K/Q3_K/IQ4_XS
10699        // keep _dp4a (the oracle/fallback). Bit-equivalent to _dp4a up to f32 reduction order.
10700        if m == 1 && self.mmvq_supports(qtype) {
10701            return self.qmatvec_mmvq(mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, mrp);
10702        }
10703        // BATCHED weight-resident matvec for the m=2-4 band (the MTP/verify forward: full_attn_verify
10704        // and decode_step_t run their projections at m=T=k=2..4). The plain _dp4a path below launches
10705        // grid.y=m INDEPENDENT blocks per output row -> the weight row is re-read m times from HBM/L2.
10706        // The _b2/_b4 kernels walk the weight ONCE and dp4a vs all m activation columns, so m tokens
10707        // cost ~1 weight read instead of m (decode is weight-BW-bound). BIT-IDENTICAL per (token,row)
10708        // to the _mmvq path (32-thread warp reduce — NOT the dp4a 128-thread reduce below).
10709        // m=2 -> mcols=2; m∈{3,4} -> mcols=4; m∈{5..8} -> mcols=8 (kernel guards c>=m).
10710        // MEMRA_NO_BATCHED forces the per-m grid.y=m path (the A/B reference); MEMRA_B8=0 keeps
10711        // m=5..8 on the old per-m path (b8-tier-only seam).
10712        // DECODE-PARITY GATE (2026-07-07): batched iff mmvq_supports — see matmul's parity note.
10713        // Without MEMRA_MMVQ, m=1 decode rides dp4a (the arm below at m=1); the verify must ride
10714        // the SAME class per column (grid.y=m dp4a = the exact m=1 dp4a program per column).
10715        if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
10716            && std::env::var("MEMRA_NO_BATCHED").is_err()
10717            && (m <= 4 || Self::b8_enabled())
10718            // b16 tier: every class routed here now has base + _rp b16 kernels (Q4_0/Q6_K
10719            // pre-existing; NVFP4/Q4_K/Q8_0-base/F8_E4M3 added lane/rp-on-st 2026-08-06), so
10720            // there is no mirror precondition left — `mrp` still selects the LAYOUT below.
10721            && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || qtype == QT_NVFP4
10722                || qtype == QT_Q4_K || qtype == QT_Q5_K || qtype == QT_F8_E4M3 || qtype == QT_Q8_0)
10723        {
10724            let mcols = Self::batched_mcols(m);
10725            return self.qmatvec_mmvq_batched(
10726                mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, mrp,
10727            );
10728        }
10729        // F8-E4M3 catch-all (m=9..15 / batched-disabled seams): grid.y=m e4m3 mmvq — this dtype
10730        // has NO _dp4a twin, and per (token,row) the mmvq body is the exact m=1 decode program.
10731        // Q4_0 joins the catch-all (2026-07-11): adaptive-K cap 8 makes verify t=9 reachable
10732        // for the first time (past the b8 tier) and Q4_0 has no dp4a twin either. The mirror
10733        // (mbytes/mrp) keeps the rp layout consistent with the m=1 decode program.
10734        if qtype == QT_F8_E4M3 || qtype == QT_Q4_0 {
10735            let (b2, r2) = if qtype == QT_Q4_0 {
10736                (mbytes, mrp)
10737            } else {
10738                (bytes, rp)
10739            };
10740            return self.qmatvec_mmvq(b2, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, r2);
10741        }
10742        let name = match qtype {
10743            QT_Q8_0 => "qmatvec_q8_0_dp4a",
10744            QT_Q4_K => "qmatvec_q4_K_dp4a",
10745            QT_Q6_K => "qmatvec_q6_K_dp4a",
10746            QT_Q5_K => "qmatvec_q5_K_dp4a",
10747            QT_Q3_K => "qmatvec_q3_K_dp4a",
10748            QT_NVFP4 => {
10749                if rp {
10750                    "qmatvec_nvfp4_dp4a_rp"
10751                } else {
10752                    "qmatvec_nvfp4_dp4a"
10753                }
10754            }
10755            QT_IQ4_XS => "qmatvec_iq4_XS_dp4a",
10756            _ => unreachable!(),
10757        };
10758        let f = self.func(name);
10759        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output: skip memset
10760        let cfg = LaunchConfig {
10761            grid_dim: (out_f as u32, m as u32, 1),
10762            block_dim: (128, 1, 1),
10763            shared_mem_bytes: 0,
10764        };
10765        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
10766        let __s_b = self.gpu.stream();
10767        let mut b = __s_b.launch_builder(&f);
10768        b.arg(bytes)
10769            .arg(aq)
10770            .arg(ad)
10771            .arg(&mut y)
10772            .arg(&inf)
10773            .arg(&outf)
10774            .arg(&mi)
10775            .arg(&rb);
10776        unsafe {
10777            b.launch(cfg)?;
10778        }
10779        if scale != 1.0 {
10780            self.scale_inplace(&mut y, scale, m * out_f)?;
10781        }
10782        Ok(y)
10783    }
10784
10785    /// DECODE-EXACT matmul at any m: guarantees the SAME warp-per-row (MMVQ, 32-thread) FP
10786    /// accumulation order as the T=1 decode path for EVERY token row. The spec-decode verify MUST
10787    /// use this for linear-attn projections to be bit-identical to greedy decode. The dp4a kernel
10788    /// (128 threads, two-level reduction) used by `matmul`/`matmul_pre` at m>=5 has a different
10789    /// shfl-tree shape that produces ULP differences propagating through gdn_scan into argmax flips.
10790    /// The MMVQ kernel with grid.y=m already processes each row independently (same 32-thread warp
10791    /// reduce as m=1); this method just forces that path unconditionally.
10792    pub fn matmul_decode_exact(
10793        &self,
10794        w: &crate::model::GpuTensor,
10795        x: &CudaSlice<f32>,
10796        m: usize,
10797    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10798        use crate::model::GpuTensor;
10799        // FLOAT tensors (35B ssm_beta/ssm_alpha on every linear layer, F32 ne=[2048,32]): the
10800        // generic path is cuBLASLt, whose reduction splits are n-DEPENDENT — m=1 vs m=2 col-0
10801        // outputs differ in every bit (probe 2026-07-06: 32/32 bit-diff, maxdiff 3.5e-3), which
10802        // shifted 35B verify logits 0.26-0.56 vs eager and flipped greedy at tight margins (the
10803        // p3 spec FAIL). Decode-exact contract: per-COLUMN m=1 cuBLASLt calls — each column's
10804        // reduction is the exact kernel the T=1 decode path runs, so verify==decode bit-for-bit.
10805        // m<=10 here (K+2 verify tier), so the extra launches are a handful of 4us gemvs.
10806        if let GpuTensor::Float { data, .. } = w {
10807            return self.linear_decode_exact(x, data, m, w.in_features(), w.out_features());
10808        }
10809        // MEMRA_FULL_PREC bf16-resident weight: dequant-on-use, then the per-column decode-exact
10810        // float linear (same n-independent reduction contract as the Float arm above).
10811        if let GpuTensor::FloatBf16 { data, .. } = w {
10812            let (in_f, out_f) = (w.in_features(), w.out_features());
10813            return self.linear_bf16_chunked(x, data, m, in_f, out_f, true);
10814        }
10815        if !self.uses_q8_1_fast(w) {
10816            return self.matmul(w, x, m);
10817        }
10818        let in_f = w.in_features();
10819        let out_f = w.out_features();
10820        let (bytes, qtype, row_bytes, scale, rp) = match w {
10821            GpuTensor::Quant {
10822                bytes,
10823                qtype,
10824                row_bytes,
10825                scale,
10826                rp,
10827                ..
10828            } => (bytes, *qtype, *row_bytes, *scale, *rp),
10829            _ => return self.matmul(w, x, m),
10830        };
10831        // Q4_0 split-plane mirror for the mmvq/batched arms below (dp4a tail = matmul_pre,
10832        // which does its own mirror pick).
10833        let (bytes, rp) = match w {
10834            GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
10835            _ => (bytes, rp),
10836        };
10837        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10838        // BLOCK-128 e4m3 (QT_F8_E4M3_BLK): the same single kernel every other entry dispatches, so
10839        // the decode-exact contract needs nothing special — grid.y=m runs the m=1 program per
10840        // (token,row) by construction, which is exactly what this method exists to guarantee.
10841        if let Some(y) = self.try_e4m3_blk_pre(w, &aq, &ad, m)? {
10842            return Ok(y);
10843        }
10844        // Batched weight-resident matvec for m=2-8: BIT-IDENTICAL per (token,row) to MMVQ (exact
10845        // integer dp4a, same warp reduce — kernel-check gate rel=0.00e0), one weight read for m
10846        // tokens. The dispatch the divergence fix must avoid is dp4a's 128-thread two-level
10847        // reduce, NOT this. m=5..8 is the K=4..7 spec-verify tier (b8): pre-b8 T=5 fell to the
10848        // grid.y=m per-row MMVQ below = 5 full weight reads/launch — the measured 27B K=4 cliff.
10849        // DECODE-PARITY GATE (2026-07-07): batched (MMVQ-class order) only when the m=1 decode
10850        // chain rides MMVQ too — without MEMRA_MMVQ decode is dp4a, so the exact-contract here
10851        // must be per-column dp4a (matmul_pre fallthrough), not the MMVQ order.
10852        if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
10853            && std::env::var("MEMRA_NO_BATCHED").is_err()
10854            && (m <= 4 || Self::b8_enabled())
10855            // Every b16 class has base + _rp twins after lane/rp-on-st (see matmul_pre's note):
10856            // no mirror precondition, `rp` selects the layout only.
10857            && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || qtype == QT_F8_E4M3
10858                || qtype == QT_NVFP4 || qtype == QT_Q4_K || qtype == QT_Q5_K || qtype == QT_Q8_0)
10859        {
10860            let mcols = Self::batched_mcols(m);
10861            return self.qmatvec_mmvq_batched(
10862                bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, rp,
10863            );
10864        }
10865        if self.mmvq_supports(qtype) {
10866            // MMVQ at grid.y=m: each row is processed by its own warp independently — same 32-thread
10867            // accumulation + warp_reduce_sum as m=1 decode. Bit-identical per row.
10868            return self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, scale, rp);
10869        }
10870        // Fallback for non-MMVQ quant types (Q5_K, Q3_K): use dp4a (the only available kernel).
10871        // These types are not used in the 27B's linear-attn NVFP4+Q4_K layers.
10872        self.matmul_pre(w, &aq, &ad, x, m)
10873    }
10874
10875    /// DECODE-EXACT matmul from a PRE-QUANTIZED q8_1 activation (batched-verify epilogue
10876    /// re-fuse, lane/vt-fixes fix 2, 2026-08-03): the EXACT `matmul_decode_exact` dispatch for
10877    /// q8_1-fast Quant tensors, with the caller's (aq, ad) replacing the internal
10878    /// `quantize_q8_1`. quantize_q8_1 is deterministic (same input bytes -> same q8 bytes), so
10879    /// sharing one quantize across sibling matmuls of the same activation — or consuming the
10880    /// q8 emitted by a fused epilogue (rms_norm_q8_1 / add_rms_norm_q8_1 /
10881    /// silu_mul_scaled_q8_1 / gated_rmsnorm_q8_1, all kernel-check-pinned bit-identical to
10882    /// their unfused chains) — cannot change any dispatched kernel's input bytes.
10883    /// Caller MUST guarantee `uses_q8_1_fast(w)` (the fused epilogues only exist on that path).
10884    pub fn matmul_decode_exact_pre(
10885        &self,
10886        w: &crate::model::GpuTensor,
10887        aq: &CudaSlice<i8>,
10888        ad: &CudaSlice<f32>,
10889        m: usize,
10890    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10891        use crate::model::GpuTensor;
10892        debug_assert!(
10893            self.uses_q8_1_fast(w),
10894            "matmul_decode_exact_pre: caller must guarantee q8_1-fast"
10895        );
10896        // BLOCK-128 e4m3: same single kernel, all m — see matmul_decode_exact's note.
10897        if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
10898            return Ok(y);
10899        }
10900        let in_f = w.in_features();
10901        let out_f = w.out_features();
10902        let (bytes, qtype, row_bytes, scale, rp) = match w {
10903            GpuTensor::Quant {
10904                bytes,
10905                qtype,
10906                row_bytes,
10907                scale,
10908                rp,
10909                ..
10910            } => (bytes, *qtype, *row_bytes, *scale, *rp),
10911            _ => {
10912                return Err(
10913                    "matmul_decode_exact_pre: Quant tensor required (q8_1-fast contract)".into(),
10914                );
10915            }
10916        };
10917        // Q4_0 split-plane mirror — same pick as matmul_decode_exact.
10918        let (bytes, rp) = match w {
10919            GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
10920            _ => (bytes, rp),
10921        };
10922        // Dispatch mirror of matmul_decode_exact's q8_1-fast tail, condition for condition.
10923        if (2..=16).contains(&m)
10924            && self.batched_supports(qtype)
10925            && self.mmvq_supports(qtype)
10926            && std::env::var("MEMRA_NO_BATCHED").is_err()
10927            && (m <= 4 || Self::b8_enabled())
10928            && (m <= 8
10929                || qtype == QT_Q4_0
10930                || qtype == QT_Q6_K
10931                || qtype == QT_F8_E4M3
10932                || qtype == QT_NVFP4
10933                || qtype == QT_Q4_K
10934                || qtype == QT_Q5_K
10935                || qtype == QT_Q8_0)
10936        {
10937            let mcols = Self::batched_mcols(m);
10938            return self.qmatvec_mmvq_batched(
10939                bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, rp,
10940            );
10941        }
10942        if self.mmvq_supports(qtype) {
10943            return self.qmatvec_mmvq(bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, rp);
10944        }
10945        // Non-MMVQ quant types (Q5_K/Q3_K under MEMRA_MMVQ=0): dp4a via matmul_pre — the same
10946        // fallback matmul_decode_exact takes. m <= 16 on the verify tier never reads x_fallback.
10947        let x0 = self.zeros(0)?;
10948        self.matmul_pre(w, aq, ad, &x0, m)
10949    }
10950
10951    /// DUAL gate+up batched matvec from a PRE-QUANTIZED activation, macro-scales DEFERRED
10952    /// (lane/vt-fixes fix 2): same eligibility as `matmul_decode_exact_dual`, but the caller's
10953    /// (aq, ad) replaces the internal quantize and the NVFP4 per-tensor scales are RETURNED
10954    /// instead of applied via two `scale_inplace` launches — the fused SwiGLU epilogue
10955    /// (`silu_mul_scaled_q8_1`) folds them, exactly like the m=1 decode chain does. Deferring
10956    /// is value-exact: `y[i]*s` inline in the epilogue is the same IEEE multiply scale_inplace
10957    /// would store (f32 store/load round-trips are exact). None -> caller falls back to the
10958    /// per-tensor path.
10959    pub fn matmul_decode_exact_dual_pre(
10960        &self,
10961        w0: &crate::model::GpuTensor,
10962        w1: &crate::model::GpuTensor,
10963        aq: &CudaSlice<i8>,
10964        ad: &CudaSlice<f32>,
10965        m: usize,
10966    ) -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn std::error::Error>>
10967    {
10968        use crate::model::GpuTensor;
10969        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
10970        let on = *ON.get_or_init(|| {
10971            std::env::var("MEMRA_SPEC_DUAL_T")
10972                .map(|v| v != "0")
10973                .unwrap_or(true)
10974        });
10975        if !on
10976            || !(2..=7).contains(&m)
10977            || std::env::var("MEMRA_NO_BATCHED").is_ok()
10978            || !self.uses_q8_1_fast(w0)
10979            || !self.uses_q8_1_fast(w1)
10980        {
10981            return Ok(None);
10982        }
10983        // DECODE-PARITY GATE (lane/nvfp4-strict, 2026-08-05): batched iff MMVQ — the dual
10984        // kernels are the MMVQ warp-reduce family, and without MEMRA_MMVQ the m=1 decode
10985        // chain this verify must match bit-for-bit rides dp4a (see matmul_decode_exact's
10986        // note). The singles enforce this via `mmvq_supports`; the dual door skipped it.
10987        if !self.mmvq_supports(QT_NVFP4) {
10988            return Ok(None);
10989        }
10990        let (in_f, out_f) = (w0.in_features(), w0.out_features());
10991        if w1.in_features() != in_f || w1.out_features() != out_f {
10992            return Ok(None);
10993        }
10994        let (b0, b1, row_bytes, s0, s1, rp) = match (w0, w1) {
10995            (
10996                GpuTensor::Quant {
10997                    bytes: b0,
10998                    qtype: q0,
10999                    row_bytes: rb0,
11000                    scale: s0,
11001                    rp: rp0,
11002                    rp4: None,
11003                    ..
11004                },
11005                GpuTensor::Quant {
11006                    bytes: b1,
11007                    qtype: q1,
11008                    row_bytes: rb1,
11009                    scale: s1,
11010                    rp: rp1,
11011                    rp4: None,
11012                    ..
11013                },
11014            ) if *q0 == QT_NVFP4 && *q1 == QT_NVFP4 && rb0 == rb1 && rp0 == rp1 => {
11015                (b0, b1, *rb0, *s0, *s1, *rp0)
11016            }
11017            _ => return Ok(None),
11018        };
11019        // m=5..7: only the exact-width rp duals exist (vt-fixes fix 1b); GGUF layout keeps
11020        // the singles. The b8 dual (MCOLS=8 at m=5..8) measured FLAT and stays dead.
11021        if m > 4 && !(rp && Self::b8_enabled() && std::env::var("MEMRA_B567").as_deref() != Ok("0"))
11022        {
11023            return Ok(None);
11024        }
11025        let (y0, y1) =
11026            self.qmatvec_batched_dual_raw(b0, b1, aq, ad, m, in_f, out_f, row_bytes, rp)?;
11027        Ok(Some(((y0, s0), (y1, s1))))
11028    }
11029
11030    /// GROUP-4 GDN-tuple BATCHED matvec (trunk-kernels slice C): the qwen35 linear-layer
11031    /// in-projection 4-tuple (wqkv / wqkv_gate / ssm_beta / ssm_alpha) from ONE pre-quantized
11032    /// activation in ONE launch. Blocks map to the concatenated row space; every out_f must be
11033    /// a multiple of 8 (rows_per_block) so each warp's row pair resolves to one tensor; per
11034    /// (tensor, token, row) the kernel body is `nvfp4_mmvq_batched_rp` VERBATIM with the
11035    /// tensor's macro-scale fused at the write (== the conditional scale_inplace pass,
11036    /// bit-identical) -> BIT-IDENTICAL to the four single launches. Split-plane rp NVFP4 only,
11037    /// m=2..8 (exact-width MCOLS at m=5..7 mirroring the B567 law; m=8 requires b8_enabled
11038    /// like the singles). None -> caller runs the four singles. MEMRA_TK_GDN_GROUP=0 rollback.
11039    pub fn matmul_decode_exact_group4_pre(
11040        &self,
11041        ws: [&crate::model::GpuTensor; 4],
11042        aq: &CudaSlice<i8>,
11043        ad: &CudaSlice<f32>,
11044        m: usize,
11045    ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
11046        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11047        let on = *ON.get_or_init(|| {
11048            std::env::var("MEMRA_TK_GDN_GROUP")
11049                .map(|v| v != "0")
11050                .unwrap_or(true)
11051        });
11052        self.matmul_decode_exact_group_pre(&ws, aq, ad, m, on, "GDN group4")
11053    }
11054
11055    /// GROUP-3 twin for the qwen35 full-attention q/k/v triple (trunk-kernels slice D):
11056    /// the SAME group4 kernels with n3=0 (blocks never reach the fourth range; W3/y3 are
11057    /// never dereferenced) — per (tensor, token, row) bit-identical to the three singles
11058    /// exactly as the group4 door is to its four. MEMRA_TK_FA_GROUP=0 rollback.
11059    pub fn matmul_decode_exact_group3_pre(
11060        &self,
11061        ws: [&crate::model::GpuTensor; 3],
11062        aq: &CudaSlice<i8>,
11063        ad: &CudaSlice<f32>,
11064        m: usize,
11065    ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
11066        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11067        let on = *ON.get_or_init(|| {
11068            std::env::var("MEMRA_TK_FA_GROUP")
11069                .map(|v| v != "0")
11070                .unwrap_or(true)
11071        });
11072        self.matmul_decode_exact_group_pre(&ws, aq, ad, m, on, "FA group3")
11073    }
11074
11075    /// Shared core of the group3/group4 doors: eligibility mirror of the singles' batched
11076    /// dispatch, then ONE `qmatvec_nvfp4_mmvq_group4_b*_rp` launch over the concatenated
11077    /// row space (3-tensor callers ride n3=0). Returns one output per input tensor.
11078    fn matmul_decode_exact_group_pre(
11079        &self,
11080        ws: &[&crate::model::GpuTensor],
11081        aq: &CudaSlice<i8>,
11082        ad: &CudaSlice<f32>,
11083        m: usize,
11084        on: bool,
11085        tag: &'static str,
11086    ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
11087        use crate::model::GpuTensor;
11088        if !on
11089            || !(2..=8).contains(&m)
11090            || std::env::var("MEMRA_NO_BATCHED").is_ok()
11091            || (m > 4 && !Self::b8_enabled())
11092            || !self.mmvq_supports(QT_NVFP4)
11093            || !self.batched_supports(QT_NVFP4)
11094        {
11095            return Ok(None);
11096        }
11097        let in_f = ws[0].in_features();
11098        let mut parts: Vec<(&CudaSlice<u8>, usize, f32)> = Vec::with_capacity(4);
11099        for w in ws {
11100            if !self.uses_q8_1_fast(w) || w.in_features() != in_f {
11101                return Ok(None);
11102            }
11103            match w {
11104                GpuTensor::Quant {
11105                    bytes,
11106                    qtype,
11107                    scale,
11108                    rp: true,
11109                    rp4: None,
11110                    ..
11111                } if *qtype == QT_NVFP4 && w.out_features() % 8 == 0 => {
11112                    parts.push((bytes, w.out_features(), *scale));
11113                }
11114                _ => return Ok(None),
11115            }
11116        }
11117        // MCOLS tier mirrors the singles: batched_mcols + the B567 exact-width law at m=5..7.
11118        static B567: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11119        let b567 = *B567.get_or_init(|| std::env::var("MEMRA_B567").as_deref() != Ok("0"));
11120        let mcols = if (5..=7).contains(&m) && b567 {
11121            m
11122        } else {
11123            Self::batched_mcols(m)
11124        };
11125        let kname: &'static str = match mcols {
11126            2 => "qmatvec_nvfp4_mmvq_group4_b2_rp",
11127            4 => "qmatvec_nvfp4_mmvq_group4_b4_rp",
11128            5 => "qmatvec_nvfp4_mmvq_group4_b5_rp",
11129            6 => "qmatvec_nvfp4_mmvq_group4_b6_rp",
11130            7 => "qmatvec_nvfp4_mmvq_group4_b7_rp",
11131            8 => "qmatvec_nvfp4_mmvq_group4_b8_rp",
11132            _ => return Ok(None),
11133        };
11134        // Engagement receipt PER DOOR (dead-arm lesson): one shared Once here suppressed
11135        // the second door's print on the slice-D battery — key the once-set by tag.
11136        if std::env::var("MEMRA_DEBUG").is_ok() {
11137            use std::sync::Mutex;
11138            static SEEN: Mutex<Vec<&'static str>> = Mutex::new(Vec::new());
11139            let mut seen = SEEN.lock().unwrap();
11140            if !seen.contains(&tag) {
11141                seen.push(tag);
11142                eprintln!("[memra] {tag} batched ENGAGED (m={m})");
11143            }
11144        }
11145        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
11146        let rows_per_block = ROWS_PER_BLOCK * 2; // WROWS=2 in the group kernel
11147        let total: usize = parts.iter().map(|p| p.1).sum();
11148        let three = parts.len() == 3;
11149        let mut y0 = self.alloc_uninit::<f32>(m * parts[0].1)?;
11150        let mut y1 = self.alloc_uninit::<f32>(m * parts[1].1)?;
11151        let mut y2 = self.alloc_uninit::<f32>(m * parts[2].1)?;
11152        // 3-tensor callers: n3=0 means no block ever resolves to the fourth range — W3/y3
11153        // are never dereferenced; a 1-element dummy keeps the launch ABI without aliasing y0.
11154        let mut y3 = self.alloc_uninit::<f32>(if three { 1 } else { m * parts[3].1 })?;
11155        let cfg = LaunchConfig {
11156            grid_dim: ((total as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
11157            block_dim: (32, ROWS_PER_BLOCK, 1),
11158            shared_mem_bytes: 0,
11159        };
11160        let (inf, mi) = (in_f as i32, m as i32);
11161        let (n0, n1, n2) = (parts[0].1 as i32, parts[1].1 as i32, parts[2].1 as i32);
11162        let n3 = if three { 0i32 } else { parts[3].1 as i32 };
11163        let (s0, s1, s2) = (parts[0].2, parts[1].2, parts[2].2);
11164        let s3 = if three { 1.0f32 } else { parts[3].2 };
11165        let w3 = if three { parts[0].0 } else { parts[3].0 };
11166        let f = self.func(kname);
11167        let __s_b = self.gpu.stream();
11168        let mut b = __s_b.launch_builder(&f);
11169        b.arg(parts[0].0)
11170            .arg(parts[1].0)
11171            .arg(parts[2].0)
11172            .arg(w3)
11173            .arg(aq)
11174            .arg(ad)
11175            .arg(&mut y0)
11176            .arg(&mut y1)
11177            .arg(&mut y2)
11178            .arg(&mut y3)
11179            .arg(&inf)
11180            .arg(&n0)
11181            .arg(&n1)
11182            .arg(&n2)
11183            .arg(&n3)
11184            .arg(&mi)
11185            .arg(&s0)
11186            .arg(&s1)
11187            .arg(&s2)
11188            .arg(&s3);
11189        unsafe {
11190            b.launch(cfg)?;
11191        }
11192        Ok(Some(if three {
11193            vec![y0, y1, y2]
11194        } else {
11195            vec![y0, y1, y2, y3]
11196        }))
11197    }
11198
11199    /// DUAL gate+up BATCHED matvec at verify t=2..8 (lane/verify-economics, 2026-08-02): ONE
11200    /// launch computes both FFN projections of a verify batch — same activation, same shape,
11201    /// blockIdx.y selects the tensor. Per (tensor, token, row) the kernel body is the single
11202    /// batched program on the SAME layout (split-plane rp: b2 rp / b4 rpr2 / b8 rpr2; GGUF:
11203    /// b2 base / b4 r2 / b8 r2) -> BIT-IDENTICAL to the two single `matmul_decode_exact`
11204    /// launches (kernel-check gates bitwise on both layouts; run-spec K=1..8 arbitrates e2e).
11205    /// The one activation quantize replaces two IDENTICAL quantizes of the same `x` (same
11206    /// kernel, same input -> same q8_1 bytes), and the two independent weight streams in one
11207    /// grid restore the memory-level parallelism the two-launch form loses to tail drain +
11208    /// launch gap (m=1 dual_mr2 precedent: DRAM 40% -> 47-50% on the 27B pair).
11209    /// `Some((y0, y1))` only when both tensors are NVFP4, the SAME layout (both rp or both
11210    /// GGUF, no rp4 mirror), identical (in_f, out_f, row_bytes), q8_1-fast, and m in 2..=4
11211    /// (the b2/b4 tiers = verify T for K=1..3, the profitable-K window — the b8 dual measured
11212    /// FLAT vs the rpsc singles x3 interleaved, research/verify-economics-20260802, and was
11213    /// killed per doctrine). None -> caller runs the two singles. MEMRA_SPEC_DUAL_T=0 rollback.
11214    pub fn matmul_decode_exact_dual(
11215        &self,
11216        w0: &crate::model::GpuTensor,
11217        w1: &crate::model::GpuTensor,
11218        x: &CudaSlice<f32>,
11219        m: usize,
11220    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
11221        use crate::model::GpuTensor;
11222        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11223        let on = *ON.get_or_init(|| {
11224            std::env::var("MEMRA_SPEC_DUAL_T")
11225                .map(|v| v != "0")
11226                .unwrap_or(true)
11227        });
11228        if !on
11229            || !(2..=4).contains(&m)
11230            || std::env::var("MEMRA_NO_BATCHED").is_ok()
11231            || !self.uses_q8_1_fast(w0)
11232            || !self.uses_q8_1_fast(w1)
11233        {
11234            return Ok(None);
11235        }
11236        // DECODE-PARITY GATE (lane/nvfp4-strict, 2026-08-05): batched iff MMVQ — same law as
11237        // the singles' `batched_supports && mmvq_supports` check in matmul_decode_exact,
11238        // which this dual door bypassed. Without MEMRA_MMVQ the m=1 decode is dp4a; the
11239        // verify must ride the per-column dp4a class, not the MMVQ-family dual.
11240        if !self.mmvq_supports(QT_NVFP4) {
11241            return Ok(None);
11242        }
11243        let (in_f, out_f) = (w0.in_features(), w0.out_features());
11244        if w1.in_features() != in_f || w1.out_features() != out_f {
11245            return Ok(None);
11246        }
11247        let (b0, b1, row_bytes, s0, s1, rp) = match (w0, w1) {
11248            (
11249                GpuTensor::Quant {
11250                    bytes: b0,
11251                    qtype: q0,
11252                    row_bytes: rb0,
11253                    scale: s0,
11254                    rp: rp0,
11255                    rp4: None,
11256                    ..
11257                },
11258                GpuTensor::Quant {
11259                    bytes: b1,
11260                    qtype: q1,
11261                    row_bytes: rb1,
11262                    scale: s1,
11263                    rp: rp1,
11264                    rp4: None,
11265                    ..
11266                },
11267            ) if *q0 == QT_NVFP4 && *q1 == QT_NVFP4 && rb0 == rb1 && rp0 == rp1 => {
11268                (b0, b1, *rb0, *s0, *s1, *rp0)
11269            }
11270            _ => return Ok(None),
11271        };
11272        // Engagement receipt (MEMRA_DEBUG=1): the first dead-arm A/B lesson — a `rp: false`
11273        // gate silently no-op'd the whole experiment; prove the arm is live in the log.
11274        if std::env::var("MEMRA_DEBUG").is_ok() {
11275            static ONCE: std::sync::Once = std::sync::Once::new();
11276            ONCE.call_once(|| eprintln!("[memra] dual gate+up batched ENGAGED (m={m} rp={rp})"));
11277        }
11278        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
11279        let (y0, y1) =
11280            self.qmatvec_batched_dual_raw(b0, b1, &aq, &ad, m, in_f, out_f, row_bytes, rp)?;
11281        let mut y0 = y0;
11282        let mut y1 = y1;
11283        if s0 != 1.0 {
11284            self.scale_inplace(&mut y0, s0, m * out_f)?;
11285        }
11286        if s1 != 1.0 {
11287            self.scale_inplace(&mut y1, s1, m * out_f)?;
11288        }
11289        Ok(Some((y0, y1)))
11290    }
11291
11292    /// Launch body of the dual batched twins from raw NVFP4 weight bytes + a pre-quantized q8_1
11293    /// activation (kernel-check's bit-equivalence entry; matmul_decode_exact_dual's core).
11294    /// mcols tier = batched_mcols(m); macro-scale NOT applied. `rp` selects the split-plane
11295    /// twins (both buffers must be the repacked layout).
11296    #[allow(clippy::too_many_arguments)]
11297    pub fn qmatvec_batched_dual_raw(
11298        &self,
11299        b0: &CudaSlice<u8>,
11300        b1: &CudaSlice<u8>,
11301        aq: &CudaSlice<i8>,
11302        ad: &CudaSlice<f32>,
11303        m: usize,
11304        in_f: usize,
11305        out_f: usize,
11306        row_bytes: usize,
11307        rp: bool,
11308    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11309        const ROWS_PER_BLOCK: u32 = 4;
11310        let mcols = Self::batched_mcols(m);
11311        // EXACT-WIDTH duals at m=5..7 (vt-fixes fix 1b): rp-only; bit-identical to the two
11312        // b5/b6/b7 singles (blockIdx.y selects the tensor, same template body).
11313        let tiny_rp1 = rp
11314            && mcols == 4
11315            && out_f <= 128
11316            && std::env::var("MEMRA_NVFP4_AUX_DUAL").as_deref() != Ok("0");
11317        let (name, rows_per_block) = if tiny_rp1 {
11318            ("qmatvec_nvfp4_mmvq_dual_b4_rp", ROWS_PER_BLOCK)
11319        } else {
11320            match (mcols, rp, m) {
11321                (2, false, _) => ("qmatvec_nvfp4_mmvq_dual_b2", ROWS_PER_BLOCK),
11322                (4, false, _) => ("qmatvec_nvfp4_mmvq_dual_b4_r2", ROWS_PER_BLOCK * 2),
11323                (2, true, _) => ("qmatvec_nvfp4_mmvq_dual_b2_rp", ROWS_PER_BLOCK),
11324                (4, true, _) => ("qmatvec_nvfp4_mmvq_dual_b4_rpr2", ROWS_PER_BLOCK * 2),
11325                (8, true, 5) => ("qmatvec_nvfp4_mmvq_dual_b5_rpr2", ROWS_PER_BLOCK * 2),
11326                (8, true, 6) => ("qmatvec_nvfp4_mmvq_dual_b6_rpr2", ROWS_PER_BLOCK * 2),
11327                (8, true, 7) => ("qmatvec_nvfp4_mmvq_dual_b7_rpr2", ROWS_PER_BLOCK * 2),
11328                _ => {
11329                    return Err(
11330                        format!("qmatvec_batched_dual_raw: no dual kernel for m {m}").into(),
11331                    );
11332                }
11333            }
11334        };
11335        let f = self.func(name);
11336        let mut y0 = self.alloc_uninit::<f32>(m * out_f)?;
11337        let mut y1 = self.alloc_uninit::<f32>(m * out_f)?;
11338        let cfg = LaunchConfig {
11339            grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 2, 1),
11340            block_dim: (32, ROWS_PER_BLOCK, 1),
11341            shared_mem_bytes: 0,
11342        };
11343        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
11344        let __s_b = self.gpu.stream();
11345        let mut b = __s_b.launch_builder(&f);
11346        b.arg(b0)
11347            .arg(b1)
11348            .arg(aq)
11349            .arg(ad)
11350            .arg(&mut y0)
11351            .arg(&mut y1)
11352            .arg(&inf)
11353            .arg(&outf)
11354            .arg(&mi)
11355            .arg(&rb);
11356        unsafe {
11357            b.launch(cfg)?;
11358        }
11359        Ok((y0, y1))
11360    }
11361
11362    /// Like `matmul_pre` but RETURNS THE RAW (un-macro-scaled) matmul output together with the
11363    /// per-tensor NVFP4 scale, instead of applying `scale_inplace` internally. Used by the fused
11364    /// SwiGLU epilogue (RANK3 LEVER 2) so the gate/up scales fold into one `silu_mul_scaled` launch.
11365    /// `Some((y_raw, scale))` only on the m==1 decode fast path (mmvq / dp4a) where the scale is a
11366    /// separate post-launch op we can defer; returns `None` for every other path (prefill GEMM, FP4
11367    /// GEMM, Stage-A, Float) so the caller falls back to the scaled `matmul_pre` + `silu_mul`.
11368    /// DUAL gate+up NVFP4 matvec (mm-fusion): ONE launch computes both projections (same
11369    /// activation, same shape) — grid.y selects the tensor. Bit-identical per element to two
11370    /// mr2 launches at m=1. Returns (gate_raw, up_raw) un-scaled (caller folds the two macro
11371    /// scales into the SwiGLU epilogue, same as the matmul_pre_noscale contract). None unless
11372    /// both tensors are NVFP4 q8_1-fast with identical (in_f, out_f, row_bytes) and m==1.
11373    pub fn matmul_pre_dual_noscale(
11374        &self,
11375        w0: &crate::model::GpuTensor,
11376        w1: &crate::model::GpuTensor,
11377        aq: &CudaSlice<i8>,
11378        ad: &CudaSlice<f32>,
11379        m: usize,
11380    ) -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn std::error::Error>>
11381    {
11382        use crate::model::GpuTensor;
11383        if m != 1 || !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) {
11384            return Ok(None);
11385        }
11386        // FP-ORDER LAW (lane/nvfp4-strict, 2026-08-05): every kernel this door can dispatch
11387        // (q8_0 fused2, nvfp4 dual_mr2) is the MMVQ family — 32-thread warp reduce. Without
11388        // MEMRA_MMVQ the m=1 singles ride dp4a (128-thread two-level reduce), so fusing here
11389        // would mix dispatch families across the pair — the exact class `q8_fused_params`
11390        // already refuses for Q8_0. The NVFP4 arm lacked this check, which is why
11391        // decode-batch-gate `--mode strict`'s equalizing env (MEMRA_MMVQ=0) never pinned
11392        // NVFP4 models: decode_step_h kept riding dual_mr2 while the batched body fell to
11393        // dp4a (gate1 maxdiff 1.639e-1 / gate2 step-8 divergence at the 2026-08-05 train
11394        // HEAD, research/nvfp4-strict-20260805/). Default env (MMVQ on) is dispatch-unchanged.
11395        if !self.mmvq_supports(QT_NVFP4) {
11396            return Ok(None);
11397        }
11398        let (in_f, out_f) = (w0.in_features(), w0.out_features());
11399        if w1.in_features() != in_f || w1.out_features() != out_f {
11400            return Ok(None);
11401        }
11402        // Q8_0 ARM (lane/q27-deepdive, 2026-08-05): the dense-FFN gate+up pair on a Q8_0 trunk fell
11403        // through this NVFP4-only gate to two `matmul_pre_noscale` launches — measured 128 of the
11404        // 1015 launches/token on q27-Q8_0 decode, the single largest un-fused class in the tick
11405        // (nsys `research/q27-deepdive-20260805/nsys/`). `q8_fused2_core` already serves the same
11406        // pair shape for the shared-expert gate/up, and its kernel body is `qmatvec_q8_0_mmvq`
11407        // VERBATIM per (tensor,row) -> BIT-IDENTICAL to the two separate launches. Q8_0 carries no
11408        // macro-scale (q8_fused_params requires scale==1.0), so the noscale contract is satisfied
11409        // by returning 1.0 for both: the SwiGLU epilogue's fold becomes the identity it already is
11410        // on this dtype today. Seam: MEMRA_Q8_FFN_FUSE2=0 rolls back to the two-launch pair.
11411        // rp4 guard: with MEMRA_Q8RP the singles route to the `_rp` split-plane twin over the
11412        // mirror buffer; the fused2 kernel has no `_rp` form, so fusing there would swap
11413        // dispatch families mid-model. Bail and let the two singles run (mirror lane unchanged).
11414        let no_mirror =
11415            |w: &crate::model::GpuTensor| !matches!(w, GpuTensor::Quant { rp4: Some(_), .. });
11416        if self.q8_ffn_fuse2_on()
11417            && no_mirror(w0)
11418            && no_mirror(w1)
11419            && let Some([p0, p1]) = self.q8_fused_params(&[w0, w1])
11420        {
11421            let (y0, y1) = self.q8_fused2_core(p0.0, p1.0, aq, ad, in_f, p0.1, p1.1, p0.2)?;
11422            return Ok(Some(((y0, 1.0), (y1, 1.0))));
11423        }
11424        // F8-E4M3 ARM (lane/fp8-decode-v1, 2026-08-05): with native e4m3 residency the FFN gate+up
11425        // pair (and the ssm beta+alpha dual, which routes through this same entry) fell through
11426        // both the NVFP4 gate below and the Q8_0 arm above to two `matmul_pre_noscale` launches —
11427        // native residency was UN-FUSING the trunk relative to the Q8_0 slab it replaces. The
11428        // fused2 kernel body is `qmatvec_e4m3_mmvq` VERBATIM per (tensor,row). Contract match:
11429        // `matmul_pre_noscale` on e4m3 launches with scale 1.0 and RETURNS the per-tensor
11430        // weight_scale for the caller to fold, so we pass ws=1.0 here and return (s0,s1) — same
11431        // bits, and the two macro-scale multiplies still fold into the SwiGLU epilogue.
11432        // MEMRA_E4M3_DUAL=0 rolls back to the two-launch pair.
11433        if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
11434            let (y0, y1) =
11435                self.e4m3_fused2_core(p0.0, p1.0, aq, ad, in_f, p0.1, p1.1, p0.2, 1.0, 1.0)?;
11436            return Ok(Some(((y0, p0.3), (y1, p1.3))));
11437        }
11438        let (b0, q0, rb0, s0, rp0) = match w0 {
11439            GpuTensor::Quant {
11440                bytes,
11441                qtype,
11442                row_bytes,
11443                scale,
11444                rp,
11445                ..
11446            } => (bytes, *qtype, *row_bytes, *scale, *rp),
11447            _ => return Ok(None),
11448        };
11449        let (b1, q1, rb1, s1, rp1) = match w1 {
11450            GpuTensor::Quant {
11451                bytes,
11452                qtype,
11453                row_bytes,
11454                scale,
11455                rp,
11456                ..
11457            } => (bytes, *qtype, *row_bytes, *scale, *rp),
11458            _ => return Ok(None),
11459        };
11460        if q0 != QT_NVFP4 || q1 != QT_NVFP4 || rb0 != rb1 || rp0 != rp1 {
11461            return Ok(None);
11462        }
11463        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
11464        const RPW: u32 = 2;
11465        let rows_per_block = ROWS_PER_BLOCK * RPW;
11466        let f = self.func(if rp0 {
11467            "qmatvec_nvfp4_mmvq_dual_mr2_rp"
11468        } else {
11469            "qmatvec_nvfp4_mmvq_dual_mr2"
11470        });
11471        let mut y0 = self.alloc_uninit::<f32>(out_f)?;
11472        let mut y1 = self.alloc_uninit::<f32>(out_f)?;
11473        let cfg = LaunchConfig {
11474            grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 2, 1),
11475            block_dim: (32, ROWS_PER_BLOCK, 1),
11476            shared_mem_bytes: 0,
11477        };
11478        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, 1i32, rb0 as i64);
11479        // noscale contract: the caller folds s0/s1 into the SwiGLU epilogue — the kernel's fused
11480        // yscale args stay 1.0 here (they exist for the single-tensor callers).
11481        let one = 1.0f32;
11482        let __s_b = self.gpu.stream();
11483        let mut b = __s_b.launch_builder(&f);
11484        b.arg(b0)
11485            .arg(b1)
11486            .arg(aq)
11487            .arg(ad)
11488            .arg(&mut y0)
11489            .arg(&mut y1)
11490            .arg(&inf)
11491            .arg(&outf)
11492            .arg(&mi)
11493            .arg(&rb)
11494            .arg(&one)
11495            .arg(&one);
11496        unsafe {
11497            b.launch(cfg)?;
11498        }
11499        Ok(Some(((y0, s0), (y1, s1))))
11500    }
11501
11502    /// FUSED NVFP4 matvec TRIPLE with unequal out_f (rig-native decode increment 1,
11503    /// lane/rig-native-nvfp4): wq+wk+wv in ONE launch via the q8_0 fused2 block-offset
11504    /// recipe. Per (tensor,row,t) the kernel body is nvfp4_mmvq_multirow_rp VERBATIM ->
11505    /// bit-identical to three separate `matmul_pre` launches; yscales fold in-kernel exactly
11506    /// as the singles do. grid.y = m (the t-parallel verify rows ride the same launch).
11507    /// None when ineligible (not all rp NVFP4 / in_f mismatch / mmvq off) — callers fall
11508    /// back to the three singles.
11509    #[allow(clippy::too_many_arguments)]
11510    pub fn matmul_nvfp4_fused3(
11511        &self,
11512        w0: &crate::model::GpuTensor,
11513        w1: &crate::model::GpuTensor,
11514        w2: &crate::model::GpuTensor,
11515        aq: &CudaSlice<i8>,
11516        ad: &CudaSlice<f32>,
11517        m: usize,
11518    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
11519    {
11520        use crate::model::GpuTensor;
11521        // m == 1 rides the original fused kernel; m = 2..=8 rides the BATCHED fused twin
11522        // (lane/samplat — the fused4_b8 pattern): the seg body is nvfp4_mmvq_batched_rp_sc
11523        // verbatim, weight rows read once for all m columns, bit-identical per
11524        // (tensor,row,column) to the three bN_rpsc singles. The old "at m>1 the fused
11525        // segments would re-read the weight per row" note described the grid.y=m lift,
11526        // which this twin deliberately is NOT.
11527        if !(1..=8).contains(&m)
11528            || !self.mmvq_supports(QT_NVFP4)
11529            || !self.uses_q8_1_fast(w0)
11530            || !self.uses_q8_1_fast(w1)
11531            || !self.uses_q8_1_fast(w2)
11532        {
11533            return Ok(None);
11534        }
11535        if m > 1 {
11536            let in_f = w0.in_features();
11537            if std::env::var("MEMRA_NVFP4_FUSED3B").as_deref() == Ok("0")
11538                || !self.batched_supports(QT_NVFP4)
11539                || std::env::var("MEMRA_NO_BATCHED").is_ok()
11540                || (m > 4 && !Self::b8_enabled())
11541                || in_f % 512 != 0
11542                || in_f / 64 > 272
11543            {
11544                return Ok(None);
11545            }
11546        }
11547        let unpack = |w: &crate::model::GpuTensor| match w {
11548            GpuTensor::Quant {
11549                bytes,
11550                qtype,
11551                scale,
11552                rp,
11553                ..
11554            } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
11555            _ => None,
11556        };
11557        let (Some(p0), Some(p1), Some(p2)) = (unpack(w0), unpack(w1), unpack(w2)) else {
11558            return Ok(None);
11559        };
11560        let in_f = w0.in_features();
11561        if w1.in_features() != in_f || w2.in_features() != in_f {
11562            return Ok(None);
11563        }
11564        let (o0, o1, o2) = (w0.out_features(), w1.out_features(), w2.out_features());
11565        const ROWS_PER_BLOCK: u32 = 4; // MEMRA_MMVQ_ROWS
11566        const RPW: u32 = 2;
11567        let rows_pb = ROWS_PER_BLOCK * RPW;
11568        let nb = |o: usize| (o as u32).div_ceil(rows_pb);
11569        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
11570        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
11571        let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
11572        let (inf, oi0, oi1, oi2, mi) = (in_f as i32, o0 as i32, o1 as i32, o2 as i32, m as i32);
11573        // SAFETY: the raw pointers come straight from the &GpuTensor borrows above and are
11574        // only dereferenced for the launch-arg build inside this call.
11575        let (b0, b1, b2) = unsafe { (&*p0.0, &*p1.0, &*p2.0) };
11576        if m > 1 {
11577            // batched twin has no in-kernel scale — refuse scale carriers (GGUF trunk = 1.0).
11578            if p0.1 != 1.0 || p1.1 != 1.0 || p2.1 != 1.0 {
11579                return Ok(None);
11580            }
11581            let f = self.func("qmatvec_nvfp4_mmvq_fused3_b8_rpsc");
11582            let cfg = LaunchConfig {
11583                grid_dim: (nb(o0) + nb(o1) + nb(o2), 1, 1),
11584                block_dim: (32, ROWS_PER_BLOCK, 1),
11585                shared_mem_bytes: 0,
11586            };
11587            let __s_b = self.gpu.stream();
11588            let mut b = __s_b.launch_builder(&f);
11589            b.arg(b0)
11590                .arg(b1)
11591                .arg(b2)
11592                .arg(aq)
11593                .arg(ad)
11594                .arg(&mut y0)
11595                .arg(&mut y1)
11596                .arg(&mut y2)
11597                .arg(&inf)
11598                .arg(&oi0)
11599                .arg(&oi1)
11600                .arg(&oi2)
11601                .arg(&mi);
11602            unsafe {
11603                b.launch(cfg)?;
11604            }
11605            return Ok(Some((y0, y1, y2)));
11606        }
11607        let f = self.func("qmatvec_nvfp4_mmvq_fused3_rp");
11608        let cfg = LaunchConfig {
11609            grid_dim: (nb(o0) + nb(o1) + nb(o2), m as u32, 1),
11610            block_dim: (32, ROWS_PER_BLOCK, 1),
11611            shared_mem_bytes: 0,
11612        };
11613        let __s_b = self.gpu.stream();
11614        let mut b = __s_b.launch_builder(&f);
11615        b.arg(b0)
11616            .arg(b1)
11617            .arg(b2)
11618            .arg(aq)
11619            .arg(ad)
11620            .arg(&mut y0)
11621            .arg(&mut y1)
11622            .arg(&mut y2)
11623            .arg(&inf)
11624            .arg(&oi0)
11625            .arg(&oi1)
11626            .arg(&oi2)
11627            .arg(&mi)
11628            .arg(&p0.1)
11629            .arg(&p1.1)
11630            .arg(&p2.1);
11631        unsafe {
11632            b.launch(cfg)?;
11633        }
11634        Ok(Some((y0, y1, y2)))
11635    }
11636
11637    /// fused2 twin of `matmul_nvfp4_fused3`, for MIXED-type weight groups: the gemma4
11638    /// dense NVFP4mix recipe keeps `attn_v` and `ffn_down` at Q8_0 (full-NVFP4 was
11639    /// measured garbage on this dense class), so its q/k/v trio and gate/up/down never
11640    /// satisfy an all-NVFP4 fused3 — the pairs that ARE uniformly NVFP4 (q,k and
11641    /// gate,up) fuse here instead. m==1 only, same law as fused3/fused4: per
11642    /// (tensor,row) the kernel seg body is VERBATIM, so the fusion is bit-identical to
11643    /// two separate launches. `MEMRA_NVFP4_FUSED2=0` is the rollback seam and the
11644    /// same-binary interleaved A/B arm.
11645    pub fn matmul_nvfp4_fused2(
11646        &self,
11647        w0: &crate::model::GpuTensor,
11648        w1: &crate::model::GpuTensor,
11649        aq: &CudaSlice<i8>,
11650        ad: &CudaSlice<f32>,
11651        m: usize,
11652    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
11653        use crate::model::GpuTensor;
11654        static FUSED2_OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11655        let off =
11656            *FUSED2_OFF.get_or_init(|| std::env::var("MEMRA_NVFP4_FUSED2").as_deref() == Ok("0"));
11657        // m==1 ONLY: at m>1 the singles ride the _b16 weight-once column program (one weight
11658        // read serves all m rows); the fused segments would re-read the weight per row.
11659        if off
11660            || m != 1
11661            || !self.mmvq_supports(QT_NVFP4)
11662            || !self.uses_q8_1_fast(w0)
11663            || !self.uses_q8_1_fast(w1)
11664        {
11665            return Ok(None);
11666        }
11667        let unpack = |w: &crate::model::GpuTensor| match w {
11668            GpuTensor::Quant {
11669                bytes,
11670                qtype,
11671                scale,
11672                rp,
11673                ..
11674            } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
11675            _ => None,
11676        };
11677        let (Some(p0), Some(p1)) = (unpack(w0), unpack(w1)) else {
11678            return Ok(None);
11679        };
11680        let in_f = w0.in_features();
11681        if w1.in_features() != in_f {
11682            return Ok(None);
11683        }
11684        let (o0, o1) = (w0.out_features(), w1.out_features());
11685        const ROWS_PER_BLOCK: u32 = 4; // MEMRA_MMVQ_ROWS
11686        const RPW: u32 = 2;
11687        let rows_pb = ROWS_PER_BLOCK * RPW;
11688        let nb = |o: usize| (o as u32).div_ceil(rows_pb);
11689        let f = self.func("qmatvec_nvfp4_mmvq_fused2_rp");
11690        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
11691        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
11692        let cfg = LaunchConfig {
11693            grid_dim: (nb(o0) + nb(o1), m as u32, 1),
11694            block_dim: (32, ROWS_PER_BLOCK, 1),
11695            shared_mem_bytes: 0,
11696        };
11697        let (inf, oi0, oi1, mi) = (in_f as i32, o0 as i32, o1 as i32, m as i32);
11698        // SAFETY: the raw pointers come straight from the &GpuTensor borrows above and are
11699        // only dereferenced for the launch-arg build inside this call.
11700        let (b0, b1) = unsafe { (&*p0.0, &*p1.0) };
11701        // PDL wave-B: the kernel carries MEMRA_PDL_ENTRY — launch overlaps the producer's
11702        // drain (quantize/norm emitting aq/ad). Same math, same order, bit-identical.
11703        if Self::pdl_on() && Self::pdl_mmvq_on() && Self::pdl_nvfp4q8_on() {
11704            {
11705                use cudarc::driver::{DevicePtr, DevicePtrMut};
11706                let s = &self.gpu.stream();
11707                let (pw0, _g0) = b0.device_ptr(s);
11708                let (pw1, _g1) = b1.device_ptr(s);
11709                let (paq, _g2) = aq.device_ptr(s);
11710                let (pad, _g3) = ad.device_ptr(s);
11711                let (py0, _g4) = y0.device_ptr_mut(s);
11712                let (py1, _g5) = y1.device_ptr_mut(s);
11713                let (s0, s1) = (p0.1, p1.1);
11714                let mut ps = [
11715                    &pw0 as *const _ as *mut std::ffi::c_void,
11716                    &pw1 as *const _ as *mut _,
11717                    &paq as *const _ as *mut _,
11718                    &pad as *const _ as *mut _,
11719                    &py0 as *const _ as *mut _,
11720                    &py1 as *const _ as *mut _,
11721                    &inf as *const _ as *mut _,
11722                    &oi0 as *const _ as *mut _,
11723                    &oi1 as *const _ as *mut _,
11724                    &mi as *const _ as *mut _,
11725                    &s0 as *const _ as *mut _,
11726                    &s1 as *const _ as *mut _,
11727                ];
11728                unsafe {
11729                    self.launch_pdl(
11730                        "qmatvec_nvfp4_mmvq_fused2_rp",
11731                        cfg.grid_dim,
11732                        cfg.block_dim,
11733                        &mut ps,
11734                    )?;
11735                }
11736            }
11737            return Ok(Some((y0, y1)));
11738        }
11739        let __s_b = self.gpu.stream();
11740        let mut b = __s_b.launch_builder(&f);
11741        b.arg(b0)
11742            .arg(b1)
11743            .arg(aq)
11744            .arg(ad)
11745            .arg(&mut y0)
11746            .arg(&mut y1)
11747            .arg(&inf)
11748            .arg(&oi0)
11749            .arg(&oi1)
11750            .arg(&mi)
11751            .arg(&p0.1)
11752            .arg(&p1.1);
11753        unsafe {
11754            b.launch(cfg)?;
11755        }
11756        Ok(Some((y0, y1)))
11757    }
11758
11759    /// Slot-fed fused2 twin (alloc-free capture lane): identical launch to
11760    /// `matmul_nvfp4_fused2`, caller-owned outputs — the gemma4 dc_slotted graph body
11761    /// needs zero mem nodes, so the allocating wrapper can't serve it. Returns false
11762    /// (decline) on any non-NVFP4/rp pair; the caller chains or refuses.
11763    pub fn matmul_nvfp4_fused2_into(
11764        &self,
11765        w0: &crate::model::GpuTensor,
11766        w1: &crate::model::GpuTensor,
11767        aq: &CudaSlice<i8>,
11768        ad: &CudaSlice<f32>,
11769        y0: &mut CudaSlice<f32>,
11770        y1: &mut CudaSlice<f32>,
11771    ) -> Result<bool, Box<dyn std::error::Error>> {
11772        use crate::model::GpuTensor;
11773        static FUSED2_OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11774        let off =
11775            *FUSED2_OFF.get_or_init(|| std::env::var("MEMRA_NVFP4_FUSED2").as_deref() == Ok("0"));
11776        if off
11777            || !self.mmvq_supports(QT_NVFP4)
11778            || !self.uses_q8_1_fast(w0)
11779            || !self.uses_q8_1_fast(w1)
11780        {
11781            return Ok(false);
11782        }
11783        let unpack = |w: &crate::model::GpuTensor| match w {
11784            GpuTensor::Quant {
11785                bytes,
11786                qtype,
11787                scale,
11788                rp,
11789                ..
11790            } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
11791            _ => None,
11792        };
11793        let (Some(p0), Some(p1)) = (unpack(w0), unpack(w1)) else {
11794            return Ok(false);
11795        };
11796        let in_f = w0.in_features();
11797        if w1.in_features() != in_f {
11798            return Ok(false);
11799        }
11800        let (o0, o1) = (w0.out_features(), w1.out_features());
11801        if y0.len() < o0 || y1.len() < o1 {
11802            return Ok(false);
11803        }
11804        const ROWS_PER_BLOCK: u32 = 4; // MEMRA_MMVQ_ROWS
11805        const RPW: u32 = 2;
11806        let rows_pb = ROWS_PER_BLOCK * RPW;
11807        let nb = |o: usize| (o as u32).div_ceil(rows_pb);
11808        let f = self.func("qmatvec_nvfp4_mmvq_fused2_rp");
11809        let cfg = LaunchConfig {
11810            grid_dim: (nb(o0) + nb(o1), 1, 1),
11811            block_dim: (32, ROWS_PER_BLOCK, 1),
11812            shared_mem_bytes: 0,
11813        };
11814        let (inf, oi0, oi1, mi) = (in_f as i32, o0 as i32, o1 as i32, 1i32);
11815        // SAFETY: the raw pointers come straight from the &GpuTensor borrows above and are
11816        // only dereferenced for the launch-arg build inside this call.
11817        let (b0, b1) = unsafe { (&*p0.0, &*p1.0) };
11818        let __s_b = self.gpu.stream();
11819        let mut b = __s_b.launch_builder(&f);
11820        b.arg(b0)
11821            .arg(b1)
11822            .arg(aq)
11823            .arg(ad)
11824            .arg(&mut *y0)
11825            .arg(&mut *y1)
11826            .arg(&inf)
11827            .arg(&oi0)
11828            .arg(&oi1)
11829            .arg(&mi)
11830            .arg(&p0.1)
11831            .arg(&p1.1);
11832        unsafe {
11833            b.launch(cfg)?;
11834        }
11835        Ok(true)
11836    }
11837
11838    /// fused4 twin of `matmul_nvfp4_fused3`: the Linear-mixer projection quartet
11839    /// (wqkv + wqkv_gate + ssm_beta + ssm_alpha) in one launch, m==1 only. Per
11840    /// (tensor,row) the kernel body is nvfp4_mmvq_multirow_rp VERBATIM — bit-identical
11841    /// to four separate launches (rig-native decode increment 2, RIG-NATIVE-DECODE.md).
11842    #[allow(clippy::type_complexity)]
11843    pub fn matmul_nvfp4_fused4(
11844        &self,
11845        w0: &crate::model::GpuTensor,
11846        w1: &crate::model::GpuTensor,
11847        w2: &crate::model::GpuTensor,
11848        w3: &crate::model::GpuTensor,
11849        aq: &CudaSlice<i8>,
11850        ad: &CudaSlice<f32>,
11851        m: usize,
11852    ) -> Result<
11853        Option<(
11854            CudaSlice<f32>,
11855            CudaSlice<f32>,
11856            CudaSlice<f32>,
11857            CudaSlice<f32>,
11858        )>,
11859        Box<dyn std::error::Error>,
11860    > {
11861        use crate::model::GpuTensor;
11862        // MEMRA_NVFP4_FUSED4=0: rollback seam + the same-binary interleaved A/B arm.
11863        // m == 1 rides the original fused kernel; m = 2..=8 rides the BATCHED fused twin
11864        // (lane/samplat, 2026-08-21): same quartet-in-one-launch shape, seg body =
11865        // nvfp4_mmvq_batched_rp_sc verbatim (weight rows read once for all m columns) —
11866        // bit-identical per (tensor,row,column) to the four bN_rpsc singles it replaces.
11867        // Admission mirrors the singles' batched gates below.
11868        if !(1..=8).contains(&m)
11869            || std::env::var("MEMRA_NVFP4_FUSED4").as_deref() == Ok("0")
11870            || !self.mmvq_supports(QT_NVFP4)
11871            || !self.uses_q8_1_fast(w0)
11872            || !self.uses_q8_1_fast(w1)
11873            || !self.uses_q8_1_fast(w2)
11874            || !self.uses_q8_1_fast(w3)
11875        {
11876            return Ok(None);
11877        }
11878        if m > 1 {
11879            // the batched-twin gates: the bN_rpsc program this must stay byte-identical to
11880            // (matmul_pre's batched arm), plus the rp-sc dispatch shape requirements.
11881            let in_f = w0.in_features();
11882            if !self.batched_supports(QT_NVFP4)
11883                || std::env::var("MEMRA_NO_BATCHED").is_ok()
11884                || (m > 4 && !Self::b8_enabled())
11885                || in_f % 512 != 0
11886                || in_f / 64 > 272
11887            {
11888                return Ok(None);
11889            }
11890        }
11891        let unpack = |w: &crate::model::GpuTensor| match w {
11892            GpuTensor::Quant {
11893                bytes,
11894                qtype,
11895                scale,
11896                rp,
11897                ..
11898            } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
11899            _ => None,
11900        };
11901        let (Some(p0), Some(p1), Some(p2), Some(p3)) =
11902            (unpack(w0), unpack(w1), unpack(w2), unpack(w3))
11903        else {
11904            return Ok(None);
11905        };
11906        let in_f = w0.in_features();
11907        if w1.in_features() != in_f || w2.in_features() != in_f || w3.in_features() != in_f {
11908            return Ok(None);
11909        }
11910        let (o0, o1, o2, o3) = (
11911            w0.out_features(),
11912            w1.out_features(),
11913            w2.out_features(),
11914            w3.out_features(),
11915        );
11916        const ROWS_PER_BLOCK: u32 = 4; // MEMRA_MMVQ_ROWS
11917        const RPW: u32 = 2;
11918        let rows_pb = ROWS_PER_BLOCK * RPW;
11919        let nb = |o: usize| (o as u32).div_ceil(rows_pb);
11920        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
11921        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
11922        let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
11923        let mut y3 = self.alloc_uninit::<f32>(m * o3)?;
11924        let (inf, oi0, oi1, oi2, oi3, mi) = (
11925            in_f as i32,
11926            o0 as i32,
11927            o1 as i32,
11928            o2 as i32,
11929            o3 as i32,
11930            m as i32,
11931        );
11932        // SAFETY: the raw pointers come straight from the &GpuTensor borrows above and are
11933        // only dereferenced for the launch-arg build inside this call.
11934        let (b0, b1, b2, b3) = unsafe { (&*p0.0, &*p1.0, &*p2.0, &*p3.0) };
11935        if m > 1 {
11936            // Batched fused twin: no in-kernel scale (the bN_rpsc program has none) — refuse
11937            // scale-carrying tensors so the singles path keeps them (GGUF trunk scales are 1.0).
11938            if p0.1 != 1.0 || p1.1 != 1.0 || p2.1 != 1.0 || p3.1 != 1.0 {
11939                return Ok(None);
11940            }
11941            let f = self.func("qmatvec_nvfp4_mmvq_fused4_b8_rpsc");
11942            let cfg = LaunchConfig {
11943                grid_dim: (nb(o0) + nb(o1) + nb(o2) + nb(o3), 1, 1),
11944                block_dim: (32, ROWS_PER_BLOCK, 1),
11945                shared_mem_bytes: 0,
11946            };
11947            let __s_b = self.gpu.stream();
11948            let mut b = __s_b.launch_builder(&f);
11949            b.arg(b0)
11950                .arg(b1)
11951                .arg(b2)
11952                .arg(b3)
11953                .arg(aq)
11954                .arg(ad)
11955                .arg(&mut y0)
11956                .arg(&mut y1)
11957                .arg(&mut y2)
11958                .arg(&mut y3)
11959                .arg(&inf)
11960                .arg(&oi0)
11961                .arg(&oi1)
11962                .arg(&oi2)
11963                .arg(&oi3)
11964                .arg(&mi);
11965            unsafe {
11966                b.launch(cfg)?;
11967            }
11968            return Ok(Some((y0, y1, y2, y3)));
11969        }
11970        let f = self.func("qmatvec_nvfp4_mmvq_fused4_rp");
11971        let cfg = LaunchConfig {
11972            grid_dim: (nb(o0) + nb(o1) + nb(o2) + nb(o3), m as u32, 1),
11973            block_dim: (32, ROWS_PER_BLOCK, 1),
11974            shared_mem_bytes: 0,
11975        };
11976        let __s_b = self.gpu.stream();
11977        let mut b = __s_b.launch_builder(&f);
11978        b.arg(b0)
11979            .arg(b1)
11980            .arg(b2)
11981            .arg(b3)
11982            .arg(aq)
11983            .arg(ad)
11984            .arg(&mut y0)
11985            .arg(&mut y1)
11986            .arg(&mut y2)
11987            .arg(&mut y3)
11988            .arg(&inf)
11989            .arg(&oi0)
11990            .arg(&oi1)
11991            .arg(&oi2)
11992            .arg(&oi3)
11993            .arg(&mi)
11994            .arg(&p0.1)
11995            .arg(&p1.1)
11996            .arg(&p2.1)
11997            .arg(&p3.1);
11998        unsafe {
11999            b.launch(cfg)?;
12000        }
12001        Ok(Some((y0, y1, y2, y3)))
12002    }
12003
12004    /// FUSED Q8_0 m=1 matvec PAIR with UNEQUAL out_f (trunk launch-fusion, 2026-07-05). Folds two
12005    /// same-input q8_0 projections (35B trunk: wqkv+wqkv_gate 8192/4096, gate_shexp+up_shexp
12006    /// 512/512) into ONE launch via a block-offset split (blocks [0,nb0) -> w0, rest -> w1) — the
12007    /// dual-mr2 recipe with the same-out_f restriction lifted. Per (tensor,row) the kernel body is
12008    /// qmatvec_q8_0_mmvq VERBATIM -> BIT-IDENTICAL to two separate m=1 launches. Returns None when
12009    /// ineligible (not both Q8_0 / in_f mismatch / MEMRA_MMVQ off / MEMRA_Q8_DUAL=0) — caller falls
12010    /// back to the per-tensor path.
12011    pub fn matmul_q8_fused2(
12012        &self,
12013        w0: &crate::model::GpuTensor,
12014        w1: &crate::model::GpuTensor,
12015        aq: &CudaSlice<i8>,
12016        ad: &CudaSlice<f32>,
12017    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
12018        // e4m3 twin (lane/fp8-decode-v1): this entry is the trunk's generic m=1 pair door
12019        // (wqkv+wqkv_gate, ssm_beta+alpha, gate_shexp+up_shexp), so admitting QT_F8_E4M3 here
12020        // fuses the NATIVE-RESIDENCY FP8 trunk at every existing call site with no call-site
12021        // change. Scale is folded in-kernel per range -> the returned buffers are already scaled,
12022        // exactly like the per-tensor `matmul_pre` e4m3 dispatch this replaces.
12023        if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
12024            return Ok(Some(self.e4m3_fused2_core(
12025                p0.0,
12026                p1.0,
12027                aq,
12028                ad,
12029                w0.in_features(),
12030                p0.1,
12031                p1.1,
12032                p0.2,
12033                p0.3,
12034                p1.3,
12035            )?));
12036        }
12037        let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
12038            return Ok(None);
12039        };
12040        Ok(Some(self.q8_fused2_core(
12041            p0.0,
12042            p1.0,
12043            aq,
12044            ad,
12045            w0.in_features(),
12046            p0.1,
12047            p1.1,
12048            p0.2,
12049        )?))
12050    }
12051
12052    #[allow(clippy::too_many_arguments)]
12053    fn q8_fused2_core(
12054        &self,
12055        b0: &CudaSlice<u8>,
12056        b1: &CudaSlice<u8>,
12057        aq: &CudaSlice<i8>,
12058        ad: &CudaSlice<f32>,
12059        in_f: usize,
12060        out0: usize,
12061        out1: usize,
12062        row_bytes: usize,
12063    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12064        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
12065        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
12066        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
12067        let f = self.func("qmatvec_q8_0_mmvq_fused2");
12068        let mut y0 = self.alloc_uninit::<f32>(out0)?;
12069        let mut y1 = self.alloc_uninit::<f32>(out1)?;
12070        let cfg = LaunchConfig {
12071            grid_dim: (nb0 + nb1, 1, 1),
12072            block_dim: (32, ROWS_PER_BLOCK, 1),
12073            shared_mem_bytes: 0,
12074        };
12075        let (inf, o0, o1, rbl) = (in_f as i32, out0 as i32, out1 as i32, row_bytes as i64);
12076        let __s_b = self.gpu.stream();
12077        let mut b = __s_b.launch_builder(&f);
12078        b.arg(b0)
12079            .arg(b1)
12080            .arg(aq)
12081            .arg(ad)
12082            .arg(&mut y0)
12083            .arg(&mut y1)
12084            .arg(&inf)
12085            .arg(&o0)
12086            .arg(&o1)
12087            .arg(&rbl);
12088        unsafe {
12089            b.launch(cfg)?;
12090        }
12091        Ok((y0, y1))
12092    }
12093
12094    /// f32-activation entry for the fused2 pair: quantizes x to q8_1 ONCE then runs the fused
12095    /// launch — replaces two `matmul(w, x, 1)` calls that would each re-quantize the same x
12096    /// (35B shared-expert gate+up per MoE layer per token). Same bits: quantize_q8_1 is
12097    /// deterministic, the fused body is the MMVQ kernel verbatim. None when ineligible (the
12098    /// callers' m==1-under-MEMRA_FAST dispatch would take MMVQ; anything else falls back).
12099    pub fn matmul_q8_fused2_x(
12100        &self,
12101        w0: &crate::model::GpuTensor,
12102        w1: &crate::model::GpuTensor,
12103        x: &CudaSlice<f32>,
12104    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
12105        if !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) {
12106            return Ok(None);
12107        }
12108        if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
12109            let (aq, ad) = self.quantize_q8_1(x, 1, w0.in_features())?;
12110            return Ok(Some(self.e4m3_fused2_core(
12111                p0.0,
12112                p1.0,
12113                &aq,
12114                &ad,
12115                w0.in_features(),
12116                p0.1,
12117                p1.1,
12118                p0.2,
12119                p0.3,
12120                p1.3,
12121            )?));
12122        }
12123        let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
12124            return Ok(None);
12125        };
12126        let (aq, ad) = self.quantize_q8_1(x, 1, w0.in_features())?;
12127        Ok(Some(self.q8_fused2_core(
12128            p0.0,
12129            p1.0,
12130            &aq,
12131            &ad,
12132            w0.in_features(),
12133            p0.1,
12134            p1.1,
12135            p0.2,
12136        )?))
12137    }
12138
12139    /// Test entry for the kernel_check gate: launch the fused2 kernel from raw weight bytes,
12140    /// quantizing the f32 activation internally (mirrors qmatvec_mmvq_raw; no env gating).
12141    #[allow(clippy::too_many_arguments)]
12142    pub fn qmatvec_q8_fused2_raw(
12143        &self,
12144        b0: &CudaSlice<u8>,
12145        b1: &CudaSlice<u8>,
12146        x: &CudaSlice<f32>,
12147        in_f: usize,
12148        out0: usize,
12149        out1: usize,
12150        row_bytes: usize,
12151    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12152        let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
12153        self.q8_fused2_core(b0, b1, &aq, &ad, in_f, out0, out1, row_bytes)
12154    }
12155
12156    /// FUSED Q8_0 m=1 matvec TRIPLE (wq+wk+wv on the 35B full-attn layers: out_f 8192/512/512).
12157    /// Same block-offset recipe as `matmul_q8_fused2` with three ranges. BIT-IDENTICAL per
12158    /// (tensor,row) to three separate m=1 MMVQ launches.
12159    /// FUSED Q4_0 m=1 TRIPLE (gemma q/k/v — same quantized input; per (tensor,row) chain
12160    /// identical to the mr2 kernel). Returns None unless all three are Q4_0 with equal in_f.
12161    pub fn matmul_q4_fused3(
12162        &self,
12163        w0: &crate::model::GpuTensor,
12164        w1: &crate::model::GpuTensor,
12165        w2: &crate::model::GpuTensor,
12166        aq: &CudaSlice<i8>,
12167        ad: &CudaSlice<f32>,
12168    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
12169    {
12170        use crate::model::GpuTensor;
12171        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
12172            match w {
12173                GpuTensor::Quant {
12174                    qtype, row_bytes, ..
12175                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
12176                _ => None,
12177            }
12178        };
12179        let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (q4(w0), q4(w1), q4(w2)) else {
12180            return Ok(None);
12181        };
12182        if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
12183            return Ok(None);
12184        }
12185        // Effective (bytes, rp) per tensor: mirror (rp4) OR the in-place swap (rp flag,
12186        // bytes already split). Mixed layouts cannot share one fused launch -> fall back to
12187        // the separate matvecs (each routes its own rp).
12188        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
12189            match w {
12190                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
12191                    Some(m) => (m, true),
12192                    None => (bytes, *rp),
12193                },
12194                _ => unreachable!(),
12195            }
12196        }
12197        let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
12198        if rp0 != rp1 || rp1 != rp2 {
12199            return Ok(None);
12200        }
12201        let rp = rp0;
12202        let rpb: u32 = 4;
12203        // mr1 (one row/warp, 2026-07-14): follows the singles' MEMRA_Q40_MR default — the
12204        // fused t=1 kernels were left on mr2 when the singles flipped (DRAM-duty map:
12205        // fused3 57% / fused2 86%; small qkv segments starve under mr2's half grid).
12206        let mr1 = rp && Self::q40_mr1_on();
12207        let nb = |o: usize| {
12208            if mr1 {
12209                (o as u32).div_ceil(rpb)
12210            } else {
12211                (o as u32).div_ceil(2).div_ceil(rpb)
12212            }
12213        };
12214        let grid = nb(o0) + nb(o1) + nb(o2);
12215        let mut y0 = self.alloc_uninit::<f32>(o0)?;
12216        let mut y1 = self.alloc_uninit::<f32>(o1)?;
12217        let mut y2 = self.alloc_uninit::<f32>(o2)?;
12218        let f = self.func(if mr1 {
12219            "qmatvec_q4_0_mmvq_fused3_mr1_rp"
12220        } else if rp {
12221            "qmatvec_q4_0_mmvq_fused3_rp"
12222        } else {
12223            "qmatvec_q4_0_mmvq_fused3"
12224        });
12225        let cfg = LaunchConfig {
12226            grid_dim: (grid, 1, 1),
12227            block_dim: (32, rpb, 1),
12228            shared_mem_bytes: 0,
12229        };
12230        let inf = w0.in_features() as i32;
12231        let (oo0, oo1, oo2) = (o0 as i32, o1 as i32, o2 as i32);
12232        let (r0, r1, r2) = (rb0 as i64, rb1 as i64, rb2 as i64);
12233        // PDL wave-A (2026-07-23): the mr1 kernel carries MEMRA_PDL_ENTRY; only that
12234        // variant may take the programmatic-serialization launch.
12235        if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
12236            {
12237                use cudarc::driver::{DevicePtr, DevicePtrMut};
12238                let s = &self.gpu.stream();
12239                let (p0, _g0) = b0.device_ptr(s);
12240                let (p1, _g1) = b1.device_ptr(s);
12241                let (p2, _g2) = b2.device_ptr(s);
12242                let (paq, _g3) = aq.device_ptr(s);
12243                let (pad, _g4) = ad.device_ptr(s);
12244                let (py0, _g5) = y0.device_ptr_mut(s);
12245                let (py1, _g6) = y1.device_ptr_mut(s);
12246                let (py2, _g7) = y2.device_ptr_mut(s);
12247                let mut ps = [
12248                    &p0 as *const _ as *mut std::ffi::c_void,
12249                    &p1 as *const _ as *mut _,
12250                    &p2 as *const _ as *mut _,
12251                    &paq as *const _ as *mut _,
12252                    &pad as *const _ as *mut _,
12253                    &py0 as *const _ as *mut _,
12254                    &py1 as *const _ as *mut _,
12255                    &py2 as *const _ as *mut _,
12256                    &inf as *const _ as *mut _,
12257                    &oo0 as *const _ as *mut _,
12258                    &oo1 as *const _ as *mut _,
12259                    &oo2 as *const _ as *mut _,
12260                    &r0 as *const _ as *mut _,
12261                    &r1 as *const _ as *mut _,
12262                    &r2 as *const _ as *mut _,
12263                ];
12264                unsafe {
12265                    self.launch_pdl(
12266                        "qmatvec_q4_0_mmvq_fused3_mr1_rp",
12267                        (grid, 1, 1),
12268                        (32, rpb, 1),
12269                        &mut ps,
12270                    )?;
12271                }
12272            }
12273            return Ok(Some((y0, y1, y2)));
12274        }
12275        let __s_b = self.gpu.stream();
12276        let mut b = __s_b.launch_builder(&f);
12277        b.arg(b0)
12278            .arg(b1)
12279            .arg(b2)
12280            .arg(aq)
12281            .arg(ad)
12282            .arg(&mut y0)
12283            .arg(&mut y1)
12284            .arg(&mut y2)
12285            .arg(&inf)
12286            .arg(&oo0)
12287            .arg(&oo1)
12288            .arg(&oo2)
12289            .arg(&r0)
12290            .arg(&r1)
12291            .arg(&r2);
12292        unsafe {
12293            b.launch(cfg)?;
12294        }
12295        Ok(Some((y0, y1, y2)))
12296    }
12297
12298    /// Slot-fed fused3 twin (alloc-free capture lane): identical launch, caller-owned outputs.
12299    /// Returns Ok(false) when the fused path is unavailable (caller falls back).
12300    #[allow(clippy::too_many_arguments)]
12301    pub fn matmul_q4_fused3_into(
12302        &self,
12303        w0: &crate::model::GpuTensor,
12304        w1: &crate::model::GpuTensor,
12305        w2: &crate::model::GpuTensor,
12306        aq: &CudaSlice<i8>,
12307        ad: &CudaSlice<f32>,
12308        y0: &mut CudaSlice<f32>,
12309        y1: &mut CudaSlice<f32>,
12310        y2: &mut CudaSlice<f32>,
12311    ) -> Result<bool, Box<dyn std::error::Error>> {
12312        use crate::model::GpuTensor;
12313        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
12314            match w {
12315                GpuTensor::Quant {
12316                    qtype, row_bytes, ..
12317                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
12318                _ => None,
12319            }
12320        };
12321        let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (q4(w0), q4(w1), q4(w2)) else {
12322            return Ok(false);
12323        };
12324        if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
12325            return Ok(false);
12326        }
12327        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
12328            match w {
12329                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
12330                    Some(m) => (m, true),
12331                    None => (bytes, *rp),
12332                },
12333                _ => unreachable!(),
12334            }
12335        }
12336        let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
12337        if rp0 != rp1 || rp1 != rp2 {
12338            return Ok(false);
12339        }
12340        let rp = rp0;
12341        let rpb: u32 = 4;
12342        let mr1 = rp && Self::q40_mr1_on();
12343        let nb = |o: usize| {
12344            if mr1 {
12345                (o as u32).div_ceil(rpb)
12346            } else {
12347                (o as u32).div_ceil(2).div_ceil(rpb)
12348            }
12349        };
12350        let grid = nb(o0) + nb(o1) + nb(o2);
12351        debug_assert!(y0.len() >= o0 && y1.len() >= o1 && y2.len() >= o2);
12352        let f = self.func(if mr1 {
12353            "qmatvec_q4_0_mmvq_fused3_mr1_rp"
12354        } else if rp {
12355            "qmatvec_q4_0_mmvq_fused3_rp"
12356        } else {
12357            "qmatvec_q4_0_mmvq_fused3"
12358        });
12359        let cfg = LaunchConfig {
12360            grid_dim: (grid, 1, 1),
12361            block_dim: (32, rpb, 1),
12362            shared_mem_bytes: 0,
12363        };
12364        let inf = w0.in_features() as i32;
12365        let (oo0, oo1, oo2) = (o0 as i32, o1 as i32, o2 as i32);
12366        let (r0, r1, r2) = (rb0 as i64, rb1 as i64, rb2 as i64);
12367        // PDL wave-A: identical to the owned twin (capture-lane parity).
12368        if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
12369            use cudarc::driver::{DevicePtr, DevicePtrMut};
12370            let s = &self.gpu.stream();
12371            let (p0, _g0) = b0.device_ptr(s);
12372            let (p1, _g1) = b1.device_ptr(s);
12373            let (p2, _g2) = b2.device_ptr(s);
12374            let (paq, _g3) = aq.device_ptr(s);
12375            let (pad, _g4) = ad.device_ptr(s);
12376            let (py0, _g5) = y0.device_ptr_mut(s);
12377            let (py1, _g6) = y1.device_ptr_mut(s);
12378            let (py2, _g7) = y2.device_ptr_mut(s);
12379            let mut ps = [
12380                &p0 as *const _ as *mut std::ffi::c_void,
12381                &p1 as *const _ as *mut _,
12382                &p2 as *const _ as *mut _,
12383                &paq as *const _ as *mut _,
12384                &pad as *const _ as *mut _,
12385                &py0 as *const _ as *mut _,
12386                &py1 as *const _ as *mut _,
12387                &py2 as *const _ as *mut _,
12388                &inf as *const _ as *mut _,
12389                &oo0 as *const _ as *mut _,
12390                &oo1 as *const _ as *mut _,
12391                &oo2 as *const _ as *mut _,
12392                &r0 as *const _ as *mut _,
12393                &r1 as *const _ as *mut _,
12394                &r2 as *const _ as *mut _,
12395            ];
12396            unsafe {
12397                self.launch_pdl(
12398                    "qmatvec_q4_0_mmvq_fused3_mr1_rp",
12399                    (grid, 1, 1),
12400                    (32, rpb, 1),
12401                    &mut ps,
12402                )?;
12403            }
12404            return Ok(true);
12405        }
12406        let __s_b = self.gpu.stream();
12407        let mut b = __s_b.launch_builder(&f);
12408        b.arg(b0)
12409            .arg(b1)
12410            .arg(b2)
12411            .arg(aq)
12412            .arg(ad)
12413            .arg(&mut *y0)
12414            .arg(&mut *y1)
12415            .arg(&mut *y2)
12416            .arg(&inf)
12417            .arg(&oo0)
12418            .arg(&oo1)
12419            .arg(&oo2)
12420            .arg(&r0)
12421            .arg(&r1)
12422            .arg(&r2);
12423        unsafe {
12424            b.launch(cfg)?;
12425        }
12426        Ok(true)
12427    }
12428
12429    /// FUSED Q4_0 m=1 PAIR (gemma shared gate+up).
12430    pub fn matmul_q4_fused2(
12431        &self,
12432        w0: &crate::model::GpuTensor,
12433        w1: &crate::model::GpuTensor,
12434        aq: &CudaSlice<i8>,
12435        ad: &CudaSlice<f32>,
12436    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
12437        use crate::model::GpuTensor;
12438        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
12439            match w {
12440                GpuTensor::Quant {
12441                    qtype, row_bytes, ..
12442                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
12443                _ => None,
12444            }
12445        };
12446        let (Some((rb0, o0)), Some((rb1, o1))) = (q4(w0), q4(w1)) else {
12447            return Ok(None);
12448        };
12449        if w0.in_features() != w1.in_features() {
12450            return Ok(None);
12451        }
12452        // Effective (bytes, rp) per tensor (mirror or in-place swap); mixed -> separate matvecs.
12453        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
12454            match w {
12455                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
12456                    Some(m) => (m, true),
12457                    None => (bytes, *rp),
12458                },
12459                _ => unreachable!(),
12460            }
12461        }
12462        let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
12463        if rp0 != rp1 {
12464            return Ok(None);
12465        }
12466        let rp = rp0;
12467        let rpb: u32 = 4;
12468        // mr1 twin — see matmul_q4_fused3.
12469        let mr1 = rp && Self::q40_mr1_on();
12470        let nb = |o: usize| {
12471            if mr1 {
12472                (o as u32).div_ceil(rpb)
12473            } else {
12474                (o as u32).div_ceil(2).div_ceil(rpb)
12475            }
12476        };
12477        let grid = nb(o0) + nb(o1);
12478        let mut y0 = self.alloc_uninit::<f32>(o0)?;
12479        let mut y1 = self.alloc_uninit::<f32>(o1)?;
12480        let f = self.func(if mr1 {
12481            "qmatvec_q4_0_mmvq_fused2_mr1_rp"
12482        } else if rp {
12483            "qmatvec_q4_0_mmvq_fused2_rp"
12484        } else {
12485            "qmatvec_q4_0_mmvq_fused2"
12486        });
12487        let cfg = LaunchConfig {
12488            grid_dim: (grid, 1, 1),
12489            block_dim: (32, rpb, 1),
12490            shared_mem_bytes: 0,
12491        };
12492        let inf = w0.in_features() as i32;
12493        let (oo0, oo1) = (o0 as i32, o1 as i32);
12494        let (r0, r1) = (rb0 as i64, rb1 as i64);
12495        // PDL wave-A: mr1 kernel carries MEMRA_PDL_ENTRY.
12496        if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
12497            {
12498                use cudarc::driver::{DevicePtr, DevicePtrMut};
12499                let s = &self.gpu.stream();
12500                let (p0, _g0) = b0.device_ptr(s);
12501                let (p1, _g1) = b1.device_ptr(s);
12502                let (paq, _g2) = aq.device_ptr(s);
12503                let (pad, _g3) = ad.device_ptr(s);
12504                let (py0, _g4) = y0.device_ptr_mut(s);
12505                let (py1, _g5) = y1.device_ptr_mut(s);
12506                let mut ps = [
12507                    &p0 as *const _ as *mut std::ffi::c_void,
12508                    &p1 as *const _ as *mut _,
12509                    &paq as *const _ as *mut _,
12510                    &pad as *const _ as *mut _,
12511                    &py0 as *const _ as *mut _,
12512                    &py1 as *const _ as *mut _,
12513                    &inf as *const _ as *mut _,
12514                    &oo0 as *const _ as *mut _,
12515                    &oo1 as *const _ as *mut _,
12516                    &r0 as *const _ as *mut _,
12517                    &r1 as *const _ as *mut _,
12518                ];
12519                unsafe {
12520                    self.launch_pdl(
12521                        "qmatvec_q4_0_mmvq_fused2_mr1_rp",
12522                        (grid, 1, 1),
12523                        (32, rpb, 1),
12524                        &mut ps,
12525                    )?;
12526                }
12527            }
12528            return Ok(Some((y0, y1)));
12529        }
12530        let __s_b = self.gpu.stream();
12531        let mut b = __s_b.launch_builder(&f);
12532        b.arg(b0)
12533            .arg(b1)
12534            .arg(aq)
12535            .arg(ad)
12536            .arg(&mut y0)
12537            .arg(&mut y1)
12538            .arg(&inf)
12539            .arg(&oo0)
12540            .arg(&oo1)
12541            .arg(&r0)
12542            .arg(&r1);
12543        unsafe {
12544            b.launch(cfg)?;
12545        }
12546        Ok(Some((y0, y1)))
12547    }
12548
12549    /// Slot-fed fused2 twin (alloc-free capture lane): identical launch, caller-owned outputs.
12550    pub fn matmul_q4_fused2_into(
12551        &self,
12552        w0: &crate::model::GpuTensor,
12553        w1: &crate::model::GpuTensor,
12554        aq: &CudaSlice<i8>,
12555        ad: &CudaSlice<f32>,
12556        y0: &mut CudaSlice<f32>,
12557        y1: &mut CudaSlice<f32>,
12558    ) -> Result<bool, Box<dyn std::error::Error>> {
12559        use crate::model::GpuTensor;
12560        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
12561            match w {
12562                GpuTensor::Quant {
12563                    qtype, row_bytes, ..
12564                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
12565                _ => None,
12566            }
12567        };
12568        let (Some((rb0, o0)), Some((rb1, o1))) = (q4(w0), q4(w1)) else {
12569            return Ok(false);
12570        };
12571        if w0.in_features() != w1.in_features() {
12572            return Ok(false);
12573        }
12574        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
12575            match w {
12576                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
12577                    Some(m) => (m, true),
12578                    None => (bytes, *rp),
12579                },
12580                _ => unreachable!(),
12581            }
12582        }
12583        let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
12584        if rp0 != rp1 {
12585            return Ok(false);
12586        }
12587        let rp = rp0;
12588        let rpb: u32 = 4;
12589        let mr1 = rp && Self::q40_mr1_on();
12590        let nb = |o: usize| {
12591            if mr1 {
12592                (o as u32).div_ceil(rpb)
12593            } else {
12594                (o as u32).div_ceil(2).div_ceil(rpb)
12595            }
12596        };
12597        let grid = nb(o0) + nb(o1);
12598        debug_assert!(y0.len() >= o0 && y1.len() >= o1);
12599        let f = self.func(if mr1 {
12600            "qmatvec_q4_0_mmvq_fused2_mr1_rp"
12601        } else if rp {
12602            "qmatvec_q4_0_mmvq_fused2_rp"
12603        } else {
12604            "qmatvec_q4_0_mmvq_fused2"
12605        });
12606        let cfg = LaunchConfig {
12607            grid_dim: (grid, 1, 1),
12608            block_dim: (32, rpb, 1),
12609            shared_mem_bytes: 0,
12610        };
12611        let inf = w0.in_features() as i32;
12612        let (oo0, oo1) = (o0 as i32, o1 as i32);
12613        let (r0, r1) = (rb0 as i64, rb1 as i64);
12614        // PDL wave-A: identical to the owned twin (capture-lane parity).
12615        if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
12616            use cudarc::driver::{DevicePtr, DevicePtrMut};
12617            let s = &self.gpu.stream();
12618            let (p0, _g0) = b0.device_ptr(s);
12619            let (p1, _g1) = b1.device_ptr(s);
12620            let (paq, _g2) = aq.device_ptr(s);
12621            let (pad, _g3) = ad.device_ptr(s);
12622            let (py0, _g4) = y0.device_ptr_mut(s);
12623            let (py1, _g5) = y1.device_ptr_mut(s);
12624            let mut ps = [
12625                &p0 as *const _ as *mut std::ffi::c_void,
12626                &p1 as *const _ as *mut _,
12627                &paq as *const _ as *mut _,
12628                &pad as *const _ as *mut _,
12629                &py0 as *const _ as *mut _,
12630                &py1 as *const _ as *mut _,
12631                &inf as *const _ as *mut _,
12632                &oo0 as *const _ as *mut _,
12633                &oo1 as *const _ as *mut _,
12634                &r0 as *const _ as *mut _,
12635                &r1 as *const _ as *mut _,
12636            ];
12637            unsafe {
12638                self.launch_pdl(
12639                    "qmatvec_q4_0_mmvq_fused2_mr1_rp",
12640                    (grid, 1, 1),
12641                    (32, rpb, 1),
12642                    &mut ps,
12643                )?;
12644            }
12645            return Ok(true);
12646        }
12647        let __s_b = self.gpu.stream();
12648        let mut b = __s_b.launch_builder(&f);
12649        b.arg(b0)
12650            .arg(b1)
12651            .arg(aq)
12652            .arg(ad)
12653            .arg(&mut *y0)
12654            .arg(&mut *y1)
12655            .arg(&inf)
12656            .arg(&oo0)
12657            .arg(&oo1)
12658            .arg(&r0)
12659            .arg(&r1);
12660        unsafe {
12661            b.launch(cfg)?;
12662        }
12663        Ok(true)
12664    }
12665
12666    /// BATCHED fused2 (2026-07-13, megakernel-microcosm probe): gate+up b-tier matvecs in
12667    /// ONE segmented-grid launch — the up segment fills SMs as the gate segment drains
12668    /// (the per-launch tail waves behind the 6x-falsified b-tier plateau). Bit-identical
12669    /// per row to two mr2_rp launches. rp layout required; m in 2..=8 (b16 has no twin).
12670    pub fn matmul_q4_fused2_batched(
12671        &self,
12672        w0: &crate::model::GpuTensor,
12673        w1: &crate::model::GpuTensor,
12674        aq: &CudaSlice<i8>,
12675        ad: &CudaSlice<f32>,
12676        m: usize,
12677    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
12678        use crate::model::GpuTensor;
12679        if m < 2 || m > 8 {
12680            return Ok(None);
12681        }
12682        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
12683            match w {
12684                GpuTensor::Quant {
12685                    qtype, row_bytes, ..
12686                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
12687                _ => None,
12688            }
12689        };
12690        let (Some((rb0, o0)), Some((_rb1, o1))) = (q4(w0), q4(w1)) else {
12691            return Ok(None);
12692        };
12693        if w0.in_features() != w1.in_features() {
12694            return Ok(None);
12695        }
12696        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
12697            match w {
12698                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
12699                    Some(mr) => (mr, true),
12700                    None => (bytes, *rp),
12701                },
12702                _ => unreachable!(),
12703            }
12704        }
12705        let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
12706        if !rp0 || !rp1 {
12707            return Ok(None);
12708        }
12709        let mcols = Self::batched_mcols(m);
12710        let rpb: u32 = 4;
12711        let nb = |o: usize| (o as u32).div_ceil(2 * rpb);
12712        let grid = nb(o0) + nb(o1);
12713        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
12714        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
12715        let f = self.func(match mcols {
12716            2 => "qmatvec_q4_0_mmvq_b2_f2_rp",
12717            4 => "qmatvec_q4_0_mmvq_b4_f2_rp",
12718            _ => "qmatvec_q4_0_mmvq_b8_f2_rp",
12719        });
12720        let cfg = LaunchConfig {
12721            grid_dim: (grid, 1, 1),
12722            block_dim: (32, rpb, 1),
12723            shared_mem_bytes: 0,
12724        };
12725        let inf = w0.in_features() as i32;
12726        let (oo0, oo1, mi) = (o0 as i32, o1 as i32, m as i32);
12727        let rb = rb0 as i64;
12728        let __s_b = self.gpu.stream();
12729        let mut b = __s_b.launch_builder(&f);
12730        b.arg(b0)
12731            .arg(b1)
12732            .arg(aq)
12733            .arg(ad)
12734            .arg(&mut y0)
12735            .arg(&mut y1)
12736            .arg(&inf)
12737            .arg(&oo0)
12738            .arg(&oo1)
12739            .arg(&mi)
12740            .arg(&rb);
12741        unsafe {
12742            b.launch(cfg)?;
12743        }
12744        Ok(Some((y0, y1)))
12745    }
12746
12747    /// BATCHED fused3 (see matmul_q4_fused2_batched): three-segment single launch for the
12748    /// verify qkv triple. Same-in_f q4_0 rp tensors, m in 2..=8. Bit-identical per row.
12749    #[allow(clippy::too_many_arguments)]
12750    pub fn matmul_q4_fused3_batched(
12751        &self,
12752        w0: &crate::model::GpuTensor,
12753        w1: &crate::model::GpuTensor,
12754        w2: &crate::model::GpuTensor,
12755        aq: &CudaSlice<i8>,
12756        ad: &CudaSlice<f32>,
12757        m: usize,
12758    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
12759    {
12760        use crate::model::GpuTensor;
12761        if m < 2 || m > 8 {
12762            return Ok(None);
12763        }
12764        let q4 = |w: &GpuTensor| -> Option<usize> {
12765            match w {
12766                GpuTensor::Quant { qtype, .. } if *qtype == QT_Q4_0 => Some(w.out_features()),
12767                _ => None,
12768            }
12769        };
12770        let (Some(o0), Some(o1), Some(o2)) = (q4(w0), q4(w1), q4(w2)) else {
12771            return Ok(None);
12772        };
12773        if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
12774            return Ok(None);
12775        }
12776        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
12777            match w {
12778                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
12779                    Some(mr) => (mr, true),
12780                    None => (bytes, *rp),
12781                },
12782                _ => unreachable!(),
12783            }
12784        }
12785        let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
12786        if !rp0 || !rp1 || !rp2 {
12787            return Ok(None);
12788        }
12789        let mcols = Self::batched_mcols(m);
12790        let rpb: u32 = 4;
12791        let nb = |o: usize| (o as u32).div_ceil(2 * rpb);
12792        let grid = nb(o0) + nb(o1) + nb(o2);
12793        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
12794        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
12795        let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
12796        let f = self.func(match mcols {
12797            2 => "qmatvec_q4_0_mmvq_b2_f3_rp",
12798            4 => "qmatvec_q4_0_mmvq_b4_f3_rp",
12799            _ => "qmatvec_q4_0_mmvq_b8_f3_rp",
12800        });
12801        let cfg = LaunchConfig {
12802            grid_dim: (grid, 1, 1),
12803            block_dim: (32, rpb, 1),
12804            shared_mem_bytes: 0,
12805        };
12806        let inf = w0.in_features() as i32;
12807        let (oo0, oo1, oo2, mi) = (o0 as i32, o1 as i32, o2 as i32, m as i32);
12808        let rb = 0i64;
12809        let __s_b = self.gpu.stream();
12810        let mut b = __s_b.launch_builder(&f);
12811        b.arg(b0)
12812            .arg(b1)
12813            .arg(b2)
12814            .arg(aq)
12815            .arg(ad)
12816            .arg(&mut y0)
12817            .arg(&mut y1)
12818            .arg(&mut y2)
12819            .arg(&inf)
12820            .arg(&oo0)
12821            .arg(&oo1)
12822            .arg(&oo2)
12823            .arg(&mi)
12824            .arg(&rb);
12825        unsafe {
12826            b.launch(cfg)?;
12827        }
12828        Ok(Some((y0, y1, y2)))
12829    }
12830
12831    pub fn matmul_q8_fused3(
12832        &self,
12833        w0: &crate::model::GpuTensor,
12834        w1: &crate::model::GpuTensor,
12835        w2: &crate::model::GpuTensor,
12836        aq: &CudaSlice<i8>,
12837        ad: &CudaSlice<f32>,
12838    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
12839    {
12840        // e4m3 twin (lane/fp8-decode-v1): the full-attn wq/wk/wv triple — on the NV-27B those three
12841        // are per-tensor FP8, so native residency without this arm meant three separate launches.
12842        if let Some([p0, p1, p2]) = self.e4m3_fused_params(&[w0, w1, w2]) {
12843            return Ok(Some(self.e4m3_fused3_core(
12844                p0.0,
12845                p1.0,
12846                p2.0,
12847                aq,
12848                ad,
12849                w0.in_features(),
12850                p0.1,
12851                p1.1,
12852                p2.1,
12853                p0.2,
12854                p0.3,
12855                p1.3,
12856                p2.3,
12857            )?));
12858        }
12859        let Some([p0, p1, p2]) = self.q8_fused_params(&[w0, w1, w2]) else {
12860            return Ok(None);
12861        };
12862        Ok(Some(self.q8_fused3_core(
12863            p0.0,
12864            p1.0,
12865            p2.0,
12866            aq,
12867            ad,
12868            w0.in_features(),
12869            p0.1,
12870            p1.1,
12871            p2.1,
12872            p0.2,
12873        )?))
12874    }
12875
12876    #[allow(clippy::too_many_arguments)]
12877    fn q8_fused3_core(
12878        &self,
12879        b0: &CudaSlice<u8>,
12880        b1: &CudaSlice<u8>,
12881        b2: &CudaSlice<u8>,
12882        aq: &CudaSlice<i8>,
12883        ad: &CudaSlice<f32>,
12884        in_f: usize,
12885        out0: usize,
12886        out1: usize,
12887        out2: usize,
12888        row_bytes: usize,
12889    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12890        const ROWS_PER_BLOCK: u32 = 4;
12891        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
12892        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
12893        let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
12894        let f = self.func("qmatvec_q8_0_mmvq_fused3");
12895        let mut y0 = self.alloc_uninit::<f32>(out0)?;
12896        let mut y1 = self.alloc_uninit::<f32>(out1)?;
12897        let mut y2 = self.alloc_uninit::<f32>(out2)?;
12898        let cfg = LaunchConfig {
12899            grid_dim: (nb0 + nb1 + nb2, 1, 1),
12900            block_dim: (32, ROWS_PER_BLOCK, 1),
12901            shared_mem_bytes: 0,
12902        };
12903        let (inf, o0, o1, o2, rbl) = (
12904            in_f as i32,
12905            out0 as i32,
12906            out1 as i32,
12907            out2 as i32,
12908            row_bytes as i64,
12909        );
12910        let __s_b = self.gpu.stream();
12911        let mut b = __s_b.launch_builder(&f);
12912        b.arg(b0)
12913            .arg(b1)
12914            .arg(b2)
12915            .arg(aq)
12916            .arg(ad)
12917            .arg(&mut y0)
12918            .arg(&mut y1)
12919            .arg(&mut y2)
12920            .arg(&inf)
12921            .arg(&o0)
12922            .arg(&o1)
12923            .arg(&o2)
12924            .arg(&rbl);
12925        unsafe {
12926            b.launch(cfg)?;
12927        }
12928        Ok((y0, y1, y2))
12929    }
12930
12931    /// Test entry for the kernel_check gate: fused3 from raw weight bytes (internal q8_1 quant).
12932    #[allow(clippy::too_many_arguments)]
12933    pub fn qmatvec_q8_fused3_raw(
12934        &self,
12935        b0: &CudaSlice<u8>,
12936        b1: &CudaSlice<u8>,
12937        b2: &CudaSlice<u8>,
12938        x: &CudaSlice<f32>,
12939        in_f: usize,
12940        out0: usize,
12941        out1: usize,
12942        out2: usize,
12943        row_bytes: usize,
12944    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12945        let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
12946        self.q8_fused3_core(b0, b1, b2, &aq, &ad, in_f, out0, out1, out2, row_bytes)
12947    }
12948
12949    /// BATCHED twin of `matmul_q8_fused2` for the verify t=2-4 tier (MEMRA_SPEC_FUSED_T call
12950    /// sites, lane/close35b): ONE launch computes both same-input Q8_0 projections for m tokens.
12951    /// Per (tensor,token,row) the kernel body is q8_0_mmvq_batched VERBATIM with the identical
12952    /// row mapping (Q8_0's batched_variant is always "base") -> BIT-IDENTICAL to the two
12953    /// per-tensor _b2/_b4 launches `matmul_decode_exact` dispatches at m=2-4, with the caller's
12954    /// single shared q8_1 activation replacing two per-call re-quantizes (quantize_q8_1 is
12955    /// deterministic -> same bytes). None when ineligible (m outside 2..=4 / not both Q8_0 /
12956    /// in_f mismatch / MEMRA_MMVQ=0 / MEMRA_Q8_DUAL=0 / MEMRA_NO_BATCHED set — the last keeps
12957    /// dispatch parity: without batched kernels decode-exact runs grid.y=m MMVQ, and the fused
12958    /// twin must not introduce a batched program the reference path would not run).
12959    pub fn matmul_q8_fused2_t(
12960        &self,
12961        w0: &crate::model::GpuTensor,
12962        w1: &crate::model::GpuTensor,
12963        aq: &CudaSlice<i8>,
12964        ad: &CudaSlice<f32>,
12965        m: usize,
12966    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
12967        // m<=8 (lane/q27-deepdive, 2026-08-05): was 2..=4 (the verify tier's mcols 2/4). The
12968        // serving tick's mcols-8 tier now has its fused2_b8 wrapper, so c=5..8 batched decode
12969        // fuses too — same template body, still bit-identical to the two _b8 launches.
12970        if !(2..=8).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok() {
12971            return Ok(None);
12972        }
12973        // e4m3 twin: MEMRA_B8 parity — without it m=5..8 e4m3 decode runs the per-m grid.y=m path,
12974        // so the fused b8 launch would introduce a batched program the reference path would not run.
12975        if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
12976            if m > 4 && !Self::b8_enabled() {
12977                return Ok(None);
12978            }
12979            return Ok(Some(self.e4m3_fused2_t_core(
12980                p0.0,
12981                p1.0,
12982                aq,
12983                ad,
12984                m,
12985                w0.in_features(),
12986                p0.1,
12987                p1.1,
12988                p0.2,
12989                p0.3,
12990                p1.3,
12991            )?));
12992        }
12993        let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
12994            return Ok(None);
12995        };
12996        Ok(Some(self.q8_fused2_t_core(
12997            p0.0,
12998            p1.0,
12999            aq,
13000            ad,
13001            m,
13002            w0.in_features(),
13003            p0.1,
13004            p1.1,
13005            p0.2,
13006        )?))
13007    }
13008
13009    #[allow(clippy::too_many_arguments)]
13010    fn q8_fused2_t_core(
13011        &self,
13012        b0: &CudaSlice<u8>,
13013        b1: &CudaSlice<u8>,
13014        aq: &CudaSlice<i8>,
13015        ad: &CudaSlice<f32>,
13016        m: usize,
13017        in_f: usize,
13018        out0: usize,
13019        out1: usize,
13020        row_bytes: usize,
13021    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13022        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
13023        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
13024        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
13025        let f = self.func(match Self::batched_mcols(m) {
13026            2 => "qmatvec_q8_0_mmvq_fused2_b2",
13027            4 => "qmatvec_q8_0_mmvq_fused2_b4",
13028            // b8 = the SERVING tier (lane/q27-deepdive): c=5..8 batched decode.
13029            _ => "qmatvec_q8_0_mmvq_fused2_b8",
13030        });
13031        let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
13032        let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
13033        let cfg = LaunchConfig {
13034            grid_dim: (nb0 + nb1, 1, 1),
13035            block_dim: (32, ROWS_PER_BLOCK, 1),
13036            shared_mem_bytes: 0,
13037        };
13038        let (inf, o0, o1, mi, rbl) = (
13039            in_f as i32,
13040            out0 as i32,
13041            out1 as i32,
13042            m as i32,
13043            row_bytes as i64,
13044        );
13045        let __s_b = self.gpu.stream();
13046        let mut b = __s_b.launch_builder(&f);
13047        b.arg(b0)
13048            .arg(b1)
13049            .arg(aq)
13050            .arg(ad)
13051            .arg(&mut y0)
13052            .arg(&mut y1)
13053            .arg(&inf)
13054            .arg(&o0)
13055            .arg(&o1)
13056            .arg(&mi)
13057            .arg(&rbl);
13058        unsafe {
13059            b.launch(cfg)?;
13060        }
13061        Ok((y0, y1))
13062    }
13063
13064    /// Test entry for the kernel_check gate: fused2 batched from raw weight bytes (internal
13065    /// q8_1 quant of the [m, in_f] activation), no env gating.
13066    #[allow(clippy::too_many_arguments)]
13067    pub fn qmatvec_q8_fused2_t_raw(
13068        &self,
13069        b0: &CudaSlice<u8>,
13070        b1: &CudaSlice<u8>,
13071        x: &CudaSlice<f32>,
13072        m: usize,
13073        in_f: usize,
13074        out0: usize,
13075        out1: usize,
13076        row_bytes: usize,
13077    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13078        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
13079        self.q8_fused2_t_core(b0, b1, &aq, &ad, m, in_f, out0, out1, row_bytes)
13080    }
13081
13082    /// BATCHED twin of `matmul_q8_fused3` (wq+wk+wv at verify t=2-4). Same contract as
13083    /// `matmul_q8_fused2_t` with three ranges.
13084    #[allow(clippy::too_many_arguments)]
13085    pub fn matmul_q8_fused3_t(
13086        &self,
13087        w0: &crate::model::GpuTensor,
13088        w1: &crate::model::GpuTensor,
13089        w2: &crate::model::GpuTensor,
13090        aq: &CudaSlice<i8>,
13091        ad: &CudaSlice<f32>,
13092        m: usize,
13093    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
13094    {
13095        if !(2..=4).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok() {
13096            return Ok(None);
13097        }
13098        if let Some([p0, p1, p2]) = self.e4m3_fused_params(&[w0, w1, w2]) {
13099            return Ok(Some(self.e4m3_fused3_t_core(
13100                p0.0,
13101                p1.0,
13102                p2.0,
13103                aq,
13104                ad,
13105                m,
13106                w0.in_features(),
13107                p0.1,
13108                p1.1,
13109                p2.1,
13110                p0.2,
13111                p0.3,
13112                p1.3,
13113                p2.3,
13114            )?));
13115        }
13116        let Some([p0, p1, p2]) = self.q8_fused_params(&[w0, w1, w2]) else {
13117            return Ok(None);
13118        };
13119        Ok(Some(self.q8_fused3_t_core(
13120            p0.0,
13121            p1.0,
13122            p2.0,
13123            aq,
13124            ad,
13125            m,
13126            w0.in_features(),
13127            p0.1,
13128            p1.1,
13129            p2.1,
13130            p0.2,
13131        )?))
13132    }
13133
13134    #[allow(clippy::too_many_arguments)]
13135    fn q8_fused3_t_core(
13136        &self,
13137        b0: &CudaSlice<u8>,
13138        b1: &CudaSlice<u8>,
13139        b2: &CudaSlice<u8>,
13140        aq: &CudaSlice<i8>,
13141        ad: &CudaSlice<f32>,
13142        m: usize,
13143        in_f: usize,
13144        out0: usize,
13145        out1: usize,
13146        out2: usize,
13147        row_bytes: usize,
13148    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13149        const ROWS_PER_BLOCK: u32 = 4;
13150        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
13151        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
13152        let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
13153        let f = self.func(if Self::batched_mcols(m) == 2 {
13154            "qmatvec_q8_0_mmvq_fused3_b2"
13155        } else {
13156            "qmatvec_q8_0_mmvq_fused3_b4"
13157        });
13158        let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
13159        let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
13160        let mut y2 = self.alloc_uninit::<f32>(m * out2)?;
13161        let cfg = LaunchConfig {
13162            grid_dim: (nb0 + nb1 + nb2, 1, 1),
13163            block_dim: (32, ROWS_PER_BLOCK, 1),
13164            shared_mem_bytes: 0,
13165        };
13166        let (inf, o0, o1, o2, mi, rbl) = (
13167            in_f as i32,
13168            out0 as i32,
13169            out1 as i32,
13170            out2 as i32,
13171            m as i32,
13172            row_bytes as i64,
13173        );
13174        let __s_b = self.gpu.stream();
13175        let mut b = __s_b.launch_builder(&f);
13176        b.arg(b0)
13177            .arg(b1)
13178            .arg(b2)
13179            .arg(aq)
13180            .arg(ad)
13181            .arg(&mut y0)
13182            .arg(&mut y1)
13183            .arg(&mut y2)
13184            .arg(&inf)
13185            .arg(&o0)
13186            .arg(&o1)
13187            .arg(&o2)
13188            .arg(&mi)
13189            .arg(&rbl);
13190        unsafe {
13191            b.launch(cfg)?;
13192        }
13193        Ok((y0, y1, y2))
13194    }
13195
13196    /// Test entry for the kernel_check gate: fused3 batched from raw weight bytes.
13197    #[allow(clippy::too_many_arguments)]
13198    pub fn qmatvec_q8_fused3_t_raw(
13199        &self,
13200        b0: &CudaSlice<u8>,
13201        b1: &CudaSlice<u8>,
13202        b2: &CudaSlice<u8>,
13203        x: &CudaSlice<f32>,
13204        m: usize,
13205        in_f: usize,
13206        out0: usize,
13207        out1: usize,
13208        out2: usize,
13209        row_bytes: usize,
13210    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13211        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
13212        self.q8_fused3_t_core(b0, b1, b2, &aq, &ad, m, in_f, out0, out1, out2, row_bytes)
13213    }
13214
13215    /// Rollback seam for the Q8_0 dense-FFN gate+up fusion arm in `matmul_pre_dual_noscale`
13216    /// (lane/q27-deepdive, 2026-08-05). Default ON; `MEMRA_Q8_FFN_FUSE2=0` restores the
13217    /// two-`matmul_pre_noscale` pair. Read once — the dispatch must not vary within a run.
13218    pub fn q8_ffn_fuse2_on(&self) -> bool {
13219        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13220        *ON.get_or_init(|| std::env::var("MEMRA_Q8_FFN_FUSE2").as_deref() != Ok("0"))
13221    }
13222
13223    /// Eligibility + param extraction for the fused q8_0 launches: every tensor must be Quant Q8_0
13224    /// with macro-scale 1.0 (always true for GGUF q8_0; only NVFP4 carries scale) and share w[0]'s
13225    /// in_f (q8_0 row_bytes is a pure function of in_f, so equal in_f => equal row_bytes). MEMRA_MMVQ
13226    /// must be on: the fused body is the MMVQ kernel; without it decode m=1 runs dp4a and fusing
13227    /// would mix dispatch families (FP-order law). MEMRA_Q8_DUAL=0 = rollback seam.
13228    #[allow(clippy::type_complexity)]
13229    fn q8_fused_params<'w, const N: usize>(
13230        &self,
13231        ws: &[&'w crate::model::GpuTensor; N],
13232    ) -> Option<[(&'w CudaSlice<u8>, usize, usize); N]> {
13233        use crate::model::GpuTensor;
13234        if std::env::var("MEMRA_MMVQ").as_deref() == Ok("0") {
13235            return None;
13236        }
13237        if std::env::var("MEMRA_Q8_DUAL").is_ok_and(|v| v == "0") {
13238            return None;
13239        }
13240        let in_f = ws[0].in_features();
13241        let mut out: [Option<(&CudaSlice<u8>, usize, usize)>; N] = [None; N];
13242        for (i, w) in ws.iter().enumerate() {
13243            match w {
13244                GpuTensor::Quant {
13245                    bytes,
13246                    qtype,
13247                    row_bytes,
13248                    scale,
13249                    ..
13250                } if *qtype == QT_Q8_0 && *scale == 1.0 && w.in_features() == in_f => {
13251                    out[i] = Some((bytes, w.out_features(), *row_bytes))
13252                }
13253                _ => return None,
13254            }
13255        }
13256        Some(out.map(|o| o.unwrap()))
13257    }
13258
13259    /// Rollback seam for the F8-E4M3 launch-fusion arm (lane/fp8-decode-v1, 2026-08-05).
13260    /// Default ON; `MEMRA_E4M3_DUAL=0` restores the per-tensor m=1/batched launches.
13261    pub fn e4m3_dual_on(&self) -> bool {
13262        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13263        *ON.get_or_init(|| std::env::var("MEMRA_E4M3_DUAL").as_deref() != Ok("0"))
13264    }
13265
13266    /// Eligibility + param extraction for the FUSED e4m3 launches — the QT_F8_E4M3 twin of
13267    /// `q8_fused_params`. Differences that are inherent to the dtype, not policy:
13268    ///   * each tensor carries its OWN per-tensor `weight_scale` (returned as the 4th field);
13269    ///     Q8_0 hard-requires scale==1.0 because it has no macro-scale at all.
13270    ///   * no MEMRA_MMVQ gate: `mmvq_supports` exempts QT_F8_E4M3 (the e4m3 mmvq family is that
13271    ///     dtype's ONLY int8-act kernel class), so the per-tensor fallback these fused kernels
13272    ///     replace is ALWAYS the same mmvq body under every env — the FP-order law holds.
13273    ///   * `row_bytes == in_f` is asserted rather than derived: the native-residency load arm keeps
13274    ///     the checkpoint's raw [out_f, in_f] rows, and a re-encoded slab must never reach here.
13275    /// Rejects any split-plane mirror (`rp`/`rp4`): there is no `_rp` e4m3 fused form, so fusing
13276    /// there would swap dispatch families mid-model. MEMRA_E4M3_DUAL=0 = rollback seam.
13277    #[allow(clippy::type_complexity)]
13278    fn e4m3_fused_params<'w, const N: usize>(
13279        &self,
13280        ws: &[&'w crate::model::GpuTensor; N],
13281    ) -> Option<[(&'w CudaSlice<u8>, usize, usize, f32); N]> {
13282        use crate::model::GpuTensor;
13283        if !self.e4m3_dual_on() {
13284            return None;
13285        }
13286        let in_f = ws[0].in_features();
13287        let mut out: [Option<(&CudaSlice<u8>, usize, usize, f32)>; N] = [None; N];
13288        for (i, w) in ws.iter().enumerate() {
13289            match w {
13290                GpuTensor::Quant {
13291                    bytes,
13292                    qtype,
13293                    row_bytes,
13294                    scale,
13295                    rp,
13296                    rp4,
13297                    ..
13298                } if *qtype == QT_F8_E4M3
13299                    && w.in_features() == in_f
13300                    && *row_bytes == in_f
13301                    && !*rp
13302                    && rp4.is_none() =>
13303                {
13304                    out[i] = Some((bytes, w.out_features(), *row_bytes, *scale))
13305                }
13306                _ => return None,
13307            }
13308        }
13309        Some(out.map(|o| o.unwrap()))
13310    }
13311
13312    /// FUSED e4m3 m=1 PAIR. Block-offset split (`qmatvec_e4m3_mmvq_fused2`), per-tensor
13313    /// weight_scale folded at the write like the single-tensor `qmatvec_e4m3_mmvq` — so per
13314    /// (tensor,row) this is BIT-IDENTICAL to two separate m=1 launches, scale included.
13315    #[allow(clippy::too_many_arguments)]
13316    fn e4m3_fused2_core(
13317        &self,
13318        b0: &CudaSlice<u8>,
13319        b1: &CudaSlice<u8>,
13320        aq: &CudaSlice<i8>,
13321        ad: &CudaSlice<f32>,
13322        in_f: usize,
13323        out0: usize,
13324        out1: usize,
13325        row_bytes: usize,
13326        ws0: f32,
13327        ws1: f32,
13328    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13329        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
13330        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
13331        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
13332        let f = self.func("qmatvec_e4m3_mmvq_fused2");
13333        let mut y0 = self.alloc_uninit::<f32>(out0)?;
13334        let mut y1 = self.alloc_uninit::<f32>(out1)?;
13335        let cfg = LaunchConfig {
13336            grid_dim: (nb0 + nb1, 1, 1),
13337            block_dim: (32, ROWS_PER_BLOCK, 1),
13338            shared_mem_bytes: 0,
13339        };
13340        let (inf, o0, o1, rbl) = (in_f as i32, out0 as i32, out1 as i32, row_bytes as i64);
13341        let __s_b = self.gpu.stream();
13342        let mut b = __s_b.launch_builder(&f);
13343        b.arg(b0)
13344            .arg(b1)
13345            .arg(aq)
13346            .arg(ad)
13347            .arg(&mut y0)
13348            .arg(&mut y1)
13349            .arg(&inf)
13350            .arg(&o0)
13351            .arg(&o1)
13352            .arg(&rbl)
13353            .arg(&ws0)
13354            .arg(&ws1);
13355        unsafe {
13356            b.launch(cfg)?;
13357        }
13358        Ok((y0, y1))
13359    }
13360
13361    /// FUSED e4m3 m=1 TRIPLE (`qmatvec_e4m3_mmvq_fused3`). Same contract as the pair.
13362    #[allow(clippy::too_many_arguments)]
13363    fn e4m3_fused3_core(
13364        &self,
13365        b0: &CudaSlice<u8>,
13366        b1: &CudaSlice<u8>,
13367        b2: &CudaSlice<u8>,
13368        aq: &CudaSlice<i8>,
13369        ad: &CudaSlice<f32>,
13370        in_f: usize,
13371        out0: usize,
13372        out1: usize,
13373        out2: usize,
13374        row_bytes: usize,
13375        ws0: f32,
13376        ws1: f32,
13377        ws2: f32,
13378    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13379        const ROWS_PER_BLOCK: u32 = 4;
13380        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
13381        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
13382        let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
13383        let f = self.func("qmatvec_e4m3_mmvq_fused3");
13384        let mut y0 = self.alloc_uninit::<f32>(out0)?;
13385        let mut y1 = self.alloc_uninit::<f32>(out1)?;
13386        let mut y2 = self.alloc_uninit::<f32>(out2)?;
13387        let cfg = LaunchConfig {
13388            grid_dim: (nb0 + nb1 + nb2, 1, 1),
13389            block_dim: (32, ROWS_PER_BLOCK, 1),
13390            shared_mem_bytes: 0,
13391        };
13392        let (inf, o0, o1, o2, rbl) = (
13393            in_f as i32,
13394            out0 as i32,
13395            out1 as i32,
13396            out2 as i32,
13397            row_bytes as i64,
13398        );
13399        let __s_b = self.gpu.stream();
13400        let mut b = __s_b.launch_builder(&f);
13401        b.arg(b0)
13402            .arg(b1)
13403            .arg(b2)
13404            .arg(aq)
13405            .arg(ad)
13406            .arg(&mut y0)
13407            .arg(&mut y1)
13408            .arg(&mut y2)
13409            .arg(&inf)
13410            .arg(&o0)
13411            .arg(&o1)
13412            .arg(&o2)
13413            .arg(&rbl)
13414            .arg(&ws0)
13415            .arg(&ws1)
13416            .arg(&ws2);
13417        unsafe {
13418            b.launch(cfg)?;
13419        }
13420        Ok((y0, y1, y2))
13421    }
13422
13423    /// BATCHED FUSED e4m3 pair (m=2..8). The batched kernels carry no `ws` arg (every batched
13424    /// kernel in the tree is scale-free), so each output takes its own `scale_inplace` — the
13425    /// SAME post-op the per-tensor batched dispatch applies, hence still bit-identical.
13426    #[allow(clippy::too_many_arguments)]
13427    fn e4m3_fused2_t_core(
13428        &self,
13429        b0: &CudaSlice<u8>,
13430        b1: &CudaSlice<u8>,
13431        aq: &CudaSlice<i8>,
13432        ad: &CudaSlice<f32>,
13433        m: usize,
13434        in_f: usize,
13435        out0: usize,
13436        out1: usize,
13437        row_bytes: usize,
13438        ws0: f32,
13439        ws1: f32,
13440    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13441        const ROWS_PER_BLOCK: u32 = 4;
13442        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
13443        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
13444        let f = self.func(match Self::batched_mcols(m) {
13445            2 => "qmatvec_e4m3_mmvq_fused2_b2",
13446            4 => "qmatvec_e4m3_mmvq_fused2_b4",
13447            _ => "qmatvec_e4m3_mmvq_fused2_b8",
13448        });
13449        let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
13450        let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
13451        let cfg = LaunchConfig {
13452            grid_dim: (nb0 + nb1, 1, 1),
13453            block_dim: (32, ROWS_PER_BLOCK, 1),
13454            shared_mem_bytes: 0,
13455        };
13456        let (inf, o0, o1, mi, rbl) = (
13457            in_f as i32,
13458            out0 as i32,
13459            out1 as i32,
13460            m as i32,
13461            row_bytes as i64,
13462        );
13463        let __s_b = self.gpu.stream();
13464        let mut b = __s_b.launch_builder(&f);
13465        b.arg(b0)
13466            .arg(b1)
13467            .arg(aq)
13468            .arg(ad)
13469            .arg(&mut y0)
13470            .arg(&mut y1)
13471            .arg(&inf)
13472            .arg(&o0)
13473            .arg(&o1)
13474            .arg(&mi)
13475            .arg(&rbl);
13476        unsafe {
13477            b.launch(cfg)?;
13478        }
13479        if ws0 != 1.0 {
13480            self.scale_inplace(&mut y0, ws0, m * out0)?;
13481        }
13482        if ws1 != 1.0 {
13483            self.scale_inplace(&mut y1, ws1, m * out1)?;
13484        }
13485        Ok((y0, y1))
13486    }
13487
13488    /// BATCHED FUSED e4m3 triple (m=2..4). Same contract as the batched pair.
13489    #[allow(clippy::too_many_arguments)]
13490    fn e4m3_fused3_t_core(
13491        &self,
13492        b0: &CudaSlice<u8>,
13493        b1: &CudaSlice<u8>,
13494        b2: &CudaSlice<u8>,
13495        aq: &CudaSlice<i8>,
13496        ad: &CudaSlice<f32>,
13497        m: usize,
13498        in_f: usize,
13499        out0: usize,
13500        out1: usize,
13501        out2: usize,
13502        row_bytes: usize,
13503        ws0: f32,
13504        ws1: f32,
13505        ws2: f32,
13506    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13507        const ROWS_PER_BLOCK: u32 = 4;
13508        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
13509        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
13510        let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
13511        let f = self.func(if Self::batched_mcols(m) == 2 {
13512            "qmatvec_e4m3_mmvq_fused3_b2"
13513        } else {
13514            "qmatvec_e4m3_mmvq_fused3_b4"
13515        });
13516        let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
13517        let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
13518        let mut y2 = self.alloc_uninit::<f32>(m * out2)?;
13519        let cfg = LaunchConfig {
13520            grid_dim: (nb0 + nb1 + nb2, 1, 1),
13521            block_dim: (32, ROWS_PER_BLOCK, 1),
13522            shared_mem_bytes: 0,
13523        };
13524        let (inf, o0, o1, o2, mi, rbl) = (
13525            in_f as i32,
13526            out0 as i32,
13527            out1 as i32,
13528            out2 as i32,
13529            m as i32,
13530            row_bytes as i64,
13531        );
13532        let __s_b = self.gpu.stream();
13533        let mut b = __s_b.launch_builder(&f);
13534        b.arg(b0)
13535            .arg(b1)
13536            .arg(b2)
13537            .arg(aq)
13538            .arg(ad)
13539            .arg(&mut y0)
13540            .arg(&mut y1)
13541            .arg(&mut y2)
13542            .arg(&inf)
13543            .arg(&o0)
13544            .arg(&o1)
13545            .arg(&o2)
13546            .arg(&mi)
13547            .arg(&rbl);
13548        unsafe {
13549            b.launch(cfg)?;
13550        }
13551        if ws0 != 1.0 {
13552            self.scale_inplace(&mut y0, ws0, m * out0)?;
13553        }
13554        if ws1 != 1.0 {
13555            self.scale_inplace(&mut y1, ws1, m * out1)?;
13556        }
13557        if ws2 != 1.0 {
13558            self.scale_inplace(&mut y2, ws2, m * out2)?;
13559        }
13560        Ok((y0, y1, y2))
13561    }
13562
13563    /// BLOCK-128 e4m3 MMVQ launcher (`qmatvec_e4m3_blk_mmvq`, lane/fp8-blk128-decode 2026-08-05).
13564    /// The per-block-dequant twin of `qmatvec_mmvq`'s QT_F8_E4M3 arm: same grid/block decomposition
13565    /// (warp per output row, ROWS_PER_BLOCK warps per block, grid.y = m), same q8_1 activation, but
13566    /// the weight scale is a resident [rows, cols] f32 grid read per k128 block inside the kernel
13567    /// instead of one scalar folded at the write. It cannot share `qmatvec_mmvq`'s body because
13568    /// that launcher's arg list is fixed at (bytes, aq, ad, y, in_f, out_f, m, row_bytes [, scale]).
13569    ///
13570    /// `mr` and `rp` have no analogue here (no split-plane e4m3 layout exists), so there is exactly
13571    /// one kernel and no name table — a shape this cannot serve must be refused at LOAD, not here.
13572    pub fn qmatvec_e4m3_blk_mmvq(
13573        &self,
13574        bytes: &CudaSlice<u8>,
13575        aq: &CudaSlice<i8>,
13576        ad: &CudaSlice<f32>,
13577        scales: &CudaSlice<f32>,
13578        m: usize,
13579        in_f: usize,
13580        out_f: usize,
13581        row_bytes: usize,
13582        scale_cols: usize,
13583    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13584        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
13585        self.qmatvec_e4m3_blk_mmvq_into(
13586            bytes, aq, ad, scales, m, in_f, out_f, row_bytes, scale_cols, &mut y,
13587        )?;
13588        Ok(y)
13589    }
13590
13591    /// Slot-fed twin of `qmatvec_e4m3_blk_mmvq` (caller-owned output; the alloc-free capture lane).
13592    #[allow(clippy::too_many_arguments)]
13593    pub fn qmatvec_e4m3_blk_mmvq_into(
13594        &self,
13595        bytes: &CudaSlice<u8>,
13596        aq: &CudaSlice<i8>,
13597        ad: &CudaSlice<f32>,
13598        scales: &CudaSlice<f32>,
13599        m: usize,
13600        in_f: usize,
13601        out_f: usize,
13602        row_bytes: usize,
13603        scale_cols: usize,
13604        y: &mut CudaSlice<f32>,
13605    ) -> Result<(), Box<dyn std::error::Error>> {
13606        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
13607        let f = self.func("qmatvec_e4m3_blk_mmvq");
13608        let cfg = LaunchConfig {
13609            grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), m as u32, 1),
13610            block_dim: (32, ROWS_PER_BLOCK, 1), // warp-per-row
13611            shared_mem_bytes: 0,                // warp-only reduce
13612        };
13613        let (inf, outf, mi, rb, sc) = (
13614            in_f as i32,
13615            out_f as i32,
13616            m as i32,
13617            row_bytes as i64,
13618            scale_cols as i32,
13619        );
13620        let __s_b = self.gpu.stream();
13621        let mut b = __s_b.launch_builder(&f);
13622        b.arg(bytes)
13623            .arg(aq)
13624            .arg(ad)
13625            .arg(scales)
13626            .arg(&mut *y)
13627            .arg(&inf)
13628            .arg(&outf)
13629            .arg(&mi)
13630            .arg(&rb)
13631            .arg(&sc);
13632        unsafe {
13633            b.launch(cfg)?;
13634        }
13635        Ok(())
13636    }
13637
13638    /// BLOCK-128 e4m3 BATCHED matvec (lane/rp-on-st, 2026-08-06): the weight-read-once twin of
13639    /// `qmatvec_e4m3_blk_mmvq` for m=2..16. Per (token,row) BIT-IDENTICAL to the grid.y=m launch
13640    /// (same fmaf chain, same per-k32 `s * ad` fold, same warp reduce), so it inherits the
13641    /// decode-exactness contract while reading the weight ONCE for up to `mcols` columns instead
13642    /// of `m` times. `mcols` must be one of {2,4,8,16} and satisfy `mcols >= m`.
13643    #[allow(clippy::too_many_arguments)]
13644    pub fn qmatvec_e4m3_blk_mmvq_batched(
13645        &self,
13646        bytes: &CudaSlice<u8>,
13647        aq: &CudaSlice<i8>,
13648        ad: &CudaSlice<f32>,
13649        scales: &CudaSlice<f32>,
13650        m: usize,
13651        in_f: usize,
13652        out_f: usize,
13653        row_bytes: usize,
13654        scale_cols: usize,
13655        mcols: usize,
13656    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13657        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
13658        debug_assert!(mcols >= m, "blk batched: mcols {mcols} < m {m}");
13659        let name = match mcols {
13660            2 => "qmatvec_e4m3_blk_mmvq_b2",
13661            4 => "qmatvec_e4m3_blk_mmvq_b4",
13662            8 => "qmatvec_e4m3_blk_mmvq_b8",
13663            16 => "qmatvec_e4m3_blk_mmvq_b16",
13664            _ => {
13665                return Err(
13666                    format!("qmatvec_e4m3_blk_mmvq_batched: no kernel for mcols {mcols}").into(),
13667                );
13668            }
13669        };
13670        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
13671        let f = self.func(name);
13672        let cfg = LaunchConfig {
13673            grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
13674            block_dim: (32, ROWS_PER_BLOCK, 1),
13675            shared_mem_bytes: 0,
13676        };
13677        let (inf, outf, mi, rb, sc) = (
13678            in_f as i32,
13679            out_f as i32,
13680            m as i32,
13681            row_bytes as i64,
13682            scale_cols as i32,
13683        );
13684        let __s_b = self.gpu.stream();
13685        let mut b = __s_b.launch_builder(&f);
13686        b.arg(bytes)
13687            .arg(aq)
13688            .arg(ad)
13689            .arg(scales)
13690            .arg(&mut y)
13691            .arg(&inf)
13692            .arg(&outf)
13693            .arg(&mi)
13694            .arg(&rb)
13695            .arg(&sc);
13696        unsafe {
13697            b.launch(cfg)?;
13698        }
13699        Ok(y)
13700    }
13701
13702    /// Test entry for the kernel_check exactness gate: the block-128 e4m3 batched MMVQ from raw
13703    /// bytes with an internal q8_1 quantize (mirrors `qmatvec_batched_raw`).
13704    #[allow(clippy::too_many_arguments)]
13705    pub fn qmatvec_e4m3_blk_batched_raw(
13706        &self,
13707        bytes: &CudaSlice<u8>,
13708        x: &CudaSlice<f32>,
13709        scales: &CudaSlice<f32>,
13710        m: usize,
13711        in_f: usize,
13712        out_f: usize,
13713        row_bytes: usize,
13714        scale_cols: usize,
13715        mcols: usize,
13716    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13717        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
13718        self.qmatvec_e4m3_blk_mmvq_batched(
13719            bytes, &aq, &ad, scales, m, in_f, out_f, row_bytes, scale_cols, mcols,
13720        )
13721    }
13722
13723    /// Test entry for the kernel_check exactness gate: the block-128 e4m3 MMVQ from raw bytes with
13724    /// an internal q8_1 quantize (mirrors `qmatvec_mmvq_raw`).
13725    #[allow(clippy::too_many_arguments)]
13726    pub fn qmatvec_e4m3_blk_mmvq_raw(
13727        &self,
13728        bytes: &CudaSlice<u8>,
13729        x: &CudaSlice<f32>,
13730        scales: &CudaSlice<f32>,
13731        m: usize,
13732        in_f: usize,
13733        out_f: usize,
13734        row_bytes: usize,
13735        scale_cols: usize,
13736    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13737        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
13738        self.qmatvec_e4m3_blk_mmvq(
13739            bytes, &aq, &ad, scales, m, in_f, out_f, row_bytes, scale_cols,
13740        )
13741    }
13742
13743    /// Test entries for the kernel_check bit-parity gate: fused e4m3 launches from raw weight
13744    /// bytes with internal q8_1 quantize, no env gating (mirrors `qmatvec_q8_fused*_raw`).
13745    #[allow(clippy::too_many_arguments)]
13746    pub fn qmatvec_e4m3_fused2_raw(
13747        &self,
13748        b0: &CudaSlice<u8>,
13749        b1: &CudaSlice<u8>,
13750        x: &CudaSlice<f32>,
13751        in_f: usize,
13752        out0: usize,
13753        out1: usize,
13754        row_bytes: usize,
13755        ws0: f32,
13756        ws1: f32,
13757    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13758        let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
13759        self.e4m3_fused2_core(b0, b1, &aq, &ad, in_f, out0, out1, row_bytes, ws0, ws1)
13760    }
13761
13762    #[allow(clippy::too_many_arguments)]
13763    pub fn qmatvec_e4m3_fused3_raw(
13764        &self,
13765        b0: &CudaSlice<u8>,
13766        b1: &CudaSlice<u8>,
13767        b2: &CudaSlice<u8>,
13768        x: &CudaSlice<f32>,
13769        in_f: usize,
13770        out0: usize,
13771        out1: usize,
13772        out2: usize,
13773        row_bytes: usize,
13774        ws0: f32,
13775        ws1: f32,
13776        ws2: f32,
13777    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13778        let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
13779        self.e4m3_fused3_core(
13780            b0, b1, b2, &aq, &ad, in_f, out0, out1, out2, row_bytes, ws0, ws1, ws2,
13781        )
13782    }
13783
13784    #[allow(clippy::too_many_arguments)]
13785    pub fn qmatvec_e4m3_fused2_t_raw(
13786        &self,
13787        b0: &CudaSlice<u8>,
13788        b1: &CudaSlice<u8>,
13789        x: &CudaSlice<f32>,
13790        m: usize,
13791        in_f: usize,
13792        out0: usize,
13793        out1: usize,
13794        row_bytes: usize,
13795        ws0: f32,
13796        ws1: f32,
13797    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13798        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
13799        self.e4m3_fused2_t_core(b0, b1, &aq, &ad, m, in_f, out0, out1, row_bytes, ws0, ws1)
13800    }
13801
13802    #[allow(clippy::too_many_arguments)]
13803    pub fn qmatvec_e4m3_fused3_t_raw(
13804        &self,
13805        b0: &CudaSlice<u8>,
13806        b1: &CudaSlice<u8>,
13807        b2: &CudaSlice<u8>,
13808        x: &CudaSlice<f32>,
13809        m: usize,
13810        in_f: usize,
13811        out0: usize,
13812        out1: usize,
13813        out2: usize,
13814        row_bytes: usize,
13815        ws0: f32,
13816        ws1: f32,
13817        ws2: f32,
13818    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13819        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
13820        self.e4m3_fused3_t_core(
13821            b0, b1, b2, &aq, &ad, m, in_f, out0, out1, out2, row_bytes, ws0, ws1, ws2,
13822        )
13823    }
13824
13825    /// THE single dispatch point for `QT_F8_E4M3_BLK` from a PRE-QUANTIZED q8_1 activation
13826    /// (lane/fp8-blk128-decode). Every `matmul_pre`-family entry calls this first, so the block-128
13827    /// class has exactly ONE code path across `matmul`, `matmul_pre`, `matmul_pre_noscale`,
13828    /// `matmul_decode_exact` and `matmul_decode_exact_pre` — the same kernel at the same grid for
13829    /// every m, which is what makes verify == decode bit-for-bit at every tier for free.
13830    ///
13831    /// Returns None for any other qtype (the caller continues its normal dispatch). The `blk: Some`
13832    /// pattern is part of the match, not an unwrap: qtype and grid presence are set together in the
13833    /// one residency arm that builds this tensor, and a qtype-without-grid would be a construction
13834    /// bug — better to fall through and hit a loud refusal than to unwrap a None here.
13835    fn try_e4m3_blk_pre(
13836        &self,
13837        w: &crate::model::GpuTensor,
13838        aq: &CudaSlice<i8>,
13839        ad: &CudaSlice<f32>,
13840        m: usize,
13841    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
13842        use crate::model::GpuTensor;
13843        if let GpuTensor::Quant {
13844            bytes,
13845            qtype,
13846            row_bytes,
13847            blk: Some(g),
13848            ..
13849        } = w
13850        {
13851            if *qtype == QT_F8_E4M3_BLK {
13852                // BATCHED tier m=2..16 (lane/rp-on-st): weight read ONCE for up to mcols columns
13853                // instead of m grid.y re-reads. Bit-identical per (token,row) to the grid.y=m form
13854                // below, so the decode-exactness contract is preserved at every width. Gated by
13855                // the same seams the other batched families honor (MEMRA_NO_BATCHED, MEMRA_B8) so
13856                // one rollback door covers every dtype's batched tier.
13857                if (2..=16).contains(&m)
13858                    && std::env::var("MEMRA_NO_BATCHED").is_err()
13859                    && (m <= 4 || Self::b8_enabled())
13860                {
13861                    let mcols = Self::batched_mcols(m);
13862                    return Ok(Some(self.qmatvec_e4m3_blk_mmvq_batched(
13863                        bytes,
13864                        aq,
13865                        ad,
13866                        &g.scales,
13867                        m,
13868                        w.in_features(),
13869                        w.out_features(),
13870                        *row_bytes,
13871                        g.cols,
13872                        mcols,
13873                    )?));
13874                }
13875                return Ok(Some(self.qmatvec_e4m3_blk_mmvq(
13876                    bytes,
13877                    aq,
13878                    ad,
13879                    &g.scales,
13880                    m,
13881                    w.in_features(),
13882                    w.out_features(),
13883                    *row_bytes,
13884                    g.cols,
13885                )?));
13886            }
13887        }
13888        Ok(None)
13889    }
13890
13891    /// PREFILL (m >= GEMM_M_THRESHOLD) for `QT_F8_E4M3_BLK` — DEQUANT-PER-CALL to the Q8_0 slab
13892    /// this class's residency replaced, then the ordinary Q8_0 prefill dispatch on the transient.
13893    ///
13894    /// WHY THIS EXISTS AT ALL, i.e. the regression it prevents: the decode kernel is a warp-per-row
13895    /// GEMV. At grid.y=m it re-reads the whole weight once PER TOKEN, so letting a 512-token prefill
13896    /// chunk reach it would be a ~500x weight-traffic blowup on the single most bandwidth-bound part
13897    /// of the forward. Native residency is a DECODE win and must not be paid for in prefill, so
13898    /// prefill keeps the floor's arithmetic and the floor's kernels.
13899    ///
13900    /// WHY DEQUANT-PER-CALL rather than a second resident slab: a resident slab is dual residency —
13901    /// it gives back the entire 1.0-vs-1.0625 B/weight win this lane exists to capture (and then
13902    /// some, since the e4m3 copy stays too). The transient costs one linear device pass per
13903    /// (projection, prefill call) and frees immediately.
13904    ///
13905    /// NUMERICALLY IT IS THE FLOOR, EXACTLY: `fp8_blk_dequant_q8_0` is the merged ARM B' kernel,
13906    /// gate-proven BYTE-IDENTICAL to the host dequant+re-encode (kernel-check `fp8-blk-gpu`). So the
13907    /// slab these bytes form is bit-for-bit the slab the `MEMRA_ST_E4M3_BLK=0` arm makes resident,
13908    /// and every prefill kernel downstream sees identical input — prefill logits under this lane are
13909    /// bit-identical to prefill logits under the floor, which is what makes the decode A/B a clean
13910    /// single-variable comparison instead of a two-variable one.
13911    ///
13912    /// WHAT IT COSTS, MEASURED, AND WHY THAT COST IS MOSTLY STRUCTURAL (27B block-128 ckpt, pp512,
13913    /// this rig = RTX 5090 Laptop, ~896 GB/s GDDR7). This arm makes prefill move the weight THREE
13914    /// times instead of once: read 6.88 GB of e4m3, write 7.31 GB of Q8_0, then the MMQ reads that
13915    /// 7.31 GB back. The two extra passes are 14.19 GB = 15.8 ms at this card's roofline against a
13916    /// ~332 ms pp512, i.e. **~-4.5% pp is a floor no kernel tuning can remove** — only deleting the
13917    /// dequant can. Measured: the dequant kernel costs 27.9 ms/pass (nsys, 208 projections) after
13918    /// the 2026-08-05 vector rewrite (was 66.5 ms at one byte per thread), and e2e pp512 is
13919    /// 1451.4 vs the slab arm's 1541.6 tok/s = -5.8% (N=3 interleaved pairs). So ~1.3pp of the
13920    /// -5.8% is residual kernel inefficiency and ~4.5pp is the extra traffic itself.
13921    ///
13922    /// SO THE DEQUANT IS NO LONGER THE DEFAULT ROUTE — it is the FALLBACK. The per-block FP8 MMQ
13923    /// tile (`try_fp8_blk_mmq`) consumes the resident e4m3 bytes + grid DIRECTLY, deleting both extra
13924    /// passes, and since 2026-08-05 it runs FIRST and by default for the native-resident source
13925    /// (`fp8_blk_mmq_native_enabled`; `MEMRA_FP8_MMQ=0` is the seam back to this dequant). On paper
13926    /// the trade was unassumable — lane/fp8-mmq-v2 measured that tile at 0.85-1.09x the Q8_0 MMQ
13927    /// floor GEMM-only, so it swapped a -4.5% traffic cost for a 0-to-15% GEMM cost of unknown sign.
13928    /// Measured on the 27B (3 arms interleaved, N=3, research/fp8blk-20260805/VERDICT.md): slab
13929    /// 1540.5 / this dequant 1449.1 / the tile 1553.3 tok/s, min(tile) > max(slab). The tile wins
13930    /// because v2's denominator had its slab already resident while this class's floor must build it
13931    /// every call; same tile, opposite sign, because the question changed.
13932    ///
13933    /// THIS ARM STILL RUNS, and is not dead code: every `try_fp8_blk_mmq` precondition (in_f % 16,
13934    /// grid dims vs shape, per-tensor scale == 1.0, the e4m3-NaN scan) refuses by falling through to
13935    /// here, so a checkpoint the tile cannot take keeps exact prefill on the floor's own bits rather
13936    /// than losing the class. It is also what `MEMRA_FP8_MMQ=0` reverts to.
13937    fn try_e4m3_blk_prefill(
13938        &self,
13939        w: &crate::model::GpuTensor,
13940        x: &CudaSlice<f32>,
13941        m: usize,
13942    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
13943        use crate::model::GpuTensor;
13944        let GpuTensor::Quant {
13945            bytes,
13946            qtype,
13947            blk: Some(g),
13948            ..
13949        } = w
13950        else {
13951            return Ok(None);
13952        };
13953        if *qtype != QT_F8_E4M3_BLK {
13954            return Ok(None);
13955        }
13956        // NO-DEQUANT ROUTE, THE DEFAULT (MEMRA_FP8_MMQ=0 reverts): the per-block MMQ tile eats the
13957        // resident e4m3 bytes and grid as-is, so neither extra weight pass happens. Its own
13958        // preconditions (in_f % 16, grid dims, scale == 1.0, no e4m3 NaN code) can refuse — fall
13959        // through to the dequant below when they do, never silently produce nothing.
13960        if let Some(y) = self.try_fp8_blk_mmq(w, x, m)? {
13961            return Ok(Some(y));
13962        }
13963        let (in_f, out_f) = (w.in_features(), w.out_features());
13964        let slab = self.fp8_blk_dequant_q8_0_dev(bytes, &g.scales, out_f, in_f)?;
13965        let tmp = GpuTensor::Quant {
13966            bytes: slab,
13967            qtype: QT_Q8_0,
13968            row_bytes: in_f / 32 * 34,
13969            ne: vec![in_f as u64, out_f as u64],
13970            scale: 1.0,
13971            rp: false,
13972            #[cfg(memra_cutlass)]
13973            cutlass: None,
13974            fp8: None,
13975            blk: None,
13976            f16: None,
13977            rp4: None,
13978        };
13979        // Recursion terminates: `tmp` is QT_Q8_0 with `blk: None`, so it cannot re-enter this arm.
13980        Ok(Some(self.matmul(&tmp, x, m)?))
13981    }
13982
13983    pub fn matmul_pre_noscale(
13984        &self,
13985        w: &crate::model::GpuTensor,
13986        aq: &CudaSlice<i8>,
13987        ad: &CudaSlice<f32>,
13988        m: usize,
13989    ) -> Result<Option<(CudaSlice<f32>, f32)>, Box<dyn std::error::Error>> {
13990        use crate::model::GpuTensor;
13991        // BLOCK-128 e4m3: every scale factor is folded inside the kernel per k128, so the
13992        // "separable post-op scale" this entry exists to defer is 1.0 — return it explicitly
13993        // rather than let the tail below refuse and cost the caller a re-dispatch.
13994        if m == 1 {
13995            if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
13996                return Ok(Some((y, 1.0)));
13997            }
13998        }
13999        // Only the m==1 fast path applies the scale as a separable post-op; bail everywhere else.
14000        if m != 1 || !self.uses_q8_1_fast(w) {
14001            return Ok(None);
14002        }
14003        let in_f = w.in_features();
14004        let out_f = w.out_features();
14005        let (bytes, qtype, row_bytes, scale, rp) = match w {
14006            GpuTensor::Quant {
14007                bytes,
14008                qtype,
14009                row_bytes,
14010                scale,
14011                rp,
14012                ..
14013            } => (bytes, *qtype, *row_bytes, *scale, *rp),
14014            _ => return Ok(None),
14015        };
14016        // MMVQ warp-per-row (scale==1.0 passed -> kernel skips its internal scale; we return scale).
14017        if self.mmvq_supports(qtype) {
14018            // Q4_0 split-plane mirror (dp4a fallback below keeps the raw GGUF bytes).
14019            let (mbytes, mrp) = match w {
14020                GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
14021                _ => (bytes, rp),
14022            };
14023            let y = self.qmatvec_mmvq(
14024                mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, /*scale*/ 1.0, mrp,
14025            )?;
14026            return Ok(Some((y, scale)));
14027        }
14028        // dp4a fallback: same launch as matmul_pre but WITHOUT the post scale_inplace.
14029        let name = match qtype {
14030            QT_Q8_0 => "qmatvec_q8_0_dp4a",
14031            QT_Q4_K => "qmatvec_q4_K_dp4a",
14032            QT_Q6_K => "qmatvec_q6_K_dp4a",
14033            QT_Q5_K => "qmatvec_q5_K_dp4a",
14034            QT_Q3_K => "qmatvec_q3_K_dp4a",
14035            QT_NVFP4 => {
14036                if rp {
14037                    "qmatvec_nvfp4_dp4a_rp"
14038                } else {
14039                    "qmatvec_nvfp4_dp4a"
14040                }
14041            }
14042            QT_IQ4_XS => "qmatvec_iq4_XS_dp4a",
14043            _ => return Ok(None),
14044        };
14045        let f = self.func(name);
14046        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
14047        let cfg = LaunchConfig {
14048            grid_dim: (out_f as u32, m as u32, 1),
14049            block_dim: (128, 1, 1),
14050            shared_mem_bytes: 0,
14051        };
14052        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
14053        let __s_b = self.gpu.stream();
14054        let mut b = __s_b.launch_builder(&f);
14055        b.arg(bytes)
14056            .arg(aq)
14057            .arg(ad)
14058            .arg(&mut y)
14059            .arg(&inf)
14060            .arg(&outf)
14061            .arg(&mi)
14062            .arg(&rb);
14063        unsafe {
14064            b.launch(cfg)?;
14065        }
14066        Ok(Some((y, scale)))
14067    }
14068
14069    /// True if `qtype` has a warp-per-row MMVQ decode kernel AND MEMRA_MMVQ is set. Only the 4
14070    /// daily-hot dtypes (Q8_0, Q4_K, Q6_K, NVFP4) — others keep the _dp4a matvec (oracle/fallback).
14071    pub fn mmvq_supports(&self, qtype: i32) -> bool {
14072        // DEFAULT ON since 2026-07-08 (MEMRA_MMVQ=0 reverts to the _dp4a matvec class).
14073        // QT_F8_E4M3 is exempt from the MEMRA_MMVQ=0 escape: the e4m3 mmvq family is that dtype's
14074        // ONLY int8-act kernel class (there is no _dp4a twin), so its m=1/verify/batched dispatch
14075        // is a pure function of the dtype — the decode-parity law holds under every env.
14076        if qtype == QT_F8_E4M3 {
14077            return true;
14078        }
14079        if std::env::var("MEMRA_MMVQ").as_deref() == Ok("0") {
14080            return false;
14081        }
14082        matches!(
14083            qtype,
14084            QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q6_K | QT_NVFP4 | QT_Q4_0
14085        )
14086    }
14087
14088    /// PERF-3 warp-per-row MMVQ launcher (decode m=1 hot path). block=(32,ROWS_PER_BLOCK,1):
14089    /// one warp owns one output row, warp-only __shfl reduction (no smem barrier). Bit-equivalent
14090    /// to qmatvec_*_dp4a up to f32 reduction order. Pre-quantized q8_1 activation (aq,ad). NVFP4
14091    /// per-tensor macro-scale applied post (scale==1.0 for other dtypes -> no-op).
14092    pub fn qmatvec_mmvq(
14093        &self,
14094        bytes: &CudaSlice<u8>,
14095        aq: &CudaSlice<i8>,
14096        ad: &CudaSlice<f32>,
14097        m: usize,
14098        in_f: usize,
14099        out_f: usize,
14100        qtype: i32,
14101        row_bytes: usize,
14102        scale: f32,
14103        rp: bool,
14104    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14105        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output: skip memset
14106        self.qmatvec_mmvq_into(
14107            bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, rp, &mut y,
14108        )?;
14109        Ok(y)
14110    }
14111
14112    /// Slot-fed MMVQ twin (alloc-free capture lane): full policy body, caller-owned output.
14113    #[allow(clippy::too_many_arguments)]
14114    pub fn qmatvec_mmvq_into(
14115        &self,
14116        bytes: &CudaSlice<u8>,
14117        aq: &CudaSlice<i8>,
14118        ad: &CudaSlice<f32>,
14119        m: usize,
14120        in_f: usize,
14121        out_f: usize,
14122        qtype: i32,
14123        row_bytes: usize,
14124        scale: f32,
14125        rp: bool,
14126        y: &mut CudaSlice<f32>,
14127    ) -> Result<(), Box<dyn std::error::Error>> {
14128        debug_assert!(y.len() >= m * out_f);
14129        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
14130        // SMALL-SHAPE GRID FILL (H100 lane, 2026-07-26 microbench: attn qkv out_f=2048 =
14131        // 0.97 waves at the 4-warp block -> 66% of peak). The g2 twin (2 warps/block)
14132        // doubles the grid when the 4-warp launch would be sub-wave; per-row program
14133        // identical -> bit-identical. MEMRA_Q80_G2=0 reverts.
14134        if qtype == QT_Q8_0
14135            && rp
14136            && m == 1
14137            && out_f >= 64
14138            && (out_f as u32).div_ceil(ROWS_PER_BLOCK) < 4 * self.sm_count() as u32
14139            && {
14140                static G2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
14141                *G2.get_or_init(|| std::env::var("MEMRA_Q80_G2").as_deref() != Ok("0"))
14142            }
14143        {
14144            let f = self.func("qmatvec_q8_0_mmvq_rp_g2");
14145            let cfg = LaunchConfig {
14146                grid_dim: ((out_f as u32).div_ceil(2), 1, 1),
14147                block_dim: (32, 2, 1),
14148                shared_mem_bytes: 0,
14149            };
14150            let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, 1i32, row_bytes as i64);
14151            let __s_b = self.gpu.stream();
14152            let mut b = __s_b.launch_builder(&f);
14153            b.arg(bytes)
14154                .arg(aq)
14155                .arg(ad)
14156                .arg(&mut *y)
14157                .arg(&inf)
14158                .arg(&outf)
14159                .arg(&mi)
14160                .arg(&rb);
14161            unsafe {
14162                b.launch(cfg)?;
14163            }
14164            if scale != 1.0 {
14165                self.scale_inplace(y, scale, out_f)?;
14166            }
14167            return Ok(());
14168        }
14169        // Multi-row-per-warp (mr2) policy, fixed since the 2026-07 sweeps (the MEMRA_MMVQ_MR
14170        // override + mr4 kernel were retired 2026-07-08 — mr4 regressed on register pressure and
14171        // crashed under rp; q4_K/q6_K mr2 measured flat, "no gain = no change"):
14172        //   NVFP4 m=1 -> mr2 (clean +1-2% on 9B: RPW acc chains hide the weight-load latency
14173        //     that pins the single-row kernel at 30-46% DRAM). Bit-identical per row.
14174        //   Q5_K m=1 -> mr2 (2026-07-05: the FR-Spec trimmed draft head is Q5_K 32768 rows = 8%
14175        //     of the 27B p3 spec wall; latency-bound like the other k-quants pre-fix).
14176        //   Q4_K/Q6_K m=1 -> single-row (mr2 measured +0.7% / flat — weight-bandwidth-bound).
14177        let mut mr: u32 = if m == 1 && (qtype == QT_NVFP4 || qtype == QT_Q5_K) {
14178            2
14179        } else {
14180            1
14181        };
14182        // Q4_0 mr (gemma trunk): DEFAULT 1 since 2026-07-13 (MEMRA_Q40_MR=2 reverts) — the
14183        // mr1 rp twin doubles the block count and wins the tail-quantization/latency battle
14184        // on every gemma model (E4B +3.75%: 198.9 vs 191.7; 26B +0.7%; 31B +0.9%; N=2-3
14185        // valid-window interleaved, bit-identical per row — same dot program).
14186        if m == 1 && qtype == QT_Q4_0 {
14187            static Q40MR: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
14188            // shape policy PROBED NEGATIVE (2026-07-13): tall-only mr1 197.2 vs
14189            // mr1-everywhere 198.7 — mr1 wins wide-output shapes too; arm removed.
14190            mr = *Q40MR.get_or_init(|| {
14191                std::env::var("MEMRA_Q40_MR")
14192                    .ok()
14193                    .and_then(|v| v.parse().ok())
14194                    .unwrap_or(1)
14195            });
14196        }
14197        // q5issue lane (2026-07-08): MEMRA_Q5K_ISSUE swaps the q5_K m=1 mmvq kernels for the
14198        // issue-reduced `_il` bodies (uint4 header/qh/qs loads + branchless scale decode —
14199        // cuts ~34 LDG.U16 + ~5 LDG.U8 + a warp-divergent scale branch per 32-elem group-row
14200        // to 5 LDG.128). Bit-identical per (token,row) to the reference kernels.
14201        // `1` = shape-aware policy (N=3 clock-locked micro-bench, mem P0, synthetic real shapes):
14202        //   out_f <= 65536 (trunk/frspec regime): il at the default mr — mr2_il -9.5%/-10.5%
14203        //     on 4096x4096/4096x8192, -3.1% on the 32768 frspec head vs the mr2-ref default;
14204        //   out_f > 65536 (the 248320-row 27B lm_head, already ~97% of the mem wall): mr2_il
14205        //     REGRESSES +22% there but mr1_il wins -2.1% vs the mr2-ref default -> force mr=1.
14206        // `2` = force il at the current mr for EVERY shape (A/B probe seam). Default OFF.
14207        let q5_mode = std::env::var("MEMRA_Q5K_ISSUE").ok();
14208        let q5_force = q5_mode.as_deref() == Some("2");
14209        // DEFAULT ON since 2026-07-08 (MEMRA_Q5K_ISSUE=0 reverts): +1.8% 9B plain e2e N=3
14210        // (128.2 -> 130.4), 27B flat (its big head is already at the mem wall), all gates green.
14211        let q5_il = qtype == QT_Q5_K
14212            && m == 1
14213            && (q5_force || q5_mode.as_deref().map(|v| v != "0").unwrap_or(true));
14214        if q5_il && !q5_force && out_f > 65536 {
14215            mr = 1;
14216        }
14217        // Q4_0 split-plane rp: mr2 default; MEMRA_Q40_MR=1 reaches the mr1 rp twin
14218        // (2026-07-13 — the tall-input/short-output tail-quantization probe).
14219        if qtype == QT_Q4_0 && rp && mr != 1 {
14220            mr = 2;
14221        }
14222        // Q8_0 rp (H100 lane): mr1 default — the q4_0 mr2 recipe MEASURED NEGATIVE on H100
14223        // (2026-07-26 N=3: mr1 186.2 vs mr2 171.5 tok/s; halving the grid on 132 SMs costs
14224        // more than 2-row ILP buys). mr2 kernel stays behind MEMRA_Q80_MR=2 for the corpus.
14225        if qtype == QT_Q8_0 && rp {
14226            static Q80MR: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
14227            mr = *Q80MR.get_or_init(|| {
14228                std::env::var("MEMRA_Q80_MR")
14229                    .ok()
14230                    .and_then(|v| v.parse().ok())
14231                    .unwrap_or(1)
14232            });
14233        }
14234        let name = match (qtype, mr, rp) {
14235            (QT_NVFP4, 2, false) => "qmatvec_nvfp4_mmvq_mr2",
14236            (QT_NVFP4, 2, true) => "qmatvec_nvfp4_mmvq_mr2_rp",
14237            (QT_NVFP4, _, true) => "qmatvec_nvfp4_mmvq_rp",
14238            (QT_Q4_0, 1, true) => "qmatvec_q4_0_mmvq_rp",
14239            (QT_Q4_0, _, true) => "qmatvec_q4_0_mmvq_mr2_rp",
14240            (QT_Q5_K, 2, _) => {
14241                if q5_il {
14242                    "qmatvec_q5_K_mmvq_mr2_il"
14243                } else {
14244                    "qmatvec_q5_K_mmvq_mr2"
14245                }
14246            }
14247            (QT_Q8_0, 2, true) => "qmatvec_q8_0_mmvq_mr2_rp",
14248            // rpca (cp.async-staged weight ring): MEASURED NEGATIVE on H100 for Q8_0
14249            // (2026-07-26 N=3: 181.8 vs plain rp 185.5 — the smem round-trip exceeds the
14250            // latency it hides for 8-bit direct-dp4a; the NVFP4 win case overlaps table
14251            // decode with half the bytes). OPT-IN via MEMRA_Q80_CA=1 for the corpus.
14252            (QT_Q8_0, _, true)
14253                if in_f % 1024 == 0 && {
14254                    static CA: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
14255                    *CA.get_or_init(|| std::env::var("MEMRA_Q80_CA").as_deref() == Ok("1"))
14256                } =>
14257            {
14258                "qmatvec_q8_0_mmvq_rpca"
14259            }
14260            (QT_Q8_0, _, true) => "qmatvec_q8_0_mmvq_rp",
14261            (QT_Q8_0, _, _) => "qmatvec_q8_0_mmvq",
14262            // K-quant split-plane twins (H100 K-quant coalescing fix, 2026-08-01): the rp4
14263            // mirror routes here; GGUF layout keeps the plain kernels. rp bytes MUST never
14264            // reach a GGUF-layout kernel or vice versa.
14265            (QT_Q4_K, _, true) => "qmatvec_q4_K_mmvq_rp",
14266            (QT_Q6_K, _, true) => "qmatvec_q6_K_mmvq_rp",
14267            (QT_Q4_K, _, _) => "qmatvec_q4_K_mmvq",
14268            (QT_Q4_0, 2, false) => "qmatvec_q4_0_mmvq_mr2",
14269            (QT_Q4_0, _, false) => "qmatvec_q4_0_mmvq",
14270            (QT_Q5_K, _, _) => {
14271                if q5_il {
14272                    "qmatvec_q5_K_mmvq_il"
14273                } else {
14274                    "qmatvec_q5_K_mmvq"
14275                }
14276            }
14277            (QT_Q6_K, _, _) => "qmatvec_q6_K_mmvq",
14278            (QT_NVFP4, _, false) => "qmatvec_nvfp4_mmvq",
14279            (QT_F8_E4M3, _, _) => "qmatvec_e4m3_mmvq",
14280            _ => panic!("qmatvec_mmvq: qtype {qtype} has no MMVQ kernel"),
14281        };
14282        let f = self.func(name);
14283        // each block still has ROWS_PER_BLOCK warps; with mr rows/warp it covers ROWS_PER_BLOCK*mr rows.
14284        let rows_per_block = ROWS_PER_BLOCK * mr;
14285        let cfg = LaunchConfig {
14286            grid_dim: (
14287                (out_f as u32 + rows_per_block - 1) / rows_per_block,
14288                m as u32,
14289                1,
14290            ),
14291            block_dim: (32, ROWS_PER_BLOCK, 1), // warp-per-row (x mr rows each)
14292            shared_mem_bytes: 0,                // warp-only reduce at m=1
14293        };
14294        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
14295        let __s_b = self.gpu.stream();
14296        let mut b = __s_b.launch_builder(&f);
14297        // NVFP4 + e4m3 mmvq kernels take the macro-scale as a fused epilogue arg (applied at the
14298        // write — bit-identical to the old separate scale_inplace pass, minus one launch per matvec:
14299        // 53 scale launches/token on the 9B; for e4m3 the scale is the checkpoint's per-tensor f32
14300        // weight_scale). Other mmvq kernels keep the 8-arg signature.
14301        if qtype == QT_NVFP4 || qtype == QT_F8_E4M3 {
14302            // PDL wave-B: the nvfp4 mr2_rp single (gemma wo / generic rp singles) joins
14303            // the wave-A launch class — 9-arg flavor (fused macro-scale epilogue).
14304            if Self::pdl_on()
14305                && Self::pdl_mmvq_on()
14306                && Self::pdl_nvfp4q8_on()
14307                && name == "qmatvec_nvfp4_mmvq_mr2_rp"
14308            {
14309                use cudarc::driver::{DevicePtr, DevicePtrMut};
14310                let s = &self.gpu.stream();
14311                let (pw, _g0) = bytes.device_ptr(s);
14312                let (paq, _g1) = aq.device_ptr(s);
14313                let (pad, _g2) = ad.device_ptr(s);
14314                let (py, _g3) = y.device_ptr_mut(s);
14315                let mut ps = [
14316                    &pw as *const _ as *mut std::ffi::c_void,
14317                    &paq as *const _ as *mut _,
14318                    &pad as *const _ as *mut _,
14319                    &py as *const _ as *mut _,
14320                    &inf as *const _ as *mut _,
14321                    &outf as *const _ as *mut _,
14322                    &mi as *const _ as *mut _,
14323                    &rb as *const _ as *mut _,
14324                    &scale as *const _ as *mut _,
14325                ];
14326                unsafe {
14327                    self.launch_pdl(name, cfg.grid_dim, cfg.block_dim, &mut ps)?;
14328                }
14329                return Ok(());
14330            }
14331            b.arg(bytes)
14332                .arg(aq)
14333                .arg(ad)
14334                .arg(&mut *y)
14335                .arg(&inf)
14336                .arg(&outf)
14337                .arg(&mi)
14338                .arg(&rb)
14339                .arg(&scale);
14340            unsafe {
14341                b.launch(cfg)?;
14342            }
14343        } else if Self::pdl_on()
14344            && Self::pdl_mmvq_on()
14345            && (matches!(
14346                name,
14347                "qmatvec_q4_0_mmvq_rp" | "qmatvec_q6_K_mmvq" | "qmatvec_q6_K_mmvq_rp"
14348            ) || (Self::pdl_nvfp4q8_on()
14349                && matches!(name, "qmatvec_q8_0_mmvq_rp" | "qmatvec_q8_0_mmvq_mr2_rp")))
14350        {
14351            // PDL wave-A (2026-07-23): the two decode-hot single-matvec kernels carry
14352            // MEMRA_PDL_ENTRY — grid launches while the producer drains. ONLY the marked
14353            // names may take this launch (unmarked kernels would read unordered).
14354            {
14355                use cudarc::driver::{DevicePtr, DevicePtrMut};
14356                let s = &self.gpu.stream();
14357                let (pw, _g0) = bytes.device_ptr(s);
14358                let (paq, _g1) = aq.device_ptr(s);
14359                let (pad, _g2) = ad.device_ptr(s);
14360                let (py, _g3) = y.device_ptr_mut(s);
14361                let mut ps = [
14362                    &pw as *const _ as *mut std::ffi::c_void,
14363                    &paq as *const _ as *mut _,
14364                    &pad as *const _ as *mut _,
14365                    &py as *const _ as *mut _,
14366                    &inf as *const _ as *mut _,
14367                    &outf as *const _ as *mut _,
14368                    &mi as *const _ as *mut _,
14369                    &rb as *const _ as *mut _,
14370                ];
14371                unsafe {
14372                    self.launch_pdl(name, cfg.grid_dim, cfg.block_dim, &mut ps)?;
14373                }
14374            }
14375            if scale != 1.0 {
14376                self.scale_inplace(y, scale, m * out_f)?;
14377            }
14378        } else {
14379            b.arg(bytes)
14380                .arg(aq)
14381                .arg(ad)
14382                .arg(&mut *y)
14383                .arg(&inf)
14384                .arg(&outf)
14385                .arg(&mi)
14386                .arg(&rb);
14387            unsafe {
14388                b.launch(cfg)?;
14389            }
14390            if scale != 1.0 {
14391                self.scale_inplace(y, scale, m * out_f)?;
14392            }
14393        }
14394        Ok(())
14395    }
14396
14397    /// Test entry for the kernel_check bit-equivalence gate: run the warp-per-row MMVQ directly
14398    /// from raw weight bytes (quantize the f32 activation `x` to q8_1 internally). NVFP4 per-tensor
14399    /// macro-scale is NOT applied (caller compares bare, like qmatvec_*_fast). Mirrors qmatvec_gemm_raw.
14400    pub fn qmatvec_mmvq_raw(
14401        &self,
14402        bytes: &CudaSlice<u8>,
14403        x: &CudaSlice<f32>,
14404        m: usize,
14405        in_f: usize,
14406        out_f: usize,
14407        qtype: i32,
14408        row_bytes: usize,
14409        rp: bool,
14410    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14411        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
14412        self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, 1.0, rp)
14413    }
14414
14415    /// True if `qtype` has a batched weight-resident (`_b2`/`_b4`) matvec kernel. These mirror the
14416    /// `_mmvq` kernels but iterate the m token columns INSIDE one warp/row, so the weight bytes leave
14417    /// HBM/L2 once for m tokens (vs grid.y=m re-reading m times). The 5 daily-hot dtypes have them.
14418    pub fn batched_supports(&self, qtype: i32) -> bool {
14419        matches!(
14420            qtype,
14421            QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q6_K | QT_NVFP4 | QT_F8_E4M3 | QT_Q4_0
14422        )
14423    }
14424
14425    /// IQ4_XS trunk fast seam: MEMRA_IQ_FAST=0 reverts non-expert IQ4_XS matmuls to the Stage-A
14426    /// f32 oracle path. Default ON since 2026-08-02 (research/kat-anomaly-20260802/): the old
14427    /// opt-in default left every IQ4_XS-trunk artifact (KAT-Coder IQ4_XS: attn_qkv/attn_gate/
14428    /// ssm_out/shexp, ~0.52GB re-read per decode tick) on the oracle kernel — decode 106.7 ->
14429    /// 193.4 tok/s (x5 interleaved), pp512 228 -> 697, same bytes, via qmatvec_iq4_XS_dp4a. The
14430    /// supported artifacts carry IQ4_XS only in EXPERT banks (their own dispatch, not this seam),
14431    /// so this admission is dispatch-unchanged for every non-IQ4_XS-trunk model.
14432    pub fn iq_fast_enabled() -> bool {
14433        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
14434        *ON.get_or_init(|| {
14435            std::env::var("MEMRA_IQ_FAST")
14436                .map(|v| v != "0")
14437                .unwrap_or(true)
14438        })
14439    }
14440
14441    /// b8 tier seam: MEMRA_B8=0 keeps m=5..8 on the per-m grid.y=m path (m=2..4 batched dispatch
14442    /// unaffected). Default ON — the K=4..7 spec-verify weight-read-once fix.
14443    pub fn b8_enabled() -> bool {
14444        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
14445        *ON.get_or_init(|| std::env::var("MEMRA_B8").map(|v| v != "0").unwrap_or(true))
14446    }
14447
14448    /// Compile-time column batch for a runtime m: 2 -> b2, 3..4 -> b4, 5..8 -> b8.
14449    pub fn batched_mcols(m: usize) -> usize {
14450        if m == 2 {
14451            2
14452        } else if m <= 4 {
14453            4
14454        } else if m <= 8 {
14455            8
14456        } else {
14457            16
14458        }
14459    }
14460
14461    /// Kernel name for the batched matvec of `(qtype, mcols)`. mcols ∈ {2,4,8}. The b8 tier is the
14462    /// K=4..7 spec-verify fix (T=5..8): pre-b8 those T fell to grid.y=m per-row MMVQ = m full
14463    /// weight reads/launch — the measured 27B K=4 cliff (101 -> 73 tok/s at p3 despite acceptance
14464    /// holding 54%). One b8 launch reads the weight ONCE for up to 8 columns (c >= m masked).
14465    fn batched_kernel_name(qtype: i32, mcols: usize) -> Option<&'static str> {
14466        Some(match (qtype, mcols) {
14467            (QT_Q8_0, 2) => "qmatvec_q8_0_mmvq_b2",
14468            (QT_Q8_0, 4) => "qmatvec_q8_0_mmvq_b4",
14469            (QT_Q8_0, 8) => "qmatvec_q8_0_mmvq_b8",
14470            // b16 now has BOTH forms (lane/rp-on-st, 2026-08-06). It used to be rp-ONLY, which
14471            // made the q8rp mirror the exact-16 tier's admission ticket for any model carrying a
14472            // single Q8_0 matmul — measured as the FP8-ST refusal (`L0.ssm_beta qtype=0
14473            // rp4=false`, 96 t / 23.9 MiB = 0.143% of resident weight). The mirror stays a
14474            // BANDWIDTH lever on Q8_0-dominant GGUFs; it is no longer a correctness prerequisite.
14475            (QT_Q8_0, 16) => "qmatvec_q8_0_mmvq_b16",
14476            (QT_Q4_K, 2) => "qmatvec_q4_K_mmvq_b2",
14477            (QT_Q4_K, 4) => "qmatvec_q4_K_mmvq_b4",
14478            (QT_Q4_K, 8) => "qmatvec_q4_K_mmvq_b8",
14479            // b16 base + _rp (lane/rp-on-st): the 9B NVFP4 GGUF's blocker — real NVFP4 GGUFs keep
14480            // Q4_K attention next to NVFP4 MLP, and the tier's predicate is an ALL.
14481            (QT_Q4_K, 16) => "qmatvec_q4_K_mmvq_b16",
14482            (QT_Q5_K, 2) => "qmatvec_q5_K_mmvq_b2",
14483            (QT_Q5_K, 4) => "qmatvec_q5_K_mmvq_b4",
14484            (QT_Q5_K, 8) => "qmatvec_q5_K_mmvq_b8",
14485            // b16 base only (lane/rp-on-st): Q5_K has no rp twins at any width, so there is
14486            // nothing to mirror. Named by the diagnostic as `L0.wqkv_gate qtype=3` on the 9B.
14487            (QT_Q5_K, 16) => "qmatvec_q5_K_mmvq_b16",
14488            (QT_Q6_K, 2) => "qmatvec_q6_K_mmvq_b2",
14489            (QT_Q6_K, 4) => "qmatvec_q6_K_mmvq_b4",
14490            (QT_Q6_K, 8) => "qmatvec_q6_K_mmvq_b8",
14491            (QT_Q6_K, 16) => "qmatvec_q6_K_mmvq_b16",
14492            (QT_NVFP4, 2) => "qmatvec_nvfp4_mmvq_b2",
14493            (QT_NVFP4, 4) => "qmatvec_nvfp4_mmvq_b4",
14494            (QT_NVFP4, 8) => "qmatvec_nvfp4_mmvq_b8",
14495            // b16 (lane/rp-on-st): no mirror needed — NVFP4's 36 B/k32 block is already the
14496            // aligned form its own kernel walks. Unlocks the exact-16 tier for every NVFP4 model
14497            // AND for the mixed FP8-ST artifact, whose 193 NVFP4 tensors were refusing it.
14498            (QT_NVFP4, 16) => "qmatvec_nvfp4_mmvq_b16",
14499            (QT_F8_E4M3, 2) => "qmatvec_e4m3_mmvq_b2",
14500            (QT_F8_E4M3, 4) => "qmatvec_e4m3_mmvq_b4",
14501            (QT_F8_E4M3, 8) => "qmatvec_e4m3_mmvq_b8",
14502            // b16 tier (lane/rp-on-st): e4m3 needs NO split-plane mirror to reach it — its native
14503            // row-major layout is already 32B-aligned per k32 block, so the base kernel IS the
14504            // aligned form. Contrast Q8_0, whose b16 exists only as the `_rp` twin (hence q8rp).
14505            (QT_F8_E4M3, 16) => "qmatvec_e4m3_mmvq_b16",
14506            (QT_Q4_0, 2) => "qmatvec_q4_0_mmvq_b2",
14507            (QT_Q4_0, 4) => "qmatvec_q4_0_mmvq_b4",
14508            (QT_Q4_0, 8) => "qmatvec_q4_0_mmvq_b8",
14509            (QT_Q4_0, 16) => "qmatvec_q4_0_mmvq_b16",
14510            _ => return None,
14511        })
14512    }
14513
14514    /// BATCHED weight-tile-resident matvec from a PRE-QUANTIZED q8_1 activation (the m=2-8 verify/MTP
14515    /// win). One warp walks the weight row ONCE, dp4a vs all m activation columns -> weight HBM/L2
14516    /// traffic 1x for m tokens (vs grid.y=m re-reading it m times). `mcols` ∈ {2,4,8} is the
14517    /// compile-time batch; m must be <= mcols (the c >= m columns are masked in-kernel). y is
14518    /// [m, out_f] token-major. NVFP4 per-tensor macro-scale applied post
14519    /// (scale==1.0 for other dtypes -> no-op). BIT-IDENTICAL per (token,row) to qmatvec_*_mmvq.
14520    ///
14521    /// NVFP4 VARIANT DISPATCH: the batched NVFP4 kernel measured memory-LATENCY bound on the real
14522    /// 27B verify (ncu --set full, 12 steady launches: long_scoreboard 18-30 stalls/issue vs <=1.7
14523    /// for every other reason, DRAM only 41-51% active, lg_throttle 0.7, L1 hit 94% — ONE 6-LDG
14524    /// weight wavefront in flight per warp is the binding constraint, NOT bandwidth and NOT the
14525    /// column-unroll break). Two exactness-free fixes, chosen PER SHAPE from the DRAM-cold 8-copy
14526    /// msweep on all six 27B shapes (2026-07-03):
14527    ///   `pf` = next-g weight-prefetch double-buffer (48 regs, occupancy intact) — wins everywhere
14528    ///          it applies for b4 (-3..-14%), never loses;
14529    ///   `r2` = two rows/warp (67 regs -> 7 resident blocks/SM) — the bigger win (-8.5..-30%) but
14530    ///          wave-quantization-sensitive: with the grid halved to ceil(out_f/8) blocks, a
14531    ///          fractional straggler wave (waves in ~1.05-1.5) costs a full extra latency round on
14532    ///          a latency-bound kernel (27B ffn_down 640 blocks / 574 resident = 1.11 waves: +17%),
14533    ///          while <=1 wave (9B ffn_down 0.89: -30%) or >=2 waves (tail amortized; qkv 2.2:
14534    ///          -8.5%, ffn_gate 3.8: -12.5%) win. For b2, r2 wins on DEEP k-loops (in_f>=6144:
14535    ///          -8..-19%) where the 2-col body starves weight MLP hardest; pf measured negative.
14536    /// b4: r2 when waves(out_f) <= 1 (and grid fills >=half the SMs) or >= 2, else pf.
14537    /// b2: in_f>=6144 -> r2, else base.
14538    /// MEMRA_MMVQ_BV=base|pf|r2|pfr2 forces one variant everywhere (A/B + rollback seam).
14539    /// All variants BIT-IDENTICAL per (token,row): same dp4a order, scales, adg factor, reduce —
14540    /// only load issue time and the row->warp mapping change (kernel-check gates all of them).
14541    /// `rp` = the weight buffer is the A6 SPLIT-PLANE repacked layout (NVFP4 only): the same
14542    /// wave-aware auto rule applies, mapped onto the `_rp` twins (rp/rpr2/rpr2w8 mirror
14543    /// pf/r2/r2w8 — regs 44/67/64 land in the same residency classes).
14544    /// The variant the batched dispatch will pick for this (shape, m, mcols, layout) — exposed so
14545    /// gates can distinguish bit-identical variants (bit-bad==0 required) from the k-split family
14546    /// (deterministic but k-reduce-order-shifted: rel<1e-3 + run-to-run bit-identity required).
14547    /// Device SM count (cached) — grid-fill policy input.
14548    pub fn sm_count(&self) -> i32 {
14549        static SMS: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
14550        *SMS.get_or_init(|| {
14551            use cudarc::driver::sys::CUdevice_attribute_enum as A;
14552            self.gpu
14553                .ctx
14554                .attribute(A::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)
14555                .unwrap_or(82)
14556        })
14557    }
14558
14559    pub fn batched_variant(
14560        &self,
14561        _m: usize,
14562        in_f: usize,
14563        out_f: usize,
14564        qtype: i32,
14565        row_bytes: usize,
14566        mcols: usize,
14567        rp: bool,
14568    ) -> &'static str {
14569        // Q8_0 never joined the auto variant machinery (on sm_120 its only batched shapes
14570        // were tiny aux tensors). On Q8_0-trunk models the layout is the whole game: the
14571        // split-plane mirror (rp) routes to the _rp twins (H100 coalescing fix, 2026-07-26);
14572        // GGUF layout stays "base". rp bytes MUST never reach the base kernel or vice versa.
14573        if qtype == QT_Q8_0 {
14574            return if rp { "rp" } else { "base" };
14575        }
14576        static BV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
14577        let bv = *BV.get_or_init(|| match std::env::var("MEMRA_MMVQ_BV").as_deref() {
14578            Ok("base") => "base",
14579            Ok("pf") => "pf",
14580            Ok("r2") => "r2",
14581            Ok("r2w8") => "r2w8",
14582            Ok("pfr2") => "pfr2",
14583            Ok("ca") => "ca",
14584            Ok("car2") => "car2",
14585            // rp* = SPLIT-PLANE REPACKED layout kernels (A6 prototype): W must already be the
14586            // repacked buffer (msweep MSWEEP_RP harness) — never valid on GGUF-layout weights.
14587            Ok("rp") => "rp",
14588            Ok("rpr2") => "rpr2",
14589            Ok("rpr2w8") => "rpr2w8",
14590            // rpca* = cp.async software-pipelined split-plane (2026-07-05): hides the _rp
14591            // long_scoreboard load stall. rp-layout only; b4/b2 (no b8 twin).
14592            Ok("rpca") => "rpca",
14593            Ok("rpcar2") => "rpcar2",
14594            // 2026-07-06 m-small latency arc: rpsc = rpr2 + per-warp smem scale prestage (kills
14595            // the scale-plane global dependency, zero reg growth); rpms/rpmsc = m-split x2
14596            // across warp pairs (2x blocks of rpr2, column halves per warp, BIT-identical to
14597            // _rp); rpks/rpksc = k-split x2 (fastest microbench cells but k-reduce-order-shifted:
14598            // run-spec self-consistency FAILED on the 27B daily driver — verify logits must be
14599            // bit-identical to the decode path — measurement corpus ONLY, never auto).
14600            Ok("rpsc") => "rpsc",
14601            Ok("rpms") => "rpms",
14602            Ok("rpmsc") => "rpmsc",
14603            Ok("rpks") => "rpks",
14604            Ok("rpksc") => "rpksc",
14605            _ => "auto",
14606        });
14607        // cp.async ring variants need 16B-aligned rows (in_f%256==0 -> (in_f/64)*36 % 16 == 0)
14608        // and whole 32-group warp iterations (nsb%32==0 <=> in_f%1024==0). All 27B/9B trunk
14609        // shapes qualify; anything else falls back to the register variants.
14610        let ca_ok = qtype == QT_NVFP4 && (row_bytes % 16 == 0) && (in_f % 1024 == 0);
14611        // rpsc: smem scale plane fits (nsb64 <= 272) + int4-aligned staging (nsb64 % 4 == 0).
14612        // rpks/rpksc: half-plane staging alignment needs nsb64 % 8 == 0 (in_f % 512 == 0).
14613        // MEMRA_KS=0 removes the 2026-07-06 rpsc/rpks/rpksc entries from AUTO (rollback seam;
14614        // forced MEMRA_MMVQ_BV values still work).
14615        static KS_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
14616        let ks_on = *KS_ON.get_or_init(|| std::env::var("MEMRA_KS").as_deref() != Ok("0"));
14617        let sc_ok = ks_on && qtype == QT_NVFP4 && (in_f % 256 == 0) && (in_f / 64 <= 272);
14618        let ks_ok = ks_on && qtype == QT_NVFP4 && (in_f % 512 == 0) && (in_f / 64 <= 272);
14619        static SMS: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
14620        let sms = *SMS.get_or_init(|| {
14621            use cudarc::driver::sys::CUdevice_attribute_enum as A;
14622            self.gpu
14623                .ctx
14624                .attribute(A::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)
14625                .unwrap_or(82)
14626        });
14627        // k-quant r2 port (2026-07-04): q4_K/q5_K/q6_K have _r2/_r2w8 twins. ncu on the DRAM-cold
14628        // 9B msweep showed q4_K/q5_K b4 memory-latency bound like NVFP4 pre-fix (long_scoreboard
14629        // 19.6/16.4 per issue, DRAM 47.7/38.2%, L2 weight hit ~13%); q6_K lm_head is the exception
14630        // at DRAM 90-91% = wall-bound (yet r2 still wins -8%: deeper MLP raises achieved DRAM).
14631        // No _pf port (a k-quant group stages 10+ words vs NVFP4's 5 — register cost outweighs;
14632        // r2 covers the same MLP) and no rp (GGUF layout only). Q8_0 stays base: its only real
14633        // batched shapes are the tiny out_f=32 ssm_alpha/beta (8-block grids never fill one SM).
14634        // AUTO RULE = the measured winners table (differs from NVFP4's!):
14635        //   r2w8 NEVER in auto — the reg squeeze (72 -> 64 regs = stack spill) loses to unbounded
14636        //     r2 on every measured k-quant cell, incl. the wave-crossing lm_heads (q6_K 1316 vs
14637        //     r2 1258us) — kernels kept behind the force seam for the corpus;
14638        //   q4_K: r2 whenever the halved grid fills the SMs (blocks >= 4*SMs), INCLUDING the
14639        //     1.05-2.0 straggler window where NVFP4's r2 lost (qkv 1.78 waves: r2 -15% here; the
14640        //     k-quant base kernel leaves more latency on the table than a straggler wave costs);
14641        //   q5_K/q6_K: r2 only at waves >= 2 (the 248320-row lm_heads, 48+ waves: q6_K -8%, q5_K
14642        //     -2%); mid shapes measured base-or-flat (q5_K qkv 49.1 base vs 49.7 r2, attn_gate
14643        //     flat, attn_k base) — the 5/6-bit two-stream unpack makes r2's staging pricier.
14644        //   b2 same table with 8-row blocks: q4_K r2 when filled (-3..-22% all measured shapes),
14645        //     q5_K/q6_K r2 at waves >= 2 (27B lm_head -2.9%; 9B q6_K flat, harmless).
14646        let kq_r2 = matches!(qtype, QT_Q4_K | QT_Q5_K | QT_Q6_K);
14647        // MEMRA_KQ_BV=base|r2|r2w8 forces the k-quant variant WITHOUT touching the NVFP4 dispatch
14648        // (MEMRA_MMVQ_BV is global — an interleaved k-quant-only e2e A/B needs this narrower seam).
14649        static KQBV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
14650        let kq_bv = *KQBV.get_or_init(|| match std::env::var("MEMRA_KQ_BV").as_deref() {
14651            Ok("base") => "base",
14652            Ok("r2") => "r2",
14653            Ok("r2w8") => "r2w8",
14654            _ => "auto",
14655        });
14656        let variant: &'static str = if qtype == QT_Q4_0 {
14657            // Q4_0 r2 (gemma verify trunk, 2026-07-10): shared activation loads + the
14658            // row-independent ones-sum computed once per (col,group) for 2 rows. Same
14659            // fill rule as q4_K: r2 when the halved grid still fills the SMs.
14660            static Q40BV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
14661            let q40 = *Q40BV.get_or_init(|| match std::env::var("MEMRA_Q40_BV").as_deref() {
14662                // ms/sm/la = force-only measurement seams (ALL FLAT/NEGATIVE 2026-07-13,
14663                // never auto): m-split flat (nvcc keeps 72 regs); smem-slab −11% (staging
14664                // + syncs cost more than the stalls, bank-pad made no difference);
14665                // register load-ahead flat (nvcc already reorders). The b-tier limiter
14666                // is still unidentified — see the jsonl row.
14667                Ok("base") => "base",
14668                Ok("r2") => "r2",
14669                Ok("ms") => "ms",
14670                Ok("sm") => "sm",
14671                Ok("la") => "la",
14672                _ => "auto",
14673            });
14674            let v = if q40 != "auto" {
14675                q40
14676            } else if (out_f as u32).div_ceil(8) >= 4 * sms as u32 {
14677                "r2"
14678            } else {
14679                "base"
14680            };
14681            // split-plane mirror twins (2026-07-10): same fill rule, _rp names.
14682            // (m-split r2 pair twin PROBED FLAT 2026-07-13 — nvcc kept 72 regs either way
14683            // and the limiter is the per-column activation load chain (long_scoreboard
14684            // 42.5%), not occupancy; arm killed per doctrine, jsonl row is the record.)
14685            if rp {
14686                match v {
14687                    "ms" => "r2ms_rp",
14688                    "sm" => "r2sm_rp",
14689                    "la" => "r2la_rp",
14690                    "r2" => "r2_rp",
14691                    _ => "rp",
14692                }
14693            } else if matches!(v, "ms" | "sm" | "la") {
14694                "r2"
14695            } else {
14696                v
14697            }
14698        } else if qtype != QT_NVFP4 && !kq_r2 {
14699            "base"
14700        } else if kq_r2 && rp {
14701            // K-quant split-plane mirror (2026-08-01): only the plain _rp batched twins are
14702            // compiled for q4_K/q6_K — rp is a LAYOUT, it must survive every heuristic
14703            // (split-plane bytes through a GGUF-layout kernel = NaN). q5_K never mirrors.
14704            "rp"
14705        } else if kq_r2 {
14706            // k-quant r2w8 only exists at b4 (b2_r2 already 8-resident; b8 has no w8 twin) ->
14707            // mcols != 4 forced r2w8 falls to unbounded r2.
14708            if kq_bv != "auto" {
14709                if kq_bv == "r2w8" && mcols != 4 {
14710                    "r2"
14711                } else {
14712                    kq_bv
14713                }
14714            } else if bv != "auto" {
14715                match bv {
14716                    "r2" | "pfr2" | "rpr2" | "car2" => "r2",
14717                    "r2w8" | "rpr2w8" => {
14718                        if mcols != 4 {
14719                            "r2"
14720                        } else {
14721                            "r2w8"
14722                        }
14723                    }
14724                    _ => "base", // base/pf/ca/rp forced -> base (no such k-quant kernels)
14725                }
14726            } else {
14727                let blocks = (out_f + 7) / 8;
14728                let waves = blocks as f64 / (7 * sms as usize) as f64;
14729                let filled = blocks >= 4 * sms as usize;
14730                let use_r2 = if qtype == QT_Q4_K {
14731                    filled
14732                } else {
14733                    waves >= 2.0
14734                };
14735                if use_r2 { "r2" } else { "base" }
14736            }
14737        } else if bv != "auto" {
14738            // r2w8 only exists for b4/b8 (the b2_r2 kernel is already 8-blocks-resident at 60 regs).
14739            // ca/car2 need the alignment gate AND have no b8 twins; pfr2 has no b8 twin either —
14740            // unsupported (shape, mcols) combos fall back to pf/r2.
14741            // On rp buffers, forced legacy names map to their rp twins (layout law).
14742            let v = if bv == "r2w8" && mcols == 2 {
14743                "r2"
14744            } else if bv == "ca" && (!ca_ok || mcols == 8) {
14745                "pf"
14746            } else if bv == "car2" && (!ca_ok || mcols == 8) {
14747                "r2"
14748            } else if bv == "pfr2" && mcols == 8 {
14749                "r2"
14750            } else if (bv == "rpr2w8" || bv == "rpr2") && mcols == 2 {
14751                "rpr2"
14752            }
14753            // rpca* has no b8 twin (falls to rpr2w8/rpr2); needs the ca alignment gate.
14754            else if (bv == "rpca" || bv == "rpcar2") && (!ca_ok || mcols == 8) {
14755                if mcols == 8 { "rpr2w8" } else { "rpr2" }
14756            } else if bv == "rpcar2" && mcols == 2 {
14757                "rpca"
14758            }
14759            // rpsc/rpmsc/rpks* gate on smem-fit + alignment; fall to rpr2 outside it
14760            // (rpms has no smem and no alignment need — always valid on rp buffers).
14761            else if (bv == "rpsc" || bv == "rpmsc") && !sc_ok {
14762                "rpr2"
14763            } else if (bv == "rpks" || bv == "rpksc") && !ks_ok {
14764                "rpr2"
14765            } else {
14766                bv
14767            };
14768            if rp {
14769                match v {
14770                    "base" | "pf" | "ca" | "rp" => "rp",
14771                    "r2" | "pfr2" | "car2" | "rpr2" => "rpr2",
14772                    "r2w8" | "rpr2w8" => {
14773                        if mcols == 2 {
14774                            "rpr2"
14775                        } else {
14776                            "rpr2w8"
14777                        }
14778                    }
14779                    other => other, // rpca/rpcar2/rpsc/rpks/rpksc pass through (already rp-layout)
14780                }
14781            } else {
14782                v
14783            }
14784        } else if mcols == 8 {
14785            // b8 AUTO (2026-07-06 m-small latency arc, g7e DRAM-cold rp msweep m=5/6/8 all five
14786            // 27B shapes): rpsc — the rpr2w8 schedule with the warp's scale rows prestaged to
14787            // smem, leaving ONE global dependency (the quant stream) in the k-loop at zero reg
14788            // growth. BIT-identical to rpr2w8 and wins or ties EVERY b8 cell: ffn_gate m5
14789            // 50.7->46.9 m8 64.1->57.1 (-11%), qkv m8 34.6->33.0, ssm_out m8 29.7->28.8,
14790            // attn_gate m8 26.9->26.1, ffn_down m5 58.2->56.9. The faster split-grid twins are
14791            // OUT: rpksc (k-split, ffn_down m5 -21%) broke run-spec self-consistency (k-reduce
14792            // order shifts verify argmax at tie margins — verify must stay bit-identical to the
14793            // m=1 decode chain); rpmsc (m-split, bit-identical) measured NEGATIVE everywhere
14794            // (twin warp's duplicated weight stream: ffn_down m5 85.7 vs 56.9).
14795            if rp {
14796                if sc_ok { "rpsc" } else { "rpr2w8" }
14797            } else {
14798                "r2w8"
14799            }
14800        } else if mcols >= 4 {
14801            // r2 runs 7 resident blocks/SM (67 regs); its __launch_bounds__(128,8) twin `r2w8`
14802            // (64 regs) runs 8. grid = ceil(out_f/8) for both. rp twins land in the same
14803            // residency classes (rp 44 regs ~ pf-class occupancy, rpr2 67, rpr2w8 64).
14804            let blocks = (out_f + 7) / 8;
14805            let r7 = 7 * sms as usize;
14806            let r8 = 8 * sms as usize;
14807            let waves = blocks as f64 / r7 as f64;
14808            let filled = blocks >= 4 * sms as usize;
14809            // 2026-07-06 m-small latency arc: b4 keeps the wave rule (rpms/rpmsc measured
14810            // flat-to-negative at m=3/4 on every shape — the m-split twin duplicates the weight
14811            // stream; rpsc b4 also negative on r2-class picks, ffn_down m4 51.1 vs 46.5).
14812            if filled && blocks.div_ceil(r8) < blocks.div_ceil(r7) {
14813                // the extra residency drops the INTEGER wave count -> the straggler wave a
14814                // latency-bound kernel pays in full disappears (ffn_down 1.11 -> 0.98 waves:
14815                // 112.5 -> 81.6us, beats pf 90.1; qkv 2.23 -> 1.95: 58.1 -> 51.1).
14816                if rp { "rpr2w8" } else { "r2w8" }
14817            } else if waves >= 2.0 || (waves <= 1.0 && filled) {
14818                // tail amortized (>=2 waves) or single wave: unbounded r2 (no reg-squeeze tax —
14819                // gate/up 81.1 vs 83.9 bounded, attn_q 61.0 vs 63.4).
14820                if rp { "rpr2" } else { "r2" }
14821            } else {
14822                // fractional straggler-wave window with no crossing, or grid too small to fill
14823                // the SMs (tiny out_f<=1024 shapes want max row-parallelism): prefetch variant
14824                // (rp = the r1 split-plane twin — measured the attn_gate winner, 35.4 vs pf 36.4).
14825                if rp { "rp" } else { "pf" }
14826            }
14827        } else if in_f >= 6144 {
14828            // b2 deep-k (2026-07-06): every new twin measured flat-to-negative here (rpms 44.1
14829            // vs rpr2 40.8 ffn_down; rpsc 43.6; the winning rpks is banned on k-order) — rpr2
14830            // stays.
14831            if rp { "rpr2" } else { "r2" }
14832        } else if rp {
14833            // b2 shallow-k: qkv (out_f=10240, 0.97 waves at 7-resident) is the one measured cell
14834            // where the r2-schedule scale-prestage twin beats the r1 rp pick (24.7 vs 28.9us
14835            // -15%); the wider (ffn_gate 1.65 waves) and smaller (attn_gate 0.58) shapes LOSE
14836            // (41.8 vs 38.2 / 16.6 vs 14.6) — gate on the single-wave window.
14837            let waves = ((out_f + 7) / 8) as f64 / (7 * sms as usize) as f64;
14838            if sc_ok && waves >= 0.9 && waves <= 1.1 {
14839                "rpsc"
14840            } else {
14841                "rp"
14842            }
14843        } else {
14844            "base"
14845        };
14846        variant
14847    }
14848
14849    pub fn qmatvec_mmvq_batched(
14850        &self,
14851        bytes: &CudaSlice<u8>,
14852        aq: &CudaSlice<i8>,
14853        ad: &CudaSlice<f32>,
14854        m: usize,
14855        in_f: usize,
14856        out_f: usize,
14857        qtype: i32,
14858        row_bytes: usize,
14859        mcols: usize,
14860        scale: f32,
14861        rp: bool,
14862    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14863        const ROWS_PER_BLOCK: u32 = 4;
14864        // TUNE SEAM (H100 lane): MEMRA_BVAR forces the batched-variant pick for the whole
14865        // process — the auto heuristics were tuned on sm_120 (82 SMs / 858 GB/s) and the
14866        // sm_90a re-tune sweeps this seam empirically. Layout variants stay safe: an rp
14867        // weight keeps its rp-layout kernel family regardless of the override.
14868        let forced: Option<&'static str> = {
14869            static V: std::sync::OnceLock<Option<String>> = std::sync::OnceLock::new();
14870            V.get_or_init(|| std::env::var("MEMRA_BVAR").ok())
14871                .as_deref()
14872                .map(|s| Box::leak(s.to_string().into_boxed_str()) as &'static str)
14873        };
14874        let variant = match forced {
14875            Some(v) if !rp || v.contains("rp") => v,
14876            _ => self.batched_variant(m, in_f, out_f, qtype, row_bytes, mcols, rp),
14877        };
14878        let base_name = Self::batched_kernel_name(qtype, mcols).ok_or_else(|| {
14879            format!("qmatvec_mmvq_batched: no kernel for qtype {qtype} mcols {mcols}")
14880        })?;
14881        // b16 tier (t=9..16 verify): only base/_rp b16 kernels are compiled — the b2..b8
14882        // per-shape perf variants (r2/pf/...) do not apply at this width. rp is a LAYOUT,
14883        // not a perf variant: it must survive (base kernel on split-plane bytes = NaN).
14884        let variant = if mcols == 16 {
14885            if rp { "rp" } else { "base" }
14886        } else {
14887            variant
14888        };
14889        // EXACT-WIDTH b5/b6/b7 twins (lane/vt-fixes fix 1, 2026-08-03): the b8 kernels
14890        // allocate acc[WROWS][8] at ANY m, so T=5..7 verify paid the full 8-wide register
14891        // tax — the measured T=4->5 cliff. The same template at MCOLS=m runs the identical
14892        // per-(token,row) chain (columns c >= m never execute in either form) ->
14893        // BIT-IDENTICAL to the b8 launch. NVFP4 split-plane only (the sm_120 default trunk);
14894        // covers both b8 auto schedules (rpsc, rpr2w8). MEMRA_B567=0 rollback.
14895        static B567: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
14896        let b567 = *B567.get_or_init(|| std::env::var("MEMRA_B567").as_deref() != Ok("0"));
14897        if b567
14898            && qtype == QT_NVFP4
14899            && rp
14900            && mcols == 8
14901            && (5..=7).contains(&m)
14902            && matches!(variant, "rpsc" | "rpr2w8")
14903        {
14904            let f = self.func(&format!("qmatvec_nvfp4_mmvq_b{m}_{variant}"));
14905            let rows_per_block = ROWS_PER_BLOCK * 2; // r2-class schedules: 2 rows/warp
14906            let mut y = self.alloc_uninit::<f32>(m * out_f)?;
14907            let cfg = LaunchConfig {
14908                grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
14909                block_dim: (32, ROWS_PER_BLOCK, 1),
14910                shared_mem_bytes: 0,
14911            };
14912            let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
14913            let __s_b = self.gpu.stream();
14914            let mut b = __s_b.launch_builder(&f);
14915            b.arg(bytes)
14916                .arg(aq)
14917                .arg(ad)
14918                .arg(&mut y)
14919                .arg(&inf)
14920                .arg(&outf)
14921                .arg(&mi)
14922                .arg(&rb);
14923            unsafe {
14924                b.launch(cfg)?;
14925            }
14926            if scale != 1.0 {
14927                self.scale_inplace(&mut y, scale, m * out_f)?;
14928            }
14929            return Ok(y);
14930        }
14931        let (name, rows_per_block): (std::borrow::Cow<'static, str>, u32) = match variant {
14932            "base" => (base_name.into(), ROWS_PER_BLOCK),
14933            "pf" => (format!("{base_name}_pf").into(), ROWS_PER_BLOCK),
14934            "ca" => (format!("{base_name}_ca").into(), ROWS_PER_BLOCK),
14935            "rp" => (format!("{base_name}_rp").into(), ROWS_PER_BLOCK),
14936            "rpca" => (format!("{base_name}_rpca").into(), ROWS_PER_BLOCK), // 1 row/warp cp.async
14937            // split families: 2 warp-pairs x 2 rows = 4 rows/block (the k-range or column set
14938            // splits across the pair's two warps; grid.x doubles vs rpr2 at the same regs).
14939            "rpks" => (format!("{base_name}_rpks").into(), ROWS_PER_BLOCK),
14940            "rpksc" => (format!("{base_name}_rpksc").into(), ROWS_PER_BLOCK),
14941            "rpms" => (format!("{base_name}_rpms").into(), ROWS_PER_BLOCK),
14942            "rpmsc" => (format!("{base_name}_rpmsc").into(), ROWS_PER_BLOCK),
14943            "r2ms_rp" => (format!("{base_name}_r2ms_rp").into(), ROWS_PER_BLOCK),
14944            "r2sm_rp" => (format!("{base_name}_r2sm_rp").into(), ROWS_PER_BLOCK * 2),
14945            "r2la_rp" => (format!("{base_name}_r2la_rp").into(), ROWS_PER_BLOCK * 2),
14946            v => (format!("{base_name}_{v}").into(), ROWS_PER_BLOCK * 2), // r2-class: 2 rows/warp
14947        };
14948        debug_assert!(
14949            !rp || name.contains("_rp"),
14950            "rp weight dispatched to a GGUF-layout kernel"
14951        );
14952        let f = self.func(&name);
14953        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
14954        // r2sm_rp: [MCOLS][32 blk][8 int] activation slab + [MCOLS][32] f32 scales.
14955        let smem = if name.contains("_r2sm_rp") {
14956            (mcols * 32 * 9 * 4 + mcols * 32 * 4) as u32
14957        } else {
14958            0
14959        };
14960        let cfg = LaunchConfig {
14961            grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
14962            block_dim: (32, ROWS_PER_BLOCK, 1),
14963            shared_mem_bytes: smem,
14964        };
14965        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
14966        let __s_b = self.gpu.stream();
14967        let mut b = __s_b.launch_builder(&f);
14968        b.arg(bytes)
14969            .arg(aq)
14970            .arg(ad)
14971            .arg(&mut y)
14972            .arg(&inf)
14973            .arg(&outf)
14974            .arg(&mi)
14975            .arg(&rb);
14976        unsafe {
14977            b.launch(cfg)?;
14978        }
14979        if scale != 1.0 {
14980            self.scale_inplace(&mut y, scale, m * out_f)?;
14981        }
14982        Ok(y)
14983    }
14984
14985    /// BATCHED weight-tile-resident matvec from raw weight bytes (quantizes the f32 activation `x` to
14986    /// q8_1 internally; macro-scale NOT applied — caller compares bare, like qmatvec_*_fast). For the
14987    /// kernel_check bit-equivalence gate. `mcols` ∈ {2,4,8}. Works for Q8_0/Q4_K/Q5_K/Q6_K/NVFP4.
14988    pub fn qmatvec_batched_raw(
14989        &self,
14990        bytes: &CudaSlice<u8>,
14991        x: &CudaSlice<f32>,
14992        m: usize,
14993        in_f: usize,
14994        out_f: usize,
14995        qtype: i32,
14996        row_bytes: usize,
14997        mcols: usize,
14998        rp: bool,
14999    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15000        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15001        self.qmatvec_mmvq_batched(
15002            bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, mcols, 1.0, rp,
15003        )
15004    }
15005
15006    /// Back-compat NVFP4-only batched raw launcher (used by older gates). Delegates to the generic one.
15007    pub fn qmatvec_nvfp4_batched_raw(
15008        &self,
15009        bytes: &CudaSlice<u8>,
15010        x: &CudaSlice<f32>,
15011        m: usize,
15012        in_f: usize,
15013        out_f: usize,
15014        row_bytes: usize,
15015        mcols: usize,
15016        rp: bool,
15017    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15018        self.qmatvec_batched_raw(bytes, x, m, in_f, out_f, QT_NVFP4, row_bytes, mcols, rp)
15019    }
15020
15021    /// Stage-C FP4 gate (MEMRA_FP4): if `w` is an NVFP4 weight with in_f%64==0, run the native mxf4
15022    /// block-scale GEMM and apply the per-tensor macro-scale, returning Some(y). Else None (caller
15023    /// falls through to the int8 GEMM / dp4a). Strict opt-in over the proven int8 path; m>=16 only.
15024    fn try_fp4_gemm(
15025        &self,
15026        w: &crate::model::GpuTensor,
15027        x: &CudaSlice<f32>,
15028        m: usize,
15029        in_f: usize,
15030        out_f: usize,
15031    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
15032        use crate::model::GpuTensor;
15033        if cfg!(memra_portable_cuda) {
15034            return Ok(None);
15035        }
15036        if std::env::var("MEMRA_FP4").is_err() {
15037            return Ok(None);
15038        }
15039        // CUTLASS prefill branch (m>=128 + MEMRA_FP4_CUTLASS + a repacked CutlassWeight present): route
15040        // to the CUTLASS sm120 NVFP4 GEMM, folding the per-tensor macro-scale into the epilogue alpha
15041        // (1/scale) — no post-matmul scale_inplace. Decode (m<128) and the m∈[16,128) middle band keep
15042        // the hand-roll below: CUTLASS's 128-row M-tile wastes work under 128.
15043        // The hand-roll applies the per-tensor macro-scale as a POST-matmul MULTIPLY (scale_inplace(y,
15044        // scale)); CUTLASS's epilogue does D = alpha * (A@B^T), so alpha == scale reproduces it exactly
15045        // (NOT 1/scale — the plan sketch had this inverted; the kernel_check arm gates it). scale==1.0
15046        // for the common no-macro-scale case.
15047        #[cfg(memra_cutlass)]
15048        if m >= 128 && std::env::var("MEMRA_FP4_CUTLASS").is_ok() {
15049            if let GpuTensor::Quant {
15050                bytes,
15051                qtype,
15052                scale,
15053                row_bytes,
15054                cutlass,
15055                ..
15056            } = w
15057            {
15058                if *qtype == QT_NVFP4 && in_f % 64 == 0 {
15059                    if let Some(cw) = cutlass {
15060                        // Resident fast path: load-time-repacked B + swizzled SFB (no per-call repack).
15061                        let y = self.cutlass_fp4_gemm(
15062                            &cw.b_packed,
15063                            &cw.sfb_swizzled,
15064                            x,
15065                            *scale,
15066                            m,
15067                            out_f,
15068                            in_f,
15069                        )?;
15070                        return Ok(Some(y));
15071                    } else if std::env::var("MEMRA_FP4_CUTLASS_OTF").is_ok() {
15072                        // On-the-fly repack (MEMRA_FP4_CUTLASS_OTF): de-interleave + swizzle the B operand
15073                        // from raw bytes per prefill call. No resident doubling of the NVFP4 weight VRAM
15074                        // (the load-time repack ~doubles it) — needed for models that don't fit the
15075                        // resident path (e.g. the 27B on 24GB). Slower (per-call repack) but argmax-exact.
15076                        let (b_packed, sfb_sw) =
15077                            self.build_cutlass_weight(bytes, out_f, in_f, *row_bytes)?;
15078                        let y =
15079                            self.cutlass_fp4_gemm(&b_packed, &sfb_sw, x, *scale, m, out_f, in_f)?;
15080                        return Ok(Some(y));
15081                    }
15082                }
15083            }
15084        }
15085        if let GpuTensor::Quant {
15086            bytes,
15087            qtype,
15088            row_bytes,
15089            scale,
15090            rp,
15091            ..
15092        } = w
15093        {
15094            // A6: the hand-rolled W4A4 mxf4 GEMM reads 36B GGUF blocks — no rp port (MEMRA_FP4 is
15095            // an opt-in accuracy tradeoff); repacked tensors fall through to the int8 GEMM.
15096            if *qtype == QT_NVFP4 && in_f % 64 == 0 && !*rp {
15097                let y =
15098                    self.qmatvec_gemm_nvfp4_fp4(bytes, x, m, in_f, out_f, *row_bytes, *scale)?;
15099                return Ok(Some(y));
15100            }
15101        }
15102        Ok(None)
15103    }
15104
15105    /// rms_norm + fused fp16 twin (task #14): f32 output verbatim `rms_norm` + the fp16
15106    /// copy the f16-mirror GEMM group would otherwise produce with a standalone convert
15107    /// launch. BIT-IDENTICAL end-to-end (same reduction, same __float2half values).
15108    pub fn rms_norm_f16out(
15109        &self,
15110        x: &CudaSlice<f32>,
15111        w: &CudaSlice<f32>,
15112        dst: &mut CudaSlice<f32>,
15113        dst16: &mut CudaSlice<u8>,
15114        ncols: usize,
15115        nrows: usize,
15116        eps: f32,
15117    ) -> Result<(), Box<dyn std::error::Error>> {
15118        let f = self.func("rms_norm_f16out_f32");
15119        let cfg = LaunchConfig {
15120            grid_dim: (nrows as u32, 1, 1),
15121            block_dim: (rms_block(), 1, 1),
15122            shared_mem_bytes: 0,
15123        };
15124        let (nc, e) = (ncols as i32, eps);
15125        let __s_b = self.gpu.stream();
15126        let mut b = __s_b.launch_builder(&f);
15127        b.arg(x).arg(w).arg(dst).arg(dst16).arg(&nc).arg(&e);
15128        unsafe {
15129            b.launch(cfg)?;
15130        }
15131        Ok(())
15132    }
15133
15134    /// add+norm(+f16out) fusion for the prefill trunk (round 28; add_rms_norm precedent —
15135    /// bit-identical to add_f32 -> rms_norm_f16out). block_dim matches rms_norm_f16out's.
15136    #[allow(clippy::too_many_arguments)]
15137    pub fn add_rms_norm_f16out(
15138        &self,
15139        a: &CudaSlice<f32>,
15140        b: &CudaSlice<f32>,
15141        w: &CudaSlice<f32>,
15142        res: &mut CudaSlice<f32>,
15143        dst: &mut CudaSlice<f32>,
15144        dst16: &mut CudaSlice<u8>,
15145        ncols: usize,
15146        nrows: usize,
15147        eps: f32,
15148    ) -> Result<(), Box<dyn std::error::Error>> {
15149        let f = self.func("add_rms_norm_f16out_f32");
15150        let cfg = LaunchConfig {
15151            grid_dim: (nrows as u32, 1, 1),
15152            block_dim: (rms_block(), 1, 1),
15153            shared_mem_bytes: 0,
15154        };
15155        let (nc, e) = (ncols as i32, eps);
15156        let __s_lb = self.gpu.stream();
15157        let mut lb = __s_lb.launch_builder(&f);
15158        lb.arg(a)
15159            .arg(b)
15160            .arg(w)
15161            .arg(res)
15162            .arg(dst)
15163            .arg(dst16)
15164            .arg(&nc)
15165            .arg(&e);
15166        unsafe {
15167            lb.launch(cfg)?;
15168        }
15169        Ok(())
15170    }
15171
15172    /// matmul_group with a PRE-EMITTED fp16 activation (task #14: the producer norm fused
15173    /// the convert). Mirror-less members fall back to `matmul` on the f32 activation.
15174    pub fn matmul_group_xh(
15175        &self,
15176        ws: &[&crate::model::GpuTensor],
15177        x: &CudaSlice<f32>,
15178        xh: &CudaSlice<u8>,
15179        m: usize,
15180    ) -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
15181        let mut out = Vec::with_capacity(ws.len());
15182        let in_f = ws[0].in_features();
15183        for w in ws {
15184            if w.in_features() == in_f && m >= 16 && !self.verify_exact_on() {
15185                if let Some(y) = self.try_f16_gemm_pre(w, xh, m)? {
15186                    out.push(y);
15187                    continue;
15188                }
15189            }
15190            out.push(self.matmul(w, x, m)?);
15191        }
15192        Ok(out)
15193    }
15194
15195    /// task #14 pad-proofing: zero beta/g_log at rows >= len_d[0] (pads become identity
15196    /// GDN steps). Layouts [T, H].
15197    pub fn gdn_pad_mask(
15198        &self,
15199        beta: &mut CudaSlice<f32>,
15200        g_log: &mut CudaSlice<f32>,
15201        len_d: &CudaSlice<i32>,
15202        h: usize,
15203        t: usize,
15204    ) -> Result<(), Box<dyn std::error::Error>> {
15205        let f = self.func("gdn_pad_mask_f32");
15206        let cfg = LaunchConfig::for_num_elems((t * h) as u32);
15207        let (hi, ti) = (h as i32, t as i32);
15208        let __s_b = self.gpu.stream();
15209        let mut b = __s_b.launch_builder(&f);
15210        b.arg(beta).arg(g_log).arg(len_d).arg(&hi).arg(&ti);
15211        unsafe {
15212            b.launch(cfg)?;
15213        }
15214        Ok(())
15215    }
15216
15217    /// task #14 pad-proofing: dst[ncols] = src row (len_d[0]-1) — device-indexed last-row
15218    /// gather for the padded prime graph's h_seed/hlast.
15219    pub fn row_gather_dev(
15220        &self,
15221        src: &CudaSlice<f32>,
15222        dst: &mut CudaSlice<f32>,
15223        len_d: &CudaSlice<i32>,
15224        ncols: usize,
15225    ) -> Result<(), Box<dyn std::error::Error>> {
15226        let f = self.func("row_gather_dev_f32");
15227        let cfg = LaunchConfig::for_num_elems(ncols as u32);
15228        let nc = ncols as i32;
15229        let __s_b = self.gpu.stream();
15230        let mut b = __s_b.launch_builder(&f);
15231        b.arg(src).arg(dst).arg(len_d).arg(&nc);
15232        unsafe {
15233            b.launch(cfg)?;
15234        }
15235        Ok(())
15236    }
15237
15238    /// Grouped matmul: several weights consuming ONE activation (hybrid layers: the GDN
15239    /// 4-tuple wqkv/gate/beta/alpha, attention q/k/v, ffn gate/up). Semantics identical to
15240    /// calling `matmul` per weight; the f16-mirror arm converts the activation ONCE for the
15241    /// whole group instead of once per GEMM (the standalone converts were ~250 launches/prime
15242    /// of small-kernel gap fuel — nsys 2026-07-26). Any member without a mirror (or with a
15243    /// different in_f) falls back to its own `matmul` — behavior unchanged.
15244    pub fn matmul_group(
15245        &self,
15246        ws: &[&crate::model::GpuTensor],
15247        x: &CudaSlice<f32>,
15248        m: usize,
15249    ) -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
15250        use crate::model::GpuTensor;
15251        let mut out = Vec::with_capacity(ws.len());
15252        let any_mirror = ws
15253            .iter()
15254            .any(|w| matches!(w, GpuTensor::Quant { f16: Some(_), .. }));
15255        if m >= 16 && any_mirror && !self.verify_exact_on() {
15256            let in_f = ws[0].in_features();
15257            let xh = self.f16_act(x, m * in_f, in_f)?;
15258            for w in ws {
15259                if w.in_features() == in_f {
15260                    if let Some(y) = self.try_f16_gemm_pre(w, &xh, m)? {
15261                        out.push(y);
15262                        continue;
15263                    }
15264                }
15265                out.push(self.matmul(w, x, m)?);
15266            }
15267            return Ok(out);
15268        }
15269        for w in ws {
15270            out.push(self.matmul(w, x, m)?);
15271        }
15272        Ok(out)
15273    }
15274
15275    /// Cross-request grouped matmul (task #13): run ONE projection group over the
15276    /// CONCATENATION of several sequences' activations (m = sum of per-seq rows — the
15277    /// GEMM-batch win vLLM gets from continuous batching), then split each output back
15278    /// into per-seq buffers. Zero view plumbing: gather/scatter are stream-ordered D2D
15279    /// copies (~us at prime sizes). NUMERIC CONFIG NOTE: a GEMM at m=sum tiles K
15280    /// differently than per-seq GEMMs — argmax-gated like every prefill GEMM change.
15281    pub fn matmul_group_multi(
15282        &self,
15283        ws: &[&crate::model::GpuTensor],
15284        xs: &[&CudaSlice<f32>],
15285        ms: &[usize],
15286    ) -> Result<Vec<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
15287        assert_eq!(xs.len(), ms.len());
15288        let in_f = ws[0].in_features();
15289        let total: usize = ms.iter().sum();
15290        let mut xcat = self.uninit(total * in_f)?;
15291        let mut off = 0usize;
15292        for (x, &m) in xs.iter().zip(ms) {
15293            self.copy_into(&mut xcat, off * in_f, x, m * in_f)?;
15294            off += m;
15295        }
15296        let ys = self.matmul_group(ws, &xcat, total)?;
15297        let mut out: Vec<Vec<CudaSlice<f32>>> = (0..xs.len()).map(|_| Vec::new()).collect();
15298        for (w, y) in ws.iter().zip(ys) {
15299            let out_f = w.out_features();
15300            let mut off = 0usize;
15301            for (s, &m) in ms.iter().enumerate() {
15302                let mut ys_s = self.uninit(m * out_f)?;
15303                let src = y.slice(off * out_f..(off + m) * out_f);
15304                self.gpu.stream().memcpy_dtod(&src, &mut ys_s)?;
15305                out[s].push(ys_s);
15306                off += m;
15307            }
15308        }
15309        Ok(out)
15310    }
15311
15312    /// True if `w`'s qtype has a batched tensor-core GEMM kernel (the prefill T>1 root fix).
15313    /// Only the 4 daily-hot dtypes: Q8_0, Q4_K, Q6_K, NVFP4. NVFP4 needs in_f % 64 == 0.
15314    /// DEFAULT-ON (2026-06-28): measured pp512 9B-NVFP4 = 1413 tok/s WITH this GEMM vs 298 with the
15315    /// dp4a fallback (4.7x) AND MORE accurate (prefill logit maxdiff 0.159 vs dp4a 0.55, both argmax
15316    /// MATCH). The int8 tensor-core GEMM is unconditional (its historical MEMRA_GEMM opt-in gate
15317    /// shipped with Phase 0 — mma + smem swizzle + cp.async — and was removed). Prefill-only
15318    /// (m>=GEMM_M_THRESHOLD); m=1 decode keeps dp4a/MMVQ (this returns true but matmul only calls it
15319    /// at m>=threshold). Portable CUDA targets always use the correctness fallback; on sm_120a,
15320    /// MEMRA_NO_GEMM forces that same dp4a fallback (the bit-reference).
15321    pub fn gemm_supports(&self, w: &crate::model::GpuTensor) -> bool {
15322        use crate::model::GpuTensor;
15323        if !legacy_quant_gemm_allowed(
15324            cfg!(memra_portable_cuda),
15325            cfg!(memra_hopper_mma),
15326            std::env::var_os("MEMRA_NO_GEMM").is_some(),
15327        ) {
15328            return false;
15329        }
15330        match w {
15331            GpuTensor::Quant { qtype, .. } => {
15332                matches!(*qtype, QT_Q8_0 | QT_Q4_K | QT_Q6_K | QT_Q5_K | QT_Q4_0)
15333                    || (*qtype == QT_NVFP4 && w.in_features() % 64 == 0)
15334            }
15335            GpuTensor::Float { .. } | GpuTensor::FloatBf16 { .. } => false,
15336        }
15337    }
15338
15339    /// Batched tensor-core int8 GEMM with a PRE-QUANTIZED q8_1 activation (aq,ad). The prefill
15340    /// (T>1) root fix: decode each weight 32-block to int8 in shared memory ONCE per (row-tile,
15341    /// K-step) and reuse it across all BN tokens via mma.sync.m16n8k32.s8 — amortizing the weight
15342    /// read/decode N-fold (vs the dp4a matvec's per-token re-read). s32 accumulate is exact vs
15343    /// dp4a; only the final f32 block-scale rounding differs. Caller MUST have checked
15344    /// `gemm_supports(w)`. y[m,out] token-major. NVFP4 per-tensor macro-scale applied post.
15345    pub fn qmatvec_gemm(
15346        &self,
15347        w: &crate::model::GpuTensor,
15348        aq: &CudaSlice<i8>,
15349        ad: &CudaSlice<f32>,
15350        m: usize,
15351    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15352        use crate::model::GpuTensor;
15353        let in_f = w.in_features();
15354        let out_f = w.out_features();
15355        let (bytes, qtype, row_bytes, scale, rp) = match w {
15356            GpuTensor::Quant {
15357                bytes,
15358                qtype,
15359                row_bytes,
15360                scale,
15361                rp,
15362                ..
15363            } => (bytes, *qtype, *row_bytes, *scale, *rp),
15364            _ => unreachable!("gemm_supports guaranteed Quant"),
15365        };
15366        // wgmma arm (sm_90a, task 8): the m64n64k32 warpgroup kernel reads the rp4 split-plane
15367        // mirror AS-IS (qplane rows = its A operand, the half dplane its scales) and the same
15368        // (aq, ad) activation planes. Same numeric class as the mma kernel below (exact s32 per
15369        // 32-block, one f32 scale fold per block, ascending K) — argmax/tolerance gated like
15370        // every prefill GEMM, not bit-gated. MEMRA_WGMMA=0 restores the portable kernel.
15371        if cfg!(memra_hopper_mma) && qtype == QT_Q8_0 && out_f % 64 == 0 && wgmma_gemm_enabled() {
15372            if let GpuTensor::Quant { rp4: Some(m4), .. } = w {
15373                let mut y = self.qmatvec_gemm_q8_0_wgmma_raw(m4, aq, ad, m, in_f, out_f)?;
15374                if scale != 1.0 {
15375                    self.scale_inplace(&mut y, scale, m * out_f)?;
15376                }
15377                return Ok(y);
15378            }
15379        }
15380        let name = match qtype {
15381            QT_Q8_0 => "qmatvec_gemm_q8_0",
15382            QT_Q4_K => "qmatvec_gemm_q4_K",
15383            QT_Q4_0 => {
15384                if rp {
15385                    "qmatvec_gemm_q4_0_rp"
15386                } else {
15387                    "qmatvec_gemm_q4_0"
15388                }
15389            }
15390            QT_Q5_K => "qmatvec_gemm_q5_K",
15391            QT_Q6_K => "qmatvec_gemm_q6_K",
15392            QT_NVFP4 => {
15393                if rp {
15394                    "qmatvec_gemm_nvfp4_rp"
15395                } else {
15396                    "qmatvec_gemm_nvfp4"
15397                }
15398            }
15399            _ => unreachable!(),
15400        };
15401        let f = self.func(name);
15402        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output: skip memset
15403        // CTA tile MUST match the .cu per-kernel tile. MMQ-PORT: kernel1 (Q8_0/Q4_K/Q5_K) runs llama's
15404        // 128x128 SQUARE tile (K1_BM=128 x K1_BN=128, 8 warps); kernel2 (Q6_K/NVFP4) keeps 64x256, 4 warps
15405        // (the macro BM/BN in the .cu). Grid dims are selected by qtype so each launches its own tile.
15406        let is_k1 = matches!(qtype, QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q4_0);
15407        // TUNE SEAM: MEMRA_GEMM_K1_LAUNCH overrides kernel1's launch tile to match a -D-swept fatbin.
15408        let k1_tile = if is_k1 {
15409            k1_launch_override().unwrap_or((128, 128, 8))
15410        } else {
15411            (128, 128, 8)
15412        };
15413        let (bm, bn): (u32, u32) = if is_k1 {
15414            (k1_tile.0, k1_tile.1)
15415        } else {
15416            (64, 256)
15417        };
15418        let warps: u32 = if is_k1 {
15419            k1_tile.2
15420        } else {
15421            match qtype {
15422                QT_NVFP4 => 8,
15423                _ => 4,
15424            }
15425        };
15426        let cfg = LaunchConfig {
15427            grid_dim: ((out_f as u32 + bm - 1) / bm, (m as u32 + bn - 1) / bn, 1),
15428            block_dim: (32, warps, 1),
15429            shared_mem_bytes: 0,
15430        };
15431        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
15432        let __s_b = self.gpu.stream();
15433        let mut b = __s_b.launch_builder(&f);
15434        b.arg(bytes)
15435            .arg(aq)
15436            .arg(ad)
15437            .arg(&mut y)
15438            .arg(&inf)
15439            .arg(&outf)
15440            .arg(&mi)
15441            .arg(&rb);
15442        unsafe {
15443            b.launch(cfg)?;
15444        }
15445        if scale != 1.0 {
15446            self.scale_inplace(&mut y, scale, m * out_f)?;
15447        }
15448        Ok(y)
15449    }
15450
15451    /// Test entry: run the GEMM directly from raw weight bytes + qtype (no GpuTensor). Quantizes
15452    /// the f32 activation `x` to q8_1 internally then launches the tensor-core GEMM. NVFP4 per-tensor
15453    /// macro-scale is NOT applied here (caller passes it separately, like the dp4a path). Used by
15454    /// kernel_check for the bit-equivalence gate vs qmatvec_*_dp4a.
15455    pub fn qmatvec_gemm_raw(
15456        &self,
15457        bytes: &CudaSlice<u8>,
15458        x: &CudaSlice<f32>,
15459        m: usize,
15460        in_f: usize,
15461        out_f: usize,
15462        qtype: i32,
15463        row_bytes: usize,
15464    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15465        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15466        let name = match qtype {
15467            QT_Q8_0 => "qmatvec_gemm_q8_0",
15468            QT_Q4_K => "qmatvec_gemm_q4_K",
15469            QT_Q4_0 => "qmatvec_gemm_q4_0",
15470            QT_Q5_K => "qmatvec_gemm_q5_K",
15471            QT_Q6_K => "qmatvec_gemm_q6_K",
15472            QT_NVFP4 => "qmatvec_gemm_nvfp4",
15473            QT_NVFP4_RP => "qmatvec_gemm_nvfp4_rp",
15474            _ => panic!("qmatvec_gemm_raw: qtype {qtype} has no GEMM kernel"),
15475        };
15476        let f = self.func(name);
15477        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output: skip memset
15478        // MMQ-PORT: kernel1 (Q8_0/Q4_K/Q5_K) = llama 128x128 tile, 8 warps; kernel2 (Q6_K/NVFP4) = 64x256,
15479        // 4/8 warps. Grid tile per qtype (must match the .cu K1_BM/K1_BN vs BM/BN). KEEP IN SYNC w/ qmatvec_gemm.
15480        let is_k1 = matches!(qtype, QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q4_0);
15481        // TUNE SEAM: MEMRA_GEMM_K1_LAUNCH overrides kernel1's launch tile to match a -D-swept fatbin.
15482        let k1_tile = if is_k1 {
15483            k1_launch_override().unwrap_or((128, 128, 8))
15484        } else {
15485            (128, 128, 8)
15486        };
15487        let (bm, bn): (u32, u32) = if is_k1 {
15488            (k1_tile.0, k1_tile.1)
15489        } else {
15490            (64, 256)
15491        };
15492        let warps: u32 = if is_k1 {
15493            k1_tile.2
15494        } else {
15495            match qtype {
15496                QT_NVFP4 | QT_NVFP4_RP => 8,
15497                _ => 4,
15498            }
15499        };
15500        let cfg = LaunchConfig {
15501            grid_dim: ((out_f as u32 + bm - 1) / bm, (m as u32 + bn - 1) / bn, 1),
15502            block_dim: (32, warps, 1),
15503            shared_mem_bytes: 0,
15504        };
15505        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
15506        let __s_b = self.gpu.stream();
15507        let mut b = __s_b.launch_builder(&f);
15508        b.arg(bytes)
15509            .arg(&aq)
15510            .arg(&ad)
15511            .arg(&mut y)
15512            .arg(&inf)
15513            .arg(&outf)
15514            .arg(&mi)
15515            .arg(&rb);
15516        unsafe {
15517            b.launch(cfg)?;
15518        }
15519        Ok(y)
15520    }
15521
15522    /// H100 warpgroup GEMM raw entry (task 8): launch `qmatvec_gemm_q8_0_wgmma` on an rp4
15523    /// split-plane mirror + pre-quantized (aq, ad) activation planes. One warpgroup (128 thr)
15524    /// owns a 64x64 C tile; grid (out_f/64, ceil(m/64)). out_f % 64 == 0 REQUIRED (row loads
15525    /// and dplane scale reads are unguarded); the token edge is guarded in-kernel.
15526    /// Standalone harness verdict (tools/bench_q8_gemm_wgmma.cu, 4096x4096x512): rel 1.6e-05
15527    /// vs CPU ref, 179us vs the portable mma kernel's 688us (3.84x, unpipelined).
15528    pub fn qmatvec_gemm_q8_0_wgmma_raw(
15529        &self,
15530        rp4: &CudaSlice<u8>,
15531        aq: &CudaSlice<i8>,
15532        ad: &CudaSlice<f32>,
15533        m: usize,
15534        in_f: usize,
15535        out_f: usize,
15536    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15537        assert!(
15538            out_f % 64 == 0 && in_f % 32 == 0,
15539            "wgmma GEMM needs out_f%64==0, in_f%32==0"
15540        );
15541        let f = self.func("qmatvec_gemm_q8_0_wgmma");
15542        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output
15543        let cfg = LaunchConfig {
15544            grid_dim: ((out_f / 64) as u32, (m as u32).div_ceil(64), 1),
15545            block_dim: (128, 1, 1),
15546            shared_mem_bytes: 0,
15547        };
15548        let (inf, outf, mi) = (in_f as i32, out_f as i32, m as i32);
15549        let __s_b = self.gpu.stream();
15550        let mut b = __s_b.launch_builder(&f);
15551        b.arg(rp4)
15552            .arg(aq)
15553            .arg(ad)
15554            .arg(&mut y)
15555            .arg(&inf)
15556            .arg(&outf)
15557            .arg(&mi);
15558        unsafe {
15559            b.launch(cfg)?;
15560        }
15561        Ok(y)
15562    }
15563
15564    /// y[i] *= s. NVFP4 per-tensor macro-scale broadcast over the whole output.
15565    pub fn scale_inplace(
15566        &self,
15567        y: &mut CudaSlice<f32>,
15568        s: f32,
15569        n: usize,
15570    ) -> Result<(), Box<dyn std::error::Error>> {
15571        let f = self.func("scale_f32");
15572        let cfg = LaunchConfig::for_num_elems(n as u32);
15573        let (sf, ni) = (s, n as i32);
15574        let __s_b = self.gpu.stream();
15575        let mut b = __s_b.launch_builder(&f);
15576        b.arg(y).arg(&sf).arg(&ni);
15577        unsafe {
15578            b.launch(cfg)?;
15579        }
15580        Ok(())
15581    }
15582
15583    /// MEMRA_FULL_PREC dequant-on-use: expand a bf16-resident weight (`GpuTensor::FloatBf16`, raw
15584    /// bf16 bytes) to a transient f32 scratch of `n` elements, which then feeds the existing f32
15585    /// cuBLASLt GEMV. The scratch is freed when the caller drops it, so peak VRAM = resident bf16
15586    /// weights + ONE (largest) weight's f32 expansion + activations. SLOW IS FINE (research mode).
15587    pub fn bf16_to_f32(
15588        &self,
15589        data: &cudarc::driver::CudaView<'_, u8>,
15590        n: usize,
15591    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15592        let mut out = self.alloc_uninit::<f32>(n)?;
15593        let f = self.func("bf16_to_f32");
15594        let cfg = LaunchConfig::for_num_elems(n as u32);
15595        let ni = n as i32;
15596        let __s_b = self.gpu.stream();
15597        let mut b = __s_b.launch_builder(&f);
15598        b.arg(data).arg(&mut out).arg(&ni);
15599        unsafe {
15600            b.launch(cfg)?;
15601        }
15602        Ok(out)
15603    }
15604
15605    /// Chunked bf16 linear (MEMRA_FULL_PREC): y[m,out] = x @ W_bf16^T with the f32 dequant scratch
15606    /// bounded to CHUNK_ROWS rows (256MB at in_f=4096) instead of the whole weight — the 4GB
15607    /// lm_head expansion OOM'd the 24GB budget. Row-chunking partitions OUTPUT rows; each row's
15608    /// dot is computed by the identical kernel on identical bytes, so per-(token,row) results are
15609    /// bit-identical to the unchunked form. `exact` selects linear_decode_exact (per-column m=1
15610    /// calls, the spec-verify contract) vs plain linear.
15611    fn linear_bf16_chunked(
15612        &self,
15613        x: &CudaSlice<f32>,
15614        data: &CudaSlice<u8>,
15615        m: usize,
15616        in_f: usize,
15617        out_f: usize,
15618        exact: bool,
15619    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15620        const CHUNK_BYTES: usize = 256 << 20;
15621        let chunk_rows = (CHUNK_BYTES / (in_f * 4)).max(1).min(out_f);
15622        if chunk_rows >= out_f {
15623            let wf32 = self.bf16_to_f32(&data.slice(0..in_f * out_f * 2), in_f * out_f)?;
15624            return if exact {
15625                self.linear_decode_exact(x, &wf32, m, in_f, out_f)
15626            } else {
15627                self.linear(x, &wf32, m, in_f, out_f)
15628            };
15629        }
15630        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
15631        let mut r0 = 0usize;
15632        while r0 < out_f {
15633            let rows = chunk_rows.min(out_f - r0);
15634            let wslice = data.slice(r0 * in_f * 2..(r0 + rows) * in_f * 2);
15635            let wf32 = self.bf16_to_f32(&wslice, in_f * rows)?;
15636            let yc = if exact {
15637                self.linear_decode_exact(x, &wf32, m, in_f, rows)?
15638            } else {
15639                self.linear(x, &wf32, m, in_f, rows)?
15640            };
15641            // scatter [m, rows] into y[m, out_f] at column offset r0 (m is tiny in decode/verify)
15642            for mi in 0..m {
15643                let src = yc.slice(mi * rows..(mi + 1) * rows);
15644                let mut dst = y.slice_mut(mi * out_f + r0..mi * out_f + r0 + rows);
15645                self.gpu.stream().memcpy_dtod(&src, &mut dst)?;
15646            }
15647            r0 += rows;
15648        }
15649        Ok(y)
15650    }
15651
15652    /// On-device linear: y[m,out] = x[m,in] @ W[out,in]^T, weights row-major [out,in] (ggml).
15653    /// cuBLASLt col-major mapping (see memra_runtime::Gpu::linear_f32 for the derivation).
15654    /// DECODE-EXACT float linear: per-column m=1 cuBLASLt calls. cuBLASLt's reduction split is
15655    /// n-dependent (lt_ndep probe: m=1 vs m=2 col0 differs every bit), so spec-verify batches
15656    /// must not batch float matmuls the T=1 decode chain runs at m=1. Used by the small-t MoE
15657    /// router/shexp sites and matmul_decode_exact's Float arm.
15658    pub fn linear_decode_exact(
15659        &self,
15660        x: &CudaSlice<f32>,
15661        w: &CudaSlice<f32>,
15662        m_tokens: usize,
15663        in_f: usize,
15664        out_f: usize,
15665    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15666        if m_tokens == 1 {
15667            return self.linear(x, w, 1, in_f, out_f);
15668        }
15669        let xv = self.view(x, m_tokens * in_f);
15670        let mut y = self.alloc_uninit::<f32>(m_tokens * out_f)?;
15671        for t in 0..m_tokens {
15672            let row = xv.slice(t * in_f..(t + 1) * in_f);
15673            let mut xr = self.alloc_uninit::<f32>(in_f)?;
15674            self.copy_view_into(&mut xr, 0, &row, in_f)?;
15675            let yr = self.linear(&xr, w, 1, in_f, out_f)?;
15676            self.copy_into(&mut y, t * out_f, &yr, out_f)?;
15677        }
15678        Ok(y)
15679    }
15680
15681    pub fn linear(
15682        &self,
15683        x: &CudaSlice<f32>,
15684        w: &CudaSlice<f32>,
15685        m_tokens: usize,
15686        in_f: usize,
15687        out_f: usize,
15688    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15689        use cudarc::cublaslt::{Matmul, MatmulConfig};
15690        let mut c = self.alloc_uninit::<f32>(m_tokens * out_f)?; // cuBLASLt beta=0: C fully written
15691        let cfg = MatmulConfig {
15692            transa: true,
15693            transb: false,
15694            transc: false,
15695            m: out_f as u64,
15696            n: m_tokens as u64,
15697            k: in_f as u64,
15698            alpha: 1.0,
15699            lda: in_f as i64,
15700            ldb: in_f as i64,
15701            beta: 0.0,
15702            ldc: out_f as i64,
15703            stride_a: None,
15704            stride_b: None,
15705            stride_c: None,
15706            stride_bias: None,
15707            batch_size: None,
15708        };
15709        unsafe {
15710            self.gpu.blas.matmul(cfg, w, x, &mut c, None, None)?;
15711        }
15712        Ok(c)
15713    }
15714
15715    /// Naive SDPA. Q:[head_dim,n_head,T], K/V:[head_dim,n_head_kv,T_kv] -> O:[head_dim,n_head,T].
15716    pub fn sdpa_naive(
15717        &self,
15718        q: &CudaSlice<f32>,
15719        k: &CudaSlice<f32>,
15720        v: &CudaSlice<f32>,
15721        o: &mut CudaSlice<f32>,
15722        head_dim: usize,
15723        n_head: usize,
15724        n_head_kv: usize,
15725        t: usize,
15726        t_kv: usize,
15727        scale: f32,
15728        causal: bool,
15729    ) -> Result<(), Box<dyn std::error::Error>> {
15730        let f = self.func("sdpa_naive_f32");
15731        let cfg = LaunchConfig {
15732            grid_dim: (n_head as u32, t as u32, 1),
15733            block_dim: (128, 1, 1),
15734            shared_mem_bytes: (t_kv * 4) as u32,
15735        };
15736        let (hd, nh, nhkv, ti, tkvi, cz) = (
15737            head_dim as i32,
15738            n_head as i32,
15739            n_head_kv as i32,
15740            t as i32,
15741            t_kv as i32,
15742            causal as i32,
15743        );
15744        let __s_b = self.gpu.stream();
15745        let mut b = __s_b.launch_builder(&f);
15746        b.arg(q)
15747            .arg(k)
15748            .arg(v)
15749            .arg(o)
15750            .arg(&hd)
15751            .arg(&nh)
15752            .arg(&nhkv)
15753            .arg(&ti)
15754            .arg(&tkvi)
15755            .arg(&scale)
15756            .arg(&cz);
15757        unsafe {
15758            b.launch(cfg)?;
15759        }
15760        Ok(())
15761    }
15762
15763    /// Island twin (lane/gemma-vision): causal + sliding-window attention with
15764    /// bidirectional image islands. `span_id` labels each absolute kv position
15765    /// (-1 text, >=0 island id); same-island keys are visible unconditionally,
15766    /// reproducing the reference's non-causal image batch. window 0 = no window.
15767    #[allow(clippy::too_many_arguments)]
15768    pub fn sdpa_naive_island(
15769        &self,
15770        q: &CudaSlice<f32>,
15771        k: &CudaSlice<f32>,
15772        v: &CudaSlice<f32>,
15773        o: &mut CudaSlice<f32>,
15774        span_id: &CudaSlice<i32>,
15775        head_dim: usize,
15776        n_head: usize,
15777        n_head_kv: usize,
15778        t: usize,
15779        t_kv: usize,
15780        scale: f32,
15781        window: usize,
15782    ) -> Result<(), Box<dyn std::error::Error>> {
15783        let f = self.func("sdpa_naive_island_f32");
15784        let cfg = LaunchConfig {
15785            grid_dim: (n_head as u32, t as u32, 1),
15786            block_dim: (128, 1, 1),
15787            shared_mem_bytes: (t_kv * 4) as u32,
15788        };
15789        let (hd, nh, nhkv, ti, tkvi, wi) = (
15790            head_dim as i32,
15791            n_head as i32,
15792            n_head_kv as i32,
15793            t as i32,
15794            t_kv as i32,
15795            window as i32,
15796        );
15797        let __s_b = self.gpu.stream();
15798        let mut b = __s_b.launch_builder(&f);
15799        b.arg(q)
15800            .arg(k)
15801            .arg(v)
15802            .arg(o)
15803            .arg(span_id)
15804            .arg(&hd)
15805            .arg(&nh)
15806            .arg(&nhkv)
15807            .arg(&ti)
15808            .arg(&tkvi)
15809            .arg(&scale)
15810            .arg(&wi);
15811        unsafe {
15812            b.launch(cfg)?;
15813        }
15814        Ok(())
15815    }
15816
15817    /// Windowed sdpa_naive twin (gemma4 R6): masks keys older than q_pos-(window-1).
15818    #[allow(clippy::too_many_arguments)]
15819    pub fn sdpa_naive_w(
15820        &self,
15821        q: &CudaSlice<f32>,
15822        k: &CudaSlice<f32>,
15823        v: &CudaSlice<f32>,
15824        o: &mut CudaSlice<f32>,
15825        head_dim: usize,
15826        n_head: usize,
15827        n_head_kv: usize,
15828        t: usize,
15829        t_kv: usize,
15830        scale: f32,
15831        causal: bool,
15832        window: usize,
15833    ) -> Result<(), Box<dyn std::error::Error>> {
15834        let f = self.func("sdpa_naive_w_f32");
15835        let cfg = LaunchConfig {
15836            grid_dim: (n_head as u32, t as u32, 1),
15837            block_dim: (128, 1, 1),
15838            shared_mem_bytes: (t_kv * 4) as u32,
15839        };
15840        let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
15841            head_dim as i32,
15842            n_head as i32,
15843            n_head_kv as i32,
15844            t as i32,
15845            t_kv as i32,
15846            causal 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(&hd)
15856            .arg(&nh)
15857            .arg(&nhkv)
15858            .arg(&ti)
15859            .arg(&tkvi)
15860            .arg(&scale)
15861            .arg(&cz)
15862            .arg(&wi);
15863        unsafe {
15864            b.launch(cfg)?;
15865        }
15866        Ok(())
15867    }
15868
15869    /// Lo-clipped windowed sdpa_naive twin (lane/dflash2-longctx, DFLASH2-EVAL §10.6(c)).
15870    /// Same mask law as `sdpa_naive_w`, but keys below every query's window floor are never
15871    /// read: kv_lo = max(0, (t_kv - t) + 1 - window) — the oldest key visible to the OLDEST
15872    /// query row (q_pos = t_kv - t). Dynamic shared memory shrinks from t_kv*4 bytes (which
15873    /// blows the 48KB launch bound at ~12k rows — the B2 ctx crash) to (t_kv - kv_lo)*4 =
15874    /// (window - 1 + t)*4, and the key scan drops from O(t_kv) to O(window + t). Output is
15875    /// byte-identical to `sdpa_naive_w` (masked keys contribute exact zeros to same-order
15876    /// reductions; kernel_check `sdpa_naive_w_lo` pins it). window == 0 (no window) keeps
15877    /// kv_lo = 0 and is then shape-identical to the legacy kernel, including its bound.
15878    #[allow(clippy::too_many_arguments)]
15879    pub fn sdpa_naive_w_lo(
15880        &self,
15881        q: &CudaSlice<f32>,
15882        k: &CudaSlice<f32>,
15883        v: &CudaSlice<f32>,
15884        o: &mut CudaSlice<f32>,
15885        head_dim: usize,
15886        n_head: usize,
15887        n_head_kv: usize,
15888        t: usize,
15889        t_kv: usize,
15890        scale: f32,
15891        causal: bool,
15892        window: usize,
15893    ) -> Result<(), Box<dyn std::error::Error>> {
15894        let kv_lo = if window > 0 {
15895            (t_kv - t + 1).saturating_sub(window)
15896        } else {
15897            0
15898        };
15899        let smem = (t_kv - kv_lo) * 4;
15900        if smem > 48 * 1024 {
15901            return Err(format!(
15902                "sdpa_naive_w_lo: window {window} + T {t} rows need {smem} bytes of dynamic \
15903                 shared memory (> 48KB launch bound) — this kernel clips the OLD side only; \
15904                 a window this wide needs the multi-pass long-ctx kernel"
15905            )
15906            .into());
15907        }
15908        let f = self.func("sdpa_naive_w_lo_f32");
15909        let cfg = LaunchConfig {
15910            grid_dim: (n_head as u32, t as u32, 1),
15911            block_dim: (128, 1, 1),
15912            shared_mem_bytes: smem as u32,
15913        };
15914        let (hd, nh, nhkv, ti, tkvi, cz, wi, lo) = (
15915            head_dim as i32,
15916            n_head as i32,
15917            n_head_kv as i32,
15918            t as i32,
15919            t_kv as i32,
15920            causal as i32,
15921            window as i32,
15922            kv_lo as i32,
15923        );
15924        let __s_b = self.gpu.stream();
15925        let mut b = __s_b.launch_builder(&f);
15926        b.arg(q)
15927            .arg(k)
15928            .arg(v)
15929            .arg(o)
15930            .arg(&hd)
15931            .arg(&nh)
15932            .arg(&nhkv)
15933            .arg(&ti)
15934            .arg(&tkvi)
15935            .arg(&scale)
15936            .arg(&cz)
15937            .arg(&wi)
15938            .arg(&lo);
15939        unsafe {
15940            b.launch(cfg)?;
15941        }
15942        Ok(())
15943    }
15944
15945    /// SDPA where K/V are CudaViews into a resident KV cache (decode hot path, no host round-trip).
15946    pub fn sdpa_naive_view(
15947        &self,
15948        q: &CudaSlice<f32>,
15949        k: &cudarc::driver::CudaView<f32>,
15950        v: &cudarc::driver::CudaView<f32>,
15951        o: &mut CudaSlice<f32>,
15952        head_dim: usize,
15953        n_head: usize,
15954        n_head_kv: usize,
15955        t: usize,
15956        t_kv: usize,
15957        scale: f32,
15958        causal: bool,
15959    ) -> Result<(), Box<dyn std::error::Error>> {
15960        let f = self.func("sdpa_naive_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: (t_kv * 4) as u32,
15965        };
15966        let (hd, nh, nhkv, ti, tkvi, cz) = (
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        );
15974        let __s_b = self.gpu.stream();
15975        let mut b = __s_b.launch_builder(&f);
15976        b.arg(q)
15977            .arg(k)
15978            .arg(v)
15979            .arg(o)
15980            .arg(&hd)
15981            .arg(&nh)
15982            .arg(&nhkv)
15983            .arg(&ti)
15984            .arg(&tkvi)
15985            .arg(&scale)
15986            .arg(&cz);
15987        unsafe {
15988            b.launch(cfg)?;
15989        }
15990        Ok(())
15991    }
15992
15993    /// Correctness fallback for quantized resident K/V views. Dequantizes K and V once into f32
15994    /// workspaces, then calls `sdpa_naive`. This is an explicit API: the optimized prefill view
15995    /// dispatch remains unchanged, so callers can use it as a reference or compatibility path.
15996    /// Dequant a quantized KV view into caller-owned f32 buffers (one grid-stride launch).
15997    /// `g` picks the kf8vf8-module stamp for e4m3 caches (same flag contract as fa_decode/
15998    /// fa_prefill_view). Used by the E4B shared-KV prefill arms (2026-07-31) to feed the
15999    /// f32 fa_prefill_w / fa_prefill_hd512 twins from the target layer's quantized rows.
16000    #[allow(clippy::too_many_arguments)]
16001    pub fn fa_dequant_kv_view_f32(
16002        &self,
16003        k: &cudarc::driver::CudaView<u8>,
16004        v: &cudarc::driver::CudaView<u8>,
16005        kf: &mut CudaSlice<f32>,
16006        vf: &mut CudaSlice<f32>,
16007        kv_dim_k: usize,
16008        kv_dim_v: usize,
16009        t_kv: usize,
16010        k_tok_bytes: usize,
16011        v_tok_bytes: usize,
16012        g: bool,
16013    ) -> Result<(), Box<dyn std::error::Error>> {
16014        let f = if g {
16015            self.func_g("fa_dequant_kv_ws_f32")
16016        } else {
16017            self.func("fa_dequant_kv_ws_f32")
16018        };
16019        let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
16020        let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
16021        let cfg = LaunchConfig {
16022            grid_dim: (nblk.max(1), 1, 1),
16023            block_dim: (256, 1, 1),
16024            shared_mem_bytes: 0,
16025        };
16026        let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
16027        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
16028        let __s_b = self.gpu.stream();
16029        let mut b = __s_b.launch_builder(&f);
16030        b.arg(k)
16031            .arg(v)
16032            .arg(&mut *kf)
16033            .arg(&mut *vf)
16034            .arg(&kdk)
16035            .arg(&kdv)
16036            .arg(&tkvi)
16037            .arg(&ktb)
16038            .arg(&vtb);
16039        unsafe {
16040            b.launch(cfg)?;
16041        }
16042        Ok(())
16043    }
16044
16045    #[allow(clippy::too_many_arguments)]
16046    pub fn sdpa_naive_quantized_view(
16047        &self,
16048        q: &CudaSlice<f32>,
16049        k: &cudarc::driver::CudaView<u8>,
16050        v: &cudarc::driver::CudaView<u8>,
16051        o: &mut CudaSlice<f32>,
16052        head_dim: usize,
16053        n_head: usize,
16054        n_head_kv: usize,
16055        t: usize,
16056        t_kv: usize,
16057        scale: f32,
16058        causal: bool,
16059        k_tok_bytes: usize,
16060        v_tok_bytes: usize,
16061    ) -> Result<(), Box<dyn std::error::Error>> {
16062        let kv_dim = n_head_kv * head_dim;
16063        let mut kf = self.uninit(t_kv * kv_dim)?;
16064        let mut vf = self.uninit(t_kv * kv_dim)?;
16065        let f = self.func("fa_dequant_kv_ws_f32");
16066        let total = (2 * t_kv * kv_dim) as u64;
16067        let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
16068        let cfg = LaunchConfig {
16069            grid_dim: (nblk.max(1), 1, 1),
16070            block_dim: (256, 1, 1),
16071            shared_mem_bytes: 0,
16072        };
16073        let (kv_dim_i, t_kv_i) = (kv_dim as i32, t_kv as i32);
16074        let (k_tok_bytes_i, v_tok_bytes_i) = (k_tok_bytes as i64, v_tok_bytes as i64);
16075        let __s_b = self.gpu.stream();
16076        let mut b = __s_b.launch_builder(&f);
16077        b.arg(k)
16078            .arg(v)
16079            .arg(&mut kf)
16080            .arg(&mut vf)
16081            .arg(&kv_dim_i)
16082            .arg(&kv_dim_i)
16083            .arg(&t_kv_i)
16084            .arg(&k_tok_bytes_i)
16085            .arg(&v_tok_bytes_i);
16086        unsafe { b.launch(cfg)? };
16087        self.sdpa_naive(
16088            q, &kf, &vf, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
16089        )
16090    }
16091
16092    /// WINDOWED twin of `sdpa_naive_quantized_view` (step35 SWA prefill): dequant the KV byte
16093    /// view into f32 workspaces with the SAME `fa_dequant_kv_ws_f32` launch, then run
16094    /// `sdpa_naive_w` instead of `sdpa_naive`. `window == 0` is the unwindowed form (the kernel
16095    /// treats a non-positive window as "no window mask"), so this is a strict superset of the
16096    /// unwindowed function above and produces bit-identical output at window == 0.
16097    ///
16098    /// Why this exists: EVERY windowed FlashAttention stamp in flash_attn.cu is head_dim-256
16099    /// only (`fa_prefill_w_f32` == `fa_prefill_f32_body<256>`, and the quantized-view windowed
16100    /// twins likewise), while step35 is head_dim 128. Its SWA layers therefore have no windowed
16101    /// FA path and take this f32 floor in v0 — same cache bytes, same numeric class as the
16102    /// unwindowed quantized-view fallback, so the chunk-invariance contract holds on both.
16103    #[allow(clippy::too_many_arguments)]
16104    pub fn sdpa_naive_w_quantized_view(
16105        &self,
16106        q: &CudaSlice<f32>,
16107        k: &cudarc::driver::CudaView<u8>,
16108        v: &cudarc::driver::CudaView<u8>,
16109        o: &mut CudaSlice<f32>,
16110        head_dim: usize,
16111        n_head: usize,
16112        n_head_kv: usize,
16113        t: usize,
16114        t_kv: usize,
16115        scale: f32,
16116        causal: bool,
16117        window: usize,
16118        k_tok_bytes: usize,
16119        v_tok_bytes: usize,
16120    ) -> Result<(), Box<dyn std::error::Error>> {
16121        let kv_dim = n_head_kv * head_dim;
16122        let mut kf = self.uninit(t_kv * kv_dim)?;
16123        let mut vf = self.uninit(t_kv * kv_dim)?;
16124        let f = self.func("fa_dequant_kv_ws_f32");
16125        let total = (2 * t_kv * kv_dim) as u64;
16126        let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
16127        let cfg = LaunchConfig {
16128            grid_dim: (nblk.max(1), 1, 1),
16129            block_dim: (256, 1, 1),
16130            shared_mem_bytes: 0,
16131        };
16132        let (kv_dim_i, t_kv_i) = (kv_dim as i32, t_kv as i32);
16133        let (k_tok_bytes_i, v_tok_bytes_i) = (k_tok_bytes as i64, v_tok_bytes as i64);
16134        let __s_b = self.gpu.stream();
16135        let mut b = __s_b.launch_builder(&f);
16136        b.arg(k)
16137            .arg(v)
16138            .arg(&mut kf)
16139            .arg(&mut vf)
16140            .arg(&kv_dim_i)
16141            .arg(&kv_dim_i)
16142            .arg(&t_kv_i)
16143            .arg(&k_tok_bytes_i)
16144            .arg(&v_tok_bytes_i);
16145        unsafe { b.launch(cfg)? };
16146        self.sdpa_naive_w(
16147            q, &kf, &vf, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal, window,
16148        )
16149    }
16150
16151    /// Hand-written FlashAttention prefill (sm_120, FA-2 online softmax on validated mma.sync,
16152    /// head_dim 256 or 128 (template-stamped twins), GQA, causal). Replaces sdpa_naive for T>1.
16153    /// Q/K/V/O [head_dim, n_head(_kv), T].
16154    pub fn fa_prefill(
16155        &self,
16156        q: &CudaSlice<f32>,
16157        k: &CudaSlice<f32>,
16158        v: &CudaSlice<f32>,
16159        o: &mut CudaSlice<f32>,
16160        head_dim: usize,
16161        n_head: usize,
16162        n_head_kv: usize,
16163        t: usize,
16164        t_kv: usize,
16165        scale: f32,
16166        causal: bool,
16167    ) -> Result<(), Box<dyn std::error::Error>> {
16168        if portable_mma_gated() {
16169            return self.sdpa_naive(
16170                q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
16171            );
16172        }
16173        // FA3 v10 arm (task #20, OPT-IN MEMRA_FA3=1 — harness-proven 883us vs the shipped
16174        // kernel's 993us at T=2048): TMA-swizzled wgmma FA, fresh causal hd256 only.
16175        // NEW NUMERIC CONFIG (GDN-mma precedent): online softmax / bf16-P class — the
16176        // run-gen argmax + greedy-stream batteries arbitrate; not bit-paired.
16177        // PROMOTED default-ON hopper (2026-07-27): 3-seed 2048-prime -> 128-decode
16178        // streams MATCH vs mma, full battery green, lane interleaved 5/5 (+2.4%).
16179        // MEMRA_FA3=0 reverts; kernel-check pins the mma config regardless.
16180        let fa3_on = head_dim == 256
16181            && causal
16182            && t == t_kv
16183            && match std::env::var("MEMRA_FA3").as_deref() {
16184                Ok("0") => false,
16185                Ok("1") => true,
16186                _ => cfg!(memra_hopper_mma),
16187            };
16188        if fa3_on {
16189            let n = t * n_head * head_dim;
16190            let nkv = t * n_head_kv * head_dim;
16191            let mut q16 = self.alloc_u8_uninit(n * 2)?;
16192            let mut k16 = self.alloc_u8_uninit(nkv * 2)?;
16193            let mut v16 = self.alloc_u8_uninit(nkv * 2)?;
16194            self.f32_to_bf16_into(q, &mut q16, n)?;
16195            self.f32_to_bf16_into(k, &mut k16, nkv)?;
16196            self.f32_to_bf16_into(v, &mut v16, nkv)?;
16197            let rc = {
16198                use cudarc::driver::{DevicePtr, DevicePtrMut};
16199                let stream = self.gpu.stream();
16200                let (qp, _g1) = q16.device_ptr(&stream);
16201                let (kp, _g2) = k16.device_ptr(&stream);
16202                let (vp, _g3) = v16.device_ptr(&stream);
16203                let (op, _g4) = o.device_ptr_mut(&stream);
16204                unsafe {
16205                    memra_fa3_prefill(
16206                        qp as *const core::ffi::c_void,
16207                        kp as *const core::ffi::c_void,
16208                        vp as *const core::ffi::c_void,
16209                        op as *mut f32,
16210                        t as i32,
16211                        n_head as i32,
16212                        n_head_kv as i32,
16213                        head_dim as i32,
16214                        scale,
16215                        stream.cu_stream() as *mut core::ffi::c_void,
16216                    )
16217                }
16218            };
16219            if rc != 0 {
16220                return Err(format!("memra_fa3_prefill rc={rc}").into());
16221            }
16222            return Ok(());
16223        }
16224        // FLOOR PORT (P2+P0a+P0b+P1): 4 warps/CTA, BLOCK_Q=64 query rows, BK=32 KV tile,
16225        // Q-in-reg + register-O, grid.y=n_head_kv (4 Q-heads share staged K/V).
16226        // P1 plain arm (MEMRA_FA_P1=1 opt-in until the qwen battery): the engine-study body
16227        // (FA2 schedule + boundary split + swizzle) on the non-windowed lane. bf16 pre-convert.
16228        static FA_P1: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16229        let fa_p1 = *FA_P1.get_or_init(|| std::env::var("MEMRA_FA_P1").as_deref() == Ok("1"));
16230        if fa_p1 && head_dim == 256 && !std::env::var("MEMRA_FA_FLOOR").is_ok() {
16231            const BLOCK_Q: usize = 64;
16232            const BKX: usize = 32;
16233            let f = self.func("fa_prefill_bf16_p1");
16234            let shmem = (2 * (2 * BKX * head_dim + BLOCK_Q * BKX)
16235                + 4 * (BLOCK_Q * BKX + 2 * BLOCK_Q)) as u32;
16236            use cudarc::driver::sys::CUfunction_attribute_enum as A;
16237            f.set_attribute(
16238                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16239                shmem as i32,
16240            )?;
16241            let cfg = LaunchConfig {
16242                grid_dim: (
16243                    (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
16244                    n_head as u32,
16245                    1,
16246                ),
16247                block_dim: (32, 4, 1),
16248                shared_mem_bytes: shmem,
16249            };
16250            let (hd, nh, nhkv, ti, tkvi, cz) = (
16251                head_dim as i32,
16252                n_head as i32,
16253                n_head_kv as i32,
16254                t as i32,
16255                t_kv as i32,
16256                causal as i32,
16257            );
16258            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
16259            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
16260            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
16261            let __s_b = self.gpu.stream();
16262            let mut b = __s_b.launch_builder(&f);
16263            b.arg(&qb)
16264                .arg(&kb)
16265                .arg(&vb)
16266                .arg(o)
16267                .arg(&hd)
16268                .arg(&nh)
16269                .arg(&nhkv)
16270                .arg(&ti)
16271                .arg(&tkvi)
16272                .arg(&scale)
16273                .arg(&cz);
16274            unsafe {
16275                b.launch(cfg)?;
16276            }
16277            return Ok(());
16278        }
16279        // Edge 5a (DEFAULT): fa_prefill_f32_pp — register-resident softmax (no sSw smem
16280        // round-trip), the FA3 softmax-GEMM overlap variant. ncu (pp512): short_scoreboard
16281        // 4.32->3.47, wait 1.99->1.45, per-call ~577us->~440us (1.31x) at flat 12.1% warps /
16282        // 255 regs / 2 CTAs (occupancy preserved). Bit-safe: 9B+27B argmax MATCH, rel 2.55e-3
16283        // vs floor 3.03e-3. MEMRA_FA_FLOOR reverts to the serialized-softmax floor kernel.
16284        const BK: usize = 32;
16285        // W2 lane (MEMRA_FA_PP_W2=1, ncu 2026-07-26): 2-warp/32-row CTA tile doubles grid.x —
16286        // bit-identical per-row math, pure coverage trade for the 6.25%-occupancy starvation.
16287        let w2 = std::env::var("MEMRA_FA_PP_W2").as_deref() == Ok("1");
16288        let (block_q, warps, w2_sfx): (usize, u32, &str) =
16289            if w2 { (32, 2, "_w2") } else { (64, 4, "") };
16290        // hd128 twins (2026-07-07): the prefill kernels are template-stamped at 256 (original
16291        // names, dispatch unchanged) and 128 (`_hd128`, the MiniMax-M3 class). Callers gate
16292        // other head_dims to sdpa_naive before reaching here.
16293        let hd_sfx = fa_hd_suffix(head_dim)?;
16294        let floor = std::env::var("MEMRA_FA_FLOOR").is_ok();
16295        // BF16-KV staging lane (2026-07-26, default ON): the kernel converts K/V to bf16
16296        // during staging anyway — pre-converting to bf16 mirrors is BIT-IDENTICAL (same
16297        // __float2bfloat16 values into the same mma) and turns the 67%-of-stalls scalar
16298        // staging into int4 vector copies. MEMRA_FA_BF16KV=0 reverts.
16299        let bf16kv = !floor && !w2 && std::env::var("MEMRA_FA_BF16KV").as_deref() != Ok("0");
16300        let (kb16, vb16) = if bf16kv {
16301            let n = t_kv * n_head_kv * head_dim;
16302            let mut kb = self.alloc_u8_uninit(n * 2)?;
16303            let mut vb = self.alloc_u8_uninit(n * 2)?;
16304            let fcv = self.func("f32_to_bf16_bulk");
16305            let ni = n as i64;
16306            let cfgc = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
16307            let __s_b = self.gpu.stream();
16308            let mut b = __s_b.launch_builder(&fcv);
16309            b.arg(k).arg(&mut kb).arg(&ni);
16310            unsafe {
16311                b.launch(cfgc)?;
16312            }
16313            let __s_b = self.gpu.stream();
16314            let mut b = __s_b.launch_builder(&fcv);
16315            b.arg(v).arg(&mut vb).arg(&ni);
16316            unsafe {
16317                b.launch(cfgc)?;
16318            }
16319            (Some(kb), Some(vb))
16320        } else {
16321            (None, None)
16322        };
16323        let f = self.func(&if bf16kv {
16324            format!("fa_prefill_bf16kv_pp{hd_sfx}")
16325        } else {
16326            format!(
16327                "fa_prefill_f32{}{}{hd_sfx}",
16328                if floor { "" } else { "_pp" },
16329                if floor { "" } else { w2_sfx }
16330            )
16331        });
16332        // persistent smem: bf16*(KV_STAGES*(sK + sV) + sP) + f32*(sS + sM + sL);
16333        // the bf16kv ring doubles the K/V stages (KV_STAGES=2).
16334        let kv_stages = if bf16kv { 2 } else { 1 };
16335        let shmem = (2 * (kv_stages * 2 * BK * head_dim + block_q * BK)
16336            + 4 * (block_q * BK + 2 * block_q)) as u32;
16337        use cudarc::driver::sys::CUfunction_attribute_enum as A;
16338        f.set_attribute(
16339            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16340            shmem as i32,
16341        )?;
16342        let cfg = LaunchConfig {
16343            grid_dim: (
16344                (t as u32 + block_q as u32 - 1) / block_q as u32,
16345                n_head as u32,
16346                1,
16347            ),
16348            block_dim: (32, warps, 1),
16349            shared_mem_bytes: shmem,
16350        };
16351        let (hd, nh, nhkv, ti, tkvi, cz) = (
16352            head_dim as i32,
16353            n_head as i32,
16354            n_head_kv as i32,
16355            t as i32,
16356            t_kv as i32,
16357            causal as i32,
16358        );
16359        let __s_b = self.gpu.stream();
16360        let mut b = __s_b.launch_builder(&f);
16361        b.arg(q);
16362        match (&kb16, &vb16) {
16363            (Some(kb), Some(vb)) => {
16364                b.arg(kb).arg(vb);
16365            }
16366            _ => {
16367                b.arg(k).arg(v);
16368            }
16369        }
16370        b.arg(o)
16371            .arg(&hd)
16372            .arg(&nh)
16373            .arg(&nhkv)
16374            .arg(&ti)
16375            .arg(&tkvi)
16376            .arg(&scale)
16377            .arg(&cz);
16378        unsafe {
16379            b.launch(cfg)?;
16380        }
16381        Ok(())
16382    }
16383
16384    /// Windowed FA prefill (gemma4 SWA layers past the sliding window, hd256): fa_prefill's
16385    /// exact dispatch (pp default, MEMRA_FA_FLOOR seam) with the sliding-window mask + tile
16386    /// skip in-kernel. Replaces the O(T*T_kv) scalar sdpa_naive_w on the prime path.
16387    #[allow(clippy::too_many_arguments)]
16388    pub fn fa_prefill_w(
16389        &self,
16390        q: &CudaSlice<f32>,
16391        k: &CudaSlice<f32>,
16392        v: &CudaSlice<f32>,
16393        o: &mut CudaSlice<f32>,
16394        head_dim: usize,
16395        n_head: usize,
16396        n_head_kv: usize,
16397        t: usize,
16398        t_kv: usize,
16399        scale: f32,
16400        causal: bool,
16401        window: usize,
16402    ) -> Result<(), Box<dyn std::error::Error>> {
16403        // sm_90a rides the mma twins (portable_mma_gated, 2026-07-31 — the raw
16404        // portable_cuda gate was stale-conservative on Hopper; fa_prefill already flipped).
16405        if portable_mma_gated() {
16406            return self.sdpa_naive_w(
16407                q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal, window,
16408            );
16409        }
16410        // Default: bf16-prestaged twin (same treatment as hd512 — Q/K/V pre-converted once,
16411        // int4 stage copies; bit-identical, kernel_check-gated). MEMRA_FAW_STAGE=f32 reverts;
16412        // MEMRA_FA_FLOOR keeps the f32 floor stamp untouched.
16413        static FAW_F32: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16414        let faw_f32 =
16415            *FAW_F32.get_or_init(|| std::env::var("MEMRA_FAW_STAGE").as_deref() == Ok("f32"));
16416        let floor = std::env::var("MEMRA_FA_FLOOR").is_ok();
16417        self.fa_prefill_w_arm(
16418            q,
16419            k,
16420            v,
16421            o,
16422            head_dim,
16423            n_head,
16424            n_head_kv,
16425            t,
16426            t_kv,
16427            scale,
16428            causal,
16429            window,
16430            floor || faw_f32,
16431            floor,
16432        )
16433    }
16434
16435    /// Windowed FA prefill with PRE-CONVERTED bf16 operands (producer-emitted; 31B glue lane).
16436    /// Launches the P1 stamp directly — callers guarantee qb/kb/vb hold the exact bf16 of q/k/v.
16437    #[allow(clippy::too_many_arguments)]
16438    pub fn fa_prefill_w_pre(
16439        &self,
16440        qb: &CudaSlice<u8>,
16441        kb: &CudaSlice<u8>,
16442        vb: &CudaSlice<u8>,
16443        o: &mut CudaSlice<f32>,
16444        head_dim: usize,
16445        n_head: usize,
16446        n_head_kv: usize,
16447        t: usize,
16448        t_kv: usize,
16449        scale: f32,
16450        causal: bool,
16451        window: usize,
16452        v_f16: bool,
16453    ) -> Result<(), Box<dyn std::error::Error>> {
16454        const BLOCK_Q: usize = 64;
16455        const BK: usize = 32;
16456        debug_assert_eq!(head_dim, 256);
16457        let hp = fa_f16pv_on() && faw_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
16458        debug_assert!(!v_f16 || hp, "f16 V emitted but the SWA hp arm is off");
16459        if hp {
16460            const BLOCK_QH: usize = 32;
16461            // V bytes must be f16 for the h2 stamp; producer normally emits f16 (v_f16),
16462            // else re-encode through the pooled scratch (stream-ordered reuse).
16463            let mut vguard = self.fa_vf16_scratch.lock().unwrap();
16464            let vh: &CudaSlice<u8> = if v_f16 {
16465                vb
16466            } else {
16467                let n = t_kv * n_head_kv * head_dim;
16468                if vguard.as_ref().map(|b| b.len() < n * 2).unwrap_or(true) {
16469                    *vguard = Some(self.alloc_uninit::<u8>(n * 2)?);
16470                }
16471                self.bf16_to_f16_into(vb, n, vguard.as_mut().unwrap())?;
16472                vguard.as_ref().unwrap()
16473            };
16474            let f = self.func("fa_prefill_w_bf16_p1h2");
16475            let shmem = (2 * (2 * BK * head_dim + 2 * BLOCK_QH * BK) + 4 * (2 * BLOCK_QH)) as u32;
16476            use cudarc::driver::sys::CUfunction_attribute_enum as A;
16477            f.set_attribute(
16478                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16479                shmem as i32,
16480            )?;
16481            let cfg = LaunchConfig {
16482                grid_dim: ((t as u32).div_ceil(BLOCK_QH as u32), (n_head / 2) as u32, 1),
16483                block_dim: (32, 4, 1),
16484                shared_mem_bytes: shmem,
16485            };
16486            let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
16487                head_dim as i32,
16488                n_head as i32,
16489                n_head_kv as i32,
16490                t as i32,
16491                t_kv as i32,
16492                causal as i32,
16493                window as i32,
16494            );
16495            let __s_b = self.gpu.stream();
16496            let mut b = __s_b.launch_builder(&f);
16497            b.arg(qb)
16498                .arg(kb)
16499                .arg(vh)
16500                .arg(o)
16501                .arg(&hd)
16502                .arg(&nh)
16503                .arg(&nhkv)
16504                .arg(&ti)
16505                .arg(&tkvi)
16506                .arg(&scale)
16507                .arg(&cz)
16508                .arg(&wi);
16509            unsafe {
16510                b.launch(cfg)?;
16511            }
16512            return Ok(());
16513        }
16514        let f = self.func("fa_prefill_w_bf16_p1");
16515        let shmem =
16516            (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
16517        use cudarc::driver::sys::CUfunction_attribute_enum as A;
16518        f.set_attribute(
16519            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16520            shmem as i32,
16521        )?;
16522        let cfg = LaunchConfig {
16523            grid_dim: (
16524                (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
16525                n_head as u32,
16526                1,
16527            ),
16528            block_dim: (32, 4, 1),
16529            shared_mem_bytes: shmem,
16530        };
16531        let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
16532            head_dim as i32,
16533            n_head as i32,
16534            n_head_kv as i32,
16535            t as i32,
16536            t_kv as i32,
16537            causal as i32,
16538            window as i32,
16539        );
16540        let __s_b = self.gpu.stream();
16541        let mut b = __s_b.launch_builder(&f);
16542        b.arg(qb)
16543            .arg(kb)
16544            .arg(vb)
16545            .arg(o)
16546            .arg(&hd)
16547            .arg(&nh)
16548            .arg(&nhkv)
16549            .arg(&ti)
16550            .arg(&tkvi)
16551            .arg(&scale)
16552            .arg(&cz)
16553            .arg(&wi);
16554        unsafe {
16555            b.launch(cfg)?;
16556        }
16557        Ok(())
16558    }
16559
16560    /// Windowed FA prefill with the stage arm FORCED — the kernel_check bit-identity entry.
16561    #[allow(clippy::too_many_arguments)]
16562    pub fn fa_prefill_w_arm(
16563        &self,
16564        q: &CudaSlice<f32>,
16565        k: &CudaSlice<f32>,
16566        v: &CudaSlice<f32>,
16567        o: &mut CudaSlice<f32>,
16568        head_dim: usize,
16569        n_head: usize,
16570        n_head_kv: usize,
16571        t: usize,
16572        t_kv: usize,
16573        scale: f32,
16574        causal: bool,
16575        window: usize,
16576        f32_stage: bool,
16577        floor: bool,
16578    ) -> Result<(), Box<dyn std::error::Error>> {
16579        const BLOCK_Q: usize = 64;
16580        const BK: usize = 32;
16581        debug_assert_eq!(head_dim, 256, "fa_prefill_w is stamped hd256 only");
16582        // P1 (2026-07-22 engine study): per-head Br=64 stamp with the FA2 schedule (V-copy
16583        // over GEMM0, next-K over softmax+GEMM1) + boundary/interior mask split. FP order
16584        // preserved -> bit-identical (gated). MEMRA_FAW_P1=0 reverts to the g4/o2 arms.
16585        static P1_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16586        let p1 = !floor
16587            && !f32_stage
16588            && *P1_ON.get_or_init(|| {
16589                std::env::var("MEMRA_FAW_P1")
16590                    .map(|v| v != "0")
16591                    .unwrap_or(true)
16592            });
16593        let hp =
16594            p1 && fa_f16pv_on() && faw_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
16595        if hp {
16596            const BLOCK_QH: usize = 32;
16597            let f = self.func("fa_prefill_w_bf16_p1h2");
16598            let shmem = (2 * (2 * BK * head_dim + 2 * BLOCK_QH * BK) + 4 * (2 * BLOCK_QH)) as u32;
16599            use cudarc::driver::sys::CUfunction_attribute_enum as A;
16600            f.set_attribute(
16601                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16602                shmem as i32,
16603            )?;
16604            let cfg = LaunchConfig {
16605                grid_dim: ((t as u32).div_ceil(BLOCK_QH as u32), (n_head / 2) as u32, 1),
16606                block_dim: (32, 4, 1),
16607                shared_mem_bytes: shmem,
16608            };
16609            let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
16610                head_dim as i32,
16611                n_head as i32,
16612                n_head_kv as i32,
16613                t as i32,
16614                t_kv as i32,
16615                causal as i32,
16616                window as i32,
16617            );
16618            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
16619            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
16620            let vh = self.f32_to_f16(v, t_kv * n_head_kv * head_dim)?;
16621            let __s_b = self.gpu.stream();
16622            let mut b = __s_b.launch_builder(&f);
16623            b.arg(&qb)
16624                .arg(&kb)
16625                .arg(&vh)
16626                .arg(o)
16627                .arg(&hd)
16628                .arg(&nh)
16629                .arg(&nhkv)
16630                .arg(&ti)
16631                .arg(&tkvi)
16632                .arg(&scale)
16633                .arg(&cz)
16634                .arg(&wi);
16635            unsafe {
16636                b.launch(cfg)?;
16637            }
16638            return Ok(());
16639        }
16640        if p1 {
16641            let f = self.func("fa_prefill_w_bf16_p1");
16642            let shmem =
16643                (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
16644            use cudarc::driver::sys::CUfunction_attribute_enum as A;
16645            f.set_attribute(
16646                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16647                shmem as i32,
16648            )?;
16649            let cfg = LaunchConfig {
16650                grid_dim: (
16651                    (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
16652                    n_head as u32,
16653                    1,
16654                ),
16655                block_dim: (32, 4, 1),
16656                shared_mem_bytes: shmem,
16657            };
16658            let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
16659                head_dim as i32,
16660                n_head as i32,
16661                n_head_kv as i32,
16662                t as i32,
16663                t_kv as i32,
16664                causal as i32,
16665                window as i32,
16666            );
16667            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
16668            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
16669            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
16670            let __s_b = self.gpu.stream();
16671            let mut b = __s_b.launch_builder(&f);
16672            b.arg(&qb)
16673                .arg(&kb)
16674                .arg(&vb)
16675                .arg(o)
16676                .arg(&hd)
16677                .arg(&nh)
16678                .arg(&nhkv)
16679                .arg(&ti)
16680                .arg(&tkvi)
16681                .arg(&scale)
16682                .arg(&cz)
16683                .arg(&wi);
16684            unsafe {
16685                b.launch(cfg)?;
16686            }
16687            return Ok(());
16688        }
16689        // MQA head-grouping (MEMRA_FAW_G4=0 reverts): 4 heads/CTA share the staged K/V —
16690        // per-(head,row) FP chain identical to the per-head stamp -> bit-identical (gated).
16691        static G4_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16692        let g4 = !floor
16693            && !f32_stage
16694            && n_head_kv == 1
16695            && n_head % 4 == 0
16696            && *G4_ON.get_or_init(|| {
16697                std::env::var("MEMRA_FAW_G4")
16698                    .map(|v| v != "0")
16699                    .unwrap_or(true)
16700            });
16701        if g4 {
16702            const SP_M: usize = 16;
16703            // Occupancy-2 twin (MEMRA_FAW_O2=0 reverts): one shared K/V buffer inside the dead
16704            // Q-stage region -> ~36.5KB smem, 2 CTA/SM (the llama hd256 mechanism). Bit-identical.
16705            static O2_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16706            let o2 = *O2_ON.get_or_init(|| {
16707                std::env::var("MEMRA_FAW_O2")
16708                    .map(|v| v != "0")
16709                    .unwrap_or(true)
16710            });
16711            let f = self.func(if o2 {
16712                "fa_prefill_w_bf16_g4o2"
16713            } else {
16714                "fa_prefill_w_bf16_g4"
16715            });
16716            let shmem = if o2 {
16717                (2 * (4 * SP_M * head_dim + 4 * SP_M * BK) + 4 * (4 * SP_M)) as u32
16718            } else {
16719                (2 * (2 * BK * head_dim + 4 * SP_M * head_dim + 4 * SP_M * BK) + 4 * (4 * SP_M))
16720                    as u32
16721            };
16722            use cudarc::driver::sys::CUfunction_attribute_enum as A;
16723            f.set_attribute(
16724                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16725                shmem as i32,
16726            )?;
16727            let cfg = LaunchConfig {
16728                grid_dim: ((t as u32).div_ceil(SP_M as u32), (n_head / 4) as u32, 1),
16729                block_dim: (32, 4, 1),
16730                shared_mem_bytes: shmem,
16731            };
16732            let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
16733                head_dim as i32,
16734                n_head as i32,
16735                n_head_kv as i32,
16736                t as i32,
16737                t_kv as i32,
16738                causal as i32,
16739                window as i32,
16740            );
16741            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
16742            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
16743            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
16744            let __s_b = self.gpu.stream();
16745            let mut b = __s_b.launch_builder(&f);
16746            b.arg(&qb)
16747                .arg(&kb)
16748                .arg(&vb)
16749                .arg(o)
16750                .arg(&hd)
16751                .arg(&nh)
16752                .arg(&nhkv)
16753                .arg(&ti)
16754                .arg(&tkvi)
16755                .arg(&scale)
16756                .arg(&cz)
16757                .arg(&wi);
16758            unsafe {
16759                b.launch(cfg)?;
16760            }
16761            return Ok(());
16762        }
16763        let f = self.func(if floor {
16764            "fa_prefill_w_f32"
16765        } else if f32_stage {
16766            "fa_prefill_w_f32_pp"
16767        } else {
16768            "fa_prefill_w_bf16_pp"
16769        });
16770        let shmem =
16771            (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
16772        use cudarc::driver::sys::CUfunction_attribute_enum as A;
16773        f.set_attribute(
16774            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16775            shmem as i32,
16776        )?;
16777        let cfg = LaunchConfig {
16778            grid_dim: (
16779                (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
16780                n_head as u32,
16781                1,
16782            ),
16783            block_dim: (32, 4, 1),
16784            shared_mem_bytes: shmem,
16785        };
16786        let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
16787            head_dim as i32,
16788            n_head as i32,
16789            n_head_kv as i32,
16790            t as i32,
16791            t_kv as i32,
16792            causal as i32,
16793            window as i32,
16794        );
16795        if f32_stage {
16796            let __s_b = self.gpu.stream();
16797            let mut b = __s_b.launch_builder(&f);
16798            b.arg(q)
16799                .arg(k)
16800                .arg(v)
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        } else {
16814            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
16815            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
16816            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
16817            let __s_b = self.gpu.stream();
16818            let mut b = __s_b.launch_builder(&f);
16819            b.arg(&qb)
16820                .arg(&kb)
16821                .arg(&vb)
16822                .arg(o)
16823                .arg(&hd)
16824                .arg(&nh)
16825                .arg(&nhkv)
16826                .arg(&ti)
16827                .arg(&tkvi)
16828                .arg(&scale)
16829                .arg(&cz)
16830                .arg(&wi);
16831            unsafe {
16832                b.launch(cfg)?;
16833            }
16834        }
16835        Ok(())
16836    }
16837
16838    /// hd512 FA prefill (gemma4 GLOBAL layers): BLOCK_Q=32 x 2 warps, Q staged in smem,
16839    /// grid.z = 2 O-halves (each CTA computes the full 512-dim scores, accumulates half the
16840    /// V dims). Replaces the scalar sdpa_naive on the prime path's globals.
16841    #[allow(clippy::too_many_arguments)]
16842    pub fn fa_prefill_hd512(
16843        &self,
16844        q: &CudaSlice<f32>,
16845        k: &CudaSlice<f32>,
16846        v: &CudaSlice<f32>,
16847        o: &mut CudaSlice<f32>,
16848        head_dim: usize,
16849        n_head: usize,
16850        n_head_kv: usize,
16851        t: usize,
16852        t_kv: usize,
16853        scale: f32,
16854        causal: bool,
16855    ) -> Result<(), Box<dyn std::error::Error>> {
16856        // sm_90a rides the mma twins (portable_mma_gated, 2026-07-31 — same flip as _w).
16857        if portable_mma_gated() {
16858            return self.sdpa_naive(
16859                q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
16860            );
16861        }
16862        // Default: pre-convert Q/K/V to bf16 once and stage int4 (8 bf16/copy) — at 1 CTA/SM the
16863        // synchronous stage serializes with compute and MQA re-stages the same K/V per head CTA;
16864        // pre-converting halves staged bytes and cuts stage instructions 8x. BIT-IDENTICAL to the
16865        // f32-staged kernel (the converter applies the same __float2bfloat16 the stage applied;
16866        // kernel_check gates the identity). MEMRA_FA512_STAGE=f32 = rollback to the f32 kernel.
16867        static F32_STAGE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16868        let f32_stage =
16869            *F32_STAGE.get_or_init(|| std::env::var("MEMRA_FA512_STAGE").as_deref() == Ok("f32"));
16870        // Single-pass arm (MEMRA_FA512_SP=0 reverts to the z=2 bf16 kernel): GEMM0 split-K across
16871        // the 2 warps instead of recomputed per O-half CTA — the 2026-07-22 kernel-diff excess.
16872        // Own numeric config (partial-sum order) — battery-gated.
16873        static SP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16874        let sp = !f32_stage
16875            && *SP_ON.get_or_init(|| {
16876                std::env::var("MEMRA_FA512_SP")
16877                    .map(|v| v != "0")
16878                    .unwrap_or(true)
16879            });
16880        self.fa_prefill_hd512_arm(
16881            q,
16882            k,
16883            v,
16884            o,
16885            head_dim,
16886            n_head,
16887            n_head_kv,
16888            t,
16889            t_kv,
16890            scale,
16891            causal,
16892            f32_stage,
16893            sp,
16894            sp && fa_f16pv_on(),
16895        )
16896    }
16897
16898    /// hd512 single-pass FA with PRE-CONVERTED bf16 operands (producer-emitted).
16899    #[allow(clippy::too_many_arguments)]
16900    pub fn fa_prefill_hd512_pre(
16901        &self,
16902        qb: &CudaSlice<u8>,
16903        kb: &CudaSlice<u8>,
16904        vb: &CudaSlice<u8>,
16905        o: &mut CudaSlice<f32>,
16906        head_dim: usize,
16907        n_head: usize,
16908        n_head_kv: usize,
16909        t: usize,
16910        t_kv: usize,
16911        scale: f32,
16912        causal: bool,
16913        v_f16: bool,
16914    ) -> Result<(), Box<dyn std::error::Error>> {
16915        debug_assert_eq!(head_dim, 512);
16916        const SP_M: usize = 16;
16917        const BKS: usize = 32;
16918        // f16-P/V door (MEMRA_FA_F16PV=1): P and the P@V accumulation in f16 (llama's fa=1 VKQ
16919        // class); KQ/softmax/rescale-band/final-normalize stay f32. Own numeric config —
16920        // battery-gated. V bytes must be f16 for the sp16 kernel (stage/ldmatrix are typeless).
16921        let f16pv = fa_f16pv_on();
16922        let nw = if f16pv { fa512_wide_warps() } else { 2 };
16923        let hp = f16pv && fa512_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
16924        debug_assert!(!v_f16 || f16pv, "f16 V emitted without the door on");
16925        let mut vguard = self.fa_vf16_scratch.lock().unwrap();
16926        let vref: &CudaSlice<u8> = if f16pv && !v_f16 {
16927            // Fallback re-encode (producer emitted bf16); the emit lane normally hands f16.
16928            let n = t_kv * n_head_kv * head_dim;
16929            let need = n * 2;
16930            if vguard.as_ref().map(|b| b.len() < need).unwrap_or(true) {
16931                *vguard = Some(self.alloc_uninit::<u8>(need)?);
16932            }
16933            let dst = vguard.as_mut().unwrap();
16934            self.bf16_to_f16_into(vb, n, dst)?;
16935            vguard.as_ref().unwrap()
16936        } else {
16937            vb
16938        };
16939        let f = self.func(if hp {
16940            "fa_prefill_bf16_hd512_sp16h2"
16941        } else {
16942            match (f16pv, nw) {
16943                (true, 4) => "fa_prefill_bf16_hd512_sp16w4",
16944                (true, _) => "fa_prefill_bf16_hd512_sp16",
16945                _ => "fa_prefill_bf16_hd512_sp",
16946            }
16947        });
16948        let (nwarp, npart) = if hp {
16949            (4usize, 4usize)
16950        } else if nw > 2 {
16951            (nw, nw)
16952        } else {
16953            (2, 1)
16954        };
16955        // h2 drops sQ (Q register-resident) and doubles sP/sS/sL for the head pair.
16956        let shmem = if hp {
16957            (2 * (2 * BKS * head_dim + 2 * SP_M * BKS) + 4 * (2 * npart * SP_M * BKS + 2 * SP_M))
16958                as u32
16959        } else {
16960            (2 * (SP_M * head_dim + 2 * BKS * head_dim + SP_M * BKS)
16961                + 4 * (npart * SP_M * BKS + SP_M)) as u32
16962        };
16963        use cudarc::driver::sys::CUfunction_attribute_enum as A;
16964        f.set_attribute(
16965            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16966            shmem as i32,
16967        )?;
16968        let grid_y = if hp {
16969            (n_head / 2) as u32
16970        } else {
16971            n_head as u32
16972        };
16973        let cfg = LaunchConfig {
16974            grid_dim: ((t as u32).div_ceil(SP_M as u32), grid_y, 1),
16975            block_dim: (32, nwarp as u32, 1),
16976            shared_mem_bytes: shmem,
16977        };
16978        let (hd, nh, nhkv, ti, tkvi, cz) = (
16979            head_dim as i32,
16980            n_head as i32,
16981            n_head_kv as i32,
16982            t as i32,
16983            t_kv as i32,
16984            causal as i32,
16985        );
16986        let __s_b = self.gpu.stream();
16987        let mut b = __s_b.launch_builder(&f);
16988        b.arg(qb)
16989            .arg(kb)
16990            .arg(vref)
16991            .arg(o)
16992            .arg(&hd)
16993            .arg(&nh)
16994            .arg(&nhkv)
16995            .arg(&ti)
16996            .arg(&tkvi)
16997            .arg(&scale)
16998            .arg(&cz);
16999        unsafe {
17000            b.launch(cfg)?;
17001        }
17002        Ok(())
17003    }
17004
17005    /// hd512 FA prefill with the stage/sp arms FORCED — the kernel_check gate entry
17006    /// (`fa_prefill_hd512` picks the arms from MEMRA_FA512_STAGE / MEMRA_FA512_SP).
17007    #[allow(clippy::too_many_arguments)]
17008    pub fn fa_prefill_hd512_arm(
17009        &self,
17010        q: &CudaSlice<f32>,
17011        k: &CudaSlice<f32>,
17012        v: &CudaSlice<f32>,
17013        o: &mut CudaSlice<f32>,
17014        head_dim: usize,
17015        n_head: usize,
17016        n_head_kv: usize,
17017        t: usize,
17018        t_kv: usize,
17019        scale: f32,
17020        causal: bool,
17021        f32_stage: bool,
17022        sp: bool,
17023        f16pv: bool,
17024    ) -> Result<(), Box<dyn std::error::Error>> {
17025        debug_assert_eq!(head_dim, 512, "fa_prefill_hd512 is hd512 only");
17026        if sp && !f32_stage {
17027            // Single-pass: 16 q-rows/CTA, 2 warps, grid (ceil(T/16), n_head, 1).
17028            // smem: sQ[16][512] + sK[32][512] + sV[32][512] + sP[16][32] (bf16) + sS[16][32]+sL f32.
17029            // f16pv: sp16 kernel — f16 P + f16 P@V accum, V operand encoded f16.
17030            const SP_M: usize = 16;
17031            const BKS: usize = 32;
17032            let nw = if f16pv { fa512_wide_warps() } else { 2 };
17033            let hp = f16pv && fa512_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
17034            let f = self.func(if hp {
17035                "fa_prefill_bf16_hd512_sp16h2"
17036            } else {
17037                match (f16pv, nw) {
17038                    (true, 4) => "fa_prefill_bf16_hd512_sp16w4",
17039                    (true, _) => "fa_prefill_bf16_hd512_sp16",
17040                    _ => "fa_prefill_bf16_hd512_sp",
17041                }
17042            });
17043            let (nwarp, npart) = if hp {
17044                (4usize, 4usize)
17045            } else if nw > 2 {
17046                (nw, nw)
17047            } else {
17048                (2, 1)
17049            };
17050            let shmem = if hp {
17051                (2 * (2 * BKS * head_dim + 2 * SP_M * BKS)
17052                    + 4 * (2 * npart * SP_M * BKS + 2 * SP_M)) as u32
17053            } else {
17054                (2 * (SP_M * head_dim + 2 * BKS * head_dim + SP_M * BKS)
17055                    + 4 * (npart * SP_M * BKS + SP_M)) as u32
17056            };
17057            use cudarc::driver::sys::CUfunction_attribute_enum as A;
17058            f.set_attribute(
17059                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
17060                shmem as i32,
17061            )?;
17062            let grid_y = if hp {
17063                (n_head / 2) as u32
17064            } else {
17065                n_head as u32
17066            };
17067            let cfg = LaunchConfig {
17068                grid_dim: ((t as u32).div_ceil(SP_M as u32), grid_y, 1),
17069                block_dim: (32, nwarp as u32, 1),
17070                shared_mem_bytes: shmem,
17071            };
17072            let (hd, nh, nhkv, ti, tkvi, cz) = (
17073                head_dim as i32,
17074                n_head as i32,
17075                n_head_kv as i32,
17076                t as i32,
17077                t_kv as i32,
17078                causal as i32,
17079            );
17080            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
17081            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
17082            let vb = if f16pv {
17083                self.f32_to_f16(v, t_kv * n_head_kv * head_dim)?
17084            } else {
17085                self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?
17086            };
17087            let __s_b = self.gpu.stream();
17088            let mut b = __s_b.launch_builder(&f);
17089            b.arg(&qb)
17090                .arg(&kb)
17091                .arg(&vb)
17092                .arg(o)
17093                .arg(&hd)
17094                .arg(&nh)
17095                .arg(&nhkv)
17096                .arg(&ti)
17097                .arg(&tkvi)
17098                .arg(&scale)
17099                .arg(&cz);
17100            unsafe {
17101                b.launch(cfg)?;
17102            }
17103            return Ok(());
17104        }
17105        const BLOCK_Q: usize = 32;
17106        const BK: usize = 32;
17107        const HALF: usize = 256;
17108        let f = self.func(if f32_stage {
17109            "fa_prefill_f32_hd512"
17110        } else {
17111            "fa_prefill_bf16_hd512"
17112        });
17113        // sQ[32][512] + sK[BK][512] + sV[BK][256] + sP[32][BK] (bf16) + sL[32] f32
17114        let shmem = (2 * (BLOCK_Q * head_dim + BK * head_dim + BK * HALF + BLOCK_Q * BK)
17115            + 4 * BLOCK_Q) as u32;
17116        use cudarc::driver::sys::CUfunction_attribute_enum as A;
17117        f.set_attribute(
17118            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
17119            shmem as i32,
17120        )?;
17121        let cfg = LaunchConfig {
17122            grid_dim: (
17123                (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
17124                n_head as u32,
17125                2,
17126            ),
17127            block_dim: (32, 2, 1),
17128            shared_mem_bytes: shmem,
17129        };
17130        let (hd, nh, nhkv, ti, tkvi, cz) = (
17131            head_dim as i32,
17132            n_head as i32,
17133            n_head_kv as i32,
17134            t as i32,
17135            t_kv as i32,
17136            causal as i32,
17137        );
17138        if f32_stage {
17139            let __s_b = self.gpu.stream();
17140            let mut b = __s_b.launch_builder(&f);
17141            b.arg(q)
17142                .arg(k)
17143                .arg(v)
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        } else {
17156            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
17157            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
17158            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
17159            let __s_b = self.gpu.stream();
17160            let mut b = __s_b.launch_builder(&f);
17161            b.arg(&qb)
17162                .arg(&kb)
17163                .arg(&vb)
17164                .arg(o)
17165                .arg(&hd)
17166                .arg(&nh)
17167                .arg(&nhkv)
17168                .arg(&ti)
17169                .arg(&tkvi)
17170                .arg(&scale)
17171                .arg(&cz);
17172            unsafe {
17173                b.launch(cfg)?;
17174            }
17175        }
17176        Ok(())
17177    }
17178
17179    /// rope_neox2 with bf16 EMIT (31B glue lane): identical rope math/stores plus the post-rope
17180    /// values written as bf16 — the FA q/k operands come from this launch (bit-identical to the
17181    /// separate f32_to_bf16 the FA entries would run).
17182    #[allow(clippy::too_many_arguments)]
17183    pub fn rope_neox2_bf16e(
17184        &self,
17185        q: &mut CudaSlice<f32>,
17186        k: &mut CudaSlice<f32>,
17187        qb: &mut CudaSlice<u8>,
17188        kb: &mut CudaSlice<u8>,
17189        pos: &CudaSlice<i32>,
17190        head_dim: usize,
17191        n_dims: usize,
17192        nh_q: usize,
17193        nh_k: usize,
17194        n_tokens: usize,
17195        base: f32,
17196        freq_scale: f32,
17197        ff: Option<&CudaSlice<f32>>,
17198    ) -> Result<(), Box<dyn std::error::Error>> {
17199        let f = self.func("rope_neox2_bf16e_f32");
17200        let rows = ((nh_q + nh_k) * n_tokens) as u32;
17201        let cfg = LaunchConfig {
17202            grid_dim: (rows, 1, 1),
17203            block_dim: ((head_dim / 2) as u32, 1, 1),
17204            shared_mem_bytes: 0,
17205        };
17206        let theta_scale = base.powf(-2.0 / n_dims as f32);
17207        let (hd, nd, nhq, nhk, nt) = (
17208            head_dim as i32,
17209            n_dims as i32,
17210            nh_q as i32,
17211            nh_k as i32,
17212            n_tokens as i32,
17213        );
17214        let __s_b = self.gpu.stream();
17215        let mut b = __s_b.launch_builder(&f);
17216        match ff {
17217            Some(t) => {
17218                b.arg(&mut *q)
17219                    .arg(&mut *k)
17220                    .arg(&mut *qb)
17221                    .arg(&mut *kb)
17222                    .arg(pos)
17223                    .arg(&hd)
17224                    .arg(&nd)
17225                    .arg(&nhq)
17226                    .arg(&nhk)
17227                    .arg(&nt)
17228                    .arg(&theta_scale)
17229                    .arg(&freq_scale)
17230                    .arg(t);
17231                unsafe {
17232                    b.launch(cfg)?;
17233                }
17234            }
17235            None => {
17236                let null: u64 = 0;
17237                b.arg(&mut *q)
17238                    .arg(&mut *k)
17239                    .arg(&mut *qb)
17240                    .arg(&mut *kb)
17241                    .arg(pos)
17242                    .arg(&hd)
17243                    .arg(&nd)
17244                    .arg(&nhq)
17245                    .arg(&nhk)
17246                    .arg(&nt)
17247                    .arg(&theta_scale)
17248                    .arg(&freq_scale)
17249                    .arg(&null);
17250                unsafe {
17251                    b.launch(cfg)?;
17252                }
17253            }
17254        }
17255        Ok(())
17256    }
17257
17258    /// Flat f32 -> bf16 conversion into a fresh scratch buffer (2 bytes/elem). `n % 4 == 0`
17259    /// (float4 in, 4x bf16 out). Feeds the bf16-staged hd512 FA prefill.
17260    pub fn f32_to_bf16(
17261        &self,
17262        x: &CudaSlice<f32>,
17263        n: usize,
17264    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
17265        assert!(n % 4 == 0, "f32_to_bf16 requires n % 4 == 0, got {n}");
17266        let mut y = self.alloc_uninit::<u8>(n * 2)?;
17267        let f = self.func("f32_to_bf16_flat");
17268        let n_i = n as i64;
17269        let cfg = LaunchConfig {
17270            grid_dim: (((n / 4) as u32).div_ceil(256), 1, 1),
17271            block_dim: (256, 1, 1),
17272            shared_mem_bytes: 0,
17273        };
17274        let __s_b = self.gpu.stream();
17275        let mut b = __s_b.launch_builder(&f);
17276        b.arg(x).arg(&mut y).arg(&n_i);
17277        unsafe {
17278            b.launch(cfg)?;
17279        }
17280        Ok(y)
17281    }
17282
17283    pub fn f32_to_f16(
17284        &self,
17285        x: &CudaSlice<f32>,
17286        n: usize,
17287    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
17288        assert!(n % 4 == 0, "f32_to_f16 requires n % 4 == 0, got {n}");
17289        let mut y = self.alloc_uninit::<u8>(n * 2)?;
17290        let f = self.func("f32_to_f16_flat");
17291        let n_i = n as i64;
17292        let cfg = LaunchConfig {
17293            grid_dim: (((n / 4) as u32).div_ceil(256), 1, 1),
17294            block_dim: (256, 1, 1),
17295            shared_mem_bytes: 0,
17296        };
17297        let __s_b = self.gpu.stream();
17298        let mut b = __s_b.launch_builder(&f);
17299        b.arg(x).arg(&mut y).arg(&n_i);
17300        unsafe {
17301            b.launch(cfg)?;
17302        }
17303        Ok(y)
17304    }
17305
17306    /// bf16 bytes -> f16 bytes, n elements (the f16-P/V door's V re-encode on the emit lane).
17307    pub fn bf16_to_f16(
17308        &self,
17309        xb: &CudaSlice<u8>,
17310        n: usize,
17311    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
17312        let mut y = self.alloc_uninit::<u8>(n * 2)?;
17313        self.bf16_to_f16_into(xb, n, &mut y)?;
17314        Ok(y)
17315    }
17316
17317    /// Same conversion into a caller-owned (pooled) buffer; `y.len() >= n*2`.
17318    pub fn bf16_to_f16_into(
17319        &self,
17320        xb: &CudaSlice<u8>,
17321        n: usize,
17322        y: &mut CudaSlice<u8>,
17323    ) -> Result<(), Box<dyn std::error::Error>> {
17324        assert!(n % 2 == 0, "bf16_to_f16 requires n % 2 == 0, got {n}");
17325        assert!(y.len() >= n * 2);
17326        let f = self.func("bf16_to_f16_flat");
17327        let n2 = (n / 2) as i64;
17328        let cfg = LaunchConfig {
17329            grid_dim: (((n / 2) as u32).div_ceil(256), 1, 1),
17330            block_dim: (256, 1, 1),
17331            shared_mem_bytes: 0,
17332        };
17333        let __s_b = self.gpu.stream();
17334        let mut b = __s_b.launch_builder(&f);
17335        b.arg(xb).arg(y).arg(&n2);
17336        unsafe {
17337            b.launch(cfg)?;
17338        }
17339        Ok(())
17340    }
17341
17342    /// task #18 (attn side): varlen FA — bf16 K/V mirrors (2 launches) + ONE
17343    /// fa_prefill_bf16kv launch for every fresh sequence. Same per-block math as the
17344    /// per-seq path (bit-gateable). Caller guarantees: fresh causal (T_kv == T),
17345    /// head_dim in {256, 128}, bf16kv lane on.
17346    #[allow(clippy::too_many_arguments)]
17347    pub fn fa_prefill_vl8(
17348        &self,
17349        seqs: &[FaSeqVl],
17350        head_dim: usize,
17351        n_head: usize,
17352        n_head_kv: usize,
17353        scale: f32,
17354    ) -> Result<(), Box<dyn std::error::Error>> {
17355        const BK: usize = 32;
17356        let b = seqs.len();
17357        assert!(b >= 1 && b <= 8);
17358        let mut packed = [FaSeqVl::default(); 8];
17359        packed[..b].copy_from_slice(seqs);
17360        let v = FaVl8(packed);
17361        let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
17362        let ept = (n_head_kv * head_dim) as i32;
17363        {
17364            let f = self.func("fa_mirror_vl");
17365            let max_n = (max_t as i64) * ept as i64;
17366            let blocks = ((max_n as u32).div_ceil(4)).div_ceil(256);
17367            for which in 0..2i32 {
17368                let cfg = LaunchConfig {
17369                    grid_dim: (blocks, 1, b as u32),
17370                    block_dim: (256, 1, 1),
17371                    shared_mem_bytes: 0,
17372                };
17373                let __s_lb = self.gpu.stream();
17374                let mut lb = __s_lb.launch_builder(&f);
17375                lb.arg(&v).arg(&ept).arg(&which);
17376                unsafe {
17377                    lb.launch(cfg)?;
17378                }
17379            }
17380        }
17381        let hd_sfx = fa_hd_suffix(head_dim)?;
17382        let f = self.func(&format!("fa_prefill_bf16kv_vl{hd_sfx}"));
17383        let block_q = 64usize;
17384        let kv_stages = 2usize;
17385        let shmem = (2 * (kv_stages * 2 * BK * head_dim + block_q * BK)
17386            + 4 * (block_q * BK + 2 * block_q)) as u32;
17387        use cudarc::driver::sys::CUfunction_attribute_enum as A;
17388        f.set_attribute(
17389            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
17390            shmem as i32,
17391        )?;
17392        let cfg = LaunchConfig {
17393            grid_dim: (max_t.div_ceil(block_q as u32), n_head as u32, b as u32),
17394            block_dim: (32, 4, 1),
17395            shared_mem_bytes: shmem,
17396        };
17397        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
17398        let __s_lb = self.gpu.stream();
17399        let mut lb = __s_lb.launch_builder(&f);
17400        lb.arg(&v).arg(&hd).arg(&nh).arg(&nhkv).arg(&scale);
17401        unsafe {
17402            lb.launch(cfg)?;
17403        }
17404        Ok(())
17405    }
17406
17407    /// task #18 (attn pre-FA): varlen split + QK-norm + RoPE + KV-append — FOUR launches
17408    /// for every fresh sequence (was 6 x B, plus the q/k/v split copies which the view
17409    /// inputs remove entirely). Fresh-only (append at t0=0, RoPE pos = token index).
17410    #[allow(clippy::too_many_arguments)]
17411    pub fn attn_pre_vl8(
17412        &self,
17413        seqs: &[AttnPreVl],
17414        wq: &CudaSlice<f32>,
17415        wk: &CudaSlice<f32>,
17416        head_dim: usize,
17417        rope_dims: usize,
17418        n_head: usize,
17419        n_head_kv: usize,
17420        eps: f32,
17421        freq_base: f32,
17422        freq_scale: f32,
17423        kv_dim_k: usize,
17424        kv_dim_v: usize,
17425        k_tok_bytes: usize,
17426        v_tok_bytes: usize,
17427    ) -> Result<(), Box<dyn std::error::Error>> {
17428        let b = seqs.len();
17429        assert!(b >= 1 && b <= 8);
17430        let mut packed = [AttnPreVl::default(); 8];
17431        packed[..b].copy_from_slice(seqs);
17432        let v = AttnPreVl8(packed);
17433        let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
17434        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
17435        {
17436            let f = self.func("q_gate_split_vl");
17437            let n = max_t * (n_head * head_dim) as u32;
17438            let cfg = LaunchConfig {
17439                grid_dim: (n.div_ceil(256), 1, b as u32),
17440                block_dim: (256, 1, 1),
17441                shared_mem_bytes: 0,
17442            };
17443            let __s_lb = self.gpu.stream();
17444            let mut lb = __s_lb.launch_builder(&f);
17445            lb.arg(&v).arg(&hd).arg(&nh);
17446            unsafe {
17447                lb.launch(cfg)?;
17448            }
17449        }
17450        {
17451            let f = self.func("attn_rms_vl");
17452            let cfg = LaunchConfig {
17453                grid_dim: (max_t * n_head as u32, 2, b as u32),
17454                block_dim: (rms_block(), 1, 1),
17455                shared_mem_bytes: 0,
17456            };
17457            let __s_lb = self.gpu.stream();
17458            let mut lb = __s_lb.launch_builder(&f);
17459            lb.arg(&v)
17460                .arg(wq)
17461                .arg(wk)
17462                .arg(&hd)
17463                .arg(&nh)
17464                .arg(&nhkv)
17465                .arg(&eps);
17466            unsafe {
17467                lb.launch(cfg)?;
17468            }
17469        }
17470        {
17471            let f = self.func("attn_rope_vl");
17472            let theta_scale = freq_base.powf(-2.0 / rope_dims as f32);
17473            let nd = rope_dims as i32;
17474            let cfg = LaunchConfig {
17475                grid_dim: (max_t * n_head as u32, 2, b as u32),
17476                block_dim: ((head_dim / 2) as u32, 1, 1),
17477                shared_mem_bytes: 0,
17478            };
17479            let __s_lb = self.gpu.stream();
17480            let mut lb = __s_lb.launch_builder(&f);
17481            lb.arg(&v)
17482                .arg(&hd)
17483                .arg(&nd)
17484                .arg(&nh)
17485                .arg(&nhkv)
17486                .arg(&theta_scale)
17487                .arg(&freq_scale);
17488            unsafe {
17489                lb.launch(cfg)?;
17490            }
17491        }
17492        {
17493            let f = self.func("append_kv_vl");
17494            let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
17495            let cfg = LaunchConfig {
17496                grid_dim: (nblk, max_t, b as u32),
17497                block_dim: (32, 1, 1),
17498                shared_mem_bytes: 0,
17499            };
17500            let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
17501            let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
17502            let __s_lb = self.gpu.stream();
17503            let mut lb = __s_lb.launch_builder(&f);
17504            lb.arg(&v).arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
17505            unsafe {
17506                lb.launch(cfg)?;
17507            }
17508        }
17509        Ok(())
17510    }
17511
17512    /// FA prefill where K/V are QUANTIZED CudaViews into the resident byte KV cache (the T=K verify
17513    /// path, MTP-PLAN §D.3). Uses `fa_prefill_q` (inline-dequant during stage-to-smem). The view's
17514    /// base+offset pointer is honored; the kernel reads [0..t_kv*tok_bytes). Q is the T fresh query
17515    /// rows; t = T, t_kv = cache len. k_tok_bytes/v_tok_bytes are the per-token byte strides.
17516    pub fn fa_prefill_view(
17517        &self,
17518        q: &CudaSlice<f32>,
17519        k: &cudarc::driver::CudaView<u8>,
17520        v: &cudarc::driver::CudaView<u8>,
17521        o: &mut CudaSlice<f32>,
17522        head_dim: usize,
17523        n_head: usize,
17524        n_head_kv: usize,
17525        t: usize,
17526        t_kv: usize,
17527        scale: f32,
17528        causal: bool,
17529        k_tok_bytes: usize,
17530        v_tok_bytes: usize,
17531        g: bool,
17532    ) -> Result<(), Box<dyn std::error::Error>> {
17533        if portable_mma_gated() {
17534            return self.sdpa_naive_quantized_view(
17535                q,
17536                k,
17537                v,
17538                o,
17539                head_dim,
17540                n_head,
17541                n_head_kv,
17542                t,
17543                t_kv,
17544                scale,
17545                causal,
17546                k_tok_bytes,
17547                v_tok_bytes,
17548            );
17549        }
17550        const BLOCK_Q: usize = 64;
17551        const BK: usize = 32;
17552        // g = e4m3 cache: the kernel parses via DQ_K_ELEM/DQ_V_ELEM (format macros) — the
17553        // kf8vf8-module stamp reads fp8 with the identical MMA/softmax/PV body.
17554        let name = format!("fa_prefill_q{}", fa_hd_suffix(head_dim)?);
17555        let f = if g {
17556            self.func_g(&name)
17557        } else {
17558            self.func(&name)
17559        };
17560        let shmem =
17561            (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
17562        use cudarc::driver::sys::CUfunction_attribute_enum as A;
17563        f.set_attribute(
17564            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
17565            shmem as i32,
17566        )?;
17567        let cfg = LaunchConfig {
17568            grid_dim: (
17569                (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
17570                n_head as u32,
17571                1,
17572            ),
17573            block_dim: (32, 4, 1),
17574            shared_mem_bytes: shmem,
17575        };
17576        let (hd, nh, nhkv, ti, tkvi, cz) = (
17577            head_dim as i32,
17578            n_head as i32,
17579            n_head_kv as i32,
17580            t as i32,
17581            t_kv as i32,
17582            causal as i32,
17583        );
17584        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
17585        let __s_b = self.gpu.stream();
17586        let mut b = __s_b.launch_builder(&f);
17587        b.arg(q)
17588            .arg(k)
17589            .arg(v)
17590            .arg(o)
17591            .arg(&hd)
17592            .arg(&nh)
17593            .arg(&nhkv)
17594            .arg(&ti)
17595            .arg(&tkvi)
17596            .arg(&scale)
17597            .arg(&cz)
17598            .arg(&ktb)
17599            .arg(&vtb);
17600        unsafe {
17601            b.launch(cfg)?;
17602        }
17603        Ok(())
17604    }
17605
17606    /// ARC B (2026-07-05): dequant-once chunk-prime FA. Same contract as `fa_prefill_view`, but
17607    /// instead of every (q-block, head) CTA re-dequanting the whole quantized KV stream inline
17608    /// (T/64 x n_head redundant at chunk prime — 30.5% of the 32k prime wall), dequant the full
17609    /// [t_kv, kv_dim] K and V ONCE into a resident bf16 workspace (fa_dequant_kv_ws_bf16), then
17610    /// run `fa_prefill_qw` (the bf16-workspace twin) over it. EXACT: the workspace holds the same
17611    /// __float2bfloat16(dq_*_elem(...)) values fa_prefill_q stages to smem, and the twin's MMA/
17612    /// softmax/PV code is byte-identical -> bit-identical O (kernel_check pins bitdiff=0).
17613    /// The workspace allocation is REUSED across layers/chunks (grown to the largest shape);
17614    /// contents are rewritten per call. MEMRA_PRIME_DEQW=0 falls back to fa_prefill_view (callers gate).
17615    #[allow(clippy::too_many_arguments)]
17616    pub fn fa_prefill_view_ws(
17617        &self,
17618        q: &CudaSlice<f32>,
17619        k: &cudarc::driver::CudaView<u8>,
17620        v: &cudarc::driver::CudaView<u8>,
17621        o: &mut CudaSlice<f32>,
17622        head_dim: usize,
17623        n_head: usize,
17624        n_head_kv: usize,
17625        t: usize,
17626        t_kv: usize,
17627        scale: f32,
17628        causal: bool,
17629        k_tok_bytes: usize,
17630        v_tok_bytes: usize,
17631        g: bool,
17632    ) -> Result<(), Box<dyn std::error::Error>> {
17633        if portable_mma_gated() {
17634            return self.sdpa_naive_quantized_view(
17635                q,
17636                k,
17637                v,
17638                o,
17639                head_dim,
17640                n_head,
17641                n_head_kv,
17642                t,
17643                t_kv,
17644                scale,
17645                causal,
17646                k_tok_bytes,
17647                v_tok_bytes,
17648            );
17649        }
17650        const BLOCK_Q: usize = 64;
17651        const BK: usize = 32;
17652        let kv_dim_k = n_head_kv * head_dim;
17653        let kv_dim_v = n_head_kv * head_dim;
17654        let k_ws_bytes = t_kv * kv_dim_k * 2; // bf16
17655        let v_ws_bytes = t_kv * kv_dim_v * 2;
17656        // Lock held across BOTH launches: enqueue-only (µs), all compute serializes on gpu.stream.
17657        let mut guard = self.prime_deqw_ws.lock().unwrap();
17658        let need_grow = match guard.as_ref() {
17659            Some((kw, vw)) => kw.len() < k_ws_bytes || vw.len() < v_ws_bytes,
17660            None => true,
17661        };
17662        if need_grow {
17663            let grow = |cur: usize, need: usize| if cur >= need { cur } else { need };
17664            let (ck, cv) = guard
17665                .as_ref()
17666                .map(|(a, b)| (a.len(), b.len()))
17667                .unwrap_or((0, 0));
17668            *guard = Some((
17669                self.alloc_u8(grow(ck, k_ws_bytes))?,
17670                self.alloc_u8(grow(cv, v_ws_bytes))?,
17671            ));
17672        }
17673        let (kw, vw) = guard.as_mut().unwrap();
17674        // pass 1: dequant K+V once into the bf16 workspace (grid-stride, 1 thread/elem)
17675        {
17676            // only THIS pass parses KV bytes — pass 2 reads the bf16 workspace (format-free).
17677            let f = if g {
17678                self.func_g("fa_dequant_kv_ws_bf16")
17679            } else {
17680                self.func("fa_dequant_kv_ws_bf16")
17681            };
17682            let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
17683            let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
17684            let cfg = LaunchConfig {
17685                grid_dim: (nblk.max(1), 1, 1),
17686                block_dim: (256, 1, 1),
17687                shared_mem_bytes: 0,
17688            };
17689            let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
17690            let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
17691            let __s_b = self.gpu.stream();
17692            let mut b = __s_b.launch_builder(&f);
17693            b.arg(k)
17694                .arg(v)
17695                .arg(&mut *kw)
17696                .arg(&mut *vw)
17697                .arg(&kdk)
17698                .arg(&kdv)
17699                .arg(&tkvi)
17700                .arg(&ktb)
17701                .arg(&vtb);
17702            unsafe {
17703                b.launch(cfg)?;
17704            }
17705        }
17706        // pass 2: the bf16-workspace prefill twin (same tile sizes/loop structure as fa_prefill_q).
17707        // DEFAULT: cp.async double-buffered staging twin (fa_prefill_qw_db, +32KB smem for the
17708        // second K/V tile pair, 1 CTA/SM): overlaps tile n+1's L2->smem copy with tile n's MMA.
17709        // Bit-identical output (staging is a pure byte copy; kernel_check pins bitdiff=0 under
17710        // both twins). A/B (27B g7e, N=3): 32k prime 17.10->16.51s, 16k 9.09->8.65s — the copy
17711        // latency hides behind the MMA pipe and beats the 2-CTA/SM occupancy of the sync twin.
17712        // MEMRA_PRIME_DEQW_DB=0 falls back to the single-buffer twin.
17713        let db = std::env::var("MEMRA_PRIME_DEQW_DB")
17714            .map(|v| v != "0")
17715            .unwrap_or(true);
17716        {
17717            let hd_sfx = fa_hd_suffix(head_dim)?;
17718            let f = self.func(&format!(
17719                "fa_prefill_qw{}{hd_sfx}",
17720                if db { "_db" } else { "" }
17721            ));
17722            let shmem = if db {
17723                // 4x KV tile buffers (bf16) + sP (bf16) + sL (f32)
17724                (2 * (4 * BK * head_dim + BLOCK_Q * BK) + 4 * BLOCK_Q) as u32
17725            } else {
17726                (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32
17727            };
17728            use cudarc::driver::sys::CUfunction_attribute_enum as A;
17729            f.set_attribute(
17730                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
17731                shmem as i32,
17732            )?;
17733            let cfg = LaunchConfig {
17734                grid_dim: (
17735                    (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
17736                    n_head as u32,
17737                    1,
17738                ),
17739                block_dim: (32, 4, 1),
17740                shared_mem_bytes: shmem,
17741            };
17742            let (hd, nh, nhkv, ti, tkvi, cz) = (
17743                head_dim as i32,
17744                n_head as i32,
17745                n_head_kv as i32,
17746                t as i32,
17747                t_kv as i32,
17748                causal as i32,
17749            );
17750            let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
17751            let __s_b = self.gpu.stream();
17752            let mut b = __s_b.launch_builder(&f);
17753            b.arg(q)
17754                .arg(&*kw)
17755                .arg(&*vw)
17756                .arg(o)
17757                .arg(&hd)
17758                .arg(&nh)
17759                .arg(&nhkv)
17760                .arg(&ti)
17761                .arg(&tkvi)
17762                .arg(&scale)
17763                .arg(&cz)
17764                .arg(&kdk)
17765                .arg(&kdv);
17766            unsafe {
17767                b.launch(cfg)?;
17768            }
17769        }
17770        Ok(())
17771    }
17772
17773    /// WINDOWED `fa_prefill_view_ws` twin at head_dim 128 (lane/pp-prefill 2026-08-07):
17774    /// step35's SWA prefill (win=512, 33 of 45 layers) previously had NO windowed FA prefill
17775    /// stamp — every windowed twin was hd256-only — and took `sdpa_naive_w_quantized_view`,
17776    /// the f32 floor, at 565 ms/layer on a pp4096 where the hd128 FA family does the harder
17777    /// causal-4096 in 3.3 ms (41% of the whole prime; research/pp-prefill-20260807 anatomy).
17778    /// Same two-pass shape as the unwindowed function: dequant K/V ONCE into the resident
17779    /// bf16 workspace, then the windowed qw kernel (`fa_prefill_qw_db_w_hd128`, cp.async
17780    /// double-buffered; MEMRA_PRIME_DEQW_DB=0 selects the single-buffer twin). The window
17781    /// mask is `fa_prefill_f32_body`'s exact predicate; `window == 0` is bit-identical to
17782    /// `fa_prefill_view_ws` by construction (default-arg body). NEW NUMERIC CLASS vs the
17783    /// f32 floor on SWA rows (bf16 MMA online-softmax vs f32 serial softmax) — adoption is
17784    /// gated by the full battery, and the class must change UNIFORMLY for a whole request
17785    /// (kernel selection keys on seq_end, never per chunk — the chunkfix law).
17786    /// hd128-only deliberately: the only windowed-prefill consumer at another head_dim is
17787    /// gemma4 (hd256), which already has `fa_prefill_w_f32`.
17788    #[allow(clippy::too_many_arguments)]
17789    pub fn fa_prefill_view_ws_w_hd128(
17790        &self,
17791        q: &CudaSlice<f32>,
17792        k: &cudarc::driver::CudaView<u8>,
17793        v: &cudarc::driver::CudaView<u8>,
17794        o: &mut CudaSlice<f32>,
17795        head_dim: usize,
17796        n_head: usize,
17797        n_head_kv: usize,
17798        t: usize,
17799        t_kv: usize,
17800        scale: f32,
17801        causal: bool,
17802        window: usize,
17803        k_tok_bytes: usize,
17804        v_tok_bytes: usize,
17805    ) -> Result<(), Box<dyn std::error::Error>> {
17806        assert_eq!(
17807            head_dim, 128,
17808            "fa_prefill_view_ws_w_hd128: only the hd128 twin is stamped"
17809        );
17810        if portable_mma_gated() {
17811            return self.sdpa_naive_w_quantized_view(
17812                q,
17813                k,
17814                v,
17815                o,
17816                head_dim,
17817                n_head,
17818                n_head_kv,
17819                t,
17820                t_kv,
17821                scale,
17822                causal,
17823                window,
17824                k_tok_bytes,
17825                v_tok_bytes,
17826            );
17827        }
17828        const BLOCK_Q: usize = 64;
17829        const BK: usize = 32;
17830        let kv_dim_k = n_head_kv * head_dim;
17831        let kv_dim_v = n_head_kv * head_dim;
17832        let k_ws_bytes = t_kv * kv_dim_k * 2; // bf16
17833        let v_ws_bytes = t_kv * kv_dim_v * 2;
17834        let mut guard = self.prime_deqw_ws.lock().unwrap();
17835        let need_grow = match guard.as_ref() {
17836            Some((kw, vw)) => kw.len() < k_ws_bytes || vw.len() < v_ws_bytes,
17837            None => true,
17838        };
17839        if need_grow {
17840            let grow = |cur: usize, need: usize| if cur >= need { cur } else { need };
17841            let (ck, cv) = guard
17842                .as_ref()
17843                .map(|(a, b)| (a.len(), b.len()))
17844                .unwrap_or((0, 0));
17845            *guard = Some((
17846                self.alloc_u8(grow(ck, k_ws_bytes))?,
17847                self.alloc_u8(grow(cv, v_ws_bytes))?,
17848            ));
17849        }
17850        let (kw, vw) = guard.as_mut().unwrap();
17851        // pass 1: dequant K+V once into the bf16 workspace (identical to fa_prefill_view_ws —
17852        // the workspace bytes are the SAME __float2bfloat16(dq(...)) values either way).
17853        {
17854            let f = self.func("fa_dequant_kv_ws_bf16");
17855            let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
17856            let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
17857            let cfg = LaunchConfig {
17858                grid_dim: (nblk.max(1), 1, 1),
17859                block_dim: (256, 1, 1),
17860                shared_mem_bytes: 0,
17861            };
17862            let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
17863            let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
17864            let __s_b = self.gpu.stream();
17865            let mut b = __s_b.launch_builder(&f);
17866            b.arg(k)
17867                .arg(v)
17868                .arg(&mut *kw)
17869                .arg(&mut *vw)
17870                .arg(&kdk)
17871                .arg(&kdv)
17872                .arg(&tkvi)
17873                .arg(&ktb)
17874                .arg(&vtb);
17875            unsafe {
17876                b.launch(cfg)?;
17877            }
17878        }
17879        // pass 2: the WINDOWED qw twin (db default, same as the unwindowed wrapper).
17880        let db = std::env::var("MEMRA_PRIME_DEQW_DB")
17881            .map(|v| v != "0")
17882            .unwrap_or(true);
17883        {
17884            let f = self.func(if db {
17885                "fa_prefill_qw_db_w_hd128"
17886            } else {
17887                "fa_prefill_qw_w_hd128"
17888            });
17889            let shmem = if db {
17890                (2 * (4 * BK * head_dim + BLOCK_Q * BK) + 4 * BLOCK_Q) as u32
17891            } else {
17892                (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32
17893            };
17894            use cudarc::driver::sys::CUfunction_attribute_enum as A;
17895            f.set_attribute(
17896                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
17897                shmem as i32,
17898            )?;
17899            let cfg = LaunchConfig {
17900                grid_dim: (
17901                    (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
17902                    n_head as u32,
17903                    1,
17904                ),
17905                block_dim: (32, 4, 1),
17906                shared_mem_bytes: shmem,
17907            };
17908            let (hd, nh, nhkv, ti, tkvi, cz) = (
17909                head_dim as i32,
17910                n_head as i32,
17911                n_head_kv as i32,
17912                t as i32,
17913                t_kv as i32,
17914                causal as i32,
17915            );
17916            let (kdk, kdv, wnd) = (kv_dim_k as i32, kv_dim_v as i32, window as i32);
17917            let __s_b = self.gpu.stream();
17918            let mut b = __s_b.launch_builder(&f);
17919            b.arg(q)
17920                .arg(&*kw)
17921                .arg(&*vw)
17922                .arg(o)
17923                .arg(&hd)
17924                .arg(&nh)
17925                .arg(&nhkv)
17926                .arg(&ti)
17927                .arg(&tkvi)
17928                .arg(&scale)
17929                .arg(&cz)
17930                .arg(&kdk)
17931                .arg(&kdv)
17932                .arg(&wnd);
17933            unsafe {
17934                b.launch(cfg)?;
17935            }
17936        }
17937        Ok(())
17938    }
17939
17940    /// FA decode (T=1 split-K) over the resident QUANTIZED KV cache (q8_0 K / q5_1 V) as u8 views.
17941    /// Replaces sdpa_naive_view for decode; inline-dequants per element. k_tok_bytes/v_tok_bytes are
17942    /// the per-token byte strides (differ: q8_0=34*nblk, q5_1=24*nblk per token).
17943    pub fn fa_decode(
17944        &self,
17945        q: &CudaSlice<f32>,
17946        k: &cudarc::driver::CudaView<u8>,
17947        v: &cudarc::driver::CudaView<u8>,
17948        o: &mut CudaSlice<f32>,
17949        head_dim: usize,
17950        n_head: usize,
17951        n_head_kv: usize,
17952        t_kv: usize,
17953        scale: f32,
17954        k_tok_bytes: usize,
17955        v_tok_bytes: usize,
17956    ) -> Result<(), Box<dyn std::error::Error>> {
17957        self.fa_decode_kvmod(
17958            q,
17959            k,
17960            v,
17961            o,
17962            head_dim,
17963            n_head,
17964            n_head_kv,
17965            t_kv,
17966            scale,
17967            k_tok_bytes,
17968            v_tok_bytes,
17969            false,
17970        )
17971    }
17972
17973    /// `fa_decode` with an explicit fp8-module flag (`g`): gemma windowed layers under
17974    /// MEMRA_GEMMA_WKV read an e4m3 cache — every kernel must come from the kf8vf8 module
17975    /// and the v4 lane (q8_0-hardcoded staging) is excluded.
17976    #[allow(clippy::too_many_arguments)]
17977    /// UNIFIED scalar decode launch (fa_decode_f32, nullable-ctr): ONE symbol for host-len
17978    /// (kvmod eager) and device-len (graph/stream) callers — the textually-identical f32_dc
17979    /// twin compiled apart and its ULP drift flipped 31B verify argmaxes (2026-07-12).
17980    #[allow(clippy::too_many_arguments)]
17981    #[allow(clippy::too_many_arguments)]
17982    fn fa_decode_scalar_unified(
17983        &self,
17984        q: &cudarc::driver::CudaView<f32>,
17985        k: &cudarc::driver::CudaView<u8>,
17986        v: &cudarc::driver::CudaView<u8>,
17987        o: &mut cudarc::driver::CudaViewMut<f32>,
17988        head_dim: usize,
17989        n_head: usize,
17990        n_head_kv: usize,
17991        t_kv_host: usize,
17992        t_kv_dev: Option<&CudaSlice<i32>>,
17993        scale: f32,
17994        n_splits: usize,
17995        split_keys: usize,
17996        k_tok_bytes: usize,
17997        v_tok_bytes: usize,
17998        g: bool,
17999        part_o: &mut CudaSlice<f32>,
18000        part_m: &mut CudaSlice<f32>,
18001        part_l: &mut CudaSlice<f32>,
18002        q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
18003    ) -> Result<(), Box<dyn std::error::Error>> {
18004        let f = if g {
18005            self.func_g("fa_decode_f32")
18006        } else {
18007            self.fa_func("fa_decode_f32", head_dim)
18008        };
18009        let cfg = LaunchConfig {
18010            grid_dim: (n_head as u32, n_splits as u32, 1),
18011            block_dim: (head_dim as u32, 1, 1),
18012            shared_mem_bytes: (4 * (head_dim + 32)) as u32,
18013        };
18014        let (hd, nh, nhkv, nsp) = (
18015            head_dim as i32,
18016            n_head as i32,
18017            n_head_kv as i32,
18018            n_splits as i32,
18019        );
18020        let (ktb, vtb, tkvi, ski) = (
18021            k_tok_bytes as i64,
18022            v_tok_bytes as i64,
18023            t_kv_host as i32,
18024            split_keys as i32,
18025        );
18026        let __s_b = self.gpu.stream();
18027        let mut b = __s_b.launch_builder(&f);
18028        match t_kv_dev {
18029            Some(d) => {
18030                b.arg(q)
18031                    .arg(k)
18032                    .arg(v)
18033                    .arg(&mut *part_o)
18034                    .arg(&mut *part_m)
18035                    .arg(&mut *part_l)
18036                    .arg(&hd)
18037                    .arg(&nh)
18038                    .arg(&nhkv)
18039                    .arg(&tkvi)
18040                    .arg(d)
18041                    .arg(&scale)
18042                    .arg(&nsp)
18043                    .arg(&ski)
18044                    .arg(&ktb)
18045                    .arg(&vtb);
18046                unsafe {
18047                    b.launch(cfg)?;
18048                }
18049            }
18050            None => {
18051                let null: u64 = 0;
18052                b.arg(q)
18053                    .arg(k)
18054                    .arg(v)
18055                    .arg(&mut *part_o)
18056                    .arg(&mut *part_m)
18057                    .arg(&mut *part_l)
18058                    .arg(&hd)
18059                    .arg(&nh)
18060                    .arg(&nhkv)
18061                    .arg(&tkvi)
18062                    .arg(&null)
18063                    .arg(&scale)
18064                    .arg(&nsp)
18065                    .arg(&ski)
18066                    .arg(&ktb)
18067                    .arg(&vtb);
18068                unsafe {
18069                    b.launch(cfg)?;
18070                }
18071            }
18072        }
18073        let cfg2 = LaunchConfig {
18074            grid_dim: (n_head as u32, 1, 1),
18075            block_dim: (head_dim as u32, 1, 1),
18076            shared_mem_bytes: 0,
18077        };
18078        if let Some((oq, od)) = q8_out {
18079            // wave-5b: q8-emitting combine — the wo matmul_pre consumes the pair directly.
18080            let fc = if g {
18081                self.func_g("fa_decode_combine_q8_1")
18082            } else {
18083                self.fa_func("fa_decode_combine_q8_1", head_dim)
18084            };
18085            let __s_b2 = self.gpu.stream();
18086            let mut b2 = __s_b2.launch_builder(&fc);
18087            b2.arg(&*part_o)
18088                .arg(&*part_m)
18089                .arg(&*part_l)
18090                .arg(oq)
18091                .arg(od)
18092                .arg(&hd)
18093                .arg(&nh)
18094                .arg(&nsp);
18095            unsafe {
18096                b2.launch(cfg2)?;
18097            }
18098            return Ok(());
18099        }
18100        let fc = if g {
18101            self.func_g("fa_decode_combine_f32")
18102        } else {
18103            self.fa_func("fa_decode_combine_f32", head_dim)
18104        };
18105        let __s_b2 = self.gpu.stream();
18106        let mut b2 = __s_b2.launch_builder(&fc);
18107        b2.arg(&*part_o)
18108            .arg(&*part_m)
18109            .arg(&*part_l)
18110            .arg(o)
18111            .arg(&hd)
18112            .arg(&nh)
18113            .arg(&nsp);
18114        unsafe {
18115            b2.launch(cfg2)?;
18116        }
18117        Ok(())
18118    }
18119
18120    pub fn fa_decode_kvmod(
18121        &self,
18122        q: &CudaSlice<f32>,
18123        k: &cudarc::driver::CudaView<u8>,
18124        v: &cudarc::driver::CudaView<u8>,
18125        o: &mut CudaSlice<f32>,
18126        head_dim: usize,
18127        n_head: usize,
18128        n_head_kv: usize,
18129        t_kv: usize,
18130        scale: f32,
18131        k_tok_bytes: usize,
18132        v_tok_bytes: usize,
18133        g: bool,
18134    ) -> Result<(), Box<dyn std::error::Error>> {
18135        let q_view = q.as_view();
18136        let mut o_view = o.as_view_mut();
18137        self.fa_decode_kvmod_view(
18138            &q_view,
18139            k,
18140            v,
18141            &mut o_view,
18142            head_dim,
18143            n_head,
18144            n_head_kv,
18145            t_kv,
18146            scale,
18147            k_tok_bytes,
18148            v_tok_bytes,
18149            g,
18150        )
18151    }
18152
18153    /// Row-view entry into `fa_decode_kvmod`. The kernel sees the selected Q/output rows as its
18154    /// base pointers, so the launch geometry and arithmetic are identical to the owned-slice entry.
18155    /// Batched fallback callers use this to avoid materializing rows around an otherwise unchanged
18156    /// per-session KV view and FA launch.
18157    #[allow(clippy::too_many_arguments)]
18158    pub fn fa_decode_kvmod_view(
18159        &self,
18160        q: &cudarc::driver::CudaView<f32>,
18161        k: &cudarc::driver::CudaView<u8>,
18162        v: &cudarc::driver::CudaView<u8>,
18163        o: &mut cudarc::driver::CudaViewMut<f32>,
18164        head_dim: usize,
18165        n_head: usize,
18166        n_head_kv: usize,
18167        t_kv: usize,
18168        scale: f32,
18169        k_tok_bytes: usize,
18170        v_tok_bytes: usize,
18171        g: bool,
18172    ) -> Result<(), Box<dyn std::error::Error>> {
18173        // PERF-4: the warp-per-token vec path replaces the scalar element-per-thread fa_decode_f32 —
18174        // warp-per-token fa_decode_vec_q (grid=(n_head_kv,n_splits), block=(32,gqa_ratio)).
18175        // The block dequants each KV tile ONCE into smem (bf16) and broadcasts to all gqa Q-head
18176        // warps -> each KV byte leaves HBM/L2 ~1x/group (vs 4x). ARGS identical; func/grid/block/
18177        // smem/n_splits differ. fa_decode_f32 stays the bit-reference fallback. Combine is shared.
18178        //
18179        // SPLIT-K: the scalar path has grid.x=n_head (32) blocks; the vec path only has
18180        // grid.x=n_head_kv (8). To avoid starving the GPU at mid ctx, the vec path splits MORE
18181        // aggressively (64 keys/split vs 256) so grid.y rises and 8*n_splits fills the SMs.
18182        // At VERY short ctx (t_kv<96) even 1 split can't fill the GPU from 8 KV heads, so the
18183        // broadcast can't beat the scalar path's 4x-more-blocks latency hiding — fall back to
18184        // scalar there (measured crossover: vec 0.68x at t_kv=64, 1.23x at t_kv=96, 2.2x at 256).
18185        // DEFAULT-ON (2026-06-28): clean clock-locked sweep proved vec beats scalar at every
18186        // t_kv>=96 and the gain WIDENS with ctx (graph decode: +9.5% @128, +11.6% @512, +11.8%
18187        // @2048) — the KV-byte-broadcast (4x fewer HBM reads/group) compounds as attention grows.
18188        // MEMRA_NO_FA_VEC forces the scalar bit-reference. Below FA_VEC_MIN_TKV the scalar path's
18189        // 4x-more-blocks (grid.x=n_head=32 vs n_head_kv=8) hides latency better, so keep scalar there.
18190        // g + no-v4: the g-module REGISTER twin mis-decodes the gemma windowed shape
18191        // (root-cause open, jsonl) — only reachable by forcing v4 off (MEMRA_FA_V4_MAX);
18192        // fall to the exact scalar there instead of the broken register arm.
18193        let mut fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && t_kv >= fa_vec_min_tkv();
18194        // hd256 under g: v4-or-scalar. hd128/other under g ride the REGISTER g-lane (the
18195        // dq_K_lane/dq_V_lane macros are format-aware; v3 is format-gated off, v2/smem
18196        // arms are excluded under g). Mirrored in kvmod / fa_decode_dc / fa_geom_eager.
18197        if g && head_dim == 256 && !fa_v4_at(t_kv) {
18198            fa_vec = false;
18199        }
18200        let sp = fa_split_keys(t_kv, n_head_kv);
18201        let n_splits = if fa_vec {
18202            ((t_kv + sp - 1) / sp).max(1)
18203        } else {
18204            ((t_kv + 255) / 256).max(1)
18205        };
18206        let o_len = n_head * n_splits * head_dim;
18207        let ml_len = n_head * n_splits;
18208        let mut part_guard = self.fa_part_pool.lock().unwrap();
18209        if part_guard
18210            .as_ref()
18211            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
18212            .unwrap_or(true)
18213        {
18214            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
18215            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
18216            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
18217            // later live allocations land at those addresses, and the next graph REPLAY writes
18218            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
18219            // output corruption began the burst after the trunk's t_kv growth first realloc'd
18220            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
18221            // the baked addresses alive (single-stream: eager writes the new buffers, replays
18222            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
18223            // (total retired < final size).
18224            let old = part_guard.take();
18225            let (co, cm) = old
18226                .as_ref()
18227                .map(|pp| (pp.0.len(), pp.1.len()))
18228                .unwrap_or((0, 0));
18229            if let Some(old) = old {
18230                self.fa_part_retired.lock().unwrap().push(old);
18231            }
18232            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
18233                eprintln!(
18234                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
18235                    co, o_len, cm, ml_len
18236                );
18237            }
18238            *part_guard = Some((
18239                self.alloc_uninit::<f32>(o_len.max(2 * co))?,
18240                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
18241                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
18242            ));
18243        }
18244        let pg = part_guard.as_mut().unwrap();
18245        self.gpu
18246            .stream()
18247            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
18248        self.gpu
18249            .stream()
18250            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
18251        self.gpu
18252            .stream()
18253            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
18254        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
18255        let (part_o, part_m, part_l) = (&mut *part_o, &mut *part_m, &mut *part_l);
18256        let (hd, nh, nhkv, tkvi, nsp) = (
18257            head_dim as i32,
18258            n_head as i32,
18259            n_head_kv as i32,
18260            t_kv as i32,
18261            n_splits as i32,
18262        );
18263        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
18264        // The vec kernel holds head_dim/32 register accumulators (FA_DEC_MAX_DPL=8 -> head_dim<=256).
18265        // All shipped models use head_dim=256; fall back to scalar for anything wider rather than
18266        // silently truncating the accumulator.
18267        let fa_vec = fa_vec && head_dim <= 512 && head_dim % 32 == 0;
18268        // hd-512 vec crossover (MEMRA_FA512_MIN, default 512): the DPL16 twin wins at depth
18269        // (82.5 -> vec at 1736) but the scalar's more-blocks latency hiding wins at tiny t_kv
18270        // (the same scalar-floor physics as hd256's old 96 floor; short-ctx plain regressed
18271        // 178.4 -> 173.7 when 512 rode vec unconditionally).
18272        let fa512_min = fa512_min_tkv();
18273        // FA-DEEP pick (bit-identical twins, see fa_deep_at): default module only — the
18274        // g-module keeps the v4 pick (its class is not the depth-decay class).
18275        let deep = fa_vec
18276            && head_dim == 256
18277            && fa_v4_at(t_kv)
18278            && !g
18279            && fa_deep_at(t_kv)
18280            && !matches!(fa_v4_mode(), "noB3" | "stage");
18281        let (f, cfg) = if fa_vec && head_dim == 512 && t_kv >= fa512_min {
18282            // gemma4 globals (hd 512): the DPL16 register twin (fa_decode_vec_q body with a
18283            // 16-slot accumulator ceiling). Scalar fallback measured 82.5us/layer at 1736 ctx.
18284            let gqa = (n_head / n_head_kv).max(1) as u32;
18285            let fv = self.fa_func("fa_decode_vec_q_dpl16", head_dim);
18286            (
18287                fv,
18288                LaunchConfig {
18289                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
18290                    block_dim: (32, gqa, 1),
18291                    shared_mem_bytes: 0,
18292                },
18293            )
18294        } else if fa_vec && head_dim <= 256 {
18295            let gqa = (n_head / n_head_kv).max(1) as u32;
18296            // DEEP-CTX smem twin (2026-07-05): the register-dequant path's GQA reuse rides L2,
18297            // which holds to ~8k ctx but dies at 40k (layer KV ~37MB) — the 4 GQA warps then
18298            // re-read every KV byte from DRAM (4x traffic). Above MEMRA_FA_SMEM_TKV (default
18299            // 1024 — the 2026-07-05 crossover re-sweep on real prompts: p3 spec 73.8->79.2 at
18300            // 2048, flat down to 512, p2 +5%, p1/9B unchanged; the ARC-A probe's synthetic
18301            // 2.1x smem-at-all-depths pointed here; 0=never) dispatch the smem-broadcast twin:
18302            // dequant each tile ONCE per block.
18303            // Bit-identical per (token,split): same bf16 round-trip, same accumulation order,
18304            // same partial layout -> same combine. Short/mid ctx keeps the register path (it won
18305            // there by 12x — latency, not bandwidth, rules small KV).
18306            static SMEM_TKV: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
18307            let smem_tkv = *SMEM_TKV.get_or_init(|| {
18308                std::env::var("MEMRA_FA_SMEM_TKV")
18309                    .ok()
18310                    .and_then(|v| v.parse().ok())
18311                    .unwrap_or_else(|| {
18312                        FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
18313                    })
18314            });
18315            if fa_v4_at(t_kv) && head_dim == 256 {
18316                // FA v4 lane (2026-07-10): key-per-lane score phase, zero shuffles per key.
18317                // NEW NUMERIC CONFIG (chunk-serial per-key dot) — battery-arbitrated.
18318                // g (fp8-windowed): the v4 staging is format-aware (2026-07-12) — kf8vf8 module.
18319                let v4name = match fa_v4_mode() {
18320                    "noB3" => "fa_decode_vec_q_v4_noB3", // phase probe (WRONG OUTPUT)
18321                    "stage" => "fa_decode_vec_q_v4_stage", // phase probe (WRONG OUTPUT)
18322                    _ if deep => "fa_decode_vec_q_v4_deep",
18323                    _ => "fa_decode_vec_q_v4",
18324                };
18325                let fv = if g {
18326                    self.func_g(v4name)
18327                } else {
18328                    self.func(v4name)
18329                };
18330                // fa_v4_smem (deep: fa_v4_deep_smem, +640B row pads) + sV (g: raw e4m3 sV
18331                // tile = 1B/elem — half the smem, 3->5 blocks/SM)
18332                let shmem = (if deep { 12160 } else { 11520 }
18333                    + 32 * head_dim * if g { 1 } else { 2 }) as u32;
18334                use cudarc::driver::sys::CUfunction_attribute_enum as A;
18335                fv.set_attribute(
18336                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18337                    shmem as i32,
18338                )?;
18339                (
18340                    fv,
18341                    LaunchConfig {
18342                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
18343                        block_dim: (32, gqa, 1),
18344                        shared_mem_bytes: shmem,
18345                    },
18346                )
18347            } else if fa_v3_active(head_dim) {
18348                // FA v3 lane: dp4a-K hybrid (register-quantized Q, raw q8_0 K, staged-V kept).
18349                // smem = sV only (half of v2's).
18350                let fv = if g {
18351                    self.func_g("fa_decode_vec_q_v3")
18352                } else {
18353                    self.func("fa_decode_vec_q_v3")
18354                };
18355                let shmem = (32 * head_dim * 2) as u32; // sV bf16 [FA_DEC_TILE=32][hd]
18356                (
18357                    fv,
18358                    LaunchConfig {
18359                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
18360                        block_dim: (32, gqa, 1),
18361                        shared_mem_bytes: shmem,
18362                    },
18363                )
18364            } else if fa_v2_on() {
18365                // FAVENDOR lane: llama fattn-vec tile-batched softmax + wide-load staging on
18366                // OUR smem KV broadcast. Replaces BOTH per-key twins when on; same grid/block/
18367                // partials; same 32KB sK+sV tile as the smem twin.
18368                let fv = if g {
18369                    self.func_g("fa_decode_vec_q_v2")
18370                } else {
18371                    self.func("fa_decode_vec_q_v2")
18372                };
18373                let shmem = (2 * 32 * head_dim * 2) as u32; // sK+sV bf16 [FA_DEC_TILE=32][hd]
18374                (
18375                    fv,
18376                    LaunchConfig {
18377                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
18378                        block_dim: (32, gqa, 1),
18379                        shared_mem_bytes: shmem,
18380                    },
18381                )
18382            } else if smem_tkv > 0 && t_kv >= smem_tkv && !g && !(head_dim == 512 && Self::gkv_on())
18383            {
18384                // (fp8 exclusions: the smem twin's V-stage is q5_1-hardcoded — neither the wkv
18385                // windowed layers (g) nor the gkv globals (hd512) may be forced onto it via
18386                // MEMRA_FA_SMEM_TKV; they fall through to the format-clean register/scalar arms.)
18387                let fv = if g {
18388                    self.func_g("fa_decode_vec_q_smem")
18389                } else {
18390                    self.func("fa_decode_vec_q_smem")
18391                };
18392                let shmem = (2 * 32 * head_dim * 2) as u32; // sK+sV bf16 [FA_DEC_TILE=32][hd]
18393                use cudarc::driver::sys::CUfunction_attribute_enum as A;
18394                fv.set_attribute(
18395                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18396                    shmem as i32,
18397                )?;
18398                (
18399                    fv,
18400                    LaunchConfig {
18401                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
18402                        block_dim: (32, gqa, 1),
18403                        shared_mem_bytes: shmem,
18404                    },
18405                )
18406            } else {
18407                // REGISTER-DEQUANT kernel (2026-07-03): per-warp direct q8_0/q5_1 register
18408                // dequant, zero dynamic shared memory.
18409                let fv = if g {
18410                    self.func_g("fa_decode_vec_q")
18411                } else {
18412                    self.func("fa_decode_vec_q")
18413                };
18414                (
18415                    fv,
18416                    LaunchConfig {
18417                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
18418                        block_dim: (32, gqa, 1),
18419                        shared_mem_bytes: 0,
18420                    },
18421                )
18422            }
18423        } else {
18424            // UNIFIED scalar (nullable-ctr symbol shared with graph/stream callers). The
18425            // split ladder value rides along so ns_eff reproduces THIS n_splits in-kernel.
18426            return self.fa_decode_scalar_unified(
18427                q,
18428                k,
18429                v,
18430                o,
18431                head_dim,
18432                n_head,
18433                n_head_kv,
18434                t_kv,
18435                None,
18436                scale,
18437                n_splits,
18438                if fa_vec { sp } else { 256 },
18439                k_tok_bytes,
18440                v_tok_bytes,
18441                g,
18442                part_o,
18443                part_m,
18444                part_l,
18445                None,
18446            );
18447        };
18448        let __s_b = self.gpu.stream();
18449        let mut b = __s_b.launch_builder(&f);
18450        b.arg(q)
18451            .arg(k)
18452            .arg(v)
18453            .arg(&mut *part_o)
18454            .arg(&mut *part_m)
18455            .arg(&mut *part_l)
18456            .arg(&hd)
18457            .arg(&nh)
18458            .arg(&nhkv)
18459            .arg(&tkvi)
18460            .arg(&scale)
18461            .arg(&nsp)
18462            .arg(&ktb)
18463            .arg(&vtb);
18464        unsafe {
18465            b.launch(cfg)?;
18466        }
18467        // (combine re-tile refuted in the fa-deep lane — flat/worse both shapes; the v4
18468        // combine stays for all arms. Receipts research/fa-decode-deep-20260802/.)
18469        let (fc, cfg2) = (
18470            if g {
18471                self.func_g("fa_decode_combine_f32")
18472            } else {
18473                self.fa_func("fa_decode_combine_f32", head_dim)
18474            },
18475            LaunchConfig {
18476                grid_dim: (n_head as u32, 1, 1),
18477                block_dim: (head_dim as u32, 1, 1),
18478                shared_mem_bytes: 0,
18479            },
18480        );
18481        let __s_b2 = self.gpu.stream();
18482        let mut b2 = __s_b2.launch_builder(&fc);
18483        b2.arg(&*part_o)
18484            .arg(&*part_m)
18485            .arg(&*part_l)
18486            .arg(o)
18487            .arg(&hd)
18488            .arg(&nh)
18489            .arg(&nsp);
18490        unsafe {
18491            b2.launch(cfg2)?;
18492        }
18493        Ok(())
18494    }
18495
18496    /// BATCHED-TICK increment 2: ONE fa_decode launch covering ALL B sequences of the
18497    /// batched decode step (blockIdx.z = sequence). Per-seq K/V cache bases ride a device
18498    /// pointer table (`kv_ptrs`, [2B] interleaved k0,v0,...); per-seq key bounds ride the
18499    /// tick's position table (`pos_seq`, T_kv = pos+1). v4-lane only: the CALLER
18500    /// (decode_batch) gates every row through `fa_seqs_eligible` AND one `fa_split_keys`
18501    /// rung (`split_keys`), so each sequence's split partition, key walk and combine order
18502    /// reproduce its per-seq eager v4 program exactly (kernel-check pins seqs-vs-loop bit
18503    /// identity; decode-batch-gate strict pins the whole tick vs decode_step_h).
18504    /// q is the stacked [B, n_head, head_dim] tick buffer read in place (no per-seq q
18505    /// copies); o is written [B, n_head, head_dim] in place (no per-seq a copies).
18506    #[allow(clippy::too_many_arguments)]
18507    pub fn fa_decode_batch_seqs_v4(
18508        &self,
18509        q: &CudaSlice<f32>,
18510        kv_ptrs: &cudarc::driver::CudaView<u64>,
18511        pos_seq: &CudaSlice<i32>,
18512        o: &mut CudaSlice<f32>,
18513        head_dim: usize,
18514        n_head: usize,
18515        n_head_kv: usize,
18516        b_n: usize,
18517        t_kv_max: usize,
18518        scale: f32,
18519        split_keys: usize,
18520        k_tok_bytes: usize,
18521        v_tok_bytes: usize,
18522    ) -> Result<(), Box<dyn std::error::Error>> {
18523        debug_assert!(head_dim == 256, "seqs twin is v4-stamped (hd256 only)");
18524        let n_splits_max = (t_kv_max + split_keys - 1) / split_keys;
18525        let o_len = b_n * n_head * n_splits_max * head_dim;
18526        let ml_len = b_n * n_head * n_splits_max;
18527        let mut part_guard = self.fa_part_pool.lock().unwrap();
18528        if part_guard
18529            .as_ref()
18530            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
18531            .unwrap_or(true)
18532        {
18533            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
18534            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
18535            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
18536            // later live allocations land at those addresses, and the next graph REPLAY writes
18537            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
18538            // output corruption began the burst after the trunk's t_kv growth first realloc'd
18539            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
18540            // the baked addresses alive (single-stream: eager writes the new buffers, replays
18541            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
18542            // (total retired < final size).
18543            let old = part_guard.take();
18544            let (co, cm) = old
18545                .as_ref()
18546                .map(|pp| (pp.0.len(), pp.1.len()))
18547                .unwrap_or((0, 0));
18548            if let Some(old) = old {
18549                self.fa_part_retired.lock().unwrap().push(old);
18550            }
18551            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
18552                eprintln!(
18553                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
18554                    co, o_len, cm, ml_len
18555                );
18556            }
18557            *part_guard = Some((
18558                self.alloc_uninit::<f32>(o_len.max(2 * co))?,
18559                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
18560                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
18561            ));
18562        }
18563        let pg = part_guard.as_mut().unwrap();
18564        self.gpu
18565            .stream()
18566            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
18567        self.gpu
18568            .stream()
18569            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
18570        self.gpu
18571            .stream()
18572            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
18573        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
18574        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
18575        let (nspm, spk) = (n_splits_max as i32, split_keys as i32);
18576        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
18577        let gqa = (n_head / n_head_kv).max(1) as u32;
18578        let f = self.func("fa_decode_vec_q_seqs_v4");
18579        // fa_v4_smem (11520B) + sV bf16 tile — the v4 eager arm's sizing on the default module.
18580        let shmem = (11520 + 32 * head_dim * 2) as u32;
18581        use cudarc::driver::sys::CUfunction_attribute_enum as A;
18582        f.set_attribute(
18583            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18584            shmem as i32,
18585        )?;
18586        let cfg = LaunchConfig {
18587            grid_dim: (n_head_kv as u32, n_splits_max as u32, b_n as u32),
18588            block_dim: (32, gqa, 1),
18589            shared_mem_bytes: shmem,
18590        };
18591        {
18592            let __s_b = self.gpu.stream();
18593            let mut b = __s_b.launch_builder(&f);
18594            b.arg(q)
18595                .arg(kv_ptrs)
18596                .arg(pos_seq)
18597                .arg(&mut *part_o)
18598                .arg(&mut *part_m)
18599                .arg(&mut *part_l)
18600                .arg(&hd)
18601                .arg(&nh)
18602                .arg(&nhkv)
18603                .arg(&scale)
18604                .arg(&nspm)
18605                .arg(&spk)
18606                .arg(&ktb)
18607                .arg(&vtb);
18608            unsafe {
18609                b.launch(cfg)?;
18610            }
18611        }
18612        let fc = self.func("fa_decode_combine_seqs");
18613        let cfg2 = LaunchConfig {
18614            grid_dim: (n_head as u32, b_n as u32, 1),
18615            block_dim: (head_dim as u32, 1, 1),
18616            shared_mem_bytes: 0,
18617        };
18618        let __s_b2 = self.gpu.stream();
18619        let mut b2 = __s_b2.launch_builder(&fc);
18620        b2.arg(&*part_o)
18621            .arg(&*part_m)
18622            .arg(&*part_l)
18623            .arg(o)
18624            .arg(&hd)
18625            .arg(&nh)
18626            .arg(pos_seq)
18627            .arg(&nspm)
18628            .arg(&spk);
18629        unsafe {
18630            b2.launch(cfg2)?;
18631        }
18632        Ok(())
18633    }
18634
18635    /// BATCHED-TICK increment 2: z-batched decode KV append — one launch appends this
18636    /// step's B rows, each into ITS OWN sequence cache at slot pos_seq[z], through the same
18637    /// [2B] interleaved pointer table the seqs FA reads. Each (block, z) warp executes the
18638    /// per-token appender's exact warp program on row z of the stacked [B, kv_dim] k/v —
18639    /// written cache bytes are BIT-IDENTICAL to the B per-seq calls it replaces
18640    /// (kernel-check pins the bytes). Default flash module only (callers exclude fp8-KV).
18641    #[allow(clippy::too_many_arguments)]
18642    pub fn append_kv_quantized_seqs(
18643        &self,
18644        k_rows: &CudaSlice<f32>,
18645        v_rows: &CudaSlice<f32>,
18646        kv_ptrs: &cudarc::driver::CudaView<u64>,
18647        pos_seq: &CudaSlice<i32>,
18648        b_n: usize,
18649        kv_dim_k: usize,
18650        kv_dim_v: usize,
18651        k_tok_bytes: usize,
18652        v_tok_bytes: usize,
18653    ) -> Result<(), Box<dyn std::error::Error>> {
18654        let f = self.func("append_quantize_kv_q8_0_q5_1_seqs");
18655        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
18656        let cfg = LaunchConfig {
18657            grid_dim: (nblk, b_n as u32, 1),
18658            block_dim: (32, 1, 1),
18659            shared_mem_bytes: 0,
18660        };
18661        let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
18662        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
18663        let __s_b = self.gpu.stream();
18664        let mut b = __s_b.launch_builder(&f);
18665        b.arg(k_rows)
18666            .arg(v_rows)
18667            .arg(kv_ptrs)
18668            .arg(pos_seq)
18669            .arg(&kdk)
18670            .arg(&kdv)
18671            .arg(&ktb)
18672            .arg(&vtb);
18673        unsafe {
18674            b.launch(cfg)?;
18675        }
18676        Ok(())
18677    }
18678
18679    /// True iff the MULTI-ROW verify FA (`fa_decode_rows`) is usable for a verify batch whose
18680    /// FIRST row attends `base_len + 1` keys: every row must take the SAME kernel eager decode
18681    /// would (the vec path) — mirrors fa_decode's gate exactly (MEMRA_NO_FA_VEC + FA_VEC_MIN_TKV +
18682    /// head_dim), evaluated at the MINIMUM row bound so no row could have picked scalar.
18683    /// MEMRA_FA_ROWS_OFF=1 is the A/B + fallback seam (per-row loop).
18684    pub fn fa_rows_eligible(&self, base_len: usize, head_dim: usize) -> bool {
18685        std::env::var("MEMRA_NO_FA_VEC").is_err()
18686            && std::env::var("MEMRA_FA_ROWS_OFF").is_err()
18687            && base_len + 1 >= fa_vec_min_tkv()
18688            && head_dim <= 256
18689            && head_dim % 32 == 0
18690    }
18691
18692    /// MULTI-ROW verify FA: run fa_decode_vec_q's EXACT per-row program for T causal query rows
18693    /// (row r attends keys [0..base_len+r+1)) in ONE kernel launch with grid.z = row, plus ONE
18694    /// row-batched combine. Replaces the T separate (fa_decode + combine) launches of the spec
18695    /// verify — same per-row split partition (n_splits_r = ceil(t_kv_r/split_keys), the
18696    /// fa_split_keys formula), same key-walk order, same reduce shapes => bit-identical outputs
18697    /// per row (kernel-check pins rows-vs-loop byte identity; run-spec is the end gate).
18698    /// Caller must have checked `fa_rows_eligible(base_len, head_dim)`.
18699    /// q is the verify's token-major [T, n_head, head_dim] stack; o is written [T, n_head, head_dim].
18700    #[allow(clippy::too_many_arguments)]
18701    pub fn fa_decode_rows(
18702        &self,
18703        q: &CudaSlice<f32>,
18704        k: &cudarc::driver::CudaView<u8>,
18705        v: &cudarc::driver::CudaView<u8>,
18706        o: &mut CudaSlice<f32>,
18707        head_dim: usize,
18708        n_head: usize,
18709        n_head_kv: usize,
18710        base_len: usize,
18711        t: usize,
18712        scale: f32,
18713        k_tok_bytes: usize,
18714        v_tok_bytes: usize,
18715        // hd512 dpl16 twin is DEVICE-LEN (graph arc): base_dev/plus feed the
18716        // kernel; host base_len keeps sizing the splits/partials. hd256 twins
18717        // keep the host arg. None is a bug for hd512 (asserted below).
18718        base_dev: Option<(&CudaSlice<i32>, i32)>,
18719        // K and V planes hold the same values (gemma globals, wv:=wk): pick
18720        // the _kv twin — V plane never read, value rides the q8_0 key dq.
18721        kv_shared: bool,
18722        // this layer's cache is e4m3 (gemma windowed under wkv): resolve the
18723        // hd256 rows kernel from the kf8vf8 module. PER-CALL — a global env
18724        // check here hijacked qwen/kernel-check hd256 rows (8 FAILs, 230ebbe).
18725        g: bool,
18726        // t=1 decode arm only: emit (int8, per-32 scales) from the dc combine
18727        // (hd512 path) — the standalone quantize launch folds away.
18728        mut q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
18729    ) -> Result<(), Box<dyn std::error::Error>> {
18730        debug_assert!(base_len + 1 >= fa_vec_min_tkv() && head_dim <= 512 && head_dim % 32 == 0);
18731        let t_kv_max = base_len + t; // LAST row's key bound
18732        let mut sp = fa_split_keys(t_kv_max, n_head_kv); // env/default — same value every row
18733        // hd512 split override (MEMRA_FA_SP512, 2026-07-11): gemma globals have n_head_kv=2 so
18734        // the grid is (2 x n_splits) — at depth ~29 splits = 58 blocks on 82 SMs (half idle,
18735        // rows_dpl16 8x off its byte floor). EVERY gemma hd512 caller shares THIS wrapper
18736        // (parity law), so the partition is freely tunable — verify and decode move together.
18737        if head_dim == 512 {
18738            static SP512: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
18739            // default 16 (2026-07-11 depth sweep, N=2: plain 155.4->156.5, depth spec
18740            // 236.9->250.4; 12/24/32 all worse). hd512 exists only on gemma globals.
18741            let v = *SP512.get_or_init(|| {
18742                std::env::var("MEMRA_FA_SP512")
18743                    .ok()
18744                    .and_then(|x| x.parse().ok())
18745                    .unwrap_or(0)
18746            });
18747            sp = if v >= 8 {
18748                v
18749            } else {
18750                FA_SP512_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
18751            };
18752        }
18753        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
18754        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
18755        let gqa = (n_head / n_head_kv).max(1) as u32;
18756        // LADDER-RUNG STRADDLE FIX (issue #10, 2026-07-13, g7e-proven): one sp for every row
18757        // diverges from eager decode when a split-ladder rung falls INSIDE the batch — row r's
18758        // eager twin used fa_split_keys(t_kv_r), the batch used fa_split_keys(t_kv_max), and
18759        // the different partition changes the combine's FP order (greedy tie flips at depth;
18760        // MEMRA_FA_SPLIT=64 pin -> PASS on the exact g7e failing config). Fix: group
18761        // consecutive rows by their OWN ladder value and launch once per group — each row then
18762        // executes the exact per-row program eager ran. Rungs land once per doubling, so this
18763        // is 1 launch in the common case and 2 on a crossing round. hd512 keeps one group (its
18764        // sp override is t_kv-independent by construction).
18765        let mut groups: Vec<(usize, usize, usize)> = Vec::new(); // (row0, t_g, sp_g)
18766        if head_dim == 512 || fa_split_keys(base_len + 1, n_head_kv) == sp {
18767            groups.push((0, t, sp));
18768        } else {
18769            let mut r0 = 0usize;
18770            while r0 < t {
18771                let sp_g = fa_split_keys(base_len + r0 + 1, n_head_kv);
18772                let mut r1 = r0 + 1;
18773                while r1 < t && fa_split_keys(base_len + r1 + 1, n_head_kv) == sp_g {
18774                    r1 += 1;
18775                }
18776                groups.push((r0, r1 - r0, sp_g));
18777                r0 = r1;
18778            }
18779        }
18780        // Deep-ctx smem twin for the VERIFY rows (2026-07-05): same threshold + rationale as
18781        // fa_decode's dispatch — at 40k the register path's GQA L2-reuse premise is dead and the
18782        // verify multiplies the 4x DRAM re-read by T rows. Bit-identical per (row,token,split).
18783        static SMEM_TKV_R: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
18784        let smem_tkv = *SMEM_TKV_R.get_or_init(|| {
18785            std::env::var("MEMRA_FA_SMEM_TKV")
18786                .ok()
18787                .and_then(|v| v.parse().ok())
18788                .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
18789        });
18790        let v4 = fa_v4_at(base_len + t) && head_dim == 256;
18791        let v3 = fa_v3_active(head_dim);
18792        let smem_rows =
18793            head_dim <= 256 && !v3 && !fa_v2_on() && smem_tkv > 0 && t_kv_max >= smem_tkv;
18794        // kv_shared twin RETIRED (2026-07-11 depth run-gen gate): the wv:=wk premise fails
18795        // POST-cache — cached K is k-normed+roped, cached V is not; the twin fed roped keys
18796        // in as values. Verify/decode/stream gates were blind (both sides shared the wrong
18797        // symbol — the parity law's blind spot); only prefill-vs-decode at depth caught it.
18798        let _ = kv_shared;
18799        // i2 twin: 2-key interleaved walk (MEMRA_FA_I2=0 reverts). i4 probed NEGATIVE
18800        // (157.3 vs 161.2 depth plain — register pressure past i2's sweet spot; jsonl).
18801        let i2 = head_dim == 512 && std::env::var("MEMRA_FA_I2").as_deref() != Ok("0");
18802        // v4-hd512 (MEMRA_FA_V512=1 opt-in, 2026-07-14): the v4 key-per-lane recipe on the
18803        // globals lane (depth profile: i2 ~4.6x off its byte floor — the v3-class
18804        // reduce-per-key latency signature). NEW NUMERIC CONFIG shared by every hd512
18805        // caller (decode+verify flip together); run-gen argmax + acceptance arbitrate.
18806        // T-BATCHED hd512 (DEFAULT ON 2026-07-14, MEMRA_FA_TB512=0 seam): one block per
18807        // (kv_head, split) stages its tile once and loops the rows over it — kills the
18808        // x t DRAM re-read of the full-ctx globals (depth cell +1.4%, plain flat, N=3
18809        // interleaved). FIXED absolute partition = NEW NUMERIC for the combine order,
18810        // shared by every hd512 caller through this wrapper (decode+verify flip together;
18811        // depth stream identical, acceptance unshifted, spec 256/256 x3 models).
18812        // Requires sp <= 32 (single staged tile; acc reused per row). The z-form v4_512
18813        // sibling (in-kernel dp4a port alone) probed FLAT — hd512 was DRAM-re-read-bound,
18814        // not unpack-bound; jsonl 2026-07-14.
18815        static TB512: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
18816        // gqa <= 16 = fa_v4_smem_512's q-array capacity; past it fall to the register twins.
18817        let tb512 = head_dim == 512
18818            && sp <= 32
18819            && n_head / n_head_kv.max(1) <= 16
18820            && *TB512.get_or_init(|| std::env::var("MEMRA_FA_TB512").as_deref() != Ok("0"));
18821        let fname = if tb512 {
18822            "fa_decode_vec_q_rows_v4_512_tb"
18823        } else if i2 {
18824            "fa_decode_vec_q_rows_dpl16_i2"
18825        } else if head_dim == 512 {
18826            "fa_decode_vec_q_rows_dpl16"
18827        }
18828        // gemma globals (parity law)
18829        else if v4 {
18830            "fa_decode_vec_q_rows_v4"
18831        } else if v3 {
18832            "fa_decode_vec_q_rows_v3"
18833        } else if fa_v2_on() {
18834            "fa_decode_vec_q_rows_v2"
18835        } else if smem_rows {
18836            "fa_decode_vec_q_rows_smem"
18837        } else {
18838            "fa_decode_vec_q_rows"
18839        };
18840        let f = if head_dim == 512 {
18841            self.fa_func(fname, head_dim)
18842        } else if g {
18843            // FP8-WINDOWED: hd256 rows over an e4m3 cache — kf8vf8 module, SAME symbol
18844            // choice as decode's kvmod dispatch (parity law: excluding v4 here paired
18845            // g-module rows against decode's g-module v4 — different programs, short-VG
18846            // maxdiff 2.0 / spec stream 0/128, 2026-07-12). rows_v4 is format-aware
18847            // since fda9790; only the smem twin stays excluded (V-stage q5_1-only).
18848            // hd128 (qwen fp8-KV) lands on the base/register rows via fname — the
18849            // dq macros are format-aware.
18850            self.func_g(if smem_rows {
18851                "fa_decode_vec_q_rows"
18852            } else {
18853                fname
18854            })
18855        } else {
18856            self.func(fname)
18857        };
18858        let shmem = if tb512 {
18859            // fa_v4_smem_512 (q 9KB gqa<=16 + k tile 18KB) + sV 32*512 (e4m3 module halves it)
18860            let gk = Self::gkv_on();
18861            let sh =
18862                (8192 + 1024 + 32 * 512 + 32 * 64 + 32 * head_dim * if gk { 1 } else { 2 }) as u32;
18863            use cudarc::driver::sys::CUfunction_attribute_enum as A;
18864            f.set_attribute(
18865                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18866                sh as i32,
18867            )?;
18868            sh
18869        } else if v4 || v3 || smem_rows || fa_v2_on() {
18870            // v4: fa_v4_smem (11.5KB) + sV; v3 stages sV only; v2/smem twins stage sK+sV.
18871            let sh = (if v4 {
18872                11520 + 32 * head_dim * if g { 1 } else { 2 }
18873            } else if v3 {
18874                32 * head_dim * 2
18875            } else {
18876                2 * 32 * head_dim * 2
18877            }) as u32;
18878            use cudarc::driver::sys::CUfunction_attribute_enum as A;
18879            f.set_attribute(
18880                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18881                sh as i32,
18882            )?;
18883            sh
18884        } else {
18885            0
18886        };
18887        // Per-GROUP launches (single group in the common case — identical to the pre-fix
18888        // single launch there): each group gets its own partials (the rows kernel indexes
18889        // partials by its LOCAL grid.z row) and q/o row-offset views.
18890        for &(r0, t_g, sp_g) in &groups {
18891            let n_splits_g = (base_len + r0 + t_g).div_ceil(sp_g);
18892            let (nspm, spk) = (n_splits_g as i32, sp_g as i32);
18893            let base_i = (base_len + r0) as i32;
18894            let o_len = t_g * n_head * n_splits_g * head_dim;
18895            let ml_len = t_g * n_head * n_splits_g;
18896            let mut part_guard = self.fa_part_pool.lock().unwrap();
18897            if part_guard
18898                .as_ref()
18899                .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
18900                .unwrap_or(true)
18901            {
18902                // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
18903                // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
18904                // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
18905                // later live allocations land at those addresses, and the next graph REPLAY writes
18906                // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
18907                // output corruption began the burst after the trunk's t_kv growth first realloc'd
18908                // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
18909                // the baked addresses alive (single-stream: eager writes the new buffers, replays
18910                // touch only the old — never concurrently). Doubling growth bounds retired VRAM
18911                // (total retired < final size).
18912                let old = part_guard.take();
18913                let (co, cm) = old
18914                    .as_ref()
18915                    .map(|pp| (pp.0.len(), pp.1.len()))
18916                    .unwrap_or((0, 0));
18917                if let Some(old) = old {
18918                    self.fa_part_retired.lock().unwrap().push(old);
18919                }
18920                if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
18921                    eprintln!(
18922                        "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
18923                        co, o_len, cm, ml_len
18924                    );
18925                }
18926                *part_guard = Some((
18927                    self.alloc_uninit::<f32>(o_len.max(2 * co))?,
18928                    self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
18929                    self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
18930                ));
18931            }
18932            let pg = part_guard.as_mut().unwrap();
18933            self.gpu
18934                .stream()
18935                .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
18936            self.gpu
18937                .stream()
18938                .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
18939            self.gpu
18940                .stream()
18941                .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
18942            let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
18943            let (part_o, part_m, part_l) = (&mut *part_o, &mut *part_m, &mut *part_l);
18944            let qv = self.view(q, t * n_head * head_dim);
18945            let q_g = qv.slice(r0 * n_head * head_dim..(r0 + t_g) * n_head * head_dim);
18946            let cfg = LaunchConfig {
18947                grid_dim: (n_head_kv as u32, n_splits_g as u32, t_g as u32),
18948                block_dim: (32, gqa, 1),
18949                shared_mem_bytes: shmem,
18950            };
18951            {
18952                let __s_b = self.gpu.stream();
18953                let mut b = __s_b.launch_builder(&f);
18954                if tb512 {
18955                    // rows-inner launch: grid.z dropped, the kernel loops n_rows itself.
18956                    let (bd, plus) =
18957                        base_dev.expect("hd512 rows twin requires a device base counter");
18958                    let plus_g = plus + r0 as i32;
18959                    let nr = t_g as i32;
18960                    if Self::pdl_on() && Self::pdl_wb_on() {
18961                        // wave-B2b: flavor mirrors fa_func(fname, 512) = gkv.
18962                        use cudarc::driver::{DevicePtr, DevicePtrMut};
18963                        let s = &self.gpu.stream();
18964                        let (pq, _b0) = q_g.device_ptr(s);
18965                        let (pk, _b1) = k.device_ptr(s);
18966                        let (pv, _b2) = v.device_ptr(s);
18967                        let (po, _b3) = part_o.device_ptr_mut(s);
18968                        let (pm, _b4) = part_m.device_ptr_mut(s);
18969                        let (pl, _b5) = part_l.device_ptr_mut(s);
18970                        let (pb, _b6) = bd.device_ptr(s);
18971                        let mut ps = [
18972                            &pq as *const _ as *mut std::ffi::c_void,
18973                            &pk as *const _ as *mut _,
18974                            &pv as *const _ as *mut _,
18975                            &po as *const _ as *mut _,
18976                            &pm as *const _ as *mut _,
18977                            &pl as *const _ as *mut _,
18978                            &hd as *const _ as *mut _,
18979                            &nh as *const _ as *mut _,
18980                            &nhkv as *const _ as *mut _,
18981                            &pb as *const _ as *mut _,
18982                            &plus_g as *const _ as *mut _,
18983                            &scale as *const _ as *mut _,
18984                            &nspm as *const _ as *mut _,
18985                            &spk as *const _ as *mut _,
18986                            &ktb as *const _ as *mut _,
18987                            &vtb as *const _ as *mut _,
18988                            &nr as *const _ as *mut _,
18989                        ];
18990                        unsafe {
18991                            self.launch_pdl_flash(
18992                                Self::gkv_on(),
18993                                "fa_decode_vec_q_rows_v4_512_tb",
18994                                (n_head_kv as u32, n_splits_g as u32, 1),
18995                                (32, gqa, 1),
18996                                shmem,
18997                                &mut ps,
18998                            )?;
18999                        }
19000                    } else {
19001                        let cfg_tb = LaunchConfig {
19002                            grid_dim: (n_head_kv as u32, n_splits_g as u32, 1),
19003                            block_dim: (32, gqa, 1),
19004                            shared_mem_bytes: shmem,
19005                        };
19006                        b.arg(&q_g)
19007                            .arg(k)
19008                            .arg(v)
19009                            .arg(&mut *part_o)
19010                            .arg(&mut *part_m)
19011                            .arg(&mut *part_l)
19012                            .arg(&hd)
19013                            .arg(&nh)
19014                            .arg(&nhkv)
19015                            .arg(bd)
19016                            .arg(&plus_g)
19017                            .arg(&scale)
19018                            .arg(&nspm)
19019                            .arg(&spk)
19020                            .arg(&ktb)
19021                            .arg(&vtb)
19022                            .arg(&nr);
19023                        unsafe {
19024                            b.launch(cfg_tb)?;
19025                        }
19026                    }
19027                } else if head_dim == 512 {
19028                    let (bd, plus) =
19029                        base_dev.expect("hd512 rows twin requires a device base counter");
19030                    let plus_g = plus + r0 as i32;
19031                    b.arg(&q_g)
19032                        .arg(k)
19033                        .arg(v)
19034                        .arg(&mut *part_o)
19035                        .arg(&mut *part_m)
19036                        .arg(&mut *part_l)
19037                        .arg(&hd)
19038                        .arg(&nh)
19039                        .arg(&nhkv)
19040                        .arg(bd)
19041                        .arg(&plus_g)
19042                        .arg(&scale)
19043                        .arg(&nspm)
19044                        .arg(&spk)
19045                        .arg(&ktb)
19046                        .arg(&vtb);
19047                    unsafe {
19048                        b.launch(cfg)?;
19049                    }
19050                } else {
19051                    b.arg(&q_g)
19052                        .arg(k)
19053                        .arg(v)
19054                        .arg(&mut *part_o)
19055                        .arg(&mut *part_m)
19056                        .arg(&mut *part_l)
19057                        .arg(&hd)
19058                        .arg(&nh)
19059                        .arg(&nhkv)
19060                        .arg(&base_i)
19061                        .arg(&scale)
19062                        .arg(&nspm)
19063                        .arg(&spk)
19064                        .arg(&ktb)
19065                        .arg(&vtb);
19066                    unsafe {
19067                        b.launch(cfg)?;
19068                    }
19069                }
19070            }
19071            let cfg2 = LaunchConfig {
19072                grid_dim: (n_head as u32, t_g as u32, 1),
19073                block_dim: (head_dim as u32, 1, 1),
19074                shared_mem_bytes: 0,
19075            };
19076            let mut o_g = o.slice_mut(r0 * n_head * head_dim..(r0 + t_g) * n_head * head_dim);
19077            if head_dim == 512 {
19078                // device-len combine (shared by verify/eager/graph — parity by symbol): the
19079                // per-row n_splits derives from the SAME counter the rows kernel read.
19080                let (bd, plus) = base_dev.unwrap();
19081                let plus_g = plus + r0 as i32;
19082                if let Some((oq, od)) = q8_out.as_mut() {
19083                    // wave-5b port (2026-07-23, t=1 decode only): q8-emitting dc combine.
19084                    debug_assert!(t == 1, "rows q8 emit is a t=1 decode arm");
19085                    if Self::pdl_on() && Self::pdl_wb_on() {
19086                        // wave-B2: flavor mirrors fa_func (hd512 + gkv → kf8vf8).
19087                        use cudarc::driver::{DevicePtr, DevicePtrMut};
19088                        let s = &self.gpu.stream();
19089                        let (po, _g0) = part_o.device_ptr(s);
19090                        let (pm, _g1) = part_m.device_ptr(s);
19091                        let (pl, _g2) = part_l.device_ptr(s);
19092                        let (pq, _g3) = oq.device_ptr_mut(s);
19093                        let (pd, _g4) = od.device_ptr_mut(s);
19094                        let (pb, _g5) = bd.device_ptr(s);
19095                        let mut ps = [
19096                            &po as *const _ as *mut std::ffi::c_void,
19097                            &pm as *const _ as *mut _,
19098                            &pl as *const _ as *mut _,
19099                            &pq as *const _ as *mut _,
19100                            &pd as *const _ as *mut _,
19101                            &hd as *const _ as *mut _,
19102                            &nh as *const _ as *mut _,
19103                            &pb as *const _ as *mut _,
19104                            &plus_g as *const _ as *mut _,
19105                            &nspm as *const _ as *mut _,
19106                            &spk as *const _ as *mut _,
19107                        ];
19108                        unsafe {
19109                            self.launch_pdl_flash(
19110                                Self::gkv_on(),
19111                                "fa_decode_combine_rows_dc_q8_1",
19112                                cfg2.grid_dim,
19113                                cfg2.block_dim,
19114                                0,
19115                                &mut ps,
19116                            )?;
19117                        }
19118                        continue;
19119                    }
19120                    let fc = self.fa_func("fa_decode_combine_rows_dc_q8_1", head_dim);
19121                    let __s_b2 = self.gpu.stream();
19122                    let mut b2 = __s_b2.launch_builder(&fc);
19123                    b2.arg(&*part_o)
19124                        .arg(&*part_m)
19125                        .arg(&*part_l)
19126                        .arg(&mut **oq)
19127                        .arg(&mut **od)
19128                        .arg(&hd)
19129                        .arg(&nh)
19130                        .arg(bd)
19131                        .arg(&plus_g)
19132                        .arg(&nspm)
19133                        .arg(&spk);
19134                    unsafe {
19135                        b2.launch(cfg2)?;
19136                    }
19137                    continue;
19138                }
19139                let fc = self.fa_func("fa_decode_combine_rows_dc", head_dim);
19140                let __s_b2 = self.gpu.stream();
19141                let mut b2 = __s_b2.launch_builder(&fc);
19142                b2.arg(&*part_o)
19143                    .arg(&*part_m)
19144                    .arg(&*part_l)
19145                    .arg(&mut o_g)
19146                    .arg(&hd)
19147                    .arg(&nh)
19148                    .arg(bd)
19149                    .arg(&plus_g)
19150                    .arg(&nspm)
19151                    .arg(&spk);
19152                unsafe {
19153                    b2.launch(cfg2)?;
19154                }
19155            } else {
19156                // q8 emit is wired for the hd512 dc-combine arm only — a Some here would
19157                // leave the caller's pair unwritten (consumer would read garbage).
19158                assert!(
19159                    q8_out.is_none(),
19160                    "rows q8 emit requires the hd512 dc combine"
19161                );
19162                let fc = self.func("fa_decode_combine_rows");
19163                let __s_b2 = self.gpu.stream();
19164                let mut b2 = __s_b2.launch_builder(&fc);
19165                b2.arg(&*part_o)
19166                    .arg(&*part_m)
19167                    .arg(&*part_l)
19168                    .arg(&mut o_g)
19169                    .arg(&hd)
19170                    .arg(&nh)
19171                    .arg(&base_i)
19172                    .arg(&nspm)
19173                    .arg(&spk);
19174                unsafe {
19175                    b2.launch(cfg2)?;
19176                }
19177            }
19178        }
19179        Ok(())
19180    }
19181
19182    /// WINDOWED verify rows (gemma R6 deep-ctx): every row attends exactly `window` keys —
19183    /// bit-identical per row to the T=1 decode's fa_decode over the window VIEW. Caller gates
19184    /// base_len + 1 >= window (no under-window rows) and head_dim == 256 (v4 stamp).
19185    #[allow(clippy::too_many_arguments)]
19186    pub fn fa_decode_rows_w(
19187        &self,
19188        q: &CudaSlice<f32>,
19189        k: &cudarc::driver::CudaView<u8>,
19190        v: &cudarc::driver::CudaView<u8>,
19191        o: &mut CudaSlice<f32>,
19192        head_dim: usize,
19193        n_head: usize,
19194        n_head_kv: usize,
19195        base_dev: &CudaSlice<i32>,
19196        base_plus: i32,
19197        t: usize,
19198        scale: f32,
19199        window: usize,
19200        k_tok_bytes: usize,
19201        v_tok_bytes: usize,
19202        q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
19203    ) -> Result<(), Box<dyn std::error::Error>> {
19204        // DEVICE-LEN (graph arc step 1, 2026-07-11): the causal base rides an i32 counter
19205        // (kernel T_kv = dev[0] + base_plus + r + 1) so depth graphs can replay with len
19206        // advancing on-device. dc paths pass kvl.len_d with plus=-1; verify/eager sync the
19207        // counter with one async set_i32_one first. Partials/splits size from `window` (host).
19208        debug_assert!(head_dim == 256);
19209        // windowed split (MEMRA_FA_SPW, default 32 — re-swept 2026-07-12 under the raw-e4m3 sV
19210        // occupancy ceiling (4 blocks/SM): t=1 decode is GRID-limited (win/sp splits x nkv
19211        // blocks), so smaller splits fill the ceiling — 1.7k 174.4/174.0 vs 48's 170.7/170.3,
19212        // 4.9k 159.8 vs 157.4 (N=2 interleaved, stable window). Spec serving prefers 64
19213        // (verify t=K+1 fills the grid via grid.z=t; depth K=7 281.3 vs 249.3 at 32) — set
19214        // MEMRA_FA_SPW=64 there, same config law as MEMRA_GEMMA_GKV=0. MUST be one value for
19215        // ALL widths: a t-keyed probe broke decode-vs-verify combine order (stream 9/128).
19216        let sp = {
19217            static SPW: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
19218            let v = *SPW.get_or_init(|| {
19219                std::env::var("MEMRA_FA_SPW")
19220                    .ok()
19221                    .and_then(|x| x.parse().ok())
19222                    .unwrap_or(0)
19223            });
19224            if v >= 8 {
19225                v
19226            } else {
19227                FA_SPW_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
19228            }
19229        };
19230        let n_splits_max = (window + sp - 1) / sp;
19231        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
19232        let (nspm, spk, wini) = (n_splits_max as i32, sp as i32, window as i32);
19233        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
19234        let gqa = (n_head / n_head_kv).max(1) as u32;
19235        let o_len = t * n_head * n_splits_max * head_dim;
19236        let ml_len = t * n_head * n_splits_max;
19237        let mut part_guard = self.fa_part_pool.lock().unwrap();
19238        if part_guard
19239            .as_ref()
19240            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
19241            .unwrap_or(true)
19242        {
19243            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
19244            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
19245            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
19246            // later live allocations land at those addresses, and the next graph REPLAY writes
19247            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
19248            // output corruption began the burst after the trunk's t_kv growth first realloc'd
19249            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
19250            // the baked addresses alive (single-stream: eager writes the new buffers, replays
19251            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
19252            // (total retired < final size).
19253            let old = part_guard.take();
19254            let (co, cm) = old
19255                .as_ref()
19256                .map(|pp| (pp.0.len(), pp.1.len()))
19257                .unwrap_or((0, 0));
19258            if let Some(old) = old {
19259                self.fa_part_retired.lock().unwrap().push(old);
19260            }
19261            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
19262                eprintln!(
19263                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
19264                    co, o_len, cm, ml_len
19265                );
19266            }
19267            *part_guard = Some((
19268                self.alloc_uninit::<f32>(o_len.max(2 * co))?,
19269                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
19270                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
19271            ));
19272        }
19273        let pg = part_guard.as_mut().unwrap();
19274        self.gpu
19275            .stream()
19276            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
19277        self.gpu
19278            .stream()
19279            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
19280        self.gpu
19281            .stream()
19282            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
19283        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
19284        // Lane pick: decode AND verify both land here in the windowed regime (parity law —
19285        // hybrid_forward verify_attn), so the pick only needs internal consistency, not
19286        // clone-of-decode bit fidelity (SASS-proven impossible for textually identical
19287        // kernels, jsonl 2026-07-10). v4 under the threshold; smem twin at/above the smem
19288        // floor (deep-ctx broadcast win); register twin between.
19289        static SMEM_TKV_W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
19290        let smem_tkv = *SMEM_TKV_W.get_or_init(|| {
19291            std::env::var("MEMRA_FA_SMEM_TKV")
19292                .ok()
19293                .and_then(|v| v.parse().ok())
19294                .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
19295        });
19296        // MULTI-ROW v4: resurrected 2026-07-14 (the '33 tok/s collapse' was a paired-map
19297        // partial-write bug, not the mechanism) and falsified HONESTLY at gqa 2: bit-exact
19298        // but −1.7% on the 31B depth cell — the sp helper warp already hides staging
19299        // in-block, and mr trades L2-cheap redundant bytes for serialized per-warp gqa
19300        // score/B3 chains. Arm deleted; jsonl row 2026-07-14 is the record.
19301        use cudarc::driver::sys::CUfunction_attribute_enum as A;
19302        // FP8-WINDOWED (wkv): the v4 family is format-aware (2026-07-12 KFMT/VFMT staging
19303        // arms) — wkv rides the SAME lane logic, resolved from the kf8vf8 module. One symbol
19304        // per (lane, format-module) keeps parity structural; the old register-i2 detour
19305        // (-33%) is retired.
19306        let wg = Self::wkv_on();
19307        // STAGING-PARALLEL v4 (MEMRA_FA_SPW2, default ON at gqa==1): warp 1 = staging helper
19308        // (v4 is 61% staging); score phases identical to v4_w. Same symbol all t.
19309        let sp2 =
19310            gqa <= 4 && fa_v4_at(window) && std::env::var("MEMRA_FA_SPW2").as_deref() != Ok("0");
19311        if sp2 {
19312            let sh = (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32;
19313            if Self::pdl_on() && Self::pdl_wb_on() {
19314                // wave-B2b: flavor mirrors wg.
19315                use cudarc::driver::{DevicePtr, DevicePtrMut};
19316                let s = &self.gpu.stream();
19317                let (pq, _b0) = q.device_ptr(s);
19318                let (pk, _b1) = k.device_ptr(s);
19319                let (pv, _b2) = v.device_ptr(s);
19320                let (po, _b3) = part_o.device_ptr_mut(s);
19321                let (pm, _b4) = part_m.device_ptr_mut(s);
19322                let (pl, _b5) = part_l.device_ptr_mut(s);
19323                let (pb, _b6) = base_dev.device_ptr(s);
19324                let mut ps = [
19325                    &pq as *const _ as *mut std::ffi::c_void,
19326                    &pk as *const _ as *mut _,
19327                    &pv as *const _ as *mut _,
19328                    &po as *const _ as *mut _,
19329                    &pm as *const _ as *mut _,
19330                    &pl as *const _ as *mut _,
19331                    &hd as *const _ as *mut _,
19332                    &nh as *const _ as *mut _,
19333                    &nhkv as *const _ as *mut _,
19334                    &pb as *const _ as *mut _,
19335                    &base_plus as *const _ as *mut _,
19336                    &scale as *const _ as *mut _,
19337                    &nspm as *const _ as *mut _,
19338                    &spk as *const _ as *mut _,
19339                    &ktb as *const _ as *mut _,
19340                    &vtb as *const _ as *mut _,
19341                    &wini as *const _ as *mut _,
19342                ];
19343                unsafe {
19344                    self.launch_pdl_flash(
19345                        wg,
19346                        "fa_decode_vec_q_rows_v4_w_sp",
19347                        (n_head_kv as u32, n_splits_max as u32, t as u32),
19348                        (32, gqa + 1, 1),
19349                        sh,
19350                        &mut ps,
19351                    )?;
19352                }
19353            } else {
19354                let f = if wg {
19355                    self.func_g("fa_decode_vec_q_rows_v4_w_sp")
19356                } else {
19357                    self.func("fa_decode_vec_q_rows_v4_w_sp")
19358                };
19359                f.set_attribute(
19360                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19361                    sh as i32,
19362                )?;
19363                let cfg = LaunchConfig {
19364                    grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
19365                    block_dim: (32, gqa + 1, 1),
19366                    shared_mem_bytes: sh,
19367                };
19368                let __s_b = self.gpu.stream();
19369                let mut b = __s_b.launch_builder(&f);
19370                b.arg(q)
19371                    .arg(k)
19372                    .arg(v)
19373                    .arg(&mut *part_o)
19374                    .arg(&mut *part_m)
19375                    .arg(&mut *part_l)
19376                    .arg(&hd)
19377                    .arg(&nh)
19378                    .arg(&nhkv)
19379                    .arg(base_dev)
19380                    .arg(&base_plus)
19381                    .arg(&scale)
19382                    .arg(&nspm)
19383                    .arg(&spk)
19384                    .arg(&ktb)
19385                    .arg(&vtb)
19386                    .arg(&wini);
19387                unsafe {
19388                    b.launch(cfg)?;
19389                }
19390            }
19391        } else {
19392            if fa_v4_at(window) && Self::pdl_on() && Self::pdl_wb_on() {
19393                // wave-B2b: the v4_w pick only (smem/reg twins stay builder-launched).
19394                let sh = (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32;
19395                use cudarc::driver::{DevicePtr, DevicePtrMut};
19396                let s = &self.gpu.stream();
19397                let (pq, _b0) = q.device_ptr(s);
19398                let (pk, _b1) = k.device_ptr(s);
19399                let (pv, _b2) = v.device_ptr(s);
19400                let (po, _b3) = part_o.device_ptr_mut(s);
19401                let (pm, _b4) = part_m.device_ptr_mut(s);
19402                let (pl, _b5) = part_l.device_ptr_mut(s);
19403                let (pb, _b6) = base_dev.device_ptr(s);
19404                let mut ps = [
19405                    &pq as *const _ as *mut std::ffi::c_void,
19406                    &pk as *const _ as *mut _,
19407                    &pv as *const _ as *mut _,
19408                    &po as *const _ as *mut _,
19409                    &pm as *const _ as *mut _,
19410                    &pl as *const _ as *mut _,
19411                    &hd as *const _ as *mut _,
19412                    &nh as *const _ as *mut _,
19413                    &nhkv as *const _ as *mut _,
19414                    &pb as *const _ as *mut _,
19415                    &base_plus as *const _ as *mut _,
19416                    &scale as *const _ as *mut _,
19417                    &nspm as *const _ as *mut _,
19418                    &spk as *const _ as *mut _,
19419                    &ktb as *const _ as *mut _,
19420                    &vtb as *const _ as *mut _,
19421                    &wini as *const _ as *mut _,
19422                ];
19423                unsafe {
19424                    self.launch_pdl_flash(
19425                        wg,
19426                        "fa_decode_vec_q_rows_v4_w",
19427                        (n_head_kv as u32, n_splits_max as u32, t as u32),
19428                        (32, gqa, 1),
19429                        sh,
19430                        &mut ps,
19431                    )?;
19432                }
19433            } else {
19434                let pick = |name: &str| {
19435                    if wg {
19436                        self.func_g(name)
19437                    } else {
19438                        self.func(name)
19439                    }
19440                };
19441                let (f, sh) = if fa_v4_at(window) {
19442                    let f = pick("fa_decode_vec_q_rows_v4_w");
19443                    (f, (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32)
19444                } else if smem_tkv > 0 && window >= smem_tkv {
19445                    // NOTE: the smem twin's V-stage is still q5_1-hardcoded — unreachable under wkv
19446                    // at the gemma window (v4 covers it); revisit if the smem floor ever drops.
19447                    (
19448                        pick("fa_decode_vec_q_rows_smem_w"),
19449                        (2 * 32 * head_dim * 2) as u32,
19450                    )
19451                } else {
19452                    (pick("fa_decode_vec_q_rows_reg_w"), 0u32)
19453                };
19454                f.set_attribute(
19455                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19456                    sh as i32,
19457                )?;
19458                let cfg = LaunchConfig {
19459                    grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
19460                    block_dim: (32, gqa, 1),
19461                    shared_mem_bytes: sh,
19462                };
19463                let __s_b = self.gpu.stream();
19464                let mut b = __s_b.launch_builder(&f);
19465                b.arg(q)
19466                    .arg(k)
19467                    .arg(v)
19468                    .arg(&mut *part_o)
19469                    .arg(&mut *part_m)
19470                    .arg(&mut *part_l)
19471                    .arg(&hd)
19472                    .arg(&nh)
19473                    .arg(&nhkv)
19474                    .arg(base_dev)
19475                    .arg(&base_plus)
19476                    .arg(&scale)
19477                    .arg(&nspm)
19478                    .arg(&spk)
19479                    .arg(&ktb)
19480                    .arg(&vtb)
19481                    .arg(&wini);
19482                unsafe {
19483                    b.launch(cfg)?;
19484                }
19485            }
19486        }
19487        let cfg2 = LaunchConfig {
19488            grid_dim: (n_head as u32, t as u32, 1),
19489            block_dim: (head_dim as u32, 1, 1),
19490            shared_mem_bytes: 0,
19491        };
19492        if let Some((oq, od)) = q8_out {
19493            // wave-5b port (2026-07-23): q8-emitting combine — the t=1 decode's wo matvec
19494            // consumes the pair directly; the standalone quantize launch folds away.
19495            if Self::pdl_on() && Self::pdl_wb_on() {
19496                // wave-B2: flavor mirrors the builder's wg choice.
19497                use cudarc::driver::{DevicePtr, DevicePtrMut};
19498                let s = &self.gpu.stream();
19499                let (po, _g0) = part_o.device_ptr(s);
19500                let (pm, _g1) = part_m.device_ptr(s);
19501                let (pl, _g2) = part_l.device_ptr(s);
19502                let (pq, _g3) = oq.device_ptr_mut(s);
19503                let (pd, _g4) = od.device_ptr_mut(s);
19504                let mut ps = [
19505                    &po as *const _ as *mut std::ffi::c_void,
19506                    &pm as *const _ as *mut _,
19507                    &pl as *const _ as *mut _,
19508                    &pq as *const _ as *mut _,
19509                    &pd as *const _ as *mut _,
19510                    &hd as *const _ as *mut _,
19511                    &nh as *const _ as *mut _,
19512                    &nspm as *const _ as *mut _,
19513                    &spk as *const _ as *mut _,
19514                    &wini as *const _ as *mut _,
19515                ];
19516                unsafe {
19517                    self.launch_pdl_flash(
19518                        wg,
19519                        "fa_decode_combine_rows_w_q8_1",
19520                        cfg2.grid_dim,
19521                        cfg2.block_dim,
19522                        0,
19523                        &mut ps,
19524                    )?;
19525                }
19526                return Ok(());
19527            }
19528            let fc = if wg {
19529                self.func_g("fa_decode_combine_rows_w_q8_1")
19530            } else {
19531                self.func("fa_decode_combine_rows_w_q8_1")
19532            };
19533            let __s_b2 = self.gpu.stream();
19534            let mut b2 = __s_b2.launch_builder(&fc);
19535            b2.arg(&*part_o)
19536                .arg(&*part_m)
19537                .arg(&*part_l)
19538                .arg(oq)
19539                .arg(od)
19540                .arg(&hd)
19541                .arg(&nh)
19542                .arg(&nspm)
19543                .arg(&spk)
19544                .arg(&wini);
19545            unsafe {
19546                b2.launch(cfg2)?;
19547            }
19548            return Ok(());
19549        }
19550        let fc = if wg {
19551            self.func_g("fa_decode_combine_rows_w")
19552        } else {
19553            self.func("fa_decode_combine_rows_w")
19554        };
19555        let __s_b2 = self.gpu.stream();
19556        let mut b2 = __s_b2.launch_builder(&fc);
19557        b2.arg(&*part_o)
19558            .arg(&*part_m)
19559            .arg(&*part_l)
19560            .arg(o)
19561            .arg(&hd)
19562            .arg(&nh)
19563            .arg(&nspm)
19564            .arg(&spk)
19565            .arg(&wini);
19566        unsafe {
19567            b2.launch(cfg2)?;
19568        }
19569        Ok(())
19570    }
19571
19572    /// ROUND-STREAM stage (c): fa rows with the causal base from a device counter. Two lanes:
19573    /// v3 (qwen stream, fa_v3_active) and v4 (gemma hd256 burst — rows_v4_dc, g-module aware);
19574    /// `t_kv_upper` sizes splits/partials — the same one-sp-for-all-rows approximation class
19575    /// the host rows path already uses (battery-arbitrated); actual per-row bounds derive
19576    /// in-kernel from the counter (+ base_plus, v4 lane only — v3's kernel has no plus arg).
19577    #[allow(clippy::too_many_arguments)]
19578    pub fn fa_decode_rows_dc(
19579        &self,
19580        q: &CudaSlice<f32>,
19581        k: &cudarc::driver::CudaView<u8>,
19582        v: &cudarc::driver::CudaView<u8>,
19583        o: &mut CudaSlice<f32>,
19584        head_dim: usize,
19585        n_head: usize,
19586        n_head_kv: usize,
19587        base_dev: &CudaSlice<i32>,
19588        t_kv_upper: usize,
19589        t: usize,
19590        scale: f32,
19591        k_tok_bytes: usize,
19592        v_tok_bytes: usize,
19593        base_plus: i32,
19594        g: bool,
19595    ) -> Result<(), Box<dyn std::error::Error>> {
19596        let v4 = head_dim == 256 && fa_v4_at(t_kv_upper);
19597        assert!(
19598            v4 || fa_v3_active(head_dim),
19599            "stream fa rows requires the v3 or v4 lane"
19600        );
19601        assert!(v4 || base_plus == 0, "v3_dc kernel takes no plus arg");
19602        if v4 {
19603            let sp = fa_split_keys(t_kv_upper, n_head_kv);
19604            let n_splits_max = (t_kv_upper + sp - 1) / sp;
19605            let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
19606            let (nspm, spk) = (n_splits_max as i32, sp as i32);
19607            let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
19608            let gqa = (n_head / n_head_kv).max(1) as u32;
19609            let o_len = t * n_head * n_splits_max * head_dim;
19610            let ml_len = t * n_head * n_splits_max;
19611            let mut part_guard = self.fa_part_pool.lock().unwrap();
19612            if part_guard
19613                .as_ref()
19614                .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
19615                .unwrap_or(true)
19616            {
19617                // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
19618                // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
19619                // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
19620                // later live allocations land at those addresses, and the next graph REPLAY writes
19621                // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
19622                // output corruption began the burst after the trunk's t_kv growth first realloc'd
19623                // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
19624                // the baked addresses alive (single-stream: eager writes the new buffers, replays
19625                // touch only the old — never concurrently). Doubling growth bounds retired VRAM
19626                // (total retired < final size).
19627                let old = part_guard.take();
19628                let (co, cm) = old
19629                    .as_ref()
19630                    .map(|pp| (pp.0.len(), pp.1.len()))
19631                    .unwrap_or((0, 0));
19632                if let Some(old) = old {
19633                    self.fa_part_retired.lock().unwrap().push(old);
19634                }
19635                if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
19636                    eprintln!(
19637                        "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
19638                        co, o_len, cm, ml_len
19639                    );
19640                }
19641                *part_guard = Some((
19642                    self.alloc_uninit::<f32>(o_len.max(2 * co))?,
19643                    self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
19644                    self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
19645                ));
19646            }
19647            let pg = part_guard.as_mut().unwrap();
19648            self.gpu
19649                .stream()
19650                .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
19651            self.gpu
19652                .stream()
19653                .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
19654            self.gpu
19655                .stream()
19656                .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
19657            let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
19658            let f = if g {
19659                self.func_g("fa_decode_vec_q_rows_v4_dc")
19660            } else {
19661                self.func("fa_decode_vec_q_rows_v4_dc")
19662            };
19663            let sh = (11520 + 32 * head_dim * if g { 1 } else { 2 }) as u32;
19664            use cudarc::driver::sys::CUfunction_attribute_enum as A;
19665            f.set_attribute(
19666                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19667                sh as i32,
19668            )?;
19669            let cfg = LaunchConfig {
19670                grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
19671                block_dim: (32, gqa, 1),
19672                shared_mem_bytes: sh,
19673            };
19674            let __s_b = self.gpu.stream();
19675            let mut b = __s_b.launch_builder(&f);
19676            b.arg(q)
19677                .arg(k)
19678                .arg(v)
19679                .arg(&mut *part_o)
19680                .arg(&mut *part_m)
19681                .arg(&mut *part_l)
19682                .arg(&hd)
19683                .arg(&nh)
19684                .arg(&nhkv)
19685                .arg(base_dev)
19686                .arg(&base_plus)
19687                .arg(&scale)
19688                .arg(&nspm)
19689                .arg(&spk)
19690                .arg(&ktb)
19691                .arg(&vtb);
19692            unsafe {
19693                b.launch(cfg)?;
19694            }
19695            let fc = self.func("fa_decode_combine_rows_dc");
19696            let cfg2 = LaunchConfig {
19697                grid_dim: (n_head as u32, t as u32, 1),
19698                block_dim: (head_dim as u32, 1, 1),
19699                shared_mem_bytes: 0,
19700            };
19701            let __s_b2 = self.gpu.stream();
19702            let mut b2 = __s_b2.launch_builder(&fc);
19703            b2.arg(&*part_o)
19704                .arg(&*part_m)
19705                .arg(&*part_l)
19706                .arg(o)
19707                .arg(&hd)
19708                .arg(&nh)
19709                .arg(base_dev)
19710                .arg(&base_plus)
19711                .arg(&nspm)
19712                .arg(&spk);
19713            unsafe {
19714                b2.launch(cfg2)?;
19715            }
19716            return Ok(());
19717        }
19718        let sp = fa_split_keys(t_kv_upper, n_head_kv);
19719        let n_splits_max = (t_kv_upper + sp - 1) / sp;
19720        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
19721        let (nspm, spk) = (n_splits_max as i32, sp as i32);
19722        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
19723        let gqa = (n_head / n_head_kv).max(1) as u32;
19724        let o_len = t * n_head * n_splits_max * head_dim;
19725        let ml_len = t * n_head * n_splits_max;
19726        let mut part_guard = self.fa_part_pool.lock().unwrap();
19727        if part_guard
19728            .as_ref()
19729            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
19730            .unwrap_or(true)
19731        {
19732            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
19733            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
19734            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
19735            // later live allocations land at those addresses, and the next graph REPLAY writes
19736            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
19737            // output corruption began the burst after the trunk's t_kv growth first realloc'd
19738            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
19739            // the baked addresses alive (single-stream: eager writes the new buffers, replays
19740            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
19741            // (total retired < final size).
19742            let old = part_guard.take();
19743            let (co, cm) = old
19744                .as_ref()
19745                .map(|pp| (pp.0.len(), pp.1.len()))
19746                .unwrap_or((0, 0));
19747            if let Some(old) = old {
19748                self.fa_part_retired.lock().unwrap().push(old);
19749            }
19750            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
19751                eprintln!(
19752                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
19753                    co, o_len, cm, ml_len
19754                );
19755            }
19756            *part_guard = Some((
19757                self.alloc_uninit::<f32>(o_len.max(2 * co))?,
19758                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
19759                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
19760            ));
19761        }
19762        let pg = part_guard.as_mut().unwrap();
19763        self.gpu
19764            .stream()
19765            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
19766        self.gpu
19767            .stream()
19768            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
19769        self.gpu
19770            .stream()
19771            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
19772        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
19773        let f = self.func("fa_decode_vec_q_rows_v3_dc");
19774        let sh = (32 * head_dim * 2) as u32;
19775        use cudarc::driver::sys::CUfunction_attribute_enum as A;
19776        f.set_attribute(
19777            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19778            sh as i32,
19779        )?;
19780        let cfg = LaunchConfig {
19781            grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
19782            block_dim: (32, gqa, 1),
19783            shared_mem_bytes: sh,
19784        };
19785        let __s_b = self.gpu.stream();
19786        let mut b = __s_b.launch_builder(&f);
19787        b.arg(q)
19788            .arg(k)
19789            .arg(v)
19790            .arg(&mut *part_o)
19791            .arg(&mut *part_m)
19792            .arg(&mut *part_l)
19793            .arg(&hd)
19794            .arg(&nh)
19795            .arg(&nhkv)
19796            .arg(base_dev)
19797            .arg(&scale)
19798            .arg(&nspm)
19799            .arg(&spk)
19800            .arg(&ktb)
19801            .arg(&vtb);
19802        unsafe {
19803            b.launch(cfg)?;
19804        }
19805        let fc = self.func("fa_decode_combine_rows_dc");
19806        let cfg2 = LaunchConfig {
19807            grid_dim: (n_head as u32, t as u32, 1),
19808            block_dim: (head_dim as u32, 1, 1),
19809            shared_mem_bytes: 0,
19810        };
19811        let plus0 = 0i32;
19812        let __s_b2 = self.gpu.stream();
19813        let mut b2 = __s_b2.launch_builder(&fc);
19814        b2.arg(&*part_o)
19815            .arg(&*part_m)
19816            .arg(&*part_l)
19817            .arg(o)
19818            .arg(&hd)
19819            .arg(&nh)
19820            .arg(base_dev)
19821            .arg(&plus0)
19822            .arg(&nspm)
19823            .arg(&spk);
19824        unsafe {
19825            b2.launch(cfg2)?;
19826        }
19827        Ok(())
19828    }
19829
19830    /// Device-counter variant of `fa_decode` (CUDA-GRAPH-PLAN Phase 2). The sequence length is read
19831    /// from `t_kv_dev[0]` (resident device i32[1]) for the attention loop bound + per-split key range;
19832    /// the GRID `n_splits` is sized for `bucket_max` (the bucket's max t_kv — baked at capture time).
19833    /// Empty splits (key range beyond the actual t_kv) write an empty partial (m=NEG_INF) so the
19834    /// shared combine skips them -> bit-correct for ANY actual t_kv <= bucket_max.
19835    ///
19836    /// BIT-IDENTITY (the gate): pass `bucket_max == actual_t_kv` and this reproduces `fa_decode`
19837    /// EXACTLY (same n_splits, same per, same split boundaries, same combine) while reading t_kv from
19838    /// device. Bucketing (bucket_max > t_kv) is for the future captured path and changes split
19839    /// grouping (different but mathematically-equal log-sum-exp merge).
19840    pub fn fa_decode_dc(
19841        &self,
19842        q: &CudaSlice<f32>,
19843        k: &cudarc::driver::CudaView<u8>,
19844        v: &cudarc::driver::CudaView<u8>,
19845        o: &mut CudaSlice<f32>,
19846        head_dim: usize,
19847        n_head: usize,
19848        n_head_kv: usize,
19849        t_kv_dev: &CudaSlice<i32>,
19850        bucket_max: usize,
19851        scale: f32,
19852        k_tok_bytes: usize,
19853        v_tok_bytes: usize,
19854        g: bool,
19855    ) -> Result<(), Box<dyn std::error::Error>> {
19856        self.fa_decode_dc_q8(
19857            q,
19858            k,
19859            v,
19860            o,
19861            head_dim,
19862            n_head,
19863            n_head_kv,
19864            t_kv_dev,
19865            bucket_max,
19866            scale,
19867            k_tok_bytes,
19868            v_tok_bytes,
19869            g,
19870            None,
19871        )
19872    }
19873
19874    /// `fa_decode_dc` with an optional q8_1 sink (wave 5b): when `q8_out` is given the
19875    /// combine emits (int8, per-32 scales) for the wo matmul_pre and skips the f32 O write.
19876    #[allow(clippy::too_many_arguments)]
19877    pub fn fa_decode_dc_q8(
19878        &self,
19879        q: &CudaSlice<f32>,
19880        k: &cudarc::driver::CudaView<u8>,
19881        v: &cudarc::driver::CudaView<u8>,
19882        o: &mut CudaSlice<f32>,
19883        head_dim: usize,
19884        n_head: usize,
19885        n_head_kv: usize,
19886        t_kv_dev: &CudaSlice<i32>,
19887        bucket_max: usize,
19888        scale: f32,
19889        k_tok_bytes: usize,
19890        v_tok_bytes: usize,
19891        g: bool,
19892        q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
19893    ) -> Result<(), Box<dyn std::error::Error>> {
19894        // The fa_vec gate + n_splits are sized from bucket_max (host, fixed at capture). The kernel
19895        // reads the ACTUAL t_kv from t_kv_dev for the per-split bound. DEFAULT-ON to MATCH the eager
19896        // `fa_decode` gate above — graph capture must mirror eager's kernel choice or the graph-vs-eager
19897        // bit-identity gate breaks. MEMRA_NO_FA_VEC forces scalar on BOTH paths in lockstep.
19898        // `g` = this layer's cache is e4m3 (gemma windowed under wkv) — every pick below must
19899        // mirror fa_decode_kvmod's g-routing or the graph diverges from eager (short/mid 1/96,
19900        // 2026-07-12).
19901        let mut fa_vec =
19902            std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
19903        if g && head_dim == 256 && !fa_v4_at(bucket_max) {
19904            fa_vec = false;
19905        } // mirror kvmod/geom
19906        let sp = fa_split_keys(bucket_max, n_head_kv);
19907        let n_splits = if fa_vec {
19908            ((bucket_max + sp - 1) / sp).max(1)
19909        } else {
19910            ((bucket_max + 255) / 256).max(1)
19911        };
19912        let o_len = n_head * n_splits * head_dim;
19913        let ml_len = n_head * n_splits;
19914        let mut part_guard = self.fa_part_pool.lock().unwrap();
19915        if part_guard
19916            .as_ref()
19917            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
19918            .unwrap_or(true)
19919        {
19920            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
19921            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
19922            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
19923            // later live allocations land at those addresses, and the next graph REPLAY writes
19924            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
19925            // output corruption began the burst after the trunk's t_kv growth first realloc'd
19926            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
19927            // the baked addresses alive (single-stream: eager writes the new buffers, replays
19928            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
19929            // (total retired < final size).
19930            let old = part_guard.take();
19931            let (co, cm) = old
19932                .as_ref()
19933                .map(|pp| (pp.0.len(), pp.1.len()))
19934                .unwrap_or((0, 0));
19935            if let Some(old) = old {
19936                self.fa_part_retired.lock().unwrap().push(old);
19937            }
19938            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
19939                eprintln!(
19940                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
19941                    co, o_len, cm, ml_len
19942                );
19943            }
19944            *part_guard = Some((
19945                self.alloc_uninit::<f32>(o_len.max(2 * co))?,
19946                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
19947                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
19948            ));
19949        }
19950        let pg = part_guard.as_mut().unwrap();
19951        self.gpu
19952            .stream()
19953            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
19954        self.gpu
19955            .stream()
19956            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
19957        self.gpu
19958            .stream()
19959            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
19960        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
19961        let (hd, nh, nhkv, nsp) = (
19962            head_dim as i32,
19963            n_head as i32,
19964            n_head_kv as i32,
19965            n_splits as i32,
19966        );
19967        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
19968        let fa_vec = fa_vec && head_dim <= 512 && head_dim % 32 == 0;
19969        // FA-DEEP pick keyed on bucket_max (the fa_v4_at precedent) — bit-identical twins,
19970        // so a threshold falling between t_kv and bucket_max cannot diverge eager-vs-graph.
19971        let deep = fa_vec
19972            && head_dim == 256
19973            && fa_v4_at(bucket_max)
19974            && !g
19975            && fa_deep_at(bucket_max)
19976            && !matches!(fa_v4_mode(), "noB3" | "stage");
19977        let (f, cfg) = if fa_vec
19978            && head_dim == 512
19979            && bucket_max >= {
19980                static FA512_MIN_DC: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
19981                *FA512_MIN_DC.get_or_init(|| {
19982                    std::env::var("MEMRA_FA512_MIN")
19983                        .ok()
19984                        .and_then(|v| v.parse().ok())
19985                        .unwrap_or(512)
19986                })
19987            } {
19988            // gemma globals dc twin (mirror the eager dpl16 pick incl the crossover floor).
19989            let gqa = (n_head / n_head_kv).max(1) as u32;
19990            (
19991                self.fa_func("fa_decode_vec_q_dpl16_dc", head_dim),
19992                LaunchConfig {
19993                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
19994                    block_dim: (32, gqa, 1),
19995                    shared_mem_bytes: 0,
19996                },
19997            )
19998        } else if fa_vec && head_dim == 512 {
19999            // under the 512 floor eager runs scalar — the SAME unified symbol, ctr non-null;
20000            // ns_eff in-kernel reproduces eager's ceil(t_kv/sp) partition for the LIVE len.
20001            let q_view = q.as_view();
20002            let mut o_view = o.as_view_mut();
20003            return self.fa_decode_scalar_unified(
20004                &q_view,
20005                k,
20006                v,
20007                &mut o_view,
20008                head_dim,
20009                n_head,
20010                n_head_kv,
20011                0,
20012                Some(t_kv_dev),
20013                scale,
20014                n_splits,
20015                sp,
20016                k_tok_bytes,
20017                v_tok_bytes,
20018                g,
20019                &mut *part_o,
20020                &mut *part_m,
20021                &mut *part_l,
20022                q8_out,
20023            );
20024        } else if fa_vec && head_dim == 256 && fa_v4_at(bucket_max) {
20025            // gemma/qwen v4 dc twin (eager default lane) — capture must mirror eager's pick,
20026            // incl the g-module route + raw-e4m3 sV sizing.
20027            let gqa = (n_head / n_head_kv).max(1) as u32;
20028            let fv = if g {
20029                self.func_g("fa_decode_vec_q_v4_dc")
20030            } else if deep {
20031                self.func("fa_decode_vec_q_v4_deep_dc")
20032            } else {
20033                self.func("fa_decode_vec_q_v4_dc")
20034            };
20035            let shmem =
20036                (if deep { 12160 } else { 11520 } + 32 * head_dim * if g { 1 } else { 2 }) as u32;
20037            use cudarc::driver::sys::CUfunction_attribute_enum as A;
20038            fv.set_attribute(
20039                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20040                shmem as i32,
20041            )?;
20042            (
20043                fv,
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: shmem,
20048                },
20049            )
20050        } else if fa_vec && fa_v3_active(head_dim) {
20051            // FA v3 lane _dc twin: the captured graph must run the SAME walk body as eager
20052            // under MEMRA_FA_V3=1 (eager, rows-verify and graph switch together).
20053            let gqa = (n_head / n_head_kv).max(1) as u32;
20054            let fv = if g {
20055                self.func_g("fa_decode_vec_q_v3_dc")
20056            } else {
20057                self.func("fa_decode_vec_q_v3_dc")
20058            };
20059            let shmem = (32 * head_dim * 2) as u32; // sV bf16 [FA_DEC_TILE=32][hd]
20060            (
20061                fv,
20062                LaunchConfig {
20063                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
20064                    block_dim: (32, gqa, 1),
20065                    shared_mem_bytes: shmem,
20066                },
20067            )
20068        } else if fa_vec && fa_v2_on() {
20069            // FAVENDOR lane: v2 _dc twin — the captured graph must run the SAME walk body as
20070            // eager under MEMRA_FA_V2=1 or graph_decode_gate's bit-identity breaks (the flag is
20071            // a numeric config; eager, rows-verify and graph all switch together).
20072            let gqa = (n_head / n_head_kv).max(1) as u32;
20073            let fv = if g {
20074                self.func_g("fa_decode_vec_q_v2_dc")
20075            } else {
20076                self.func("fa_decode_vec_q_v2_dc")
20077            };
20078            let shmem = (2 * 32 * head_dim * 2) as u32; // sK+sV bf16 [FA_DEC_TILE=32][hd]
20079            (
20080                fv,
20081                LaunchConfig {
20082                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
20083                    block_dim: (32, gqa, 1),
20084                    shared_mem_bytes: shmem,
20085                },
20086            )
20087        } else if fa_vec {
20088            let gqa = (n_head / n_head_kv).max(1) as u32;
20089            // REGISTER-DEQUANT twin: zero dynamic smem (see fa_decode above).
20090            let fv = if g {
20091                self.func_g("fa_decode_vec_q_dc")
20092            } else {
20093                self.func("fa_decode_vec_q_dc")
20094            };
20095            (
20096                fv,
20097                LaunchConfig {
20098                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
20099                    block_dim: (32, gqa, 1),
20100                    shared_mem_bytes: 0,
20101                },
20102            )
20103        } else {
20104            let q_view = q.as_view();
20105            let mut o_view = o.as_view_mut();
20106            return self.fa_decode_scalar_unified(
20107                &q_view,
20108                k,
20109                v,
20110                &mut o_view,
20111                head_dim,
20112                n_head,
20113                n_head_kv,
20114                0,
20115                Some(t_kv_dev),
20116                scale,
20117                n_splits,
20118                if fa_vec { sp } else { 256 },
20119                k_tok_bytes,
20120                v_tok_bytes,
20121                g,
20122                &mut *part_o,
20123                &mut *part_m,
20124                &mut *part_l,
20125                q8_out,
20126            );
20127        };
20128        let ski = sp as i32; // one-partition law: the twins derive ns_eff from (T_kv, ski)
20129        let __s_b = self.gpu.stream();
20130        let mut b = __s_b.launch_builder(&f);
20131        b.arg(q)
20132            .arg(k)
20133            .arg(v)
20134            .arg(&mut *part_o)
20135            .arg(&mut *part_m)
20136            .arg(&mut *part_l)
20137            .arg(&hd)
20138            .arg(&nh)
20139            .arg(&nhkv)
20140            .arg(t_kv_dev)
20141            .arg(&scale)
20142            .arg(&nsp)
20143            .arg(&ski)
20144            .arg(&ktb)
20145            .arg(&vtb);
20146        unsafe {
20147            b.launch(cfg)?;
20148        }
20149        let cfg2 = LaunchConfig {
20150            grid_dim: (n_head as u32, 1, 1),
20151            block_dim: (head_dim as u32, 1, 1),
20152            shared_mem_bytes: 0,
20153        };
20154        if let Some((oq, od)) = q8_out {
20155            let fc = if g {
20156                self.func_g("fa_decode_combine_q8_1")
20157            } else {
20158                self.fa_func("fa_decode_combine_q8_1", head_dim)
20159            };
20160            let __s_b2 = self.gpu.stream();
20161            let mut b2 = __s_b2.launch_builder(&fc);
20162            b2.arg(&*part_o)
20163                .arg(&*part_m)
20164                .arg(&*part_l)
20165                .arg(oq)
20166                .arg(od)
20167                .arg(&hd)
20168                .arg(&nh)
20169                .arg(&nsp);
20170            unsafe {
20171                b2.launch(cfg2)?;
20172            }
20173            return Ok(());
20174        }
20175        let fc = if g {
20176            self.func_g("fa_decode_combine_f32")
20177        } else {
20178            self.fa_func("fa_decode_combine_f32", head_dim)
20179        };
20180        let __s_b2 = self.gpu.stream();
20181        let mut b2 = __s_b2.launch_builder(&fc);
20182        b2.arg(&*part_o)
20183            .arg(&*part_m)
20184            .arg(&*part_l)
20185            .arg(o)
20186            .arg(&hd)
20187            .arg(&nh)
20188            .arg(&nsp);
20189        unsafe {
20190            b2.launch(cfg2)?;
20191        }
20192        Ok(())
20193    }
20194
20195    /// EAGER fa_decode geometry for a given actual `t_kv` (CUDA-GRAPH-PLAN §3.3 bucketing). Returns
20196    /// `(fa_vec, n_splits)` EXACTLY as `fa_decode` computes them so the graph-capture path can key its
20197    /// bucket on the same `(kernel, n_splits)` pair and pass a `bucket_max` that reproduces eager's
20198    /// n_splits bit-for-bit. (Per = ceil(t_kv/n_splits) is then recomputed from the DEVICE t_kv inside
20199    /// the kernel and matches eager when n_splits matches — the bit-identity contract.)
20200    pub fn fa_geom_eager(
20201        &self,
20202        t_kv: usize,
20203        head_dim: usize,
20204        n_head_kv: usize,
20205        g: bool,
20206    ) -> (bool, usize) {
20207        // MUST mirror `fa_decode` / `fa_decode_dc` (default-ON 2026-06-28). This is the bucket-key
20208        // source: if it disagrees with the actual kernel pick, the graph captures the wrong path and
20209        // replay diverges from eager. All three sites read MEMRA_NO_FA_VEC in lockstep.
20210        let fa_ok = std::env::var("MEMRA_NO_FA_VEC").is_err() && t_kv >= fa_vec_min_tkv();
20211        // hd512 dpl16 vec lane (gemma globals, 2026-07-11 graph-arc fix): the original key
20212        // hardcoded vec = hd<=256, so for hd512 it bucketed by the SCALAR 256-key splits while
20213        // the dpl16/rows_dpl16 kernels split by the ladder — n_splits changed WITHIN a bucket
20214        // (mid-ctx graph mismatch at pos 19 + partials OOB at longer runs). Mirror the real
20215        // fa_decode dispatch: vec512 above the fa512 floor, vec256 as before.
20216        let vec512 = fa_ok && head_dim == 512 && t_kv >= fa512_min_tkv();
20217        let mut fa_vec = vec512 || (fa_ok && head_dim <= 256 && head_dim % 32 == 0);
20218        // g (fp8-windowed): mirror kvmod's clamp — only the v4 lane parses e4m3 in the vec
20219        // family; everything else falls to the g-module scalar.
20220        // hd256 under g: v4-or-scalar. hd128/other under g ride the REGISTER g-lane (the
20221        // dq_K_lane/dq_V_lane macros are format-aware; v3 is format-gated off, v2/smem
20222        // arms are excluded under g). Mirrored in kvmod / fa_decode_dc / fa_geom_eager.
20223        if g && head_dim == 256 && !fa_v4_at(t_kv) {
20224            fa_vec = false;
20225        }
20226        let sp = fa_split_keys(t_kv, n_head_kv);
20227        let n_splits = if fa_vec {
20228            ((t_kv + sp - 1) / sp).max(1)
20229        } else {
20230            ((t_kv + 255) / 256).max(1)
20231        };
20232        (fa_vec, n_splits)
20233    }
20234
20235    /// `bucket_max` (host t_kv to feed `fa_decode_dc` / `full_attn_decode_dc`) that makes the _dc
20236    /// kernel pick the SAME (fa_vec, n_splits) as eager would for actual `t_kv`. Because the dc
20237    /// launcher derives both from `bucket_max` via the same formulas, we just hand it `t_kv` itself:
20238    /// the n_splits is then identical, and the per-split boundaries (computed from the DEVICE t_kv in
20239    /// the kernel) match eager exactly. The bucket KEY (for the graph HashMap) is `(fa_vec, n_splits)`.
20240    pub fn fa_bucket_key(
20241        &self,
20242        t_kv: usize,
20243        head_dim: usize,
20244        n_head_kv: usize,
20245        g: bool,
20246    ) -> (bool, usize) {
20247        self.fa_geom_eager(t_kv, head_dim, n_head_kv, g)
20248    }
20249
20250    /// CUDA-graph capture wrapper (CUDA-GRAPH-PLAN §3.2, llama.cpp warmup pattern). Runs `step`
20251    /// inline TWICE (warmup — lets the caching allocator settle to stable pointers and any one-time
20252    /// kernel attribute/JIT happen outside capture), then captures a THIRD invocation on the Engine's
20253    /// decode stream (RELAXED mode) and instantiates it into a replayable `CudaGraph`. The closure
20254    /// must enqueue ONLY device work on `e.stream()` (no dtoh / no synchronize / no host branch on
20255    /// device data) — every per-step varying scalar must come from a device counter. Returns the
20256    /// instantiated graph; `CudaGraph::launch()` replays the whole step in one dispatch.
20257    /// `capture_graph` with CAPTURE-RETAIN: every Engine allocation made during the warmups
20258    /// and the capture is kept alive in the returned keeper — hold it as long as the graph
20259    /// replays (transients returning to the pool get reused by unrelated work and corrupt
20260    /// replays; the draft-graph root cause). Model-generic, next capture reuses it.
20261    pub fn capture_graph_retained<F>(
20262        &self,
20263        step: F,
20264    ) -> Result<
20265        (
20266            cudarc::driver::CudaGraph,
20267            Vec<Box<dyn std::any::Any + Send>>,
20268        ),
20269        Box<dyn std::error::Error>,
20270    >
20271    where
20272        F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
20273    {
20274        use cudarc::driver::sys::CUgraphInstantiate_flags;
20275        self.capture_graph_retained_flags(
20276            CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
20277            step,
20278        )
20279    }
20280
20281    /// Retained capture with an explicit instantiate flag. ALLOC-FREE captured graphs
20282    /// (zero mem nodes — the gemma slotted door) should pass UPLOAD instead of
20283    /// AUTO_FREE_ON_LAUNCH: the auto-free flag's launch-time mem-pool scan was measured at
20284    /// ~0.25us/node (205us on the 826-node step) even with nothing to free.
20285    pub fn capture_graph_retained_flags<F>(
20286        &self,
20287        flags: cudarc::driver::sys::CUgraphInstantiate_flags,
20288        mut step: F,
20289    ) -> Result<
20290        (
20291            cudarc::driver::CudaGraph,
20292            Vec<Box<dyn std::any::Any + Send>>,
20293        ),
20294        Box<dyn std::error::Error>,
20295    >
20296    where
20297        F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
20298    {
20299        use cudarc::driver::sys::CUstreamCaptureMode;
20300        // KEEP scope = WARMUPS ONLY (2026-07-13): keep_if_capturing retains via
20301        // CudaSlice::clone, which is a device ALLOC + D2D COPY on the stream — clones made
20302        // while the capture region is open become dead copy NODES replayed every launch
20303        // (E4B: 1440 copies = 0.74ms/token, the whole graph-vs-eager regression). The
20304        // warmup runs allocate the same transient sequence at the same pool addresses, so
20305        // retaining the warmup clones preserves the draft-graph fix without polluting the
20306        // captured graph.
20307        self.capture_keep.lock().unwrap().clear();
20308        let was_tracking = self.gpu.ctx.is_event_tracking();
20309        if was_tracking {
20310            unsafe {
20311                self.gpu.ctx.disable_event_tracking();
20312            }
20313        }
20314        let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
20315            self.capture_keep_on
20316                .store(true, std::sync::atomic::Ordering::Relaxed);
20317            let w = (|| {
20318                step(self)?;
20319                step(self)
20320            })();
20321            self.capture_keep_on
20322                .store(false, std::sync::atomic::Ordering::Relaxed);
20323            w?;
20324            self.gpu.stream().synchronize()?;
20325            self.gpu
20326                .stream()
20327                .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
20328            let r = step(self);
20329            let g = self.gpu.stream().end_capture(flags);
20330            r?;
20331            let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
20332            graph.upload()?;
20333            Ok(graph)
20334        };
20335        let result = run();
20336        self.capture_keep_on
20337            .store(false, std::sync::atomic::Ordering::Relaxed);
20338        if was_tracking {
20339            unsafe {
20340                self.gpu.ctx.enable_event_tracking();
20341            }
20342        }
20343        let keeper = std::mem::take(&mut *self.capture_keep.lock().unwrap());
20344        Ok((result?, keeper))
20345    }
20346
20347    pub fn capture_graph<F>(
20348        &self,
20349        mut step: F,
20350    ) -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>>
20351    where
20352        F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
20353    {
20354        use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
20355        // EVENT TRACKING OFF for capture. The Engine creates a 2nd stream (copy_stream) so cudarc is in
20356        // multi-stream mode and, by default, records a CudaEvent per CudaSlice alloc/use to serialize
20357        // cross-stream access. Those per-buffer event waits issue stream ops that are NOT permitted
20358        // inside a capture region (CUDA_ERROR_STREAM_CAPTURE_UNSUPPORTED). The captured decode step is
20359        // strictly SINGLE-STREAM (every kernel on gpu.stream), so this synchronization is unnecessary
20360        // here — disable it for the whole warmup+capture, re-enable after. SAFETY: the decode-dc path
20361        // touches only gpu.stream; no buffer crosses to copy_stream during capture.
20362        let was_tracking = self.gpu.ctx.is_event_tracking();
20363        if was_tracking {
20364            unsafe {
20365                self.gpu.ctx.disable_event_tracking();
20366            }
20367        }
20368        // Q1 PROBE (MEMRA_GRAPH_IFLAG): the generic capture body's cuMemAllocAsync nodes are
20369        // EXACTLY BALANCED by in-graph free nodes (measured census q27: 1589 ALLOC / 1589
20370        // FREE), so AUTO_FREE_ON_LAUNCH has nothing to reclaim at launch — it only pays its
20371        // per-node launch-time mem-pool scan. `upload` / `none` select the alternatives to
20372        // measure that scan's real cost on the generic path. Diagnostic door only; the
20373        // default stays AUTO_FREE until a measured A/B justifies moving it.
20374        let iflag = {
20375            static F: std::sync::OnceLock<CUgraphInstantiate_flags> = std::sync::OnceLock::new();
20376            *F.get_or_init(|| match std::env::var("MEMRA_GRAPH_IFLAG").as_deref() {
20377                // UPLOAD = the gemma slotted door's zero-mem-node choice; PRIORITY = the flag
20378                // hybrid_forward.rs:5935 actually ships (both drop the auto-free launch scan).
20379                Ok("upload") => CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_UPLOAD,
20380                Ok("priority") => {
20381                    CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY
20382                }
20383                _ => CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
20384            })
20385        };
20386        // MEMRA_GRAPH_CAPTIME=1 (Q1 lane): phase-resolved capture cost. Recapture is paid at
20387        // every kernel-class crossing, so it — not steady-state decode — is the quantity a
20388        // mem-node reduction could plausibly shrink. Only `instantiate` (cuStreamEndCapture +
20389        // cuGraphInstantiateWithFlags) and `upload` scale with node count; the warmups are
20390        // eager step executions and are node-count-invariant. Printing the split bounds the
20391        // refactor's ceiling instead of assuming it.
20392        let ct = {
20393            static T: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20394            *T.get_or_init(|| std::env::var("MEMRA_GRAPH_CAPTIME").as_deref() == Ok("1"))
20395        };
20396        // MEMRA_GRAPH_WARMUPS (Q1 lane; DEFAULT 1 since lane/graph-warmups 2026-08-05): the
20397        // phase split showed the eager warmups are 80% of recapture cost (q27 27.4 of 34.4 ms
20398        // pod / 42% of 52.6 ms 5090) — 3x larger than the ENTIRE mem-node ceiling the audit
20399        // chased, and node-count-invariant, so no capture-body refactor could touch it.
20400        // Warmup 2's theorized job was async-pool ADDRESS STABILITY: warmup 1's allocs may
20401        // grow/map the pool, warmup 2 re-walks the same sequence over the freed blocks so the
20402        // captured third run bakes settled addresses. That hazard is the #68 stale-baked-
20403        // address class — which the engine now guards STRUCTURALLY rather than by re-walking:
20404        // in-body transients are captured as BALANCED in-graph alloc/free node pairs (census
20405        // 1589/1589 — replays allocate for themselves; no baked transient pointers), every
20406        // externally-referenced buffer is stable-pointer by design (fa_part_pool retires-on-
20407        // grow and never frees, resident counters/scratch, cache set in place), and the
20408        // draft-graph path additionally rides capture_graph_retained (capture_keep holds all
20409        // warmup+capture allocs alive). One warmup therefore suffices for kernel-attr
20410        // settling and pool mapping. Arbitrated adversarially, not by taste:
20411        // graph-warmup-stress (pool-growth cycles large<->small x10, overlap arm, forced
20412        // recaptures over freed blocks — bit-identity vs eager + canary teeth) is GREEN at
20413        // warmups=1 on the deployment rig, plus graph-decode-gate 256-step bit-identity,
20414        // graph-session-gate, run-spec K=1..8 (receipts research/graph-warmups-5090-20260805/
20415        // + the pod's research/graph-allocfree-20260805/). Measured: recapture -38..-42% q27 /
20416        // -41% q9, decode +~1%, capture+prime -13ms. MEMRA_GRAPH_WARMUPS=2 = the rollback
20417        // seam; tools/graph-warmup-stress-gate.sh = the gate any regression re-runs.
20418        let warmups = {
20419            static W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
20420            *W.get_or_init(|| {
20421                std::env::var("MEMRA_GRAPH_WARMUPS")
20422                    .ok()
20423                    .and_then(|v| v.parse().ok())
20424                    .filter(|n| *n >= 1)
20425                    .unwrap_or(1)
20426            })
20427        };
20428        let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
20429            let t_w = std::time::Instant::now();
20430            // warmup: inline runs (no capture) so allocator pointers + kernel attrs are stable.
20431            for _ in 0..warmups {
20432                step(self)?;
20433            }
20434            self.gpu.stream().synchronize()?;
20435            let ms_warm = t_w.elapsed().as_secs_f64() * 1e3;
20436            // capture the third run.
20437            let t_c = std::time::Instant::now();
20438            self.gpu
20439                .stream()
20440                .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
20441            // If the body errors mid-capture, end the capture before propagating so the stream isn't
20442            // left in a capturing state.
20443            let r = step(self);
20444            let ms_body = t_c.elapsed().as_secs_f64() * 1e3;
20445            let t_i = std::time::Instant::now();
20446            let g = self.gpu.stream().end_capture(iflag);
20447            let ms_inst = t_i.elapsed().as_secs_f64() * 1e3;
20448            r?;
20449            let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
20450            let t_u = std::time::Instant::now();
20451            graph.upload()?;
20452            if ct {
20453                println!(
20454                    "[graph-captime] warmup2x {ms_warm:.2} ms  capture-body {ms_body:.2} ms  \
20455                          instantiate {ms_inst:.2} ms  upload {:.2} ms",
20456                    t_u.elapsed().as_secs_f64() * 1e3
20457                );
20458            }
20459            Ok(graph)
20460        };
20461        let result = run();
20462        if was_tracking {
20463            unsafe {
20464                self.gpu.ctx.enable_event_tracking();
20465            }
20466        }
20467        result
20468    }
20469
20470    /// gdn_scan variant where state_in/out are CudaViews (resident SSM state, in-place per step).
20471    pub fn gdn_scan_s128_view(
20472        &self,
20473        q: &CudaSlice<f32>,
20474        k: &CudaSlice<f32>,
20475        v: &CudaSlice<f32>,
20476        g: &CudaSlice<f32>,
20477        beta: &CudaSlice<f32>,
20478        state_in: &cudarc::driver::CudaView<f32>,
20479        state_out: &mut cudarc::driver::CudaViewMut<f32>,
20480        o: &mut CudaSlice<f32>,
20481        n_head: usize,
20482        t: usize,
20483        scale: f32,
20484    ) -> Result<(), Box<dyn std::error::Error>> {
20485        let f = self.func("gdn_scan_s128");
20486        const S_V: u32 = 128;
20487        const WARP: u32 = 32;
20488        const COLS: u32 = 4;
20489        let cfg = LaunchConfig {
20490            grid_dim: (n_head as u32, 1, S_V / COLS),
20491            block_dim: (WARP, COLS, 1),
20492            shared_mem_bytes: 0,
20493        };
20494        let (h, ti) = (n_head as i32, t as i32);
20495        let __s_b = self.gpu.stream();
20496        let mut b = __s_b.launch_builder(&f);
20497        b.arg(q)
20498            .arg(k)
20499            .arg(v)
20500            .arg(g)
20501            .arg(beta)
20502            .arg(state_in)
20503            .arg(state_out)
20504            .arg(o)
20505            .arg(&h)
20506            .arg(&ti)
20507            .arg(&scale);
20508        unsafe {
20509            b.launch(cfg)?;
20510        }
20511        Ok(())
20512    }
20513
20514    /// conv1d where the input is a CudaView (resident conv state assembled in place).
20515    pub fn ssm_conv1d_view(
20516        &self,
20517        x: &cudarc::driver::CudaView<f32>,
20518        w: &CudaSlice<f32>,
20519        y: &mut CudaSlice<f32>,
20520        conv_dim: usize,
20521        t: usize,
20522        d_conv: usize,
20523        silu: bool,
20524    ) -> Result<(), Box<dyn std::error::Error>> {
20525        let f = self.func("ssm_conv1d_silu_f32");
20526        // grid.x = channel, grid.y = T-tiles (block 256 strides over T) — parallel over both axes.
20527        let cfg = LaunchConfig {
20528            grid_dim: (conv_dim as u32, ((t as u32 + 255) / 256).max(1), 1),
20529            block_dim: (256, 1, 1),
20530            shared_mem_bytes: 0,
20531        };
20532        let (cd, ti, dc, s) = (conv_dim as i32, t as i32, d_conv as i32, silu as i32);
20533        let __s_b = self.gpu.stream();
20534        let mut b = __s_b.launch_builder(&f);
20535        b.arg(x).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc).arg(&s);
20536        unsafe {
20537            b.launch(cfg)?;
20538        }
20539        Ok(())
20540    }
20541
20542    /// Depthwise causal conv1d + optional SiLU.
20543    /// x:[conv_dim, T+d_conv-1] channel-major (first d_conv-1 cols = carried state),
20544    /// w:[d_conv, conv_dim] kernel-major, y:[conv_dim, T] channel-major.
20545    /// FUSED prefill conv (token-major input, zero left-state): replaces
20546    /// transpose + zeros + conv_left_pad + ssm_conv1d with ONE launch reading the matmul output
20547    /// directly. Output channel-major [conv_dim, T], SiLU applied. BIT-IDENTICAL accumulation.
20548    pub fn ssm_conv1d_tm(
20549        &self,
20550        qkv_tm: &CudaSlice<f32>,
20551        w: &CudaSlice<f32>,
20552        y: &mut CudaSlice<f32>,
20553        conv_dim: usize,
20554        t: usize,
20555        d_conv: usize,
20556    ) -> Result<(), Box<dyn std::error::Error>> {
20557        let f = self.func("ssm_conv1d_tm_f32");
20558        let cfg = LaunchConfig {
20559            grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
20560            block_dim: (256, 1, 1),
20561            shared_mem_bytes: 0,
20562        };
20563        let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
20564        let __s_b = self.gpu.stream();
20565        let mut b = __s_b.launch_builder(&f);
20566        b.arg(qkv_tm).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc);
20567        unsafe {
20568            b.launch(cfg)?;
20569        }
20570        Ok(())
20571    }
20572
20573    /// BATCHED verify conv (T>1, carried state): window reads the resident conv ring for
20574    /// negative rows; separate ring-update launch afterwards. BIT-IDENTICAL per value to the
20575    /// T=1 chain. T >= pad rides the pure input-column ring update (unchanged legacy path);
20576    /// T < pad (the MEMRA_SPEC_M2 t=2 verify arm) needs old-ring sources for the roll — the
20577    /// update kernel would race reading the ring it rewrites, so that arm clones the ring
20578    /// (dtod) and rolls via ssm_conv_ring_rebuild (PURE COPIES: the ring stores raw input
20579    /// columns; the final ring == what T sequential decode ring rolls leave).
20580    pub fn ssm_conv1d_tm_state(
20581        &self,
20582        qkv_tm: &CudaSlice<f32>,
20583        conv_state: &mut CudaSlice<f32>,
20584        w: &CudaSlice<f32>,
20585        y: &mut CudaSlice<f32>,
20586        conv_dim: usize,
20587        t: usize,
20588        d_conv: usize,
20589    ) -> Result<(), Box<dyn std::error::Error>> {
20590        self.ssm_conv1d_tm_state_pad(qkv_tm, conv_state, w, y, conv_dim, t, d_conv, None)
20591    }
20592
20593    /// task #14: `pad_len` = device true length for PADDED prime graphs — the ring update
20594    /// reads rows [len-pad, len) instead of the pad tail. None = the classic host-T path.
20595    #[allow(clippy::too_many_arguments)]
20596    pub fn ssm_conv1d_tm_state_pad(
20597        &self,
20598        qkv_tm: &CudaSlice<f32>,
20599        conv_state: &mut CudaSlice<f32>,
20600        w: &CudaSlice<f32>,
20601        y: &mut CudaSlice<f32>,
20602        conv_dim: usize,
20603        t: usize,
20604        d_conv: usize,
20605        pad_len: Option<&CudaSlice<i32>>,
20606    ) -> Result<(), Box<dyn std::error::Error>> {
20607        assert!(t >= 1, "ssm_conv1d_tm_state requires T >= 1");
20608        // clone BEFORE the window kernel is issued is not required (stream-ordered: the dtod and
20609        // the window kernel both read the pre-roll ring; the roll launches after both) — but
20610        // cloning first keeps the ordering trivially correct under any future stream split.
20611        let ring_old = if t < d_conv - 1 {
20612            Some(self.clone_dtod(conv_state)?)
20613        } else {
20614            None
20615        };
20616        {
20617            let f = self.func("ssm_conv1d_tm_state_f32");
20618            let cfg = LaunchConfig {
20619                grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
20620                block_dim: (256, 1, 1),
20621                shared_mem_bytes: 0,
20622            };
20623            let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
20624            let __s_b = self.gpu.stream();
20625            let mut b = __s_b.launch_builder(&f);
20626            b.arg(qkv_tm)
20627                .arg(&*conv_state)
20628                .arg(w)
20629                .arg(y)
20630                .arg(&cd)
20631                .arg(&ti)
20632                .arg(&dc);
20633            unsafe {
20634                b.launch(cfg)?;
20635            }
20636        }
20637        match (ring_old, pad_len) {
20638            (None, Some(len_d)) => {
20639                let f = self.func("ssm_conv_ring_update_dev_f32");
20640                let n = conv_dim * (d_conv - 1);
20641                let cfg = LaunchConfig::for_num_elems(n as u32);
20642                let (cd, dc) = (conv_dim as i32, d_conv as i32);
20643                let __s_b = self.gpu.stream();
20644                let mut b = __s_b.launch_builder(&f);
20645                b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
20646                unsafe {
20647                    b.launch(cfg)?;
20648                }
20649            }
20650            (None, None) => {
20651                let f = self.func("ssm_conv_ring_update_f32");
20652                let n = conv_dim * (d_conv - 1);
20653                let cfg = LaunchConfig::for_num_elems(n as u32);
20654                let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
20655                let __s_b = self.gpu.stream();
20656                let mut b = __s_b.launch_builder(&f);
20657                b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
20658                unsafe {
20659                    b.launch(cfg)?;
20660                }
20661            }
20662            (Some(old), _) => {
20663                self.ssm_conv_ring_rebuild(qkv_tm, &old, conv_state, conv_dim, t, d_conv)?
20664            }
20665        }
20666        Ok(())
20667    }
20668
20669    /// qkv-view twin (task #16): batched prime reads the concat GEMM output directly.
20670    pub fn ssm_conv1d_tm_state_pad_v(
20671        &self,
20672        qkv_tm: &cudarc::driver::CudaView<f32>,
20673        conv_state: &mut CudaSlice<f32>,
20674        w: &CudaSlice<f32>,
20675        y: &mut CudaSlice<f32>,
20676        conv_dim: usize,
20677        t: usize,
20678        d_conv: usize,
20679        pad_len: Option<&CudaSlice<i32>>,
20680    ) -> Result<(), Box<dyn std::error::Error>> {
20681        assert!(t >= 1, "ssm_conv1d_tm_state requires T >= 1");
20682        // clone BEFORE the window kernel is issued is not required (stream-ordered: the dtod and
20683        // the window kernel both read the pre-roll ring; the roll launches after both) — but
20684        // cloning first keeps the ordering trivially correct under any future stream split.
20685        let ring_old = if t < d_conv - 1 {
20686            Some(self.clone_dtod(conv_state)?)
20687        } else {
20688            None
20689        };
20690        {
20691            let f = self.func("ssm_conv1d_tm_state_f32");
20692            let cfg = LaunchConfig {
20693                grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
20694                block_dim: (256, 1, 1),
20695                shared_mem_bytes: 0,
20696            };
20697            let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
20698            let __s_b = self.gpu.stream();
20699            let mut b = __s_b.launch_builder(&f);
20700            b.arg(qkv_tm)
20701                .arg(&*conv_state)
20702                .arg(w)
20703                .arg(y)
20704                .arg(&cd)
20705                .arg(&ti)
20706                .arg(&dc);
20707            unsafe {
20708                b.launch(cfg)?;
20709            }
20710        }
20711        match (ring_old, pad_len) {
20712            (None, Some(len_d)) => {
20713                let f = self.func("ssm_conv_ring_update_dev_f32");
20714                let n = conv_dim * (d_conv - 1);
20715                let cfg = LaunchConfig::for_num_elems(n as u32);
20716                let (cd, dc) = (conv_dim as i32, d_conv as i32);
20717                let __s_b = self.gpu.stream();
20718                let mut b = __s_b.launch_builder(&f);
20719                b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
20720                unsafe {
20721                    b.launch(cfg)?;
20722                }
20723            }
20724            (None, None) => {
20725                let f = self.func("ssm_conv_ring_update_f32");
20726                let n = conv_dim * (d_conv - 1);
20727                let cfg = LaunchConfig::for_num_elems(n as u32);
20728                let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
20729                let __s_b = self.gpu.stream();
20730                let mut b = __s_b.launch_builder(&f);
20731                b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
20732                unsafe {
20733                    b.launch(cfg)?;
20734                }
20735            }
20736            (Some(_), _) => unreachable!(
20737                "ssm_conv1d_tm_state_pad_v: T < d_conv-1 has no view path (PRIME_MIN_T gates it)"
20738            ),
20739        }
20740        Ok(())
20741    }
20742
20743    /// PREFIX conv-ring rebuild (spec REPLAY-FREE partial accept): overwrite the resident ring
20744    /// with the state a T=1 chain holds after only the FIRST `tc` columns of `qkv_tm` — the last
20745    /// `pad` entries of [ring_old | cols 0..tc-1]. PURE COPIES (the ring stores raw inputs; no
20746    /// arithmetic, cannot perturb FP order). `ring_old` = the pre-round snapshot ring.
20747    pub fn ssm_conv_ring_rebuild(
20748        &self,
20749        qkv_tm: &CudaSlice<f32>,
20750        ring_old: &CudaSlice<f32>,
20751        conv_state: &mut CudaSlice<f32>,
20752        conv_dim: usize,
20753        tc: usize,
20754        d_conv: usize,
20755    ) -> Result<(), Box<dyn std::error::Error>> {
20756        let f = self.func("ssm_conv_ring_rebuild_f32");
20757        let n = conv_dim * (d_conv - 1);
20758        let cfg = LaunchConfig::for_num_elems(n as u32);
20759        let (cd, ti, dc) = (conv_dim as i32, tc as i32, d_conv as i32);
20760        let __s_b = self.gpu.stream();
20761        let mut b = __s_b.launch_builder(&f);
20762        b.arg(qkv_tm)
20763            .arg(ring_old)
20764            .arg(conv_state)
20765            .arg(&cd)
20766            .arg(&ti)
20767            .arg(&dc);
20768        unsafe {
20769            b.launch(cfg)?;
20770        }
20771        Ok(())
20772    }
20773
20774    /// FUSED decode GDN prep (T=1): repack + q/k L2-norm + beta sigmoid + g_log in one launch.
20775    /// Replaces 5 tiny serialized kernels on the decode critical path. L2 reduce runs as a 32-lane
20776    /// warp tree (vs l2_norm_f32's 256-thread two-level tree) — same math, different FP sum order;
20777    /// the argmax + run-spec gates are the authority.
20778    #[allow(clippy::too_many_arguments)]
20779    pub fn gdn_prep_decode(
20780        &self,
20781        conv_out: &CudaSlice<f32>,
20782        beta_raw: &CudaSlice<f32>,
20783        alpha: &CudaSlice<f32>,
20784        dt_bias: &CudaSlice<f32>,
20785        a: &CudaSlice<f32>,
20786        q_l2: &mut CudaSlice<f32>,
20787        k_l2: &mut CudaSlice<f32>,
20788        v_g: &mut CudaSlice<f32>,
20789        beta: &mut CudaSlice<f32>,
20790        g_log: &mut CudaSlice<f32>,
20791        d_state: usize,
20792        num_v: usize,
20793        num_k: usize,
20794        key_dim: usize,
20795        eps: f32,
20796    ) -> Result<(), Box<dyn std::error::Error>> {
20797        let f = self.func("gdn_prep_decode_f32");
20798        let cfg = LaunchConfig {
20799            grid_dim: (num_v as u32, 1, 1),
20800            block_dim: (32, 4, 1),
20801            shared_mem_bytes: 0,
20802        };
20803        let (ds, nv, nk, kd) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
20804        let __s_b = self.gpu.stream();
20805        let mut b = __s_b.launch_builder(&f);
20806        b.arg(conv_out)
20807            .arg(beta_raw)
20808            .arg(alpha)
20809            .arg(dt_bias)
20810            .arg(a)
20811            .arg(q_l2)
20812            .arg(k_l2)
20813            .arg(v_g)
20814            .arg(beta)
20815            .arg(g_log)
20816            .arg(&ds)
20817            .arg(&nv)
20818            .arg(&nk)
20819            .arg(&kd)
20820            .arg(&eps);
20821        unsafe {
20822            b.launch(cfg)?;
20823        }
20824        Ok(())
20825    }
20826
20827    /// FUSED prefill conv + GDN repack: token-major qkv -> q_g/k_g/v_g in ONE launch (no conv_out
20828    /// materialization, no qkv_to_gdn_repack pass). BIT-IDENTICAL values; scatter matches
20829    /// qkv_to_gdn_repack's modulo head-repeat mapping exactly.
20830    #[allow(clippy::too_many_arguments)]
20831    pub fn ssm_conv1d_gdn(
20832        &self,
20833        qkv_tm: &CudaSlice<f32>,
20834        w: &CudaSlice<f32>,
20835        q_g: &mut CudaSlice<f32>,
20836        k_g: &mut CudaSlice<f32>,
20837        v_g: &mut CudaSlice<f32>,
20838        conv_dim: usize,
20839        t: usize,
20840        d_conv: usize,
20841        d_state: usize,
20842        num_v: usize,
20843        num_k: usize,
20844        key_dim: usize,
20845    ) -> Result<(), Box<dyn std::error::Error>> {
20846        let f = self.func("ssm_conv1d_gdn_f32");
20847        let cfg = LaunchConfig {
20848            grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
20849            block_dim: (256, 1, 1),
20850            shared_mem_bytes: 0,
20851        };
20852        let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
20853        let (ds, nv, nk, kd) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
20854        let __s_b = self.gpu.stream();
20855        let mut b = __s_b.launch_builder(&f);
20856        b.arg(qkv_tm)
20857            .arg(w)
20858            .arg(q_g)
20859            .arg(k_g)
20860            .arg(v_g)
20861            .arg(&cd)
20862            .arg(&ti)
20863            .arg(&dc)
20864            .arg(&ds)
20865            .arg(&nv)
20866            .arg(&nk)
20867            .arg(&kd);
20868        unsafe {
20869            b.launch(cfg)?;
20870        }
20871        Ok(())
20872    }
20873
20874    pub fn ssm_conv1d(
20875        &self,
20876        x: &CudaSlice<f32>,
20877        w: &CudaSlice<f32>,
20878        y: &mut CudaSlice<f32>,
20879        conv_dim: usize,
20880        t: usize,
20881        d_conv: usize,
20882        silu: bool,
20883    ) -> Result<(), Box<dyn std::error::Error>> {
20884        let f = self.func("ssm_conv1d_silu_f32");
20885        let cfg = LaunchConfig {
20886            grid_dim: (conv_dim as u32, ((t as u32 + 255) / 256).max(1), 1),
20887            block_dim: (256, 1, 1),
20888            shared_mem_bytes: 0,
20889        };
20890        let (cd, ti, dc, s) = (conv_dim as i32, t as i32, d_conv as i32, silu as i32);
20891        let __s_b = self.gpu.stream();
20892        let mut b = __s_b.launch_builder(&f);
20893        b.arg(x).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc).arg(&s);
20894        unsafe {
20895            b.launch(cfg)?;
20896        }
20897        Ok(())
20898    }
20899
20900    /// Gated DeltaNet scan, S_v=128. q,k,v:[128,H,T]; g,beta:[H,T]; state:[128,128,H] transposed;
20901    /// o:[128,H,T]. Single sequence.
20902    pub fn gdn_scan_s128(
20903        &self,
20904        q: &CudaSlice<f32>,
20905        k: &CudaSlice<f32>,
20906        v: &CudaSlice<f32>,
20907        g: &CudaSlice<f32>,
20908        beta: &CudaSlice<f32>,
20909        state_in: &CudaSlice<f32>,
20910        state_out: &mut CudaSlice<f32>,
20911        o: &mut CudaSlice<f32>,
20912        n_head: usize,
20913        t: usize,
20914        scale: f32,
20915    ) -> Result<(), Box<dyn std::error::Error>> {
20916        let f = self.func("gdn_scan_s128");
20917        const S_V: u32 = 128;
20918        const WARP: u32 = 32;
20919        const COLS_PER_BLOCK: u32 = 4;
20920        let cfg = LaunchConfig {
20921            grid_dim: (n_head as u32, 1, S_V / COLS_PER_BLOCK),
20922            block_dim: (WARP, COLS_PER_BLOCK, 1),
20923            shared_mem_bytes: 0,
20924        };
20925        let (h, ti) = (n_head as i32, t as i32);
20926        let __s_b = self.gpu.stream();
20927        let mut b = __s_b.launch_builder(&f);
20928        b.arg(q)
20929            .arg(k)
20930            .arg(v)
20931            .arg(g)
20932            .arg(beta)
20933            .arg(state_in)
20934            .arg(state_out)
20935            .arg(o)
20936            .arg(&h)
20937            .arg(&ti)
20938            .arg(&scale);
20939        unsafe {
20940            b.launch(cfg)?;
20941        }
20942        Ok(())
20943    }
20944
20945    // ==== B2' batched decode state ops (decode_batch.rs) ====
20946    // Per-seq state pointers ride device u64 arrays (views into the per-step pointer table).
20947    // Bodies are the single-seq kernels per sequence — bit-identical per row.
20948
20949    #[allow(clippy::too_many_arguments)]
20950    pub fn ssm_conv1d_fused_decode_b(
20951        &self,
20952        qkv_cols: &CudaSlice<f32>,
20953        conv_state_ptrs: &cudarc::driver::CudaView<u64>,
20954        w: &CudaSlice<f32>,
20955        conv_outs: &mut CudaSlice<f32>,
20956        conv_dim: usize,
20957        d_conv: usize,
20958        b_n: usize,
20959    ) -> Result<(), Box<dyn std::error::Error>> {
20960        let f = self.func("ssm_conv1d_fused_decode_b_f32");
20961        let cfg = LaunchConfig {
20962            grid_dim: (((conv_dim + 255) / 256) as u32, 1, b_n as u32),
20963            block_dim: (256, 1, 1),
20964            shared_mem_bytes: 0,
20965        };
20966        let (cd, dc) = (conv_dim as i32, d_conv as i32);
20967        let __s_b = self.gpu.stream();
20968        let mut b = __s_b.launch_builder(&f);
20969        b.arg(qkv_cols)
20970            .arg(conv_state_ptrs)
20971            .arg(w)
20972            .arg(conv_outs)
20973            .arg(&cd)
20974            .arg(&dc);
20975        unsafe {
20976            b.launch(cfg)?;
20977        }
20978        Ok(())
20979    }
20980
20981    #[allow(clippy::too_many_arguments)]
20982    pub fn gdn_prep_decode_b(
20983        &self,
20984        conv_outs: &CudaSlice<f32>,
20985        beta_raws: &CudaSlice<f32>,
20986        alphas: &CudaSlice<f32>,
20987        dt_bias: &CudaSlice<f32>,
20988        a: &CudaSlice<f32>,
20989        q_l2: &mut CudaSlice<f32>,
20990        k_l2: &mut CudaSlice<f32>,
20991        v_g: &mut CudaSlice<f32>,
20992        beta: &mut CudaSlice<f32>,
20993        g_log: &mut CudaSlice<f32>,
20994        d_state: usize,
20995        num_v: usize,
20996        num_k: usize,
20997        key_dim: usize,
20998        eps: f32,
20999        conv_dim: usize,
21000        b_n: usize,
21001    ) -> Result<(), Box<dyn std::error::Error>> {
21002        let f = self.func("gdn_prep_decode_b_f32");
21003        let cfg = LaunchConfig {
21004            grid_dim: (num_v as u32, 1, b_n as u32),
21005            block_dim: (32, 4, 1),
21006            shared_mem_bytes: 0,
21007        };
21008        let (ds, nv, nk, kd, cd) = (
21009            d_state as i32,
21010            num_v as i32,
21011            num_k as i32,
21012            key_dim as i32,
21013            conv_dim as i32,
21014        );
21015        let __s_b = self.gpu.stream();
21016        let mut b = __s_b.launch_builder(&f);
21017        b.arg(conv_outs)
21018            .arg(beta_raws)
21019            .arg(alphas)
21020            .arg(dt_bias)
21021            .arg(a)
21022            .arg(q_l2)
21023            .arg(k_l2)
21024            .arg(v_g)
21025            .arg(beta)
21026            .arg(g_log)
21027            .arg(&ds)
21028            .arg(&nv)
21029            .arg(&nk)
21030            .arg(&kd)
21031            .arg(&eps)
21032            .arg(&cd);
21033        unsafe {
21034            b.launch(cfg)?;
21035        }
21036        Ok(())
21037    }
21038
21039    #[allow(clippy::too_many_arguments)]
21040    pub fn gdn_scan_s128_batched(
21041        &self,
21042        q: &CudaSlice<f32>,
21043        k: &CudaSlice<f32>,
21044        v: &CudaSlice<f32>,
21045        g: &CudaSlice<f32>,
21046        beta: &CudaSlice<f32>,
21047        state_in_ptrs: &cudarc::driver::CudaView<u64>,
21048        state_out_ptrs: &cudarc::driver::CudaView<u64>,
21049        o: &mut CudaSlice<f32>,
21050        n_head: usize,
21051        b_n: usize,
21052        scale: f32,
21053    ) -> Result<(), Box<dyn std::error::Error>> {
21054        let f = self.func("gdn_scan_s128_b");
21055        const S_V: u32 = 128;
21056        const WARP: u32 = 32;
21057        const COLS_PER_BLOCK: u32 = 4;
21058        let cfg = LaunchConfig {
21059            grid_dim: (n_head as u32, b_n as u32, S_V / COLS_PER_BLOCK),
21060            block_dim: (WARP, COLS_PER_BLOCK, 1),
21061            shared_mem_bytes: 0,
21062        };
21063        let h = n_head as i32;
21064        let __s_b = self.gpu.stream();
21065        let mut b = __s_b.launch_builder(&f);
21066        b.arg(q)
21067            .arg(k)
21068            .arg(v)
21069            .arg(g)
21070            .arg(beta)
21071            .arg(state_in_ptrs)
21072            .arg(state_out_ptrs)
21073            .arg(o)
21074            .arg(&h)
21075            .arg(&scale);
21076        unsafe {
21077            b.launch(cfg)?;
21078        }
21079        Ok(())
21080    }
21081
21082    /// VIEW twins of the three GDN decode state wrappers (lane/verify-launchslim): identical
21083    /// launches, row args as views into the caller's packed [T, ...] buffers — the t-parallel
21084    /// verify's per-row loop passes slices instead of paying an arithmetic-free dtod per row
21085    /// per kernel (48 layers x T rows x 4 copies/round on the money path). Same kernels, same
21086    /// numeric class; only the pointer arithmetic moved host-side.
21087    #[allow(clippy::too_many_arguments)]
21088    pub fn ssm_conv1d_fused_decode_b_view(
21089        &self,
21090        qkv_cols: &cudarc::driver::CudaView<f32>,
21091        conv_state_ptrs: &cudarc::driver::CudaView<u64>,
21092        w: &CudaSlice<f32>,
21093        conv_outs: &mut CudaSlice<f32>,
21094        conv_dim: usize,
21095        d_conv: usize,
21096        b_n: usize,
21097    ) -> Result<(), Box<dyn std::error::Error>> {
21098        let f = self.func("ssm_conv1d_fused_decode_b_f32");
21099        let cfg = LaunchConfig {
21100            grid_dim: (((conv_dim + 255) / 256) as u32, 1, b_n as u32),
21101            block_dim: (256, 1, 1),
21102            shared_mem_bytes: 0,
21103        };
21104        let (cd, dc) = (conv_dim as i32, d_conv as i32);
21105        let __s_b = self.gpu.stream();
21106        let mut b = __s_b.launch_builder(&f);
21107        b.arg(qkv_cols)
21108            .arg(conv_state_ptrs)
21109            .arg(w)
21110            .arg(conv_outs)
21111            .arg(&cd)
21112            .arg(&dc);
21113        unsafe {
21114            b.launch(cfg)?;
21115        }
21116        Ok(())
21117    }
21118
21119    #[allow(clippy::too_many_arguments)]
21120    pub fn gdn_prep_decode_b_view(
21121        &self,
21122        conv_outs: &CudaSlice<f32>,
21123        beta_raws: &cudarc::driver::CudaView<f32>,
21124        alphas: &cudarc::driver::CudaView<f32>,
21125        dt_bias: &CudaSlice<f32>,
21126        a: &CudaSlice<f32>,
21127        q_l2: &mut CudaSlice<f32>,
21128        k_l2: &mut CudaSlice<f32>,
21129        v_g: &mut CudaSlice<f32>,
21130        beta: &mut CudaSlice<f32>,
21131        g_log: &mut CudaSlice<f32>,
21132        d_state: usize,
21133        num_v: usize,
21134        num_k: usize,
21135        key_dim: usize,
21136        eps: f32,
21137        conv_dim: usize,
21138        b_n: usize,
21139    ) -> Result<(), Box<dyn std::error::Error>> {
21140        let f = self.func("gdn_prep_decode_b_f32");
21141        let cfg = LaunchConfig {
21142            grid_dim: (num_v as u32, 1, b_n as u32),
21143            block_dim: (32, 4, 1),
21144            shared_mem_bytes: 0,
21145        };
21146        let (ds, nv, nk, kd, cd) = (
21147            d_state as i32,
21148            num_v as i32,
21149            num_k as i32,
21150            key_dim as i32,
21151            conv_dim as i32,
21152        );
21153        let __s_b = self.gpu.stream();
21154        let mut b = __s_b.launch_builder(&f);
21155        b.arg(conv_outs)
21156            .arg(beta_raws)
21157            .arg(alphas)
21158            .arg(dt_bias)
21159            .arg(a)
21160            .arg(q_l2)
21161            .arg(k_l2)
21162            .arg(v_g)
21163            .arg(beta)
21164            .arg(g_log)
21165            .arg(&ds)
21166            .arg(&nv)
21167            .arg(&nk)
21168            .arg(&kd)
21169            .arg(&eps)
21170            .arg(&cd);
21171        unsafe {
21172            b.launch(cfg)?;
21173        }
21174        Ok(())
21175    }
21176
21177    #[allow(clippy::too_many_arguments)]
21178    pub fn gdn_scan_s128_batched_view(
21179        &self,
21180        q: &CudaSlice<f32>,
21181        k: &CudaSlice<f32>,
21182        v: &CudaSlice<f32>,
21183        g: &CudaSlice<f32>,
21184        beta: &CudaSlice<f32>,
21185        state_in_ptrs: &cudarc::driver::CudaView<u64>,
21186        state_out_ptrs: &cudarc::driver::CudaView<u64>,
21187        o: &mut cudarc::driver::CudaViewMut<f32>,
21188        n_head: usize,
21189        b_n: usize,
21190        scale: f32,
21191    ) -> Result<(), Box<dyn std::error::Error>> {
21192        let f = self.func("gdn_scan_s128_b");
21193        const S_V: u32 = 128;
21194        const WARP: u32 = 32;
21195        const COLS_PER_BLOCK: u32 = 4;
21196        let cfg = LaunchConfig {
21197            grid_dim: (n_head as u32, b_n as u32, S_V / COLS_PER_BLOCK),
21198            block_dim: (WARP, COLS_PER_BLOCK, 1),
21199            shared_mem_bytes: 0,
21200        };
21201        let h = n_head as i32;
21202        let __s_b = self.gpu.stream();
21203        let mut b = __s_b.launch_builder(&f);
21204        b.arg(q)
21205            .arg(k)
21206            .arg(v)
21207            .arg(g)
21208            .arg(beta)
21209            .arg(state_in_ptrs)
21210            .arg(state_out_ptrs)
21211            .arg(o)
21212            .arg(&h)
21213            .arg(&scale);
21214        unsafe {
21215            b.launch(cfg)?;
21216        }
21217        Ok(())
21218    }
21219
21220    /// A4 seam: chunked WY GDN prefill. DEFAULT ON (`MEMRA_GDN_CHUNKED=0` = rollback to the
21221    /// sequential scan). Flipped 2026-07-04 with the full battery green: kernel-check ALL
21222    /// GREEN x {9B, 27B} incl the f64-truth chunk gates; run-gen argmax 82==82 both models
21223    /// on AND off (24/24 sweep runs); run-spec K={1,2,3,4,6,8} PASS x {9B synth, 9B text,
21224    /// 27B p2, 27B p3}; e2e first-16-token agreement 6/6 (full-256 drifts at index 47-125
21225    /// on 5/6 prompts — accepted cache-state-FP class, batched-prime precedent).
21226    /// PREFILL-ONLY: decode + spec verify never route here (decode==verify dispatch
21227    /// identity law); prime_cache/forward/forward_last are the only callers.
21228    pub fn gdn_chunked_enabled() -> bool {
21229        static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
21230        *E.get_or_init(|| {
21231            std::env::var("MEMRA_GDN_CHUNKED")
21232                .map(|v| v != "0")
21233                .unwrap_or(true)
21234        })
21235    }
21236
21237    /// A4 chunk size (MEMRA_GDN_CHUNK, default 32 — the sweep winner: the O(T*C) chunk
21238    /// matrices grow with C while the sequential state pass is C-flat, so smaller chunks
21239    /// win; C=32/64 also get the register-history solve template). Clamped to multiples
21240    /// of 32 in [32, 128] (kernel row mappings require it).
21241    pub fn gdn_chunk_size() -> usize {
21242        static C: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
21243        *C.get_or_init(|| {
21244            let c: usize = std::env::var("MEMRA_GDN_CHUNK")
21245                .ok()
21246                .and_then(|v| v.parse().ok())
21247                .unwrap_or(32);
21248            c.clamp(32, 128) / 32 * 32
21249        })
21250    }
21251
21252    /// A4: chunked WY / blockwise-inverse GDN prefill (see cu/hybrid.cu K1-K5 header for the
21253    /// math). Same contract as `gdn_scan_s128` (layouts, state ping-pong) but chunk-parallel:
21254    /// NOT bit-identical to the sequential scan (chunked FP accumulation order); run-gen
21255    /// argmax + run-spec batteries are the accuracy authority. PREFILL callers only.
21256    #[allow(clippy::too_many_arguments)]
21257    /// task #18: K1-K3 of the chunked WY scan (shared by the per-seq path and the
21258    /// batched-prime varlen path). Returns (gcum, P, U, W); `A` is K3-internal.
21259    #[allow(clippy::too_many_arguments, clippy::type_complexity)]
21260    #[allow(clippy::too_many_arguments)]
21261    pub fn gdn_chunk_k123(
21262        &self,
21263        q: &CudaSlice<f32>,
21264        k: &CudaSlice<f32>,
21265        v: &CudaSlice<f32>,
21266        g: &CudaSlice<f32>,
21267        beta: &CudaSlice<f32>,
21268        wb16: Option<&mut CudaSlice<u8>>,
21269        n_head: usize,
21270        t: usize,
21271        c: usize,
21272        hk: usize,
21273        k2w: Option<(&CudaSlice<u8>, &CudaSlice<u8>, &mut CudaSlice<u8>)>,
21274    ) -> Result<
21275        (
21276            CudaSlice<f32>,
21277            CudaSlice<f32>,
21278            CudaSlice<f32>,
21279            CudaSlice<f32>,
21280        ),
21281        Box<dyn std::error::Error>,
21282    > {
21283        const D: usize = 128;
21284        let h = n_head;
21285        let nc = (t + c - 1) / c;
21286        let (hi, ti, ci) = (h as i32, t as i32, c as i32);
21287        let mut gcum = self.uninit(t * h)?;
21288        let mut a = self.uninit(nc * h * c * c)?;
21289        let mut p = self.uninit(nc * h * c * c)?;
21290        let mut u = self.uninit(nc * h * c * D)?;
21291        let mut w = self.uninit(nc * h * c * D)?;
21292        {
21293            // K1
21294            let f = self.func("gdn_chunk_cumgate_f32");
21295            let cfg = LaunchConfig {
21296                grid_dim: (nc as u32, h as u32, 1),
21297                block_dim: (32, 1, 1),
21298                shared_mem_bytes: 0,
21299            };
21300            let __s_b = self.gpu.stream();
21301            let mut b = __s_b.launch_builder(&f);
21302            b.arg(g).arg(&mut gcum).arg(&hi).arg(&ti).arg(&ci);
21303            unsafe {
21304                b.launch(cfg)?;
21305            }
21306        }
21307        if let Some((qb, kb, pb)) = k2w {
21308            // K2-wgmma (MEMRA_GDN_WGMMA path, c==32): A + pre-masked Pb16 in one kernel;
21309            // the P f32 buffer stays UNWRITTEN (its only wgmma-path consumer is Pb16).
21310            assert!(c == 32, "gdn_k2_wgmma is a C==32 tile");
21311            let f = self.func("gdn_k2_wgmma");
21312            let cfg = LaunchConfig {
21313                grid_dim: (nc as u32, h as u32, 1),
21314                block_dim: (128, 1, 1),
21315                shared_mem_bytes: 0,
21316            };
21317            let hki = hk as i32;
21318            let __s_b = self.gpu.stream();
21319            let mut b = __s_b.launch_builder(&f);
21320            b.arg(qb)
21321                .arg(kb)
21322                .arg(&gcum)
21323                .arg(beta)
21324                .arg(&mut a)
21325                .arg(&mut *pb)
21326                .arg(&hi)
21327                .arg(&ti)
21328                .arg(&ci)
21329                .arg(&hki);
21330            unsafe {
21331                b.launch(cfg)?;
21332            }
21333        } else if c <= 64 && !portable_mma_gated() {
21334            // K2 register-tiled (2x2 outputs/thread, whole-chunk smem k tile)
21335            let f = self.func("gdn_chunk_attn_f32");
21336            let jt = ((c + 31) / 32) as u32;
21337            let cfg = LaunchConfig {
21338                grid_dim: (nc as u32, h as u32, jt),
21339                block_dim: (256, 1, 1),
21340                shared_mem_bytes: 0,
21341            };
21342            let hki = hk as i32;
21343            let __s_b = self.gpu.stream();
21344            let mut b = __s_b.launch_builder(&f);
21345            b.arg(q)
21346                .arg(k)
21347                .arg(&gcum)
21348                .arg(beta)
21349                .arg(&mut a)
21350                .arg(&mut p)
21351                .arg(&hi)
21352                .arg(&ti)
21353                .arg(&ci)
21354                .arg(&hki);
21355            unsafe {
21356                b.launch(cfg)?;
21357            }
21358        } else {
21359            // K2 generic (C = 128, or the portable target's low-smem fallback)
21360            assert!(
21361                hk == h,
21362                "generic K2 is broadcast-only (de-broadcast rides C==32)"
21363            );
21364            let f = self.func("gdn_chunk_attn_g_f32");
21365            let cfg = LaunchConfig {
21366                grid_dim: (nc as u32, h as u32, 1),
21367                block_dim: (32, 8, 1),
21368                shared_mem_bytes: 0,
21369            };
21370            let __s_b = self.gpu.stream();
21371            let mut b = __s_b.launch_builder(&f);
21372            b.arg(q)
21373                .arg(k)
21374                .arg(&gcum)
21375                .arg(beta)
21376                .arg(&mut a)
21377                .arg(&mut p)
21378                .arg(&hi)
21379                .arg(&ti)
21380                .arg(&ci);
21381            unsafe {
21382                b.launch(cfg)?;
21383            }
21384        }
21385        {
21386            // K3 (register-history templates for C=32/64; local-memory generic otherwise)
21387            let cfg = LaunchConfig {
21388                grid_dim: (nc as u32, h as u32, 1),
21389                block_dim: (256, 1, 1),
21390                shared_mem_bytes: 0,
21391            };
21392            match c {
21393                32 | 64 => {
21394                    let f = self.func(if c == 32 {
21395                        "gdn_chunk_solve32_f32"
21396                    } else {
21397                        "gdn_chunk_solve64_f32"
21398                    });
21399                    // mirror-fold: W's bf16 twin emitted on store (0 = skip)
21400                    let wb: u64 = match wb16 {
21401                        Some(d) => self.addr_u8(d),
21402                        None => 0,
21403                    };
21404                    let hki = hk as i32;
21405                    let __s_b = self.gpu.stream();
21406                    let mut b = __s_b.launch_builder(&f);
21407                    b.arg(v)
21408                        .arg(k)
21409                        .arg(&a)
21410                        .arg(&gcum)
21411                        .arg(&mut u)
21412                        .arg(&mut w)
21413                        .arg(&wb)
21414                        .arg(&hi)
21415                        .arg(&ti)
21416                        .arg(&hki);
21417                    unsafe {
21418                        b.launch(cfg)?;
21419                    }
21420                }
21421                _ => {
21422                    assert!(hk == h, "generic K3 is broadcast-only");
21423                    let f = self.func("gdn_chunk_solve_f32");
21424                    let __s_b = self.gpu.stream();
21425                    let mut b = __s_b.launch_builder(&f);
21426                    b.arg(v)
21427                        .arg(k)
21428                        .arg(&a)
21429                        .arg(&gcum)
21430                        .arg(&mut u)
21431                        .arg(&mut w)
21432                        .arg(&hi)
21433                        .arg(&ti)
21434                        .arg(&ci);
21435                    unsafe {
21436                        b.launch(cfg)?;
21437                    }
21438                }
21439            }
21440        }
21441        Ok((gcum, p, u, w))
21442    }
21443
21444    /// task #21 de-broadcast seam: q/k stored at num_k distinct GQA heads instead of
21445    /// the num_v broadcast. MEMRA_GDN_DB=0 reverts. Only the chunked prefill path
21446    /// consumes the compact layout (hk plumbed; hk == H reproduces broadcast exactly).
21447    pub fn gdn_db_on() -> bool {
21448        std::env::var("MEMRA_GDN_DB").as_deref() != Ok("0")
21449    }
21450
21451    /// Whether the K4/K5 mma pair serves at chunk size `c` (mirrors gdn_scan_chunked's
21452    /// seam read — env re-read per call ON PURPOSE, kernel-check pins both configs).
21453    /// DEFAULT ON for sm_120a builds too (lane/moeprime-nvfp4-direct, 2026-08-21): the pair
21454    /// was qualified on 90a only and left env-opt-in elsewhere; measured on Blackwell it
21455    /// wins on BOTH rigs — one RTX PRO 6000 (ornith15 pp14715 12,036 -> 12,751/12,957,
21456    /// +6-8%, both orders) and the local 5090 (q38-27b pp6435 1,397/1,429 -> 1,427/1,446,
21457    /// both orders) — with kernel-check/run-gen/margin-gate/run-spec green under the flag.
21458    /// bf16 HMMA (m16n8k16) is sm_80-class PTX; only the wgmma nest stays Hopper-gated.
21459    pub fn gdn_mma_enabled(&self, c: usize) -> bool {
21460        !portable_mma_gated()
21461            && c == 32
21462            && match std::env::var("MEMRA_GDN_MMA").as_deref() {
21463                Ok("1") => true,
21464                Ok("0") => false,
21465                _ => gdn_mma_default_on(),
21466            }
21467    }
21468
21469    /// task #22: whether the fused K4+K5 (+K2) wgmma path serves (nested inside the
21470    /// mma config; same per-call env read discipline). HARD-gated to the Hopper build:
21471    /// the wgmma asm bodies exist only at __CUDA_ARCH__ == 900 (MEMRA_K45_REAL,
21472    /// wgmma_common.cuh) — on every other arch the kernel compiles EMPTY, so an env
21473    /// force would silently produce garbage. Required since the sm_120a mma default
21474    /// flip made MEMRA_GDN_WGMMA=1 alone reach this branch there.
21475    pub fn gdn_wgmma_on(&self, c: usize) -> bool {
21476        cfg!(memra_hopper_mma)
21477            && self.gdn_mma_enabled(c)
21478            && std::env::var("MEMRA_GDN_WGMMA").as_deref() != Ok("0")
21479    }
21480
21481    /// task #18 conv-fuse: carried-ring conv + SiLU + GDN repack in ONE pass (the
21482    /// conv_out intermediate and its transposed re-read disappear — 11.8ms of the
21483    /// T=2048 prime). Ring update stays the separate follow-up launch (pad-aware).
21484    /// BIT-IDENTICAL values to ssm_conv1d_tm_state_pad + qkv_to_gdn_repack.
21485    #[allow(clippy::too_many_arguments)]
21486    pub fn ssm_conv1d_gdn_state_pad(
21487        &self,
21488        qkv_tm: &cudarc::driver::CudaView<f32>,
21489        conv_state: &mut CudaSlice<f32>,
21490        w: &CudaSlice<f32>,
21491        q_g: &mut CudaSlice<f32>,
21492        k_g: &mut CudaSlice<f32>,
21493        v_g: &mut CudaSlice<f32>,
21494        conv_dim: usize,
21495        t: usize,
21496        d_conv: usize,
21497        d_state: usize,
21498        num_v: usize,
21499        num_k: usize,
21500        key_dim: usize,
21501        hk: usize,
21502        pad_len: Option<&CudaSlice<i32>>,
21503    ) -> Result<(), Box<dyn std::error::Error>> {
21504        assert!(
21505            t >= d_conv - 1,
21506            "fused state conv requires T >= pad (PRIME_MIN_T gates)"
21507        );
21508        {
21509            let f = self.func("ssm_conv1d_gdn_state_f32");
21510            let cfg = LaunchConfig {
21511                grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
21512                block_dim: (256, 1, 1),
21513                shared_mem_bytes: 0,
21514            };
21515            let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
21516            let (ds, nv, nk, kd, hki) = (
21517                d_state as i32,
21518                num_v as i32,
21519                num_k as i32,
21520                key_dim as i32,
21521                hk as i32,
21522            );
21523            let __s_b = self.gpu.stream();
21524            let mut b = __s_b.launch_builder(&f);
21525            b.arg(qkv_tm)
21526                .arg(&*conv_state)
21527                .arg(w)
21528                .arg(q_g)
21529                .arg(k_g)
21530                .arg(v_g)
21531                .arg(&cd)
21532                .arg(&ti)
21533                .arg(&dc)
21534                .arg(&ds)
21535                .arg(&nv)
21536                .arg(&nk)
21537                .arg(&kd)
21538                .arg(&hki);
21539            unsafe {
21540                b.launch(cfg)?;
21541            }
21542        }
21543        match pad_len {
21544            Some(len_d) => {
21545                let f = self.func("ssm_conv_ring_update_dev_f32");
21546                let n = conv_dim * (d_conv - 1);
21547                let cfg = LaunchConfig::for_num_elems(n as u32);
21548                let (cd, dc) = (conv_dim as i32, d_conv as i32);
21549                let __s_b = self.gpu.stream();
21550                let mut b = __s_b.launch_builder(&f);
21551                b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
21552                unsafe {
21553                    b.launch(cfg)?;
21554                }
21555            }
21556            None => {
21557                let f = self.func("ssm_conv_ring_update_f32");
21558                let n = conv_dim * (d_conv - 1);
21559                let cfg = LaunchConfig::for_num_elems(n as u32);
21560                let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
21561                let __s_b = self.gpu.stream();
21562                let mut b = __s_b.launch_builder(&f);
21563                b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
21564                unsafe {
21565                    b.launch(cfg)?;
21566                }
21567            }
21568        }
21569        Ok(())
21570    }
21571
21572    /// task #18 increment 2: allocate ONE sequence's chunk buffers (no launches) —
21573    /// K1-K5 all run varlen afterwards. `a`/`w` become struct members so the varlen
21574    /// K2/K3 can write them.
21575    pub fn gdn_chunk_alloc(
21576        &self,
21577        n_head: usize,
21578        t: usize,
21579        c: usize,
21580        hk: usize,
21581    ) -> Result<GdnChunkBufs, Box<dyn std::error::Error>> {
21582        const D: usize = 128;
21583        assert!(
21584            c == 32,
21585            "gdn_chunk_alloc: varlen chain is the C==32 mma pair"
21586        );
21587        let h = n_head;
21588        let nc = (t + c - 1) / c;
21589        Ok(GdnChunkBufs {
21590            gcum: self.uninit(t * h)?,
21591            a: self.uninit(nc * h * c * c)?,
21592            p: self.uninit(nc * h * c * c)?,
21593            u: self.uninit(nc * h * c * D)?,
21594            w: self.uninit(nc * h * c * D)?,
21595            kb16: self.alloc_u8_uninit(t * hk * D * 2)?,
21596            wb16: self.alloc_u8_uninit(nc * h * c * D * 2)?,
21597            y16: self.alloc_u8_uninit(nc * h * c * D * 2)?,
21598            ssnap16: self.alloc_u8_uninit(nc * h * D * D * 2)?,
21599            qb16: self.alloc_u8_uninit(t * hk * D * 2)?,
21600            pb16: self.alloc_u8_uninit(nc * h * c * c * 2)?,
21601            o: self.uninit(D * h * t)?,
21602            t,
21603            nc,
21604        })
21605    }
21606
21607    /// view-source twin of f32_to_bf16 (the batched FA3 v mirror reads a concat view).
21608    pub fn f32_to_bf16_v(
21609        &self,
21610        x: &cudarc::driver::CudaView<f32>,
21611        dst: &mut CudaSlice<u8>,
21612        n: usize,
21613    ) -> Result<(), Box<dyn std::error::Error>> {
21614        let f = self.func("f32_to_bf16_bulk");
21615        let ni = n as i64;
21616        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
21617        let __s_b = self.gpu.stream();
21618        let mut b = __s_b.launch_builder(&f);
21619        b.arg(x).arg(dst).arg(&ni);
21620        unsafe {
21621            b.launch(cfg)?;
21622        }
21623        Ok(())
21624    }
21625
21626    /// f32 -> bf16 bulk mirror into a caller buffer (the K4/K5 operand mirrors).
21627    pub fn f32_to_bf16_into(
21628        &self,
21629        x: &CudaSlice<f32>,
21630        dst: &mut CudaSlice<u8>,
21631        n: usize,
21632    ) -> Result<(), Box<dyn std::error::Error>> {
21633        let f = self.func("f32_to_bf16_bulk");
21634        let ni = n as i64;
21635        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
21636        let __s_b = self.gpu.stream();
21637        let mut b = __s_b.launch_builder(&f);
21638        b.arg(x).arg(dst).arg(&ni);
21639        unsafe {
21640            b.launch(cfg)?;
21641        }
21642        Ok(())
21643    }
21644
21645    /// task #18 increment 2: varlen K1+K2+K3 — three launches run every sequence's
21646    /// cumgate/attn/solve (per-block math identical to the per-seq kernels).
21647    pub fn gdn_chunk_k123_vl8(
21648        &self,
21649        seqs: &[GdnSeqVl],
21650        n_head: usize,
21651        hk: usize,
21652        wq: Option<&GdnWVl8>,
21653    ) -> Result<(), Box<dyn std::error::Error>> {
21654        let b = seqs.len();
21655        assert!(b >= 1 && b <= 8, "gdn_chunk_k123_vl8: 1..=8 sequences");
21656        let mut packed = [GdnSeqVl::default(); 8];
21657        packed[..b].copy_from_slice(seqs);
21658        let v = GdnVl8(packed);
21659        let (hi, ci) = (n_head as i32, 32i32);
21660        let max_nc = seqs.iter().map(|a| a.nc).max().unwrap() as u32;
21661        {
21662            let f = self.func("gdn_chunk_cumgate_vl");
21663            let cfg = LaunchConfig {
21664                grid_dim: (max_nc, n_head as u32, b as u32),
21665                block_dim: (32, 1, 1),
21666                shared_mem_bytes: 0,
21667            };
21668            let __s_lb = self.gpu.stream();
21669            let mut lb = __s_lb.launch_builder(&f);
21670            lb.arg(&v).arg(&hi).arg(&ci);
21671            unsafe {
21672                lb.launch(cfg)?;
21673            }
21674        }
21675        let hki = hk as i32;
21676        if let Some(w) = wq {
21677            // K2-wgmma vl twin (writes A + pre-masked Pb16)
21678            let f = self.func("gdn_k2_wgmma_vl");
21679            let cfg = LaunchConfig {
21680                grid_dim: (max_nc, n_head as u32, b as u32),
21681                block_dim: (128, 1, 1),
21682                shared_mem_bytes: 0,
21683            };
21684            let __s_lb = self.gpu.stream();
21685            let mut lb = __s_lb.launch_builder(&f);
21686            lb.arg(&v).arg(w).arg(&hi).arg(&ci).arg(&hki);
21687            unsafe {
21688                lb.launch(cfg)?;
21689            }
21690        } else {
21691            let f = self.func("gdn_chunk_attn_vl");
21692            let cfg = LaunchConfig {
21693                grid_dim: (max_nc, n_head as u32, b as u32),
21694                block_dim: (256, 1, 1),
21695                shared_mem_bytes: 0,
21696            };
21697            let __s_lb = self.gpu.stream();
21698            let mut lb = __s_lb.launch_builder(&f);
21699            lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
21700            unsafe {
21701                lb.launch(cfg)?;
21702            }
21703        }
21704        {
21705            let f = self.func("gdn_chunk_solve32_vl");
21706            let cfg = LaunchConfig {
21707                grid_dim: (max_nc, n_head as u32, b as u32),
21708                block_dim: (256, 1, 1),
21709                shared_mem_bytes: 0,
21710            };
21711            let __s_lb = self.gpu.stream();
21712            let mut lb = __s_lb.launch_builder(&f);
21713            lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
21714            unsafe {
21715                lb.launch(cfg)?;
21716            }
21717        }
21718        Ok(())
21719    }
21720
21721    /// task #18 increment 3: varlen PREP chain — conv(+ring) / repack / fused-l2 /
21722    /// fused gate-prep, 5 launches for every sequence (per-element math identical
21723    /// to the per-seq kernels; l2/gate fusions write disjoint outputs).
21724    #[allow(clippy::too_many_arguments)]
21725    pub fn gdn_prep_vl8(
21726        &self,
21727        seqs: &[GdnPrepVl],
21728        conv_w: &CudaSlice<f32>,
21729        dt_bias: &CudaSlice<f32>,
21730        a: &CudaSlice<f32>,
21731        conv_dim: usize,
21732        d_conv: usize,
21733        d_state: usize,
21734        num_v: usize,
21735        num_k: usize,
21736        key_dim: usize,
21737        hk: usize,
21738        eps: f32,
21739    ) -> Result<(), Box<dyn std::error::Error>> {
21740        let b = seqs.len();
21741        assert!(b >= 1 && b <= 8);
21742        let mut packed = [GdnPrepVl::default(); 8];
21743        packed[..b].copy_from_slice(seqs);
21744        let v = GdnPrepVl8(packed);
21745        let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
21746        let (cdi, dci) = (conv_dim as i32, d_conv as i32);
21747        let conv_fuse = std::env::var("MEMRA_CONV_FUSE").as_deref() != Ok("0");
21748        assert!(
21749            conv_fuse || hk == num_v,
21750            "de-broadcast requires the fused conv"
21751        );
21752        if conv_fuse {
21753            let f = self.func("ssm_conv1d_gdn_state_vl");
21754            let cfg = LaunchConfig {
21755                grid_dim: ((conv_dim as u32).div_ceil(256), max_t, b as u32),
21756                block_dim: (256, 1, 1),
21757                shared_mem_bytes: 0,
21758            };
21759            let (dsi, nvi, nki, kdi, hki) = (
21760                d_state as i32,
21761                num_v as i32,
21762                num_k as i32,
21763                key_dim as i32,
21764                hk as i32,
21765            );
21766            let __s_lb = self.gpu.stream();
21767            let mut lb = __s_lb.launch_builder(&f);
21768            lb.arg(&v)
21769                .arg(conv_w)
21770                .arg(&cdi)
21771                .arg(&dci)
21772                .arg(&dsi)
21773                .arg(&nvi)
21774                .arg(&nki)
21775                .arg(&kdi)
21776                .arg(&hki);
21777            unsafe {
21778                lb.launch(cfg)?;
21779            }
21780        } else {
21781            let f = self.func("ssm_conv1d_tm_state_vl");
21782            let cfg = LaunchConfig {
21783                grid_dim: ((conv_dim as u32).div_ceil(256), max_t, b as u32),
21784                block_dim: (256, 1, 1),
21785                shared_mem_bytes: 0,
21786            };
21787            let __s_lb = self.gpu.stream();
21788            let mut lb = __s_lb.launch_builder(&f);
21789            lb.arg(&v).arg(conv_w).arg(&cdi).arg(&dci);
21790            unsafe {
21791                lb.launch(cfg)?;
21792            }
21793        }
21794        {
21795            let f = self.func("ssm_conv_ring_update_vl");
21796            let n = (conv_dim * (d_conv - 1)) as u32;
21797            let cfg = LaunchConfig {
21798                grid_dim: (n.div_ceil(256), 1, b as u32),
21799                block_dim: (256, 1, 1),
21800                shared_mem_bytes: 0,
21801            };
21802            let __s_lb = self.gpu.stream();
21803            let mut lb = __s_lb.launch_builder(&f);
21804            lb.arg(&v).arg(&cdi).arg(&dci);
21805            unsafe {
21806                lb.launch(cfg)?;
21807            }
21808        }
21809        if !conv_fuse {
21810            let f = self.func("qkv_to_gdn_repack_vl");
21811            let n = max_t * (num_v * d_state) as u32;
21812            let cfg = LaunchConfig {
21813                grid_dim: (n.div_ceil(256), 1, b as u32),
21814                block_dim: (256, 1, 1),
21815                shared_mem_bytes: 0,
21816            };
21817            let (dsi, nvi, nki, kdi) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
21818            let __s_lb = self.gpu.stream();
21819            let mut lb = __s_lb.launch_builder(&f);
21820            lb.arg(&v).arg(&dsi).arg(&nvi).arg(&nki).arg(&kdi);
21821            unsafe {
21822                lb.launch(cfg)?;
21823            }
21824        }
21825        if Self::l2_v2_on(d_state) {
21826            let f = self.func("gdn_l2_v2_vl");
21827            let cfg = LaunchConfig {
21828                grid_dim: ((max_t * hk as u32).div_ceil(8), 2, b as u32),
21829                block_dim: (256, 1, 1),
21830                shared_mem_bytes: 0,
21831            };
21832            let (dsi, nvi) = (d_state as i32, hk as i32);
21833            let __s_lb = self.gpu.stream();
21834            let mut lb = __s_lb.launch_builder(&f);
21835            lb.arg(&v).arg(&dsi).arg(&nvi).arg(&eps);
21836            unsafe {
21837                lb.launch(cfg)?;
21838            }
21839        } else {
21840            let f = self.func("gdn_l2_vl");
21841            let cfg = LaunchConfig {
21842                grid_dim: (max_t * hk as u32, 2, b as u32),
21843                block_dim: (256, 1, 1),
21844                shared_mem_bytes: 0,
21845            };
21846            let (dsi, nvi) = (d_state as i32, hk as i32);
21847            let __s_lb = self.gpu.stream();
21848            let mut lb = __s_lb.launch_builder(&f);
21849            lb.arg(&v).arg(&dsi).arg(&nvi).arg(&eps);
21850            unsafe {
21851                lb.launch(cfg)?;
21852            }
21853        }
21854        {
21855            let f = self.func("gdn_gate_prep_vl");
21856            let n = max_t * num_v as u32;
21857            let cfg = LaunchConfig {
21858                grid_dim: (n.div_ceil(256), 1, b as u32),
21859                block_dim: (256, 1, 1),
21860                shared_mem_bytes: 0,
21861            };
21862            let nvi = num_v as i32;
21863            let __s_lb = self.gpu.stream();
21864            let mut lb = __s_lb.launch_builder(&f);
21865            lb.arg(&v).arg(dt_bias).arg(a).arg(&nvi);
21866            unsafe {
21867                lb.launch(cfg)?;
21868            }
21869        }
21870        Ok(())
21871    }
21872
21873    /// varlen bf16 mirrors over the gdnseq_t table (which: 0 = k_l2 -> kb16, 1 = w -> wb16).
21874    pub fn gdn_mirror_vl8(
21875        &self,
21876        seqs: &[GdnSeqVl],
21877        n_head: usize,
21878        which: i32,
21879        hk: usize,
21880    ) -> Result<(), Box<dyn std::error::Error>> {
21881        let b = seqs.len();
21882        assert!(b >= 1 && b <= 8);
21883        let mut packed = [GdnSeqVl::default(); 8];
21884        packed[..b].copy_from_slice(seqs);
21885        let v = GdnVl8(packed);
21886        let ept = (if which == 0 { hk } else { n_head } * 128) as i32;
21887        let max_n = seqs
21888            .iter()
21889            .map(|s| {
21890                if which == 0 {
21891                    s.t as i64 * ept as i64
21892                } else {
21893                    s.nc as i64 * ept as i64 * 32
21894                }
21895            })
21896            .max()
21897            .unwrap();
21898        let f = self.func("gdn_mirror_vl");
21899        let blocks = ((max_n as u32).div_ceil(4)).div_ceil(256);
21900        let cfg = LaunchConfig {
21901            grid_dim: (blocks, 1, b as u32),
21902            block_dim: (256, 1, 1),
21903            shared_mem_bytes: 0,
21904        };
21905        let __s_lb = self.gpu.stream();
21906        let mut lb = __s_lb.launch_builder(&f);
21907        lb.arg(&v).arg(&ept).arg(&which);
21908        unsafe {
21909            lb.launch(cfg)?;
21910        }
21911        Ok(())
21912    }
21913
21914    /// varlen gated-norm tail (+f16out) — one launch replaces B gated_rmsnorm calls.
21915    pub fn gdn_tail_vl8(
21916        &self,
21917        seqs: &[GdnPrepVl],
21918        norm_w: &CudaSlice<f32>,
21919        d_state: usize,
21920        num_v: usize,
21921        eps: f32,
21922    ) -> Result<(), Box<dyn std::error::Error>> {
21923        let b = seqs.len();
21924        assert!(b >= 1 && b <= 8);
21925        let mut packed = [GdnPrepVl::default(); 8];
21926        packed[..b].copy_from_slice(seqs);
21927        let v = GdnPrepVl8(packed);
21928        let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
21929        let f = self.func("gated_rmsnorm_f16out_vl");
21930        // block_dim MUST match gated_rmsnorm's (128): the reduction tree order pins the scale
21931        let cfg = LaunchConfig {
21932            grid_dim: (max_t * num_v as u32, 1, b as u32),
21933            block_dim: (128, 1, 1),
21934            shared_mem_bytes: 0,
21935        };
21936        let (dsi, nvi) = (d_state as i32, num_v as i32);
21937        let __s_lb = self.gpu.stream();
21938        let mut lb = __s_lb.launch_builder(&f);
21939        lb.arg(&v).arg(norm_w).arg(&dsi).arg(&nvi).arg(&eps);
21940        unsafe {
21941            lb.launch(cfg)?;
21942        }
21943        Ok(())
21944    }
21945
21946    /// Raw device address helpers for the varlen by-value arg struct (single-stream
21947    /// launches; every buffer outlives the call — the f16 FFI discipline).
21948    pub fn addr_f32(&self, x: &CudaSlice<f32>) -> u64 {
21949        use cudarc::driver::DevicePtr;
21950        let s = self.gpu.stream();
21951        let (p, _g) = x.device_ptr(&s);
21952        p as u64
21953    }
21954    pub fn addr_f32_mut(&self, x: &mut CudaSlice<f32>) -> u64 {
21955        use cudarc::driver::DevicePtrMut;
21956        let s = self.gpu.stream();
21957        let (p, _g) = x.device_ptr_mut(&s);
21958        p as u64
21959    }
21960    pub fn addr_f32v(&self, x: &cudarc::driver::CudaView<f32>) -> u64 {
21961        use cudarc::driver::DevicePtr;
21962        let s = self.gpu.stream();
21963        let (p, _g) = x.device_ptr(&s);
21964        p as u64
21965    }
21966    pub fn addr_u8(&self, x: &CudaSlice<u8>) -> u64 {
21967        use cudarc::driver::DevicePtr;
21968        let s = self.gpu.stream();
21969        let (p, _g) = x.device_ptr(&s);
21970        p as u64
21971    }
21972
21973    /// task #18: the varlen K4+K5 pair — TWO launches run every sequence's state pass
21974    /// and output pass (grid gains a seq dim; per-block math identical to the per-seq
21975    /// launches, so this is strictly bit-gateable against them).
21976    pub fn gdn_chunk_vl8(
21977        &self,
21978        seqs: &[GdnSeqVl],
21979        n_head: usize,
21980        scale: f32,
21981        hk: usize,
21982        wq: Option<&GdnWVl8>,
21983    ) -> Result<(), Box<dyn std::error::Error>> {
21984        const NSPLIT: u32 = 4;
21985        let b = seqs.len();
21986        assert!(b >= 1 && b <= 8, "gdn_chunk_vl8: 1..=8 sequences");
21987        let mut packed = [GdnSeqVl::default(); 8];
21988        packed[..b].copy_from_slice(seqs);
21989        let v = GdnVl8(packed);
21990        let (hi, ci) = (n_head as i32, 32i32);
21991        let max_nc = seqs.iter().map(|a| a.nc).max().unwrap() as u32;
21992        let hki = hk as i32;
21993        if let Some(w) = wq {
21994            // K4+K5 fused wgmma vl twin: one launch, Y/Ssnap never materialized.
21995            let f = self.func("gdn_k45_wgmma_vl");
21996            let cfg = LaunchConfig {
21997                grid_dim: (n_head as u32, NSPLIT, b as u32),
21998                block_dim: (256, 1, 1),
21999                shared_mem_bytes: 0,
22000            };
22001            let __s_lb = self.gpu.stream();
22002            let mut lb = __s_lb.launch_builder(&f);
22003            lb.arg(&v).arg(w).arg(&scale).arg(&hi).arg(&ci).arg(&hki);
22004            unsafe {
22005                lb.launch(cfg)?;
22006            }
22007            let _ = max_nc;
22008            return Ok(());
22009        }
22010        {
22011            let f = self.func("gdn_chunk_state_mma_vl");
22012            let cfg = LaunchConfig {
22013                grid_dim: (n_head as u32, NSPLIT, b as u32),
22014                block_dim: (256, 1, 1),
22015                shared_mem_bytes: 0,
22016            };
22017            let __s_lb = self.gpu.stream();
22018            let mut lb = __s_lb.launch_builder(&f);
22019            lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
22020            unsafe {
22021                lb.launch(cfg)?;
22022            }
22023        }
22024        {
22025            let f = self.func("gdn_chunk_output_mma_vl");
22026            let cfg = LaunchConfig {
22027                grid_dim: (max_nc, n_head as u32, b as u32),
22028                block_dim: (256, 1, 1),
22029                shared_mem_bytes: 0,
22030            };
22031            let __s_lb = self.gpu.stream();
22032            let mut lb = __s_lb.launch_builder(&f);
22033            lb.arg(&v).arg(&hi).arg(&ci).arg(&scale).arg(&hki);
22034            unsafe {
22035                lb.launch(cfg)?;
22036            }
22037        }
22038        Ok(())
22039    }
22040    pub fn gdn_scan_chunked(
22041        &self,
22042        q: &CudaSlice<f32>,
22043        k: &CudaSlice<f32>,
22044        v: &CudaSlice<f32>,
22045        g: &CudaSlice<f32>,
22046        beta: &CudaSlice<f32>,
22047        kb16_pre: Option<&CudaSlice<u8>>,
22048        qb16_pre: Option<&CudaSlice<u8>>,
22049        state_in: &CudaSlice<f32>,
22050        state_out: &mut CudaSlice<f32>,
22051        o: &mut CudaSlice<f32>,
22052        n_head: usize,
22053        t: usize,
22054        scale: f32,
22055        c: usize,
22056        hk: usize,
22057    ) -> Result<(), Box<dyn std::error::Error>> {
22058        const D: usize = 128;
22059        const NSPLIT: u32 = 4;
22060        assert!(c >= 1 && c <= 128, "gdn_scan_chunked: C must be in 1..=128");
22061        let h = n_head;
22062        let nc = (t + c - 1) / c;
22063        let (hi, ti, ci) = (h as i32, t as i32, c as i32);
22064        // mirror-fold (round 27): on the mma path W's bf16 twin is emitted by K3's store
22065        // (wb16 pre-allocated and threaded through k123) and k's by the producer l2 when
22066        // the caller hands `kb16_pre` — both standalone mirror passes disappear.
22067        // default mirrors gdn_mma_enabled (incl. the sm_120a-build ON, lane/moeprime-nvfp4-direct)
22068        let gdn_mma_pre = !portable_mma_gated()
22069            && c == 32
22070            && match std::env::var("MEMRA_GDN_MMA").as_deref() {
22071                Ok("1") => true,
22072                Ok("0") => false,
22073                _ => gdn_mma_default_on(),
22074            };
22075        let mut wb16_pre: Option<CudaSlice<u8>> = if gdn_mma_pre {
22076            Some(self.alloc_u8_uninit(nc * h * c * D * 2)?)
22077        } else {
22078            None
22079        };
22080        // K2-wgmma pre-work (MEMRA_GDN_WGMMA): the kb16/qb16 mirrors hoist ABOVE K123 so
22081        // K2 rides them via cp.async; K2 writes the pre-masked Pb16 directly (the
22082        // gdn_p_bf16_masked pass and the in-branch mirror builds disappear).
22083        // mirrors gdn_wgmma_on: hard Hopper-build gate (empty wgmma bodies elsewhere)
22084        let gdn_wgmma_pre = cfg!(memra_hopper_mma)
22085            && gdn_mma_pre
22086            && std::env::var("MEMRA_GDN_WGMMA").as_deref() != Ok("0");
22087        let nk = t * hk * D;
22088        let mut kb16_local: Option<CudaSlice<u8>> = None;
22089        if gdn_mma_pre && kb16_pre.is_none() {
22090            let mut kb = self.alloc_u8_uninit(nk * 2)?;
22091            let f = self.func("f32_to_bf16_bulk");
22092            let n2 = nk as i64;
22093            let cfg2 = LaunchConfig::for_num_elems((nk as u32).div_ceil(4));
22094            let __s_b = self.gpu.stream();
22095            let mut b = __s_b.launch_builder(&f);
22096            b.arg(k).arg(&mut kb).arg(&n2);
22097            unsafe {
22098                b.launch(cfg2)?;
22099            }
22100            kb16_local = Some(kb);
22101        }
22102        let kb16_ref0: Option<&CudaSlice<u8>> = kb16_local.as_ref().or(kb16_pre);
22103        if let Some(kb) = kb16_pre {
22104            assert!(kb.len() >= nk * 2, "kb16_pre too small");
22105        }
22106        let mut qb16: Option<CudaSlice<u8>> = None;
22107        let mut pb16: Option<CudaSlice<u8>> = None;
22108        if gdn_wgmma_pre {
22109            // mirror-fold (round 35): prep's l2 v2 emits qb16 in-epilogue (kb16 pattern);
22110            // the standalone bulk cvt only serves callers without the prep mirror.
22111            if qb16_pre.is_none() {
22112                let mut qb = self.alloc_u8_uninit(nk * 2)?;
22113                let f = self.func("f32_to_bf16_bulk");
22114                let n2 = nk as i64;
22115                let cfg2 = LaunchConfig::for_num_elems((nk as u32).div_ceil(4));
22116                let __s_b = self.gpu.stream();
22117                let mut b = __s_b.launch_builder(&f);
22118                b.arg(q).arg(&mut qb).arg(&n2);
22119                unsafe {
22120                    b.launch(cfg2)?;
22121                }
22122                qb16 = Some(qb);
22123            } else if let Some(qb) = qb16_pre {
22124                assert!(qb.len() >= nk * 2, "qb16_pre too small");
22125            }
22126            pb16 = Some(self.alloc_u8_uninit(nc * h * c * c * 2)?);
22127        }
22128        let qb16_ref0: Option<&CudaSlice<u8>> = qb16.as_ref().or(qb16_pre);
22129        let k2w = if gdn_wgmma_pre {
22130            Some((
22131                *qb16_ref0.as_ref().unwrap(),
22132                *kb16_ref0.as_ref().unwrap(),
22133                pb16.as_mut().unwrap(),
22134            ))
22135        } else {
22136            None
22137        };
22138        let (gcum, p, u, w) =
22139            self.gdn_chunk_k123(q, k, v, g, beta, wb16_pre.as_mut(), n_head, t, c, hk, k2w)?;
22140        let _ = &w;
22141        let mut y = self.uninit(nc * h * c * D)?;
22142        let mut ssnap = self.uninit(nc * h * D * D)?; // chunk-start state snapshots (K5 phase 1)
22143        // K4-MMA seam (MEMRA_GDN_MMA; harness verdict 1.75x — tools/bench_gdn_k4.cu, ledger
22144        // 2026-07-26): M in mma accumulator fragments, bf16 W/k mirrors through a cp.async
22145        // ring. C==32 only (the kernel's tile). PROMOTED default-ON on the Hopper lane
22146        // after the STATE-CARRY battery (2026-07-26): 2048-token prime (64 in-kernel state
22147        // carries) -> 256 greedy decode tokens IDENTICAL to f32 on 3 seeds, AND chunked-
22148        // continuation prime (MEMRA_PRIME_CHUNK=512, 4 cross-call carries via cache.recur)
22149        // IDENTICAL on 2 seeds; plus argmax MATCH, pp512 +3.5% (17286), oracle out
22150        // mean_rel ~1e-4. kernel-check pins BOTH configs (f32 tight band forced =0; mma
22151        // band 8e-2/8e-1 vs f64 truth). =0 reverts; portable stays f32. NOT read via
22152        // OnceLock ON PURPOSE: kernel-check toggles the env per call to pin both forms.
22153        // default mirrors gdn_mma_enabled (incl. the sm_120a-build ON) — all three read
22154        // sites must agree or the pre-work arms while the scan takes the scalar route.
22155        let gdn_mma = !portable_mma_gated()
22156            && c == 32
22157            && match std::env::var("MEMRA_GDN_MMA").as_deref() {
22158                Ok("1") => true,
22159                Ok("0") => false,
22160                _ => gdn_mma_default_on(),
22161            };
22162        if gdn_mma {
22163            let wb16 = wb16_pre
22164                .take()
22165                .expect("mma path pre-allocates wb16 (K3 store fold)");
22166            let kb16_ref: &CudaSlice<u8> = kb16_ref0.expect("mma path pre-builds kb16 above K123");
22167            // K4+K5 FUSED wgmma seam (MEMRA_GDN_WGMMA, task #22; harness verdict
22168            // tools/bench_gdn_wgmma.cu v5, ledger 1f08b997: in-band Y 1.07e-2 / state
22169            // 1.03e-2 / O 1.08e-2, 91.3us vs 70.4 K4-only at H=32 T=512). K5's output
22170            // pass runs inside the persistent-M kernel; Y and Ssnap are never
22171            // materialized. New numeric class (gk folds into k^T instead of ys) —
22172            // explicit opt-in until the state-carry battery promotes it. Env read per
22173            // call (kernel-check pins configs by toggling env, GDN_MMA precedent).
22174            // PROMOTED default-ON hopper (2026-07-27): full battery green — harness
22175            // in-band, argmax gate PASS, 3-seed greedy IDENTICAL after ~2k prime,
22176            // chunked-continuation IDENTICAL, kernel-check + decode-batch gates green,
22177            // official prefill lane +0.74% interleaved x5 (5/5 rounds). =0 reverts.
22178            if gdn_wgmma_pre {
22179                // qb16/pb16 pre-built above K123 (K2-wgmma wrote the masked Pb16).
22180                let qb16 = qb16_ref0.unwrap();
22181                let pb16 = pb16.as_ref().unwrap();
22182                {
22183                    let f = self.func("gdn_k45_wgmma");
22184                    let cfg = LaunchConfig {
22185                        grid_dim: (h as u32, 4, 1),
22186                        block_dim: (256, 1, 1),
22187                        shared_mem_bytes: 0,
22188                    };
22189                    let hki = hk as i32;
22190                    let __s_b = self.gpu.stream();
22191                    let mut b = __s_b.launch_builder(&f);
22192                    b.arg(kb16_ref)
22193                        .arg(&gcum)
22194                        .arg(beta)
22195                        .arg(&u)
22196                        .arg(&wb16)
22197                        .arg(qb16)
22198                        .arg(pb16)
22199                        .arg(o)
22200                        .arg(&scale)
22201                        .arg(state_in)
22202                        .arg(&mut *state_out)
22203                        .arg(&hi)
22204                        .arg(&ti)
22205                        .arg(&ci)
22206                        .arg(&hki);
22207                    unsafe {
22208                        b.launch(cfg)?;
22209                    }
22210                }
22211                return Ok(());
22212            }
22213            // COUPLED PAIR: K4-mma writes Y and Ssnap as bf16 (their only consumer is
22214            // K5-mma, which rounds to bf16 regardless — identical numerics, half the
22215            // traffic; harness K5 63.0 -> 35.3us). Fresh bf16 buffers replace the f32 ones.
22216            let mut y16 = self.alloc_u8_uninit(nc * h * c * D * 2)?;
22217            let mut ssnap16 = self.alloc_u8_uninit(nc * h * D * D * 2)?;
22218            {
22219                let f = self.func("gdn_chunk_state_mma");
22220                let cfg = LaunchConfig {
22221                    grid_dim: (h as u32, NSPLIT, 1),
22222                    block_dim: (256, 1, 1),
22223                    shared_mem_bytes: 0,
22224                };
22225                let hki = hk as i32;
22226                let __s_b = self.gpu.stream();
22227                let mut b = __s_b.launch_builder(&f);
22228                b.arg(kb16_ref)
22229                    .arg(&gcum)
22230                    .arg(beta)
22231                    .arg(&u)
22232                    .arg(&wb16)
22233                    .arg(&mut y16)
22234                    .arg(&mut ssnap16)
22235                    .arg(state_in)
22236                    .arg(&mut *state_out)
22237                    .arg(&hi)
22238                    .arg(&ti)
22239                    .arg(&ci)
22240                    .arg(&hki);
22241                unsafe {
22242                    b.launch(cfg)?;
22243                }
22244            }
22245            {
22246                // K5-mma (bf16 St/Y consumers)
22247                let f = self.func("gdn_chunk_output_mma");
22248                let jt = ((c + 31) / 32) as u32;
22249                let cfg = LaunchConfig {
22250                    grid_dim: (nc as u32, h as u32, jt),
22251                    block_dim: (256, 1, 1),
22252                    shared_mem_bytes: 0,
22253                };
22254                let hki = hk as i32;
22255                let __s_b = self.gpu.stream();
22256                let mut b = __s_b.launch_builder(&f);
22257                b.arg(q)
22258                    .arg(&gcum)
22259                    .arg(&p)
22260                    .arg(&y16)
22261                    .arg(&ssnap16)
22262                    .arg(o)
22263                    .arg(&hi)
22264                    .arg(&ti)
22265                    .arg(&ci)
22266                    .arg(&scale)
22267                    .arg(&hki);
22268                unsafe {
22269                    b.launch(cfg)?;
22270                }
22271            }
22272            return Ok(());
22273        }
22274        {
22275            // K4 (sequential over chunks inside; blocks col-partition the state)
22276            let f = self.func("gdn_chunk_state_f32");
22277            let cfg = LaunchConfig {
22278                grid_dim: (h as u32, NSPLIT, 1),
22279                block_dim: (256, 1, 1),
22280                shared_mem_bytes: 0,
22281            };
22282            let __s_b = self.gpu.stream();
22283            let mut b = __s_b.launch_builder(&f);
22284            b.arg(k)
22285                .arg(&gcum)
22286                .arg(beta)
22287                .arg(&u)
22288                .arg(&w)
22289                .arg(&mut y)
22290                .arg(&mut ssnap)
22291                .arg(state_in)
22292                .arg(&mut *state_out)
22293                .arg(&hi)
22294                .arg(&ti)
22295                .arg(&ci);
22296            unsafe {
22297                b.launch(cfg)?;
22298            }
22299        }
22300        {
22301            // K5 (j-blocked: grid.z = 32-row output blocks per chunk; writes o fully)
22302            let f = self.func("gdn_chunk_output_f32");
22303            let jt = ((c + 31) / 32) as u32;
22304            let cfg = LaunchConfig {
22305                grid_dim: (nc as u32, h as u32, jt),
22306                block_dim: (256, 1, 1),
22307                shared_mem_bytes: 0,
22308            };
22309            let __s_b = self.gpu.stream();
22310            let mut b = __s_b.launch_builder(&f);
22311            b.arg(q)
22312                .arg(&gcum)
22313                .arg(&p)
22314                .arg(&y)
22315                .arg(&ssnap)
22316                .arg(o)
22317                .arg(&hi)
22318                .arg(&ti)
22319                .arg(&ci)
22320                .arg(&scale);
22321            unsafe {
22322                b.launch(cfg)?;
22323            }
22324        }
22325        Ok(())
22326    }
22327
22328    /// PREFILL GDN scan dispatch (the A4 seam): chunked WY form when enabled and T is in the
22329    /// batched-prefill regime, else the sequential scan. Callers: hybrid_forward::linear_attn
22330    /// (forward/forward_last) + linear_attn_prime (prime_cache). Decode (T=1) and the spec
22331    /// verify call `gdn_scan_s128` DIRECTLY — the decode==verify dispatch identity is untouched.
22332    ///
22333    /// MEMRA_GDN_DIFF=1: numerical-oracle mode — runs BOTH forms on the same inputs, prints the
22334    /// per-call (== per-layer, in call order) output/state error distribution, and keeps the
22335    /// SEQUENTIAL results so the run stays on the shipped path (stage-1 prototype evidence).
22336    #[allow(clippy::too_many_arguments)]
22337    #[allow(clippy::too_many_arguments)]
22338    pub fn gdn_scan_prefill(
22339        &self,
22340        q: &CudaSlice<f32>,
22341        k: &CudaSlice<f32>,
22342        v: &CudaSlice<f32>,
22343        g: &CudaSlice<f32>,
22344        beta: &CudaSlice<f32>,
22345        kb16_pre: Option<&CudaSlice<u8>>,
22346        qb16_pre: Option<&CudaSlice<u8>>,
22347        state_in: &CudaSlice<f32>,
22348        state_out: &mut CudaSlice<f32>,
22349        o: &mut CudaSlice<f32>,
22350        n_head: usize,
22351        t: usize,
22352        scale: f32,
22353        hk: usize,
22354    ) -> Result<(), Box<dyn std::error::Error>> {
22355        if std::env::var("MEMRA_GDN_DIFF").is_ok() && t >= 16 {
22356            assert!(hk == n_head, "GDN_DIFF oracle is broadcast-only");
22357            return self.gdn_scan_diff(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale);
22358        }
22359        if Self::gdn_chunked_enabled() && t >= 16 {
22360            self.gdn_scan_chunked(
22361                q,
22362                k,
22363                v,
22364                g,
22365                beta,
22366                kb16_pre,
22367                qb16_pre,
22368                state_in,
22369                state_out,
22370                o,
22371                n_head,
22372                t,
22373                scale,
22374                Self::gdn_chunk_size(),
22375                hk,
22376            )
22377        } else {
22378            assert!(
22379                hk == n_head,
22380                "s128 scan is broadcast-only (prep guarantees by predicate)"
22381            );
22382            self.gdn_scan_s128(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale)
22383        }
22384    }
22385
22386    /// Stage-1 oracle: run sequential AND chunked, report per-call error stats, keep sequential.
22387    #[allow(clippy::too_many_arguments)]
22388    fn gdn_scan_diff(
22389        &self,
22390        q: &CudaSlice<f32>,
22391        k: &CudaSlice<f32>,
22392        v: &CudaSlice<f32>,
22393        g: &CudaSlice<f32>,
22394        beta: &CudaSlice<f32>,
22395        state_in: &CudaSlice<f32>,
22396        state_out: &mut CudaSlice<f32>,
22397        o: &mut CudaSlice<f32>,
22398        n_head: usize,
22399        t: usize,
22400        scale: f32,
22401    ) -> Result<(), Box<dyn std::error::Error>> {
22402        static CALL: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
22403        let call = CALL.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
22404        let mut o_c = self.uninit(o.len())?;
22405        let mut st_c = self.uninit(state_out.len())?;
22406        self.gdn_scan_chunked(
22407            q,
22408            k,
22409            v,
22410            g,
22411            beta,
22412            None,
22413            None,
22414            state_in,
22415            &mut st_c,
22416            &mut o_c,
22417            n_head,
22418            t,
22419            scale,
22420            Self::gdn_chunk_size(),
22421            n_head,
22422        )?;
22423        self.gdn_scan_s128(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale)?;
22424        let (oh_s, oh_c) = (self.dtoh(o)?, self.dtoh(&o_c)?);
22425        let (sh_s, sh_c) = (self.dtoh(state_out)?, self.dtoh(&st_c)?);
22426        let stats = |a: &[f32], b: &[f32]| -> (f32, f32, f64) {
22427            let mut max_abs = 0f32;
22428            let mut max_rel = 0f32;
22429            let mut sum_rel = 0f64;
22430            for (x, y) in a.iter().zip(b) {
22431                let ad = (x - y).abs();
22432                let rel = ad / x.abs().max(y.abs()).max(1e-3);
22433                if ad > max_abs {
22434                    max_abs = ad;
22435                }
22436                if rel > max_rel {
22437                    max_rel = rel;
22438                }
22439                sum_rel += rel as f64;
22440            }
22441            (max_abs, max_rel, sum_rel / a.len() as f64)
22442        };
22443        let (o_ma, o_mr, o_mean) = stats(&oh_s, &oh_c);
22444        let (s_ma, s_mr, s_mean) = stats(&sh_s, &sh_c);
22445        println!(
22446            "[gdn-diff call {call:3} T={t} C={}] out: max_abs={o_ma:.3e} max_rel={o_mr:.3e} mean_rel={o_mean:.3e} | \
22447                  state: max_abs={s_ma:.3e} max_rel={s_mr:.3e} mean_rel={s_mean:.3e}",
22448            Self::gdn_chunk_size()
22449        );
22450        Ok(())
22451    }
22452
22453    /// softplus-based g_log: g_log[h,t] = a[h] * softplus(alpha[h,t] + dt_bias[h]). a pre-negated.
22454    pub fn gdn_glog(
22455        &self,
22456        alpha: &CudaSlice<f32>,
22457        dt_bias: &CudaSlice<f32>,
22458        a: &CudaSlice<f32>,
22459        g_log: &mut CudaSlice<f32>,
22460        n_head: usize,
22461        t: usize,
22462    ) -> Result<(), Box<dyn std::error::Error>> {
22463        let f = self.func("gdn_glog_f32");
22464        let cfg = LaunchConfig::for_num_elems((n_head * t) as u32);
22465        let (h, ti) = (n_head as i32, t as i32);
22466        let __s_b = self.gpu.stream();
22467        let mut b = __s_b.launch_builder(&f);
22468        b.arg(alpha).arg(dt_bias).arg(a).arg(g_log).arg(&h).arg(&ti);
22469        unsafe {
22470            b.launch(cfg)?;
22471        }
22472        Ok(())
22473    }
22474
22475    /// view twins (task #16): the batched prime's GDN core reads the CONCAT projection
22476    /// buffers at row offsets (CudaView) — same kernels, same values, no split copies.
22477    pub fn sigmoid_v(
22478        &self,
22479        x: &cudarc::driver::CudaView<f32>,
22480        y: &mut CudaSlice<f32>,
22481        n: usize,
22482    ) -> Result<(), Box<dyn std::error::Error>> {
22483        let f = self.func("sigmoid_f32");
22484        let cfg = LaunchConfig::for_num_elems(n as u32);
22485        let ni = n as i32;
22486        let __s_b = self.gpu.stream();
22487        let mut b = __s_b.launch_builder(&f);
22488        b.arg(x).arg(y).arg(&ni);
22489        unsafe {
22490            b.launch(cfg)?;
22491        }
22492        Ok(())
22493    }
22494
22495    pub fn gdn_glog_v(
22496        &self,
22497        alpha: &cudarc::driver::CudaView<f32>,
22498        dt_bias: &CudaSlice<f32>,
22499        a: &CudaSlice<f32>,
22500        g_log: &mut CudaSlice<f32>,
22501        n_head: usize,
22502        t: usize,
22503    ) -> Result<(), Box<dyn std::error::Error>> {
22504        let f = self.func("gdn_glog_f32");
22505        let cfg = LaunchConfig::for_num_elems((n_head * t) as u32);
22506        let (h, ti) = (n_head as i32, t as i32);
22507        let __s_b = self.gpu.stream();
22508        let mut b = __s_b.launch_builder(&f);
22509        b.arg(alpha).arg(dt_bias).arg(a).arg(g_log).arg(&h).arg(&ti);
22510        unsafe {
22511            b.launch(cfg)?;
22512        }
22513        Ok(())
22514    }
22515
22516    pub fn sigmoid(
22517        &self,
22518        x: &CudaSlice<f32>,
22519        y: &mut CudaSlice<f32>,
22520        n: usize,
22521    ) -> Result<(), Box<dyn std::error::Error>> {
22522        let f = self.func("sigmoid_f32");
22523        let cfg = LaunchConfig::for_num_elems(n as u32);
22524        let ni = n as i32;
22525        let __s_b = self.gpu.stream();
22526        let mut b = __s_b.launch_builder(&f);
22527        b.arg(x).arg(y).arg(&ni);
22528        unsafe {
22529            b.launch(cfg)?;
22530        }
22531        Ok(())
22532    }
22533
22534    /// attn out-gate fused epilogue (task #17): dst = a * sigmoid(g) + fp16 twin, one launch
22535    /// (replaces sigmoid + mul + convert). Bit-identical class.
22536    pub fn sig_mul_f16out(
22537        &self,
22538        a: &CudaSlice<f32>,
22539        g: &CudaSlice<f32>,
22540        dst: &mut CudaSlice<f32>,
22541        dst16: &mut CudaSlice<u8>,
22542        n: usize,
22543    ) -> Result<(), Box<dyn std::error::Error>> {
22544        let f = self.func("sig_mul_f16out_f32");
22545        let cfg = LaunchConfig::for_num_elems(n as u32);
22546        let ni = n as i32;
22547        let __s_b = self.gpu.stream();
22548        let mut b = __s_b.launch_builder(&f);
22549        b.arg(a).arg(g).arg(dst).arg(dst16).arg(&ni);
22550        unsafe {
22551            b.launch(cfg)?;
22552        }
22553        Ok(())
22554    }
22555
22556    /// step35 (Step-3.7-Flash) SEPARATE head-wise attention gate: one scalar per query head,
22557    /// broadcast over head_dim. `dst = a * sigmoid(g)` where `a`/`dst` are `[head_dim, n_head, T]`
22558    /// (the `q_gate_split` layout) and `g` is the PRE-sigmoid `attn_gate` projection output in
22559    /// token-major `[T, n_head]`. `dst16` is the optional fp16 operand for wo (None -> skipped).
22560    ///
22561    /// NOT interchangeable with `sig_mul_f16out`, which gates FULL WIDTH (qwen35 packs one gate
22562    /// value per (head, dim) element inside wq). Using this for that, or that for this, silently
22563    /// applies the wrong number of distinct gate values.
22564    #[allow(clippy::too_many_arguments)]
22565    pub fn attn_head_gate(
22566        &self,
22567        a: &CudaSlice<f32>,
22568        g: &CudaSlice<f32>,
22569        dst: &mut CudaSlice<f32>,
22570        dst16: Option<&mut CudaSlice<u8>>,
22571        head_dim: usize,
22572        n_head: usize,
22573        t: usize,
22574    ) -> Result<(), Box<dyn std::error::Error>> {
22575        let f = self.func("attn_head_gate_f32");
22576        let cfg = LaunchConfig::for_num_elems((head_dim * n_head * t) as u32);
22577        let (hd, nh, ti) = (head_dim as i32, n_head as i32, t as i32);
22578        // nullable device pointer by value (0 = skip), same convention as `l2_norm_pp`.
22579        let d16: u64 = match dst16 {
22580            Some(d) => self.addr_u8(d),
22581            None => 0,
22582        };
22583        let __s_b = self.gpu.stream();
22584        let mut b = __s_b.launch_builder(&f);
22585        b.arg(a)
22586            .arg(g)
22587            .arg(dst)
22588            .arg(&d16)
22589            .arg(&hd)
22590            .arg(&nh)
22591            .arg(&ti);
22592        unsafe {
22593            b.launch(cfg)?;
22594        }
22595        Ok(())
22596    }
22597
22598    /// step35 CLAMPED SwiGLU: `dst = min(silu(gate*gs), limit) * clamp(up*us, +-limit)`.
22599    /// Verbatim from llama.cpp `llama-graph.cpp:2146-2165` (routed, `swiglu_clamp_exp`) and
22600    /// `:1751-1770` (shared, `swiglu_clamp_shexp`), non-DEEPSEEK4 branch.
22601    ///
22602    /// This is NOT `swigluoai_mul_scaled`: that one clamps the gate BEFORE swish and multiplies by
22603    /// `(1 + clamp(up))`. Caller MUST check `limit > 1e-6` (upstream's eps gate) and use the plain
22604    /// `silu_mul_scaled` path otherwise — at limit=0 this kernel would clamp every positive
22605    /// activation to zero. On Step-3.7-Flash only layers 43 (7.0) and 44 (16.0) have a live limit.
22606    #[allow(clippy::too_many_arguments)]
22607    pub fn swiglu_clamped_mul_scaled(
22608        &self,
22609        gate: &CudaSlice<f32>,
22610        up: &CudaSlice<f32>,
22611        gs: f32,
22612        us: f32,
22613        limit: f32,
22614        dst: &mut CudaSlice<f32>,
22615        n: usize,
22616    ) -> Result<(), Box<dyn std::error::Error>> {
22617        debug_assert!(
22618            limit > 1e-6,
22619            "swiglu_clamped needs a live limit; use silu_mul_scaled"
22620        );
22621        let f = self.func("swiglu_clamped_mul_scaled_f32");
22622        let cfg = LaunchConfig::for_num_elems(n as u32);
22623        let ni = n as i32;
22624        let __s_b = self.gpu.stream();
22625        let mut b = __s_b.launch_builder(&f);
22626        b.arg(gate)
22627            .arg(up)
22628            .arg(&gs)
22629            .arg(&us)
22630            .arg(&limit)
22631            .arg(dst)
22632            .arg(&ni);
22633        unsafe {
22634            b.launch(cfg)?;
22635        }
22636        Ok(())
22637    }
22638
22639    /// gated RMSNorm: dst = RMSNorm(o, w[ncols]) * silu(z), per row of ncols. nrows blocks.
22640    pub fn gated_rmsnorm(
22641        &self,
22642        o: &CudaSlice<f32>,
22643        w: &CudaSlice<f32>,
22644        z: &CudaSlice<f32>,
22645        dst: &mut CudaSlice<f32>,
22646        ncols: usize,
22647        nrows: usize,
22648        eps: f32,
22649    ) -> Result<(), Box<dyn std::error::Error>> {
22650        let f = self.func("gated_rmsnorm_f32");
22651        let cfg = LaunchConfig {
22652            grid_dim: (nrows as u32, 1, 1),
22653            block_dim: (128, 1, 1),
22654            shared_mem_bytes: 0,
22655        };
22656        let (nc, e) = (ncols as i32, eps);
22657        let __s_b = self.gpu.stream();
22658        let mut b = __s_b.launch_builder(&f);
22659        b.arg(o).arg(w).arg(z).arg(dst).arg(&nc).arg(&e);
22660        unsafe {
22661            b.launch(cfg)?;
22662        }
22663        Ok(())
22664    }
22665
22666    /// f16out twin of `gated_rmsnorm` (task #17): epilogue also emits the fp16 operand for
22667    /// the ssm_out GEMM. Bit-identical class (same floats + the cvt kernel's __float2half).
22668    pub fn gated_rmsnorm_f16out(
22669        &self,
22670        o: &CudaSlice<f32>,
22671        w: &CudaSlice<f32>,
22672        z: &CudaSlice<f32>,
22673        dst: &mut CudaSlice<f32>,
22674        dst16: &mut CudaSlice<u8>,
22675        ncols: usize,
22676        nrows: usize,
22677        eps: f32,
22678    ) -> Result<(), Box<dyn std::error::Error>> {
22679        let f = self.func("gated_rmsnorm_f16out_f32");
22680        // block_dim MUST match gated_rmsnorm's (128): the reduction tree order pins the scale
22681        let cfg = LaunchConfig {
22682            grid_dim: (nrows as u32, 1, 1),
22683            block_dim: (128, 1, 1),
22684            shared_mem_bytes: 0,
22685        };
22686        let (nc, e) = (ncols as i32, eps);
22687        let __s_b = self.gpu.stream();
22688        let mut b = __s_b.launch_builder(&f);
22689        b.arg(o).arg(w).arg(z).arg(dst).arg(dst16).arg(&nc).arg(&e);
22690        unsafe {
22691            b.launch(cfg)?;
22692        }
22693        Ok(())
22694    }
22695
22696    /// add+RMSNorm emitting the f32 normed row AND its q8_1 quantization in one launch (the MoE
22697    /// layer input: z feeds the router matmul as f32, the expert dp4a as q8_1). BIT-IDENTICAL to
22698    /// add_rms_norm + quantize_q8_1. Returns (q, d) alongside the caller-provided res/z buffers.
22699    #[allow(clippy::too_many_arguments)]
22700    pub fn add_rms_norm_zq8(
22701        &self,
22702        a: &CudaSlice<f32>,
22703        b_in: &CudaSlice<f32>,
22704        w: &CudaSlice<f32>,
22705        res: &mut CudaSlice<f32>,
22706        z: &mut CudaSlice<f32>,
22707        ncols: usize,
22708        nrows: usize,
22709        eps: f32,
22710    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
22711        assert!(ncols % 32 == 0);
22712        let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
22713        let mut d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
22714        let f = self.func("add_rms_norm_zq8");
22715        let cfg = LaunchConfig {
22716            grid_dim: (nrows as u32, 1, 1),
22717            block_dim: (1024, 1, 1),
22718            shared_mem_bytes: 0,
22719        };
22720        let (nc, ep) = (ncols as i32, eps);
22721        let __s_b = self.gpu.stream();
22722        let mut b = __s_b.launch_builder(&f);
22723        b.arg(a)
22724            .arg(b_in)
22725            .arg(w)
22726            .arg(res)
22727            .arg(z)
22728            .arg(&mut q)
22729            .arg(&mut d)
22730            .arg(&nc)
22731            .arg(&ep);
22732        unsafe {
22733            b.launch(cfg)?;
22734        }
22735        Ok((q, d))
22736    }
22737
22738    /// gated RMSNorm emitting q8_1 directly (fused quantize epilogue) — the ssm_out matvec input.
22739    /// BIT-IDENTICAL bytes to gated_rmsnorm + quantize_q8_1 (ncols % 32 == 0; blocks never straddle
22740    /// rows). Saves one launch per linear-attn layer (36/token on the 9B).
22741    /// z-view twins of gated_rmsnorm(+f16out) — task #16 batched-prime split removal.
22742    pub fn gated_rmsnorm_zv(
22743        &self,
22744        o: &CudaSlice<f32>,
22745        w: &CudaSlice<f32>,
22746        z: &cudarc::driver::CudaView<f32>,
22747        dst: &mut CudaSlice<f32>,
22748        ncols: usize,
22749        nrows: usize,
22750        eps: f32,
22751    ) -> Result<(), Box<dyn std::error::Error>> {
22752        let f = self.func("gated_rmsnorm_f32");
22753        let cfg = LaunchConfig {
22754            grid_dim: (nrows as u32, 1, 1),
22755            block_dim: (128, 1, 1),
22756            shared_mem_bytes: 0,
22757        };
22758        let (nc, e) = (ncols as i32, eps);
22759        let __s_b = self.gpu.stream();
22760        let mut b = __s_b.launch_builder(&f);
22761        b.arg(o).arg(w).arg(z).arg(dst).arg(&nc).arg(&e);
22762        unsafe {
22763            b.launch(cfg)?;
22764        }
22765        Ok(())
22766    }
22767
22768    pub fn gated_rmsnorm_f16out_zv(
22769        &self,
22770        o: &CudaSlice<f32>,
22771        w: &CudaSlice<f32>,
22772        z: &cudarc::driver::CudaView<f32>,
22773        dst: &mut CudaSlice<f32>,
22774        dst16: &mut CudaSlice<u8>,
22775        ncols: usize,
22776        nrows: usize,
22777        eps: f32,
22778    ) -> Result<(), Box<dyn std::error::Error>> {
22779        let f = self.func("gated_rmsnorm_f16out_f32");
22780        // block_dim MUST match gated_rmsnorm's (128): the reduction tree order pins the scale
22781        let cfg = LaunchConfig {
22782            grid_dim: (nrows as u32, 1, 1),
22783            block_dim: (128, 1, 1),
22784            shared_mem_bytes: 0,
22785        };
22786        let (nc, e) = (ncols as i32, eps);
22787        let __s_b = self.gpu.stream();
22788        let mut b = __s_b.launch_builder(&f);
22789        b.arg(o).arg(w).arg(z).arg(dst).arg(dst16).arg(&nc).arg(&e);
22790        unsafe {
22791            b.launch(cfg)?;
22792        }
22793        Ok(())
22794    }
22795
22796    pub fn gated_rmsnorm_q8_1(
22797        &self,
22798        o: &CudaSlice<f32>,
22799        w: &CudaSlice<f32>,
22800        z: &CudaSlice<f32>,
22801        ncols: usize,
22802        nrows: usize,
22803        eps: f32,
22804    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
22805        assert!(ncols % 32 == 0);
22806        let f = self.func("gated_rmsnorm_q8_1");
22807        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
22808        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
22809        let cfg = LaunchConfig {
22810            grid_dim: (nrows as u32, 1, 1),
22811            block_dim: (128, 1, 1),
22812            shared_mem_bytes: 0,
22813        };
22814        let (nc, ep) = (ncols as i32, eps);
22815        let __s_b = self.gpu.stream();
22816        let mut b = __s_b.launch_builder(&f);
22817        b.arg(o)
22818            .arg(w)
22819            .arg(z)
22820            .arg(&mut out_q)
22821            .arg(&mut out_d)
22822            .arg(&nc)
22823            .arg(&ep);
22824        unsafe {
22825            b.launch(cfg)?;
22826        }
22827        Ok((out_q, out_d))
22828    }
22829
22830    /// transpose [rows,cols] row-major -> [cols,rows] row-major.
22831    pub fn transpose(
22832        &self,
22833        inp: &CudaSlice<f32>,
22834        rows: usize,
22835        cols: usize,
22836    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22837        let f = self.func("transpose_f32");
22838        let mut out = self.zeros(rows * cols)?;
22839        let cfg = LaunchConfig::for_num_elems((rows * cols) as u32);
22840        let (r, c) = (rows as i32, cols as i32);
22841        let __s_b = self.gpu.stream();
22842        let mut b = __s_b.launch_builder(&f);
22843        b.arg(inp).arg(&mut out).arg(&r).arg(&c);
22844        unsafe {
22845            b.launch(cfg)?;
22846        }
22847        Ok(out)
22848    }
22849
22850    /// repeat-interleave heads: in[head_dim,n_in,T] -> out[head_dim,n_out,T].
22851    pub fn repeat_heads(
22852        &self,
22853        inp: &CudaSlice<f32>,
22854        out: &mut CudaSlice<f32>,
22855        head_dim: usize,
22856        n_in: usize,
22857        n_out: usize,
22858        t: usize,
22859    ) -> Result<(), Box<dyn std::error::Error>> {
22860        let f = self.func("repeat_heads_f32");
22861        let cfg = LaunchConfig::for_num_elems((head_dim * n_out * t) as u32);
22862        let (hd, ni, no, ti) = (head_dim as i32, n_in as i32, n_out as i32, t as i32);
22863        let __s_b = self.gpu.stream();
22864        let mut b = __s_b.launch_builder(&f);
22865        b.arg(inp).arg(out).arg(&hd).arg(&ni).arg(&no).arg(&ti);
22866        unsafe {
22867            b.launch(cfg)?;
22868        }
22869        Ok(())
22870    }
22871
22872    /// q|gate split (on-device). qf:[T, n_head*2*head_dim] -> q_out,gate_out:[head_dim,n_head,T].
22873    /// Replaces the dtoh->host-double-loop->htod in full_attn / full_attn_decode.
22874    ///
22875    /// Bounds-checked against `qf` before the launch: the kernel reads `2*head_dim*n_head*t`
22876    /// floats, and running it on a `wq` output that carries no fused gate reads 2x off the end
22877    /// (silently, on the device). A layout mismatch is a typed `FusedQGateExtent` here instead.
22878    pub fn q_gate_split(
22879        &self,
22880        qf: &CudaSlice<f32>,
22881        q_out: &mut CudaSlice<f32>,
22882        gate_out: &mut CudaSlice<f32>,
22883        head_dim: usize,
22884        n_head: usize,
22885        t: usize,
22886    ) -> Result<(), Box<dyn std::error::Error>> {
22887        memra_gguf::config::check_fused_q_gate_extent(qf.len(), head_dim, n_head, t)?;
22888        let out_need = head_dim * n_head * t;
22889        if q_out.len() < out_need || gate_out.len() < out_need {
22890            return Err(format!(
22891                "q_gate_split destinations too small: need {out_need} each, have q={} gate={}",
22892                q_out.len(),
22893                gate_out.len()
22894            )
22895            .into());
22896        }
22897        let f = self.func("q_gate_split_f32");
22898        let cfg = LaunchConfig::for_num_elems((head_dim * n_head * t) as u32);
22899        let (hd, nh, ti) = (head_dim as i32, n_head as i32, t as i32);
22900        let __s_b = self.gpu.stream();
22901        let mut b = __s_b.launch_builder(&f);
22902        b.arg(qf)
22903            .arg(q_out)
22904            .arg(gate_out)
22905            .arg(&hd)
22906            .arg(&nh)
22907            .arg(&ti);
22908        unsafe {
22909            b.launch(cfg)?;
22910        }
22911        Ok(())
22912    }
22913
22914    /// qkv->GDN repack (on-device). conv_out:[conv_dim,T] channel-major ->
22915    /// q_g/k_g/v_g:[d_state,num_v,T] with q/k head-repeat kh = vh % num_k (validated modulo mapping).
22916    /// Replaces the dtoh->host-q/k/v-repack->3x-htod in linear_attn / linear_attn_decode.
22917    pub fn qkv_to_gdn_repack(
22918        &self,
22919        conv_out: &CudaSlice<f32>,
22920        q_g: &mut CudaSlice<f32>,
22921        k_g: &mut CudaSlice<f32>,
22922        v_g: &mut CudaSlice<f32>,
22923        d_state: usize,
22924        num_v: usize,
22925        num_k: usize,
22926        key_dim: usize,
22927        t: usize,
22928    ) -> Result<(), Box<dyn std::error::Error>> {
22929        let f = self.func("qkv_to_gdn_repack_f32");
22930        let cfg = LaunchConfig::for_num_elems((d_state * num_v * t) as u32);
22931        let (ds, nv, nk, kd, ti) = (
22932            d_state as i32,
22933            num_v as i32,
22934            num_k as i32,
22935            key_dim as i32,
22936            t as i32,
22937        );
22938        let __s_b = self.gpu.stream();
22939        let mut b = __s_b.launch_builder(&f);
22940        b.arg(conv_out)
22941            .arg(q_g)
22942            .arg(k_g)
22943            .arg(v_g)
22944            .arg(&ds)
22945            .arg(&nv)
22946            .arg(&nk)
22947            .arg(&kd)
22948            .arg(&ti);
22949        unsafe {
22950            b.launch(cfg)?;
22951        }
22952        Ok(())
22953    }
22954
22955    /// conv left zero-pad (prefill from zero state). src:[conv_dim,T] -> dst:[conv_dim,T+pad],
22956    /// cols 0..pad = 0, cols pad..pad+T = src. `dst` MUST be pre-zeroed. No dtoh/host-loop/htod.
22957    pub fn conv_left_pad(
22958        &self,
22959        src: &CudaSlice<f32>,
22960        dst: &mut CudaSlice<f32>,
22961        conv_dim: usize,
22962        t: usize,
22963        pad: usize,
22964    ) -> Result<(), Box<dyn std::error::Error>> {
22965        let f = self.func("conv_left_pad_f32");
22966        let cfg = LaunchConfig::for_num_elems((conv_dim * t) as u32);
22967        let (cd, ti, p) = (conv_dim as i32, t as i32, pad as i32);
22968        let __s_b = self.gpu.stream();
22969        let mut b = __s_b.launch_builder(&f);
22970        b.arg(src).arg(dst).arg(&cd).arg(&ti).arg(&p);
22971        unsafe {
22972            b.launch(cfg)?;
22973        }
22974        Ok(())
22975    }
22976
22977    /// conv-state assemble + ring roll (decode T=1). conv_state:[conv_dim,pad] (resident),
22978    /// qkv_col:[conv_dim] -> conv_in:[conv_dim,pad+1]; AND rolls conv_state (keep last pad cols).
22979    /// Replaces the dtoh->host-conv-ring-assemble->ring-update->htod in linear_attn_decode.
22980    pub fn conv_assemble_and_roll(
22981        &self,
22982        qkv_col: &CudaSlice<f32>,
22983        conv_state: &mut CudaSlice<f32>,
22984        conv_in: &mut CudaSlice<f32>,
22985        conv_dim: usize,
22986        pad: usize,
22987    ) -> Result<(), Box<dyn std::error::Error>> {
22988        let f = self.func("conv_assemble_and_roll_f32");
22989        let cfg = LaunchConfig::for_num_elems(conv_dim as u32);
22990        let (cd, p) = (conv_dim as i32, pad as i32);
22991        let __s_b = self.gpu.stream();
22992        let mut b = __s_b.launch_builder(&f);
22993        b.arg(qkv_col).arg(conv_state).arg(conv_in).arg(&cd).arg(&p);
22994        unsafe {
22995            b.launch(cfg)?;
22996        }
22997        Ok(())
22998    }
22999
23000    /// RANK3 LEVER (conv fuse, T=1 DECODE): fused conv_assemble_and_roll + ssm_conv1d_silu in ONE
23001    /// launch. Assembles the conv window [conv_state | qkv_col] in registers, computes the depthwise
23002    /// causal conv + SiLU into `conv_out`, and rolls the ring — never materializing conv_in to HBM.
23003    /// Replaces e.conv_assemble_and_roll(...) + e.ssm_conv1d(...). BIT-IDENTICAL to that two-kernel
23004    /// sequence (same 8-wide accumulation order, same SiLU). `conv_out` is [conv_dim] (T=1).
23005    pub fn ssm_conv1d_fused_decode(
23006        &self,
23007        qkv_col: &CudaSlice<f32>,
23008        conv_state: &mut CudaSlice<f32>,
23009        w: &CudaSlice<f32>,
23010        conv_out: &mut CudaSlice<f32>,
23011        conv_dim: usize,
23012        d_conv: usize,
23013    ) -> Result<(), Box<dyn std::error::Error>> {
23014        let f = self.func("ssm_conv1d_fused_decode_f32");
23015        let cfg = LaunchConfig::for_num_elems(conv_dim as u32);
23016        let (cd, dc) = (conv_dim as i32, d_conv as i32);
23017        let __s_b = self.gpu.stream();
23018        let mut b = __s_b.launch_builder(&f);
23019        b.arg(qkv_col)
23020            .arg(conv_state)
23021            .arg(w)
23022            .arg(conv_out)
23023            .arg(&cd)
23024            .arg(&dc);
23025        unsafe {
23026            b.launch(cfg)?;
23027        }
23028        Ok(())
23029    }
23030
23031    /// Copy a contiguous range [start, start+len) out of src into a fresh slice (device→device via host).
23032    /// Used for qkv split views. Small/rare; not perf-critical in Stage 1.
23033    pub fn slice_range(
23034        &self,
23035        src: &CudaSlice<f32>,
23036        start: usize,
23037        len: usize,
23038    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
23039        let host = self.gpu.stream().clone_dtoh(src)?;
23040        self.gpu.stream().synchronize()?;
23041        Ok(self.htod(&host[start..start + len])?)
23042    }
23043}
23044
23045#[cfg(test)]
23046mod target_dispatch_tests {
23047    use super::legacy_quant_gemm_allowed;
23048
23049    #[test]
23050    fn legacy_quant_gemm_arch_policy_honors_the_escape_hatch() {
23051        // sm_120a native lane
23052        assert!(legacy_quant_gemm_allowed(false, false, false));
23053        assert!(!legacy_quant_gemm_allowed(false, false, true));
23054        // pure portable lane (sm_89): gated
23055        assert!(!legacy_quant_gemm_allowed(true, false, false));
23056        assert!(!legacy_quant_gemm_allowed(true, false, true));
23057        // Hopper-MMA lane (sm_90a): portable build, int8-MMA GEMM re-admitted
23058        assert!(legacy_quant_gemm_allowed(true, true, false));
23059        assert!(!legacy_quant_gemm_allowed(true, true, true));
23060    }
23061
23062    #[cfg(all(memra_portable_cuda, not(memra_hopper_mma)))]
23063    #[test]
23064    fn portable_build_disables_legacy_quant_gemm_without_an_env_override() {
23065        assert!(!legacy_quant_gemm_allowed(
23066            cfg!(memra_portable_cuda),
23067            cfg!(memra_hopper_mma),
23068            false
23069        ));
23070    }
23071
23072    #[cfg(memra_hopper_mma)]
23073    #[test]
23074    fn hopper_mma_build_re_admits_legacy_quant_gemm() {
23075        assert!(legacy_quant_gemm_allowed(
23076            cfg!(memra_portable_cuda),
23077            cfg!(memra_hopper_mma),
23078            false
23079        ));
23080        assert!(super::portable_mma_gated() == false);
23081    }
23082}
23083
23084/// The memra-kv device seam (Phase D): the cache's 7 ops delegate to the engine's
23085/// inherent methods (inherent methods win name resolution, so no recursion).
23086impl memra_kv::KvDev for Engine {
23087    fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
23088        Engine::zeros(self, n)
23089    }
23090    fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
23091        Engine::uninit(self, n)
23092    }
23093    fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
23094        Engine::alloc_u8(self, n)
23095    }
23096    fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
23097        Engine::htod_i32(self, v)
23098    }
23099    fn clone_dtod(
23100        &self,
23101        src: &CudaSlice<f32>,
23102    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
23103        Engine::clone_dtod(self, src)
23104    }
23105    fn copy_into(
23106        &self,
23107        dst: &mut CudaSlice<f32>,
23108        off: usize,
23109        src: &CudaSlice<f32>,
23110        len: usize,
23111    ) -> Result<(), Box<dyn std::error::Error>> {
23112        Engine::copy_into(self, dst, off, src, len)
23113    }
23114    fn set_i32_one(
23115        &self,
23116        d: &mut CudaSlice<i32>,
23117        v: i32,
23118    ) -> Result<(), Box<dyn std::error::Error>> {
23119        Engine::set_i32_one(self, d, v)
23120    }
23121}
23122
23123#[cfg(test)]
23124mod fused_gate_bounds_tests {
23125    use super::*;
23126
23127    /// The fused `[q|gate]` split's read-site guard, on the device.
23128    ///
23129    /// `q_gate_split_f32` reads `2*head_dim*n_head*T` floats out of `qf`. A checkpoint whose gate
23130    /// is a SEPARATE tensor produces a `wq` output of exactly half that, so before 2026-08-19 the
23131    /// kernel launched and read 2x past the end of the allocation — an out-of-bounds DEVICE read:
23132    /// no panic, no error, just whatever memory follows. The guard turns it into a typed
23133    /// `FusedQGateExtent` before the launch.
23134    ///
23135    /// Catch demonstration for this test (guard temporarily removed, then restored):
23136    /// `compute-sanitizer --tool memcheck` on the half-width case reported invalid `__global__`
23137    /// reads of size 4 in `q_gate_split_f32`; with the guard in place the same run is clean and
23138    /// the call returns `Err`. Receipt in the lane report.
23139    #[test]
23140    #[ignore = "requires a CUDA GPU"]
23141    fn q_gate_split_refuses_a_separate_gate_wq_instead_of_reading_past_it() {
23142        let e = Engine::new(0).unwrap();
23143        let (head_dim, n_head, t) = (8usize, 4usize, 2usize);
23144        let fused = 2 * head_dim * n_head * t;
23145        let out_n = head_dim * n_head * t;
23146
23147        // half-width `qf` = the separate-gate / ungated layout. MUST be refused.
23148        let narrow = e.htod(&vec![1.0f32; out_n]).unwrap();
23149        let mut q = e.uninit(out_n).unwrap();
23150        let mut gate = e.uninit(out_n).unwrap();
23151        let err = e
23152            .q_gate_split(&narrow, &mut q, &mut gate, head_dim, n_head, t)
23153            .expect_err("half-width wq must be refused, not read past")
23154            .to_string();
23155        assert!(err.contains("NO fused gate"), "{err}");
23156        assert!(err.contains(&format!("{fused}")), "{err}");
23157
23158        // full-width `qf` = a real qwen3.5 fused layout. MUST still run, and split correctly:
23159        // per head hh the block is [q(head_dim) | gate(head_dim)] at stride 2*head_dim.
23160        let host: Vec<f32> = (0..fused).map(|i| i as f32).collect();
23161        let wide = e.htod(&host).unwrap();
23162        e.q_gate_split(&wide, &mut q, &mut gate, head_dim, n_head, t)
23163            .expect("full-width wq splits");
23164        let (qh, gh) = (e.dtoh(&q).unwrap(), e.dtoh(&gate).unwrap());
23165        for tok in 0..t {
23166            for hh in 0..n_head {
23167                for d in 0..head_dim {
23168                    let base = tok * (n_head * 2 * head_dim) + hh * (2 * head_dim);
23169                    let idx = tok * (n_head * head_dim) + hh * head_dim + d;
23170                    assert_eq!(qh[idx], host[base + d], "q t{tok} h{hh} d{d}");
23171                    assert_eq!(gh[idx], host[base + head_dim + d], "gate t{tok} h{hh} d{d}");
23172                }
23173            }
23174        }
23175
23176        // undersized destinations are refused too (the other half of the extent contract)
23177        let mut small = e.uninit(out_n - 1).unwrap();
23178        assert!(
23179            e.q_gate_split(&wide, &mut small, &mut gate, head_dim, n_head, t)
23180                .is_err()
23181        );
23182    }
23183}
23184
23185/// FULL-WIDTH-ROPE CONTRACT on the fused rms_norm+qkv+rope kernels
23186/// (lane/graph-s-key-exactness-20260819, probe O-10 follow-up). CPU-only: the guard runs before
23187/// any launch, so the refusal is testable without a device.
23188#[cfg(test)]
23189mod fused_rope_width_tests {
23190    use super::Engine;
23191
23192    /// gemma-4: rotary width == head width on both classes (GGUF 256/256 and 512/512, and the
23193    /// safetensors route derives the same), which is why the fusion is legal there today.
23194    #[test]
23195    fn full_width_is_accepted() {
23196        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 256, 256).is_ok());
23197        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_cat", 512, 512).is_ok());
23198        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append", 128, 128).is_ok());
23199    }
23200
23201    /// The widths the gemma-4 31B OFFICIAL artifact declares, read from its own GGUF header
23202    /// (`gemma-4-31B-it-official-Q8_0-MTP.gguf`, box3, 2026-08-19):
23203    ///
23204    /// ```text
23205    /// attention.key_length     512   rope.dimension_count     512   (global class)
23206    /// attention.key_length_swa 256   rope.dimension_count_swa 256   (SWA class)
23207    /// ```
23208    ///
23209    /// Both classes satisfy `n_rot == head_dim`, which is why the fusion is legal for gemma and
23210    /// why `HybridModel::gemma4_rope_dims` can feed this guard without refusing what we serve.
23211    /// An artifact that ever declares otherwise gets a loud refusal at the first fused launch
23212    /// instead of a silently over-rotated head.
23213    #[test]
23214    fn gemma4_official_artifact_widths_pass() {
23215        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 512, 512).is_ok());
23216        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append_dc", 256, 256).is_ok());
23217    }
23218
23219    /// The bug this guard exists to make impossible: a partial-rotary arch fused onto a kernel
23220    /// with no `n_dims`, silently rotating the pass-through band.
23221    #[test]
23222    fn partial_rotary_is_refused_with_the_geometry_named() {
23223        // qwen3.5: n_rot 64 of head_dim 256 (the shape probe O-10 pinned in the split path).
23224        let err = Engine::full_width_rope_only("rms_norm_qkv_rope", 64, 256)
23225            .expect_err("partial rotary must refuse");
23226        let msg = err.to_string();
23227        assert!(msg.contains("PARTIAL ROTARY REFUSED"), "{msg}");
23228        assert!(msg.contains("n_rot 64"), "{msg}");
23229        assert!(msg.contains("head_dim 256"), "{msg}");
23230        assert!(
23231            msg.contains("64..256"),
23232            "names the band it would corrupt: {msg}"
23233        );
23234        // step35 full-attn: 64 of 128 (upstream halves n_rot_full).
23235        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append_dc", 64, 128).is_err());
23236        // and the reverse mismatch (a wider rope than the head) is not "close enough" either.
23237        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 256, 128).is_err());
23238    }
23239}