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::sys::CUfunction_attribute_enum::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES;
4use cudarc::driver::{
5    CudaContext, CudaFunction, CudaModule, CudaSlice, CudaStream, DeviceSlice, LaunchConfig,
6    PushKernelArg,
7};
8use cudarc::nvrtc::Ptx;
9use std::sync::{Arc, Mutex};
10
11const GDN_K2_DYNAMIC_SHARED_BYTES: u32 = 67_072;
12
13#[cfg(debug_assertions)]
14pub(crate) fn debug_assert_tensor_stream_device<T>(
15    tensor: &CudaSlice<T>,
16    stream: &CudaStream,
17    site: &str,
18) {
19    let tensor_dev = tensor.ordinal();
20    let stream_dev = stream.context().ordinal();
21    assert_eq!(
22        tensor_dev, stream_dev,
23        "PP cross-device tensor read at {site}: tensor on dev{tensor_dev}, stream on dev{stream_dev}"
24    );
25}
26
27fn ensure_tensor_stream_device<T>(
28    tensor: &impl DeviceSlice<T>,
29    stream: &CudaStream,
30    site: &str,
31) -> Result<(), Box<dyn std::error::Error>> {
32    let tensor_dev = tensor.stream().context().ordinal();
33    let stream_dev = stream.context().ordinal();
34    if tensor_dev != stream_dev {
35        return Err(format!(
36            "PP cross-device tensor access at {site}: tensor on dev{tensor_dev}, \
37             stream on dev{stream_dev}"
38        )
39        .into());
40    }
41    Ok(())
42}
43
44pub use memra_gguf;
45pub use memra_runtime;
46
47pub mod forward;
48pub mod hybrid;
49pub mod hybrid_forward;
50pub mod model;
51pub mod sigrouter_contract;
52pub mod vision;
53pub mod vision_gemma;
54pub mod vision_pre;
55/// The dual cache lives in the shared `memra-kv` crate (Phase D extraction); this
56/// re-export keeps every `crate::cache::` / `memra_engine::cache::` path unchanged.
57pub mod cache {
58    pub use memra_kv::*;
59}
60pub mod decode;
61pub mod decode_batch;
62pub mod dflash;
63pub mod eagle;
64pub mod gemma_spec;
65pub mod graph_update;
66/// MLA (multi-head latent attention) CPU f32 reference — GLM-5.2 bring-up lane increment 1.
67/// Naive vs absorbed decode forms + NORM/NEOX rope permutation, unit-tested; the permanent
68/// oracle for the MLA kernel family (`research/mla-bringup-20260801/DESIGN.md`). No CUDA deps.
69pub mod mla;
70pub mod moesd;
71pub mod parallel;
72pub mod plan_backend;
73pub mod pp;
74pub mod round_stream;
75pub mod spec;
76pub mod tp;
77pub use memra_sampling as sampler;
78
79/// In-house MoE router GEMV on the spec-verify small-t path (DEFAULT ON since 2026-07-10:
80/// battery green on 35B p2/p3 K=1..8, acceptance bit-identical, +2-4% spec e2e — replaces
81/// ~240 per-column cuBLAS gemv launches/round). MEMRA_ROUTER_KERNEL=0 is the rollback seam.
82/// MoE grouped f16 GEMM door (experimental until gated), f16-mirror numeric class:
83/// per-layer expert dequant to f16 + one grouped f16 GEMM over the CSR groups.
84///   MEMRA_MOE_F16G=1  cublasGemmGroupedBatchedEx (round 46 arc 2). The grouped API issues
85///                     through cublas-internal streams NOT ordered with ours — v1 pays a full
86///                     stream sync per projection (round-47 ledgered defect).
87///   MEMRA_MOE_F16G=2  single-kernel grouped GEMM on the engine stream (round 49): ordered by
88///                     construction, zero syncs, f32 C with the act row-scale folded in.
89/// DEFAULT (2026-08-01, round 49 promotion): mode 1 on the Hopper lane — with the 41/41
90/// dequant coverage fix the q35 board-2048 prime measured 5490 (MMQ) / 8380 (mode 1,
91/// +53%) / 7990 (mode 2) x3 interleaved on the H100, argmax MATCH — the last board loss
92/// flips. The 5090 measured FLAT (858GB/s makes the dequant-workspace traffic cancel the
93/// GEMM win) — but that verdict is for expert banks the int8-MMA MMQ arm can take
94/// (IQ3_S/IQ4_XS/Q4_0). MEMRA_MOE_F16G=0 kills anywhere.
95///
96/// HOPPER RE-VERDICT (2026-08-02, lane/h100-flip-full): mode 2 with full direct coverage
97/// (Q4_K/Q6_K/IQ4_XS/IQ3_S tile loaders, lane/iq-direct-loaders) + the deep tail
98/// (lane/sk-tail-form) FLIPS past cublas mode 1 on the H100 — q35 board-2048 prime
99/// 13163.6 (mode 2, cross=32) vs 8626.5 (mode 1) vs 8073.4 (round-51 sk form), +52.6%,
100/// interleaved x5 zero overlap, argmax MATCH 30/30. The round-54 NO-FLIP (8547 vs 8112)
101/// was coverage-priced at 5.2% direct; ~100% coverage kills the workspace pass and the
102/// verdict inverts. Hopper naked default -> mode 2 (this arm); the gemma (gelu) site
103/// stays env-explicit-only via moe_f16g_gemma_on (Err => closed, unaffected by this arm).
104///
105/// MODE-2 DEFAULT (sm_120a naked, 2026-08-02, lane/f16g-default-rearb): with the direct
106/// tile loaders covering Q4_K/Q6_K/IQ4_XS/IQ3_S, the sk visitor beats the int8-MMA MMQ
107/// tiles on the IQ-bank models too (q35 board-2048 +33.9%, KAT pp512 +46.7% / pp2048
108/// +30.6% — research/iq-direct-loaders-20260802 §3-5, confirmed + full battery in
109/// research/f16g-default-rearb-20260802/), so every f16g-admitted expert layer rides
110/// mode 2 naked. Decode/verify stay on dp4a (t >= 16 floor). f16-mirror numeric class
111/// for naked q35/KAT prefill+prime — new token-sha anchors stamped in the rearb lane.
112///
113/// AUTO-KQUANT (mode 3, 2026-08-02, lane/q4k-expert-prefill): the previous sm_120a
114/// default, kept reachable via MEMRA_MOE_F16G=3. The mode-2 sk form is admitted ONLY for
115/// layers the MMA MMQ arm rejects (k-quant expert projections — Q3_K/Q4_K/Q6_K), i.e.
116/// exactly where the baseline is the per-pair moe_pairs_matvec_q8_em fallback with zero
117/// token reuse (Ornith-35B Q4_K_M board-2048 1098.2 -> 3453.7, 3.14x,
118/// research/q4k-expert-prefill-20260802/). Its "IQ banks keep their measured-faster MMQ
119/// tiles" ruling was priced BEFORE the IQ direct loaders and is refuted on the 5090 —
120/// the k-quant-only admission survives as the rollback seam, not the default.
121/// The gemma (gelu) site stays env-explicit-only (moe_f16g_gemma_on).
122pub fn moe_f16g_mode() -> u8 {
123    static M: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
124    *M.get_or_init(|| match std::env::var("MEMRA_MOE_F16G").as_deref() {
125        Ok("0") => 0,
126        Ok("2") => 2,
127        Ok("3") => 3,
128        Ok(_) => 1,
129        // Both arches independently re-arbitrated to mode 2 on 2026-08-02
130        // (5090: lane/f16g-default-rearb; H100: lane/h100-flip-full) — unset = 2 everywhere.
131        Err(_) => 2,
132    })
133}
134/// Mode-2 sk kernel form policy (round 51, lane/sk-bm128): the single-kernel grouped GEMM runs
135/// as a persistent problem-visitor over the real CSR tiles with two tile forms. Returns
136/// (shape_sel, cross) for the FFI:
137///   MEMRA_F16G_SK=0    -> (-1, _): the round-49 grid-scan kernel (rollback seam).
138///   MEMRA_F16G_SK=32   -> all groups on the 32x64x32 2-stage form (cross = i32::MAX).
139///   MEMRA_F16G_SK=128  -> all groups on the 128x64x64 3-stage form (cross = 1; groups fall
140///                         back to 32x64 in-launcher when the device/in_f can't take it).
141///   unset              -> hybrid split: groups with m_e >= MEMRA_F16G_SK_CROSS ride the 128
142///                         form. Default cross = 64 (5090 sweep 2026-08-01, receipts
143///                         research/sk-bm128-20260801/; H100 re-swept on the direct+tail
144///                         form 2026-08-02, lane/h100-flip-full: {16,32,64} ->
145///                         12868/13192/13225 — 64 wins there too, the pre-direct 32
146///                         verdict was stale).
147pub fn moe_f16g_sk_params() -> (i32, i32) {
148    static P: std::sync::OnceLock<(i32, i32)> = std::sync::OnceLock::new();
149    *P.get_or_init(|| match std::env::var("MEMRA_F16G_SK").as_deref() {
150        Ok("0") => (-1, 0),
151        Ok("32") => (0, i32::MAX),
152        Ok("128") => (0, 1),
153        _ => {
154            let cross = std::env::var("MEMRA_F16G_SK_CROSS")
155                .ok()
156                .and_then(|v| v.parse().ok())
157                .unwrap_or(64);
158            (0, cross)
159        }
160    })
161}
162/// DIRECT-FROM-QUANT sk tile loaders (lane/kquant-tile-loaders, 2026-08-02; IQ classes added
163/// by lane/iq-direct-loaders): Q4_K/Q6_K/IQ4_XS/IQ3_S expert projections on the mode-2/3 sk
164/// visitor forms dequant their weight tiles in-register from the quant superblocks instead of
165/// running the per-(layer,projection) dequant pass into an f16 workspace (41.8% of Ornith-35B
166/// t=512 kernel time — the pp512 wall, research/q4k-expert-prefill-20260802 §5; the IQ classes
167/// are 94.8% of q35's bank bytes — the h100-sk-direct coverage pricing). Bit-identical to the
168/// workspace path by construction (kernel-check "f16g-kq-direct" gates it bitwise) — a
169/// data-movement change, not a numeric-class change. Default ON; MEMRA_F16G_DIRECT=0 reverts
170/// to the workspace path everywhere; MEMRA_F16G_DIRECT=kq keeps the k-quant loaders and
171/// reverts only the IQ classes (the iq-direct-loaders A/B seam — the pre-lane shipped config).
172pub fn moe_f16g_direct_on(qtype: i32) -> bool {
173    static M: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
174    let m = *M.get_or_init(|| match std::env::var("MEMRA_F16G_DIRECT").as_deref() {
175        Ok("0") => 0,
176        Ok("kq") => 1,
177        _ => 2,
178    });
179    match m {
180        0 => false,
181        1 => qtype == QT_Q4_K || qtype == QT_Q6_K,
182        _ => true,
183    }
184}
185/// DEEP-TAIL sk form (lane/sk-tail-form, 2026-08-02): groups below the visitor crossover ride
186/// a 32x64x64 3-STAGE cp.async tile instead of the round-51 32x64x32 2-stage — the same 32-row
187/// tile (zero extra padding), 2 k-blocks in flight instead of 1 and half the syncs per k. The
188/// H100 ncu pricing (research/sk-bm128-20260801) put the 2-stage tail at 31% of the sk GEMM
189/// stage under q35's routing skew. Bit-identical to every other sk form by construction
190/// (kernel-check "f16g-sk" gates all tail arms maxdiff==0); exists in both the workspace-f16
191/// and direct-from-quant variants. Default ON; MEMRA_F16G_TAIL=0 = rollback to the 2-stage
192/// tail. in_f % 64 != 0 falls back in-launcher.
193pub fn moe_f16g_tail_on() -> bool {
194    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
195    *ON.get_or_init(|| std::env::var("MEMRA_F16G_TAIL").as_deref() != Ok("0"))
196}
197
198/// Per-model door for the gemma-MoE (gelu) grouped path: round 49's Hopper default
199/// REGRESSED g26 board-2048 prefill -8.3% interleaved x5 on-box (def median 10380,
200/// wild 8.9k-11.7k spread; off 11317, ±0.13%) — the +6-15% probe verdict didn't
201/// survive the board workload (stale-verdict law, round 50). The silu/qwen class
202/// keeps the round-49 default (q35 +53% board-2048). Explicit MEMRA_MOE_F16G=1/2
203/// still opens this door for A/B.
204pub fn moe_f16g_gemma_on() -> bool {
205    static M: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
206    *M.get_or_init(|| !matches!(std::env::var("MEMRA_MOE_F16G").as_deref(), Ok("0") | Err(_)))
207}
208
209/// Fused act-epilogue (silu/gelu-mul + q8_1_mmq quantize in one launch) for the MoE prefill
210/// MMA arms. Byte-identical to the two-pass path (kernel-check gated) — default ON.
211/// MEMRA_MOE_FUSE_ACTQ=0 is the rollback/A-B seam.
212pub fn moe_fuse_actq_on() -> bool {
213    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
214    *ON.get_or_init(|| std::env::var("MEMRA_MOE_FUSE_ACTQ").as_deref() != Ok("0"))
215}
216
217/// PREFILL router m-invariance (lane/concat-prime-exact, 2026-08-02). The batched cuBLASLt
218/// router GEMM changes a row's logits when OTHER rows join the call (probed: first change at
219/// m=65 on the Ornith-35B router, 3.9e-3 — while the MMQ/f16 trunk GEMMs are bit-identical
220/// across m). Feeding a top-k discontinuity, that made a served request's expert selection a
221/// function of its CO-ARRIVALS under cross-request prime batching. The in-house router GEMV
222/// is m-invariant, so prefill uses it too and routing depends on a session's own tokens only.
223/// DEFAULT ON: it is the serving isolation contract, and it is the same kernel decode and spec
224/// verify already use (dispatch parity, one router kernel for every t).
225/// MEMRA_ROUTER_PREFILL_EXACT=0 reverts to the batched GEMM.
226pub fn router_prefill_exact_on() -> bool {
227    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
228    *ON.get_or_init(|| std::env::var("MEMRA_ROUTER_PREFILL_EXACT").as_deref() != Ok("0"))
229}
230
231pub fn router_kernel_on() -> bool {
232    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
233    *ON.get_or_init(|| {
234        let on = std::env::var("MEMRA_ROUTER_KERNEL").as_deref() != Ok("0");
235        if !on {
236            eprintln!("[memra] router kernel OFF (rollback: per-column cuBLAS gemv)");
237        }
238        on
239    })
240}
241
242/// FAST-ROUTER batch twin (lane/fast-router, 2026-08-02). The concat-prime exactness fix
243/// (router_prefill_exact_on) routes prefill through router_gemv — m-invariant, but a
244/// per-(expert,token) GEMV program with zero operand reuse, so q35 board-2048 prefill paid
245/// -10% on the 5090. router_gemv_f32_w8_batch register-tiles (8x8 expert-x-token) the same
246/// per-row FP chains (BIT-IDENTICAL per row — kernel-check sweeps m=1..2048 on real router
247/// weights), so the t crossover below is pure perf, not a numeric config. Swept on-box
248/// (research/fast-router-20260802/crossover-router*.jsonl): plain wins t<=4, batch +7-9%
249/// at t=8, 1.9x at t=16 rising to 3.45x at t=2048 — MIN_T=8. Decode t=1 and spec verify
250/// t<8 keep the plain w8 form. MEMRA_ROUTER_BATCH=0 forces plain at every t (rollback
251/// seam, perf-only: bits are equal by the kernel-check gate).
252/// Killed arms (same sweep, JSONL is the record): the 8x16 tile lost to 8x8 at every t
253/// (128-accumulator register pressure beats the halved w-traffic), and the same-shape
254/// sigmoid_dot_rows twin (out_f=1) measured 0.62-0.89x at every prefill t
255/// (launch-latency-bound, ~7us/layer at m=2048) — both bit-identity-PASSED before dying.
256pub const ROUTER_BATCH_MIN_T: usize = 8;
257pub fn router_batch_on() -> bool {
258    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
259    *ON.get_or_init(|| std::env::var("MEMRA_ROUTER_BATCH").as_deref() != Ok("0"))
260}
261mod cpu_experts;
262#[cfg(memra_cutlass)]
263pub mod cutlass_ffi;
264pub mod dsv4_ffi;
265pub mod dsv4_gpu;
266pub mod f16_ffi;
267pub mod fp8_ffi;
268pub mod mmq_ffi;
269pub mod moe_cache;
270pub mod prime_graph;
271pub mod spill;
272mod spill_pread;
273
274// Fatbins are EMBEDDED (crates-release lane, 2026-08-04): build.rs still writes them to
275// OUT_DIR, but the bytes ship inside the binary via include_bytes! and load through
276// cuModuleLoadData. Distribution contract: a prebuilt or cargo-installed binary must be
277// self-contained — the old baked OUT_DIR *paths* pointed at the builder's temp dir and
278// broke every machine that wasn't the build machine. Same bytes, same module image;
279// the runtime MEMRA_GEMM_FATBIN tune-seam override below is preserved.
280const FATBIN: &[u8] = include_bytes!(env!("MEMRA_ENGINE_FATBIN"));
281const HYBRID_FATBIN: &[u8] = include_bytes!(env!("MEMRA_HYBRID_FATBIN"));
282const QMATVEC_FATBIN: &[u8] = include_bytes!(env!("MEMRA_QMATVEC_FATBIN"));
283const FLASH_FATBIN: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN"));
284const GEMM_FATBIN: &[u8] = include_bytes!(env!("MEMRA_GEMM_FATBIN"));
285const ROUTER_FATBIN: &[u8] = include_bytes!(env!("MEMRA_ROUTER_FATBIN"));
286/// spec_sample.cu: sampled-spec primitives (Philox Gumbel-max / softmax gather / residual sampler).
287const SAMPLE_FATBIN: &[u8] = include_bytes!(env!("MEMRA_SAMPLE_FATBIN"));
288
289/// TUNE SEAM (tools/sweep): a RUNTIME `MEMRA_GEMM_FATBIN=<path>` overrides the baked-in
290/// qmatvec_gemm.cu fatbin path (build.rs bakes the same name at COMPILE time via
291/// cargo:rustc-env — that constant is the default). Lets the sweep harness swap in a
292/// `-D`-tuned fatbin per process with NO rust rebuild. Unset at runtime => the
293/// compile-time default (zero behavior change).
294fn gemm_fatbin_bytes() -> std::borrow::Cow<'static, [u8]> {
295    assert!(
296        !(portable_mma_gated() && std::env::var_os("MEMRA_GEMM_FATBIN").is_some()),
297        "MEMRA_GEMM_FATBIN overrides are not allowed in the portable CUDA lane"
298    );
299    match std::env::var("MEMRA_GEMM_FATBIN") {
300        Ok(path) => std::borrow::Cow::Owned(
301            std::fs::read(&path).unwrap_or_else(|e| panic!("MEMRA_GEMM_FATBIN read {path}: {e}")),
302        ),
303        Err(_) => std::borrow::Cow::Borrowed(GEMM_FATBIN),
304    }
305}
306
307/// Phase A (ARCHITECTURE-H100.md): sm_90a re-enables the portable-PTX tensor-core paths
308/// (int8 mma.m16n8k32/k16.s8, bf16 m16n8k16, ldmatrix, cp.async — all sm_80-class, native
309/// on Hopper) that the portable boot lane gates off. Dispatch guards that used to test
310/// `cfg!(memra_portable_cuda)` test this instead; sm_89 keeps the pure-portable behavior.
311/// The sm_120a/sm_100a-only MMA kinds (mxf4nvf4, kind::f8f6f4) are NOT covered — their
312/// launchers stay fail-closed stubs on 90a and their dispatch arms stay arch-gated.
313pub(crate) const fn portable_mma_gated() -> bool {
314    cfg!(memra_portable_cuda) && !cfg!(memra_hopper_mma)
315}
316
317/// The GDN K4/K5 mma pair's UNSET-env default — ONE definition for the three read sites
318/// (gdn_mma_enabled, the k123 pre-work, gdn_scan_chunked's dispatch). They read the env
319/// per call ON PURPOSE (kernel-check toggles it to pin both configs), so the shared part
320/// is this compile-time constant: ON for Hopper-MMA builds (the original 90a promotion)
321/// and for sm_120a builds (lane/moeprime-nvfp4-direct, 2026-08-21 — measured on one RTX
322/// PRO 6000 ornith15 pp14715 +6-8% and the local 5090 q38-27b +1-2%, both orders both
323/// rigs). A site defaulting differently from its peers arms the mma pre-work while the
324/// scan takes the scalar route — measured as a 0.8% LOSS, the drift this helper kills.
325pub(crate) const fn gdn_mma_default_on() -> bool {
326    cfg!(memra_hopper_mma) || konst_eq(env!("MEMRA_BUILT_CUDA_ARCH"), "120a")
327}
328
329/// const str-eq (std `==` on &str is not const-stable on this toolchain floor).
330const fn konst_eq(a: &str, b: &str) -> bool {
331    let (a, b) = (a.as_bytes(), b.as_bytes());
332    if a.len() != b.len() {
333        return false;
334    }
335    let mut i = 0;
336    while i < a.len() {
337        if a[i] != b[i] {
338            return false;
339        }
340        i += 1;
341    }
342    true
343}
344
345/// The legacy quantized prefill GEMMs are tuned and validated for sm_120a; sm_90a re-admits
346/// them through the Hopper-MMA lane (int8 m16n8k32.s8 is sm_80-class PTX).  Keep the policy
347/// in a pure helper so the dispatch guard can be regression-tested without constructing an
348/// Engine or allocating a GPU tensor.
349const fn legacy_quant_gemm_allowed(portable_cuda: bool, hopper_mma: bool, no_gemm: bool) -> bool {
350    (!portable_cuda || hopper_mma) && !no_gemm
351}
352
353// ---- KV-cache format selection (kvbytes lane, 2026-07-08; default OFF = daily config) ----
354// `MEMRA_KV_K` = q8_0 (default, 34 B/32elem) | fp8 (raw e4m3, 32 B — the -6% K-bytes arm)
355// `MEMRA_KV_V` = q5_1 (default, 24 B/32elem) | q4_0 (18 B, -25% V bytes) | fp8 (32 B, +33%)
356// A non-default format is a NEW NUMERIC CONFIG: its own run-gen argmax baseline is legal,
357// but the gate battery (kernel-check, run-spec self-consistency) must pass WITHIN it and
358// the choice is explicit env, never silent. flash_attn.cu is compiled once per format pair
359// (build.rs); the kernels keep their names — Engine::new just loads the matching fatbin.
360const FLASH_FATBIN_VQ4: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_VQ4"));
361const FLASH_FATBIN_VF8: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_VF8"));
362const FLASH_FATBIN_KF8: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_KF8"));
363const FLASH_FATBIN_KF8VQ4: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_KF8VQ4"));
364const FLASH_FATBIN_KF8VF8: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_KF8VF8"));
365
366/// KV format policy moved to the shared `memra-kv` crate (Phase D); re-exported so the
367/// fatbin router below and every existing `crate::kv_blk_bytes()` call site is unchanged.
368pub use memra_kv::{kv_blk_bytes, kv_cache_formats};
369
370/// The flash_attn fatbin matching the selected KV formats.
371fn flash_fatbin_bytes() -> &'static [u8] {
372    match kv_cache_formats() {
373        ("q8_0", "q5_1") => FLASH_FATBIN,
374        ("q8_0", "q4_0") => FLASH_FATBIN_VQ4,
375        ("q8_0", "fp8") => FLASH_FATBIN_VF8,
376        ("fp8", "q5_1") => FLASH_FATBIN_KF8,
377        ("fp8", "q4_0") => FLASH_FATBIN_KF8VQ4,
378        ("fp8", "fp8") => FLASH_FATBIN_KF8VF8,
379        other => unreachable!("kv_cache_formats returned {other:?}"),
380    }
381}
382
383/// TUNE SEAM (tools/sweep): kernel1 (Q8_0/Q4_K/Q5_K) launch-tile override,
384/// `MEMRA_GEMM_K1_LAUNCH="BM,BN,NWARP"`. MUST match the `-D K1_BM/K1_BN/NWARP` the swept
385/// fatbin was compiled with (the .cu tile and the host launch grid/block have to agree —
386/// the hardcoded (128,128,8) in qmatvec_gemm/qmatvec_gemm_raw is the shipped default).
387/// Kernel2 (Q6_K/NVFP4) launch is untouched. Unset or malformed => None => shipped
388/// defaults (zero behavior change).
389fn k1_launch_override() -> Option<(u32, u32, u32)> {
390    static K1: std::sync::OnceLock<Option<(u32, u32, u32)>> = std::sync::OnceLock::new();
391    *K1.get_or_init(|| {
392        let v = std::env::var("MEMRA_GEMM_K1_LAUNCH").ok()?;
393        let p: Vec<u32> = v.split(',').filter_map(|s| s.trim().parse().ok()).collect();
394        match p.as_slice() {
395            [bm, bn, w] => Some((*bm, *bn, *w)),
396            _ => None,
397        }
398    })
399}
400
401/// H100 wgmma prefill-GEMM seam (task 8, ARCHITECTURE-H100.md): OPT-IN (MEMRA_WGMMA=1).
402/// v0 verdict (2026-07-26, N=5 pp512 9B-Q8_0): wgmma 3845 tok/s vs MMQ 8692 — the
403/// standalone harness's "688us MMQ ref" was a pp2048-shape figure, so v0 (unpipelined,
404/// 64x64 tile, wait_group<0> every 32-K step) is ~3x SLOWER per launch at m=512 model
405/// shapes. Default stays MMQ until the pipelined version beats it N=5 (repo law).
406/// Correctness stays pinned regardless: kernel-check's wgmma case is cfg-gated, not env-gated.
407pub(crate) fn wgmma_gemm_enabled() -> bool {
408    static V: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
409    *V.get_or_init(|| std::env::var("MEMRA_WGMMA").as_deref() == Ok("1"))
410}
411
412/// TUNE SEAM: keys per FA-decode split (`MEMRA_FA_SPLIT` forces a fixed size; default 64). Smaller
413/// splits raise grid.y so grid = n_head_kv * n_splits fills the 82 SMs at short/mid ctx (vec path
414/// launches only n_head_kv=8 CTAs per split). Swept clock-locked 2026-07-03 (graph tg128): 32 beat
415/// 64 at ctx 128/512 (+0.5/+1.2%) and lost at 2048 (-3%) — BUT the adaptive 32/64 default BROKE the
416/// MTP spec-decode exact-match gate (run-spec K=1/2 self-consistency FAIL with 32; PASS with 64):
417/// the split count changes the combine's FP summation order, and the spec verify's batched forward
418/// only argmax-matches single-step decode under the 64-split order on real prompts. Spec exactness
419/// (the bigger lever) outranks a <=1.2% decode win -> default stays FIXED 64; sweeps use the env.
420/// Takes t_kv so eager, _dc capture, and fa_geom_eager stay signature-compatible for future
421/// adaptive retries (any retry MUST pass run-spec self-consistency first).
422/// Minimum t_kv for the warp-per-token vec FA path (below it the scalar path's 4x-more-blocks
423/// hides latency better — measured crossover, see `fa_decode`). Shared by fa_decode / fa_decode_dc /
424/// fa_geom_eager / fa_decode_rows-eligibility (spec verify) so the kernel pick NEVER diverges
425/// between eager decode and the verify (the spec-exactness law).
426pub const FA_VEC_MIN_TKV: usize = 96;
427/// Env-overridable crossover (MEMRA_FA_VEC_MIN, default FA_VEC_MIN_TKV). The 96 floor was
428/// measured on the qwen geometry (nkv=2); gemma4 SWA layers run nkv=8 = 4x the vec grid,
429/// which moves the crossover — sweep per model, adopt per the battery.
430pub fn fa_vec_min_tkv() -> usize {
431    static V: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
432    *V.get_or_init(|| {
433        std::env::var("MEMRA_FA_VEC_MIN")
434            .ok()
435            .and_then(|v| v.parse().ok())
436            .unwrap_or_else(|| FA_VEC_MIN_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
437    })
438}
439
440/// f16-P/V class (DEFAULT since 2026-07-23 stamp v4; MEMRA_FA_F16PV=0 = f32-class rollback):
441/// llama-fa=1-style f16 P + f16 P@V accumulation on the hd512/SWA prefill stamps
442/// (KQ/softmax/normalize stay f32). Laptop stamp: 12B 1.045x, 31B 0.979x vs llama.
443///
444/// SPEC-SERVING FLIP (2026-07-26, the wkv acceptance-law pattern): with MEMRA_DRAFT set the
445/// default is OFF. f16 P/V shifts the PRIME's hidden states/KV in the sub-argmax logit
446/// space the drafter feeds on — argmax gates stay MATCH while depth acceptance falls off a
447/// cliff (26B d1736 0.883 -> 0.405, -40% e2e; f16pv-off alone restores 0.846/314 tok/s —
448/// the perf-ci acceptance battery is the only gate that sees this class). Explicit
449/// MEMRA_FA_F16PV always wins; plain serving keeps the f16 prefill win.
450pub fn fa_f16pv_on() -> bool {
451    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
452    *ON.get_or_init(|| {
453        std::env::var("MEMRA_FA_F16PV")
454            .map(|v| v != "0")
455            .unwrap_or_else(|_| std::env::var("MEMRA_DRAFT").is_err())
456    })
457}
458
459/// hd512 head-pair arm (DEFAULT since stamp v4; MEMRA_FA512_HP=0 reverts to sp16): GQA
460/// ncols2=2 — 2 heads per CTA share each staged K/V tile, Q register-resident. Engages
461/// when n_head is even and the GQA group (n_head/n_head_kv) is even.
462pub fn fa512_hp_on() -> bool {
463    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
464    *ON.get_or_init(|| std::env::var("MEMRA_FA512_HP").as_deref() != Ok("0"))
465}
466
467/// SWA head-pair arm (DEFAULT since stamp v4; MEMRA_FAW_HP=0 reverts to p1): llama-class
468/// windowed geometry — 32 q-rows x 2 heads per CTA sharing staged K/V, f16 P@V
469/// accumulation. Even n_head and even GQA group required (guarded per call).
470pub fn faw_hp_on() -> bool {
471    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
472    *ON.get_or_init(|| std::env::var("MEMRA_FAW_HP").as_deref() != Ok("0"))
473}
474
475/// 4-warp sp16 experiment arm (MEMRA_FA512_W4=1, requires the f16pv door): GEMM0 split-K
476/// 4-way + GEMM1 4x128 O-dims. Own partial-sum order — oracle-band gated. Returns warp
477/// count (2 = base sp16). 8-warp arm measured NEGATIVE 2026-07-23 (jsonl) and removed.
478pub fn fa512_wide_warps() -> usize {
479    static N: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
480    *N.get_or_init(|| match std::env::var("MEMRA_FA512_W4").as_deref() {
481        Ok("1") => 4,
482        _ => 2,
483    })
484}
485
486/// hd-512 vec crossover floor (MEMRA_FA512_MIN, default 512) — shared by fa_decode dispatch
487/// and the gemma global-layer rows/parity call sites.
488pub fn fa512_min_tkv() -> usize {
489    static FA512_MIN: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
490    *FA512_MIN.get_or_init(|| {
491        std::env::var("MEMRA_FA512_MIN")
492            .ok()
493            .and_then(|v| v.parse().ok())
494            .unwrap_or(512)
495    })
496}
497/// Per-model crossover default, set at model load BEFORE the first decode (per-model
498/// numeric-config adoption law). qwen keeps the measured 96; gemma4 (nkv=8 SWA) measured
499/// vec-always fastest: 119.9 (96) / 130.0 (48) / 133.2 (1) tok/s tg128-regime, 2026-07-10.
500pub static FA_VEC_MIN_DEFAULT: std::sync::atomic::AtomicUsize =
501    std::sync::atomic::AtomicUsize::new(FA_VEC_MIN_TKV);
502/// Per-model windowed-split default (MEMRA_FA_SPW overrides): gemma MoE (26B, nkv=8) measured
503/// 32 (grid-limited t=1 under the raw-e4m3 sV ceiling, 2026-07-12); dense gemma (31B)
504/// measured 64 (37.13/37.12 vs 36.87/36.86 at 1.7k, N=2 — different attention geometry).
505pub static FA_SPW_DEFAULT: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(32);
506/// Per-model hd512 (gemma globals) split default (MEMRA_FA_SP512 overrides): 26B measured 16
507/// (2026-07-11 N=2), dense 31B measured 32 (36.86/36.93 vs 36.73/36.73 at 1.7k, 2026-07-12).
508/// fused t=1 q4_0 pair/triple row mapping: true = mr1 (one row/warp). Per-model default
509/// (dense gemma wins +1.1% short / +0.6% depth on the 31B; MoE 26B REGRESSES −1.2% —
510/// its shared-expert fused2 shapes lose to the finer grid). MEMRA_Q40_MR env still wins.
511pub static FUSED_MR1_DEFAULT: std::sync::atomic::AtomicBool =
512    std::sync::atomic::AtomicBool::new(false);
513/// Per-model router-GEMV form (2026-07-31): the 8-warp twin is +8.8% on the H100 q35
514/// decode step (router was 14.8% of it) with argmax + spec self-consistency green on
515/// qwen-class MoE both rigs. The gemma-4 26B knife-edge block (2026-07-31, single
516/// synthetic prompt) was RE-ARBITRATED 2026-08-01 on 6 real prompts — gate outcomes
517/// identical to the lone-warp arm, +13% g26 decode — so gemma4 rides the default too
518/// (research/g26-decode-20260801/). MEMRA_ROUTER_V2 env overrides either way.
519pub static ROUTER_W8_DEFAULT: std::sync::atomic::AtomicBool =
520    std::sync::atomic::AtomicBool::new(true);
521pub static FA_SP512_DEFAULT: std::sync::atomic::AtomicUsize =
522    std::sync::atomic::AtomicUsize::new(16);
523/// Per-model rms_norm block size (per-model numeric-config law: the per-thread partial-sum
524/// split changes with blockDim -> different FP order -> battery-arbitrated per model).
525/// qwen keeps the shipped 256; gemma4 adopts 1024 (single-row 2816-col norms are one-block
526/// latency-bound at 256 threads — 7us/launch measured).
527pub static RMS_BLOCK_DEFAULT: std::sync::atomic::AtomicU32 = std::sync::atomic::AtomicU32::new(256);
528/// gemma4 fa split ladder switch (set at model load; see fa_split_keys).
529pub static FA_SP_GEMMA: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
530/// Per-model stream-k entry override for SPEC serving (-1 = unset → env/default;
531/// 0 = force tiling; 1 = admit the deterministic form selector). The former timing
532/// selector made identical boots choose different fold orders; `MEMRA_MMQ_SK_FORM` is the
533/// explicit numerical-form seam. mmq_ffi reads this before the env.
534pub static MMQ_SK_FORCE: std::sync::atomic::AtomicI8 = std::sync::atomic::AtomicI8::new(-1);
535/// Per-model FP8-KV door — lives in memra-kv next to the format policy it drives
536/// (re-export keeps `crate::KV_FP8_FORCE` setters in model.rs/hybrid.rs working).
537pub use memra_kv::KV_FP8_FORCE;
538/// bf16 matvec family block size (MEMRA_MMV_BLOCK, default 128, clamped to [64, 256] and a
539/// multiple of 32 — the f32acc twin's shared reduce caps at 256). NUMERIC-CLASS knob: the
540/// per-thread stride and reduction order change with the block, same acceptance class as
541/// MEMRA_RMS_BLOCK (fresh-tape identity + battery at the pinned value).
542pub(crate) fn mmv_block() -> u32 {
543    static V: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
544    *V.get_or_init(|| {
545        std::env::var("MEMRA_MMV_BLOCK")
546            .ok()
547            .and_then(|v| v.parse().ok())
548            .filter(|&b: &u32| (64..=256).contains(&b) && b % 32 == 0)
549            .unwrap_or(128)
550    })
551}
552
553/// MEMRA_TOPK_FAST=1: barrier-lean sigmoid top-k twin (warp-local top-k + one merge).
554/// Selection and weight arithmetic identical to the round-robin kernel — a latency twin.
555/// MEMRA_SIG_EXPF_DEV=1: device-libm expf sigmoid router (numeric-class door — the
556/// host-glibc transcription is FP64-rate-bound on consumer Blackwell). New tape + battery.
557pub(crate) fn sig_expf_dev_on() -> bool {
558    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
559    *ON.get_or_init(|| std::env::var("MEMRA_SIG_EXPF_DEV").as_deref() == Ok("1"))
560}
561
562pub(crate) fn topk_fast_on() -> bool {
563    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
564    *ON.get_or_init(|| std::env::var("MEMRA_TOPK_FAST").as_deref() == Ok("1"))
565}
566
567pub(crate) fn rms_block() -> u32 {
568    static V: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
569    *V.get_or_init(|| {
570        std::env::var("MEMRA_RMS_BLOCK")
571            .ok()
572            .and_then(|v| v.parse().ok())
573            .unwrap_or_else(|| RMS_BLOCK_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
574    })
575}
576
577pub(crate) fn fa_split_keys(t_kv: usize, n_head_kv: usize) -> usize {
578    static S: std::sync::OnceLock<Option<usize>> = std::sync::OnceLock::new();
579    if let Some(forced) = *S.get_or_init(|| {
580        std::env::var("MEMRA_FA_SPLIT")
581            .ok()
582            .and_then(|v| v.parse().ok())
583            .filter(|&s: &usize| s >= 8 && s % 8 == 0)
584    }) {
585        return forced;
586    }
587    // CTX-ADAPTIVE default (2026-07-05 40k sweep: sp32 24.5 vs sp128 26.0 tok/s = +5.8% — at
588    // deep ctx the n_splits count explodes (40k/32 = 1265 splits x 8 kv-heads) and the combine
589    // + partial-buffer cost dominates; at short ctx small splits fill the SMs). Exactness: split
590    // size only changes the PARTITION of keys; the rows/combine order per split is fixed and the
591    // gate battery (kernel-check + run-spec K=1..8) arbitrates every default change.
592    //
593    // SM-AWARE SHORT-CTX RUNG (2026-07-06 rtx6000): the 32-key rung was tuned on the 82-SM 5090.
594    // On 188 SMs the vec grid (n_head_kv x n_splits CTAs) starves at short ctx — the 35B has
595    // n_head_kv=2, so ctx128/split32 = 8 CTAs on 188 SMs. Measured on rtx6000 (N=1 sweep + N=3
596    // interleaved confirm): 35B ctx128 sp16 179 vs sp32 161 (+11%), ctx512 178 vs 158, ctx2048
597    // flat, ctx>=4096 sp64 edges sp16 by ~3%; 27B ctx128 70.9 vs 66.3 (+7%); 9B 177 vs 163
598    // (+9%). Rigs <=100 SMs keep the validated 5090 ladder EXACTLY (default unchanged there —
599    // rig-divergence law: this branch is measured on 188 SMs only).
600    // gemma4 all-16 ladder probe REVERTED (2026-07-10): +1.3 plain at d1736 (157.5 vs 156.2)
601    // but depth VERIFY collapsed (spec 203.5 -> 169 — the windowed rows' per-row combine over
602    // 64 splits). The mixed default (swa nkv=8 -> 32, globals nkv=2 -> 8-ladder) stays; a
603    // caller-split policy would break row-vs-decode split parity. FA_SP_GEMMA kept as a seam.
604    if FA_SP_GEMMA.load(std::sync::atomic::Ordering::Relaxed)
605        && std::env::var("MEMRA_FA_SP16").as_deref() == Ok("1")
606    {
607        return if t_kv <= 8192 {
608            16
609        } else if t_kv <= 16384 {
610            64
611        } else {
612            128
613        };
614    }
615    let big_rig = fa_sm_count() >= 128;
616    if big_rig {
617        let _ = n_head_kv;
618        if t_kv <= 2048 {
619            16
620        } else if t_kv <= 16384 {
621            64
622        } else {
623            128
624        }
625    } else if n_head_kv <= 4 {
626        // KV-HEAD-AWARE RUNG (2026-07-08, 5090): the 8192->32 rung was validated on kv=8 models
627        // (27B/9B: 8 heads x n_splits fills 82 SMs). The 35B has n_head_kv=2 — at ctx512/sp32
628        // the vec grid is 2 x 20 = 40 CTAs on 82 SMs (half idle). Measured (35B, run-gen 128tok
629        // N=1 sweep + N=3 confirm): sp8 162.1 / sp16 161.3 / sp32 159.4 at short ctx.
630        // DEPTH TAPER (same day, the deep-ctx lesson re-learned on this rung): sp8 at d6257 =
631        // 782 splits -> combine + partial-buffer cost dominates (141.2 tok/s); the d6257 sweep
632        // says sp64 = 153.0 (sp16/32 147, sp96 147.6, sp128 141). Few-kv-head models need the
633        // taper EARLIER than kv=8 (per-split grid 4x thinner, same per-split combine cost).
634        // Crossover hunt: sp8 vs sp64 = 156.7/155.9 at d3072, 151.7/155.6 at d4096 -> boundary 3072.
635        // RUNG RE-SWEPT UNDER THE DEEP KERNEL (2026-08-02, lane/ladder-3072 — the stale-verdict
636        // law: the 3072 boundary was calibrated on the conflicted v4 core; the deep rewrite cut
637        // vec cost ~1.2-1.4x while combine scales with n_splits, so sp8's combine bill
638        // dominates far earlier). Kernel receipts (quiet-rig nsys, deep vec + combine us):
639        // d1024 sp8 17.1 vs sp64 10.6; d2048 31.0 vs 12.2; d3072 44.0 vs 18.3. e2e run-gen
640        // tg128 N=3 interleaved (KAT + q35, research/ladder-3072-20260802/): sp8 loses at
641        // EVERY depth >= 1024 (KAT d2048 182.6 vs 188.0 = -2.9%, d3072 175.9 vs 186.4 =
642        // -5.6%; q35 d4096 169.2 vs 182.6 = -7.4%); d512 flat (+-0.2%, inside noise). sp32
643        // ties sp64 within noise in the mid band and loses at d4096 -> no extra rung.
644        // Boundary 3072 -> 512: sp8 keeps only the short-ctx band it was validated on
645        // (ctx128-512); sp64 takes over where the deep kernel made combine the bill.
646        if t_kv <= 512 {
647            8
648        } else if t_kv <= 16384 {
649            64
650        } else {
651            128
652        }
653    } else {
654        if t_kv <= 8192 {
655            32
656        } else if t_kv <= 16384 {
657            64
658        } else {
659            128
660        }
661    }
662}
663
664/// SM count of device 0, cached (used by fa_split_keys' rig-size rung; primary-context query,
665/// same attribute Engine::batched_variant reads).
666fn fa_sm_count() -> i32 {
667    static N: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
668    *N.get_or_init(|| {
669        cudarc::driver::result::init().ok();
670        cudarc::driver::result::device::get(0)
671            .and_then(|d| unsafe { cudarc::driver::result::device::get_attribute(
672                d, cudarc::driver::sys::CUdevice_attribute_enum::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT) })
673            .unwrap_or(82)
674    })
675}
676
677/// FA-prefill kernel-name suffix for a head_dim (the template-stamped twins in flash_attn.cu):
678/// 256 = the original names (qwen35 class, dispatch unchanged), 128 = `_hd128` (MiniMax-M3).
679/// Any other dim errors — callers gate to sdpa_naive before dispatching FA.
680fn fa_hd_suffix(head_dim: usize) -> Result<&'static str, Box<dyn std::error::Error>> {
681    match head_dim {
682        256 => Ok(""),
683        128 => Ok("_hd128"),
684        d => Err(format!(
685            "fa_prefill: no kernel stamped for head_dim={d} (only 256/128); \
686                          callers must gate to sdpa_naive"
687        )
688        .into()),
689    }
690}
691
692/// Quant type codes matching qmatvec.cu QType enum.
693pub const QT_Q8_0: i32 = 0;
694pub const QT_Q4_K: i32 = 1;
695pub const QT_Q6_K: i32 = 2;
696pub const QT_Q5_K: i32 = 3;
697pub const QT_Q3_K: i32 = 4;
698pub const QT_IQ4_XS: i32 = 5;
699pub const QT_IQ3_S: i32 = 6;
700pub const QT_NVFP4: i32 = 7;
701/// Checkpoint-native FP8-E4M3 (MEMRA_ST_E4M3, lane e4m3dec): raw safetensors e4m3 weight bytes
702/// [out_f, in_f] row-major (row_bytes == in_f), per-tensor f32 weight_scale in GpuTensor `scale`
703/// (fused at the mmvq write / post-matmul scale_inplace). Decode = qmatvec_e4m3_mmvq (+ _b2/_b4/_b8
704/// batched twins); prefill (m>=16) = the cuBLASLt FP8 GEMM on the SAME resident bytes (fp8_ffi.rs)
705/// — ONE weight copy total, no Q8_0 re-encode duplicate.
706pub const QT_F8_E4M3: i32 = 10;
707/// Device-side tag for the A6 SPLIT-PLANE repacked NVFP4 layout (Stage-A generic kernel only;
708/// GpuTensor keeps qtype=QT_NVFP4 + an `rp` flag — this tag never lives in a GpuTensor).
709pub const QT_NVFP4_RP: i32 = 9;
710/// Unquantized f32 weight (safetensors MoE Path A: experts dequantized to f32 host-resident).
711pub const QT_F32: i32 = 8;
712pub const QT_BF16: i32 = 11;
713pub const QT_Q4_0: i32 = 12; // gemma-4 QAT GGUF weight format (18B/32: fp16 d + nibbles)
714/// GGUF Q2_K. Appended after the existing Q4_0 code so kernel ABI values do not move.
715/// Mixed-expert artifacts use the generic f32-dequant staged kernel until a target-rig-gated
716/// dp4a/MMQ implementation exists.
717pub const QT_Q2_K: i32 = 13;
718/// Checkpoint-native FP8-E4M3 with a BLOCK-128 weight-scale GRID (lane/fp8-blk128-decode,
719/// 2026-08-05) — the Qwen-official FP8 / DeepSeek-V3 scale class. Same raw e4m3 bytes as
720/// `QT_F8_E4M3` ([out_f, in_f] row-major, row_bytes == in_f), but the dequant scale is
721/// `GpuTensor::Quant.blk` (`Fp8BlockScales`, [ceil(out_f/128), ceil(in_f/128)] f32) and the
722/// scalar `scale` field is 1.0 by the layout contract.
723///
724/// WHY A DISTINCT CODE rather than `QT_F8_E4M3` + a `blk` flag: every existing QT_F8_E4M3
725/// consumer (qmatvec_e4m3_mmvq and its batched/fused twins, e4m3_fused_params,
726/// matmul_pre_dual_noscale's F8 arm, try_fp8_gemm) threads exactly ONE scalar weight scale. Under
727/// a shared code, any consumer that was not taught the grid would still MATCH and would dequant
728/// every tile at scale 1.0 — a silent numeric corruption. Under a distinct code every untaught
729/// consumer refuses loudly instead (`mmvq_supports`/`gemm_supports`/`mmq_supports` return false;
730/// the mmvq name match panics), so a missed dispatch site is a crash or a refusal receipt, never
731/// wrong numbers. Decode = `qmatvec_e4m3_blk_mmvq`; prefill (m>=16) = the per-block FP8 MMQ tile
732/// on the SAME resident bytes+grid (fp8_ffi::try_fp8_blk_mmq) — ONE weight copy total.
733pub const QT_F8_E4M3_BLK: i32 = 14;
734
735/// Engine device context: CUDA context, stream, loaded kernel modules, cuBLASLt (via runtime::Gpu).
736pub struct Engine {
737    pub gpu: memra_runtime::Gpu,
738    module: Arc<CudaModule>,
739    hybrid: Arc<CudaModule>,
740    qmatvec: Arc<CudaModule>,
741    flash: Arc<CudaModule>,
742    /// FP8-GLOBALS module (2026-07-11): the kf8vf8 fatbin loaded ALONGSIDE the default —
743    /// gemma GLOBAL layers (hd512) append + attend in e4m3 (dequant-latency arc, HANDOVER).
744    /// Lazy: loaded on first global-format use; None until then.
745    flash_g: std::sync::OnceLock<Arc<CudaModule>>,
746    gemm: Arc<CudaModule>,
747    router: Arc<CudaModule>,
748    /// Sampled-spec kernels (research/sampled-spec-impl-map.md piece A).
749    sample: Arc<CudaModule>,
750    /// EDGE-1 §B: one shared SLRU expert-residency cache, lazily built on first MoE dispatch under
751    /// MEMRA_MOE_CACHE. `Mutex` makes it multi-agent safe (§E.2); the lock covers only lookup/admit/
752    /// memcpy-issue (µs), NOT the GEMM, so streams still overlap. `None` => cache disabled.
753    moe_cache: Mutex<Option<crate::moe_cache::MoeSlotCache>>,
754    /// Exact retained expert-block lengths collected after model load. Mixed-layout models use
755    /// this inventory to preallocate fixed-address size classes instead of sizing every slot to
756    /// the single largest block. The cache still owns every address for its full lifetime.
757    moe_cache_layout: Mutex<Option<Vec<usize>>>,
758    /// CAPTURE-RETAIN mode (graph arc, 2026-07-12): while a graph capture (and its allocator
759    /// warmups) runs, every Engine allocation is ALSO kept alive here — a captured graph's
760    /// transient buffers must never return to the pool, or later allocations (e.g. the spec
761    /// verify between replays) reuse their addresses and the replay reads/writes live memory
762    /// (the draft-graph corruption root cause). Fast-path cost when off: one relaxed atomic.
763    capture_keep_on: std::sync::atomic::AtomicBool,
764    /// VERIFY-EXACT scope (dflash lane, 2026-07-13): when set, matmul/matmul_pre skip the
765    /// m>=16 prefill-GEMM branches so a t>=16 batched VERIFY rides the decode-exact b-tier
766    /// class (the parity law). The t=16 dflash verify tripped the GEMM threshold — 770us/
767    /// matmul (54% of the round) AND a different FP order than decode (issue-10 landmine).
768    verify_exact: std::sync::atomic::AtomicBool,
769    capture_keep: Mutex<Vec<Box<dyn std::any::Any + Send>>>,
770    /// EDGE-1 §C.2: dedicated H2D copy stream for async prefetch (event-synced to the compute stream).
771    pub copy_stream: Arc<CudaStream>,
772    /// Resident CUTLASS NVFP4 prefill scratch (workspace + a_packed + sfa_linear + sfa_sw + y + alpha),
773    /// allocated ONCE and grown to the largest prefill GEMM shape, then reused per-call. Removes the
774    /// 6 fresh allocations + alpha htod that `cutlass_fp4_gemm` did every prefill matmul (~200/prefill).
775    /// Safe as a single shared buffer because all GPU compute serializes on the one `gpu.stream` worker
776    /// thread (the server runs one GPU worker; no concurrent CUTLASS GEMMs share this scratch). `None`
777    /// until the first CUTLASS FP4 GEMM. Mutex guards lazy build/grow only (matches `moe_cache`).
778    #[cfg(memra_cutlass)]
779    cutlass_scratch: Mutex<Option<crate::cutlass_ffi::CutlassScratch>>,
780    /// FP8-ACT PREFILL scratch (MEMRA_PP_FP8): quantized-activation buffer + scale block + cuBLASLt
781    /// workspace, allocated once and grown to the largest prefill m*k (see fp8_ffi.rs). `None`
782    /// until the first FP8 prefill GEMM; Mutex guards lazy build/grow only (matches cutlass_scratch).
783    fp8_scratch: Mutex<Option<crate::fp8_ffi::Fp8Scratch>>,
784    /// f16-P/V door: pooled V re-encode buffer (bf16->f16) for the hd512 _pre path. Lazy-grow;
785    /// per-call cudaMalloc was a laptop-regression suspect (VRAM pressure, 31B nkv=4 = 4x bytes).
786    fa_vf16_scratch: Mutex<Option<CudaSlice<u8>>>,
787    /// Pooled fa-decode split partials (part_o, part_m, part_l): per-call zeros() was 3
788    /// alloc+memset pairs per fa launch (~144 mem nodes per decode token — the graph door's
789    /// residual launch tax) — lazy-grow, memset-prefix per use, stream-ordered reuse.
790    fa_part_pool: Mutex<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>>,
791    /// Retired fa-part pool generations (#68): old buffers whose addresses captured graphs may
792    /// have baked — kept alive for the Engine's lifetime instead of returning to the async pool
793    /// (see the RETIRE-ON-GROW comment at the realloc sites). Doubling growth bounds the total.
794    fa_part_retired: Mutex<Vec<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>>,
795    /// name -> resolved CudaFunction (capture-safe lookups; see `func`).
796    fn_cache: Mutex<std::collections::HashMap<String, CudaFunction>>,
797    f16_scratch: Mutex<Option<crate::f16_ffi::F16Scratch>>,
798    /// RANK1 LEVER (parallel argmax): resident pass-1 partials scratch (part_v[NB] f32, part_i[NB] i32),
799    /// allocated ONCE on first parallel-argmax call and reused. Stable pointers so the 2-pass argmax
800    /// is CUDA-graph-capturable (the buffer is referenced by both captured passes; lazy-allocated
801    /// before capture under the generate_graph tracking-off window so it carries no events).
802    argmax_partials: Mutex<Option<(CudaSlice<f32>, CudaSlice<i32>)>>,
803    /// ARC B (chunk-prime dequant-once): resident bf16 K/V workspace for `fa_prefill_view_ws`
804    /// ((K bytes, V bytes) u8 buffers holding [t_kv, kv_dim] bf16). Grown lazily to the largest
805    /// (t_kv, kv_dim) seen, REUSED across layers/chunks/calls (contents rewritten per launch —
806    /// safe because all compute serializes on the one gpu.stream). ~82MB at 40k ctx on the 27B.
807    prime_deqw_ws: Mutex<Option<(CudaSlice<u8>, CudaSlice<u8>)>>,
808    /// LAUNCH-STRUCTURE STAGE 1: persistent PINNED (cacheable, flags=0) host staging buffer for the
809    /// fused-router sel/w readback — one async DtoH pair + ONE sync instead of two synced dtohs.
810    /// Grown lazily; reused every MoE layer (single-threaded decode serializes on the sync).
811    router_stage: Mutex<Option<PinnedStage>>,
812}
813
814/// FAVENDOR lane env gate (2026-07-08): MEMRA_FA_V2=1 dispatches the llama-fattn-vec-mechanism
815/// decode kernels (fa_decode_vec_q_v2 / fa_decode_vec_q_rows_v2 / fa_decode_vec_q_v2_dc):
816/// tile-batched online softmax (one alpha rescale per 32-key tile instead of per key) + wide-load
817/// block dequant in the staging phase. NOTE rev2: llama's register streaming (no smem) was ALSO
818/// tried and measured 2x WORSE at depth in our gqa-warps frame — the smem KV-tile broadcast stays
819/// (see the kernel comment). NEW NUMERIC CONFIG (tile-level softmax regrouping changes FP order vs
820/// the per-key twins) — own argmax baseline; eager decode, the spec-verify rows path AND the
821/// graph _dc path switch TOGETHER (the spec-exactness law). Default OFF. Read per call (not
822/// OnceLock) so the gate battery can A/B within one process, matching the MEMRA_NO_FA_VEC pattern.
823fn fa_v2_on() -> bool {
824    // DEFAULT ON since 2026-07-08 (MEMRA_FA_V2=0 reverts): tile-batched online softmax, e2e
825    // measured across every model x depth — 35B 168.7->173.4 (d512) / 153.1->158.5 (d6257),
826    // 9B 131.2->132.7 / 108.4->124.5 (+15% — the engine-wide depth-slope fix), 27B 47.2->47.7 /
827    // 42.2->44.9. One-time numeric-config change; kernel-check + argmax + spec self-consistency
828    // + graph bit-identity green on all three models.
829    std::env::var("MEMRA_FA_V2")
830        .map(|v| v != "0")
831        .unwrap_or(true)
832}
833
834/// FA v3 gate (default ON since 2026-07-09; MEMRA_FA_V3=0 reverts to v2 — research/fa/fa_v3_design.md):
835/// HYBRID decode twins (fa_decode_vec_q_v3 / _rows_v3 / _v3_dc): llama's int8-dp4a K.Q with
836/// register-quantized Q (no K dequant, no K smem) + OUR CTA-shared staged bf16 V tile + OUR
837/// split partition/combine. NEW NUMERIC CONFIG (int8 Q quantization changes the K.Q accumulation
838/// vs the bf16-roundtrip FMA chain) — own argmax baseline; eager decode, the spec-verify rows
839/// path AND the graph _dc path switch TOGETHER (the spec-exactness law). Read per call so the
840/// gate battery can A/B within one process (the MEMRA_FA_V2 pattern).
841fn fa_v3_on() -> bool {
842    // DEFAULT ON since 2026-07-09 (MEMRA_FA_V3=0 reverts to v2): dp4a-K hybrid FA decode —
843    // fa kernel -21-23% at depth (micro), 35B spec p3 +5% (190->200, the last spec cell),
844    // d6257 +1.7%. Own numeric config; full battery green on 35B+9B incl graph bit-identity.
845    std::env::var("MEMRA_FA_V3")
846        .map(|v| v != "0")
847        .unwrap_or(true)
848}
849
850/// The v3 dp4a K path reads RAW q8_0 bytes (34B blocks) and stages q5_1 V verbatim — it is only
851/// correct on the DEFAULT KV formats — and needs dpl % 4 == 0 consecutive quants per lane
852/// (head_dim % 128 == 0; both daily models are hd256). All three dispatch sites share this
853/// predicate so the twins can never diverge.
854fn fa_v4_mode() -> &'static str {
855    static M: std::sync::OnceLock<String> = std::sync::OnceLock::new();
856    M.get_or_init(|| std::env::var("MEMRA_FA_V4").unwrap_or_default())
857}
858fn fa_v4_on() -> bool {
859    fa_v4_mode() != "0"
860} // DEFAULT ON 2026-07-10 (MEMRA_FA_V4=0 rollback)
861/// t_kv-conditional v4 pick (gemma depth lesson 2026-07-10: v4's key-per-lane pipeline starves
862/// at the 1024-window with short splits — MEMRA_FA_V4=0 measured depth plain 158.0 vs 156.7).
863/// Threshold MEMRA_FA_V4_MAX (default usize::MAX = unchanged behavior; gemma sets 1024 at load
864/// via FA_V4_MAX_DEFAULT). Applied at EVERY dispatch site (eager, rows, rows_w, dc) so verify
865/// stays kernel-family-identical to decode at the same t_kv.
866/// Per-model deep-ctx smem floor default (MEMRA_FA_SMEM_TKV env overrides): gemma pushes it
867/// above the 1024 window so the windowed decode + verify rows share the REGISTER family.
868pub static FA_SMEM_TKV_DEFAULT: std::sync::atomic::AtomicUsize =
869    std::sync::atomic::AtomicUsize::new(1024);
870pub static FA_V4_MAX_DEFAULT: std::sync::atomic::AtomicUsize =
871    std::sync::atomic::AtomicUsize::new(usize::MAX);
872pub fn fa_v4_at_pub(t_kv: usize) -> bool {
873    fa_v4_at(t_kv)
874}
875fn fa_v4_at(t_kv: usize) -> bool {
876    static M: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
877    let mx = *M.get_or_init(|| {
878        std::env::var("MEMRA_FA_V4_MAX")
879            .ok()
880            .and_then(|v| v.parse().ok())
881            .unwrap_or_else(|| FA_V4_MAX_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
882    });
883    fa_v4_on() && t_kv < mx
884}
885/// FA-DEEP gate (2026-08-02, lane fa-decode-deep): deep-ctx v4 twins
886/// (fa_decode_vec_q_v4_deep / _deep_dc) — the depth-decode lane's priced fix. Unlike
887/// v2/v3/v4 this is NOT a numeric config: the deep twins run the v4 program VERBATIM
888/// (same split partition, same softmax/accumulation order, same partials/combine) and only
889/// move the smem physical layout (bank de-conflict row pads) + the load schedule (next-tile
890/// L2 prefetch) — kernel-check pins bitdiff==0 vs the v4 twins across depths, so eager /
891/// rows-verify / graph / seqs stay mutually bit-identical wherever the threshold falls.
892/// Engages at t_kv >= MEMRA_FA_DEEP_MIN. The swept floor is 0 = ALWAYS ON where v4 ran
893/// (fa-deep-bench fine grid 96..6144, 2026-08-02: deep flat-or-better at EVERY depth,
894/// 1.01-1.26x, no losing cell — so there is no engagement boundary and no new
895/// capture-recapture edge; the env stays as a sweep/diagnostic seam only).
896/// MEMRA_FA_DEEP=0 is the rollback seam. Read per call so the battery + bench can A/B
897/// within one process (the v2/v3 pattern).
898pub const FA_DEEP_MIN_DEFAULT: usize = 0;
899fn fa_deep_at(t_kv: usize) -> bool {
900    if std::env::var("MEMRA_FA_DEEP").as_deref() == Ok("0") {
901        return false;
902    }
903    let min = std::env::var("MEMRA_FA_DEEP_MIN")
904        .ok()
905        .and_then(|v| v.parse().ok())
906        .unwrap_or(FA_DEEP_MIN_DEFAULT);
907    t_kv >= min
908}
909/// Public twin (kernel-check builds the deep-vs-v4 bit pin; bench sweeps the floor).
910pub fn fa_deep_at_pub(t_kv: usize) -> bool {
911    fa_deep_at(t_kv)
912}
913
914fn fa_v3_active(head_dim: usize) -> bool {
915    // v3's dp4a-K walk reads raw q8_0 K bytes — no e4m3 arm; the fp8-KV arm (MEMRA_KV_FP8)
916    // must fall back like any non-default KV format (the rows_dc stream path asserts on it).
917    fa_v3_on()
918        && head_dim % 128 == 0
919        && kv_cache_formats() == ("q8_0", "q5_1")
920        && !Engine::kv_fp8_on()
921}
922
923/// BATCHED-TICK increment 2 (2026-08-01): true iff a row at this t_kv would take the v4
924/// eager arm in `fa_decode_kvmod`'s dispatch — the exact precondition for the z-batched
925/// `fa_decode_vec_q_seqs_v4` twin to reproduce its per-seq program bit-identically.
926/// Mirrors the kvmod predicates: vec on + above the vec floor + hd256 + inside the v4
927/// window + the PRODUCTION v4 body (the noB3/stage phase probes are wrong-output) + the
928/// default flash module (no fp8-KV g-module). Callers must ALSO group rows on one
929/// `fa_split_keys` rung (the rows-twins' straddle law) before batching.
930pub fn fa_seqs_eligible(t_kv: usize, head_dim: usize) -> bool {
931    std::env::var("MEMRA_NO_FA_VEC").is_err()
932        && t_kv >= fa_vec_min_tkv()
933        && head_dim == 256
934        && fa_v4_at(t_kv)
935        && !matches!(fa_v4_mode(), "noB3" | "stage")
936        && !Engine::kv_fp8_on()
937}
938/// Public twin of the crate-private split ladder (kernel-check builds the seqs-vs-loop pin).
939pub fn fa_split_keys_pub(t_kv: usize, n_head_kv: usize) -> usize {
940    fa_split_keys(t_kv, n_head_kv)
941}
942
943/// A raw pinned (page-locked, CACHEABLE — flags=0, not write-combined) host allocation for
944/// DtoH staging. cudarc's `alloc_pinned` uses CU_MEMHOSTALLOC_WRITECOMBINED, which is right for
945/// HtoD streams but pathologically slow for host READS — the router readback is host-read-heavy,
946/// so we allocate through `result::malloc_host` with flags=0 directly.
947struct PinnedStage {
948    ptr: *mut u8,
949    cap: usize,
950}
951unsafe impl Send for PinnedStage {}
952impl PinnedStage {
953    fn new(cap: usize) -> Result<Self, Box<dyn std::error::Error>> {
954        let ptr = unsafe { cudarc::driver::result::malloc_host(cap, 0)? } as *mut u8;
955        Ok(PinnedStage { ptr, cap })
956    }
957}
958impl Drop for PinnedStage {
959    fn drop(&mut self) {
960        let _ = unsafe { cudarc::driver::result::free_host(self.ptr as _) };
961    }
962}
963
964/// Number of pass-1 blocks for the parallel argmax (fan-out across SMs to saturate HBM). 256 blocks
965/// x 256 threads = 65536 threads covering the 248K-vocab scan in ~4 strided loads/thread.
966pub const ARGMAX_NB: usize = 256;
967
968/// crate-visible alias for the batched FA3 shim entry (hybrid_forward's batch arm).
969pub(crate) use memra_fa3_vl as fa3_vl_raw;
970
971unsafe extern "C" {
972    /// FA3 v10 shim (cu/fa3_prefill.cu): TMA-swizzled wgmma FA, fresh causal hd256.
973    fn memra_fa3_prefill(
974        q16: *const core::ffi::c_void,
975        k16: *const core::ffi::c_void,
976        v16: *const core::ffi::c_void,
977        o: *mut f32,
978        t: i32,
979        h: i32,
980        hkv: i32,
981        d: i32,
982        scale: f32,
983        stream: *mut core::ffi::c_void,
984    ) -> i32;
985    /// batched varlen twin: host arrays of device pointers per seq (B <= 8).
986    pub(crate) fn memra_fa3_vl(
987        q16s: *const *const core::ffi::c_void,
988        k16s: *const *const core::ffi::c_void,
989        v16s: *const *const core::ffi::c_void,
990        os: *const *mut f32,
991        ts: *const i32,
992        b: i32,
993        h: i32,
994        hkv: i32,
995        d: i32,
996        scale: f32,
997        stream: *mut core::ffi::c_void,
998    ) -> i32;
999}
1000
1001/// STAGE-2 GROUPED DECODE: 8 expert weight-block device pointers passed BY VALUE as one kernel
1002/// param (matches the CUDA `wptr8_t` struct: 8x 64-bit pointers, `#[repr(C)]` => identical
1003/// layout). The pointers are SLRU cache-slot base addresses — fixed for the engine's lifetime
1004/// (slots are never re-allocated), so passing raw values is stable across the launch.
1005#[repr(C)]
1006#[derive(Clone, Copy)]
1007pub struct WPtr8(pub [u64; 8]);
1008unsafe impl cudarc::driver::DeviceRepr for WPtr8 {}
1009
1010/// task #18 varlen GDN: per-seq args for gdn_chunk_{state,output}_mma_vl — one launch
1011/// runs all B<=8 sequences' K4/K5 (CUDA `gdnseq_t`/`gdnvl_t`, layout-identical repr(C)).
1012/// Raw addresses are valid for the launch: every referenced buffer outlives the call and
1013/// all work is on the single compute stream (same discipline as the f16 GEMM FFI).
1014#[repr(C)]
1015#[derive(Clone, Copy, Default)]
1016pub struct GdnSeqVl {
1017    pub kb16: u64,
1018    pub gcum: u64,
1019    pub beta: u64,
1020    pub u: u64,
1021    pub wb16: u64,
1022    pub y: u64,
1023    pub ssnap: u64,
1024    pub state_in: u64,
1025    pub state_out: u64,
1026    pub q: u64,
1027    pub p: u64,
1028    pub o: u64,
1029    pub k: u64,
1030    pub v: u64,
1031    pub g: u64,
1032    pub a: u64,
1033    pub w: u64,
1034    pub t: i32,
1035    pub nc: i32,
1036}
1037unsafe impl cudarc::driver::DeviceRepr for GdnSeqVl {}
1038#[repr(C)]
1039#[derive(Clone, Copy)]
1040pub struct GdnVl8(pub [GdnSeqVl; 8]);
1041unsafe impl cudarc::driver::DeviceRepr for GdnVl8 {}
1042
1043/// task #22: per-seq wgmma-fused extras (CUDA `gdnw_t`/`gdnwvl_t`) — qb16 mirror +
1044/// pre-masked Pb16, riding NEXT TO GdnSeqVl so the base struct stays untouched.
1045#[repr(C)]
1046#[derive(Clone, Copy, Default)]
1047pub struct GdnWVl {
1048    pub qb16: u64,
1049    pub pb16: u64,
1050}
1051unsafe impl cudarc::driver::DeviceRepr for GdnWVl {}
1052#[repr(C)]
1053#[derive(Clone, Copy)]
1054pub struct GdnWVl8(pub [GdnWVl; 8]);
1055unsafe impl cudarc::driver::DeviceRepr for GdnWVl8 {}
1056
1057/// task #18 increment 3: per-seq PREP/TAIL args (CUDA `gdnprep_t`/`gdnprepvl_t`).
1058#[repr(C)]
1059#[derive(Clone, Copy, Default)]
1060pub struct GdnPrepVl {
1061    pub qkv: u64,
1062    pub conv_state: u64,
1063    pub conv_out: u64,
1064    pub q_g: u64,
1065    pub k_g: u64,
1066    pub v_g: u64,
1067    pub q_l2: u64,
1068    pub k_l2: u64,
1069    pub beta_raw: u64,
1070    pub alpha: u64,
1071    pub beta: u64,
1072    pub g_log: u64,
1073    pub o: u64,
1074    pub z: u64,
1075    pub gn: u64,
1076    pub gn16: u64,
1077    pub kb16: u64,
1078    pub qb16: u64,
1079    pub t: i32,
1080    pub pad: i32,
1081}
1082unsafe impl cudarc::driver::DeviceRepr for GdnPrepVl {}
1083#[repr(C)]
1084#[derive(Clone, Copy)]
1085pub struct GdnPrepVl8(pub [GdnPrepVl; 8]);
1086unsafe impl cudarc::driver::DeviceRepr for GdnPrepVl8 {}
1087
1088/// task #18 (attn side): per-seq varlen FA args (CUDA `faseq_t`/`favl_t`).
1089#[repr(C)]
1090#[derive(Clone, Copy, Default)]
1091pub struct FaSeqVl {
1092    pub q: u64,
1093    pub k16: u64,
1094    pub v16: u64,
1095    pub o: u64,
1096    pub kf: u64,
1097    pub vf: u64,
1098    pub t: i32,
1099    pub pad: i32,
1100}
1101unsafe impl cudarc::driver::DeviceRepr for FaSeqVl {}
1102#[repr(C)]
1103#[derive(Clone, Copy)]
1104pub struct FaVl8(pub [FaSeqVl; 8]);
1105unsafe impl cudarc::driver::DeviceRepr for FaVl8 {}
1106
1107/// task #18 (attn pre-FA): per-seq split/norm/rope/append args (CUDA `attnpre_t`).
1108#[repr(C)]
1109#[derive(Clone, Copy, Default)]
1110pub struct AttnPreVl {
1111    pub qf: u64,
1112    pub kf: u64,
1113    pub vf: u64,
1114    pub q: u64,
1115    pub gate: u64,
1116    pub qn: u64,
1117    pub kn: u64,
1118    pub kc: u64,
1119    pub vc: u64,
1120    pub t: i32,
1121    pub pad: i32,
1122}
1123unsafe impl cudarc::driver::DeviceRepr for AttnPreVl {}
1124#[repr(C)]
1125#[derive(Clone, Copy)]
1126pub struct AttnPreVl8(pub [AttnPreVl; 8]);
1127unsafe impl cudarc::driver::DeviceRepr for AttnPreVl8 {}
1128
1129/// task #18 increment 2: one sequence's FULL chunk-buffer set (alloc-only; the
1130/// varlen K1-K5 chain fills them).
1131pub struct GdnChunkBufs {
1132    pub gcum: CudaSlice<f32>,
1133    pub a: CudaSlice<f32>,
1134    pub p: CudaSlice<f32>,
1135    pub u: CudaSlice<f32>,
1136    pub w: CudaSlice<f32>,
1137    pub kb16: CudaSlice<u8>,
1138    pub wb16: CudaSlice<u8>,
1139    pub y16: CudaSlice<u8>,
1140    pub ssnap16: CudaSlice<u8>,
1141    pub qb16: CudaSlice<u8>,
1142    pub pb16: CudaSlice<u8>,
1143    pub o: CudaSlice<f32>,
1144    pub t: usize,
1145    pub nc: usize,
1146}
1147
1148/// STAGE-2 GROUPED DECODE: the 8 routed-expert weights by value (CUDA `f32x8_t`).
1149#[repr(C)]
1150#[derive(Clone, Copy)]
1151pub struct F32x8(pub [f32; 8]);
1152unsafe impl cudarc::driver::DeviceRepr for F32x8 {}
1153
1154/// Harness timing contract: wall nanos of the LAST generate/generate_spec prompt prime on this
1155/// process. Bench binaries read it right after the call to print gen-only throughput without the
1156/// prime-subtraction hack (which amplifies prime jitter into the gen number at long prompts).
1157pub static PRIME_NANOS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1158
1159impl Engine {
1160    pub fn new(ordinal: usize) -> Result<Self, Box<dyn std::error::Error>> {
1161        let gpu = memra_runtime::Gpu::new(ordinal)?;
1162        // ARCH GUARD (unified dual-arch engine): the fatbins carry single-arch SASS, so a
1163        // binary/device mismatch otherwise dies at first module load with an opaque CUDA
1164        // error. Fail early with the rebuild hint instead. MEMRA_ARCH_CHECK=0 skips.
1165        if std::env::var("MEMRA_ARCH_CHECK").as_deref() != Ok("0") {
1166            use cudarc::driver::sys::CUdevice_attribute_enum as A;
1167            let (maj, min) = cudarc::driver::result::device::get(ordinal as i32)
1168                .and_then(|d| unsafe {
1169                    Ok((
1170                        cudarc::driver::result::device::get_attribute(
1171                            d,
1172                            A::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR,
1173                        )?,
1174                        cudarc::driver::result::device::get_attribute(
1175                            d,
1176                            A::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR,
1177                        )?,
1178                    ))
1179                })
1180                .unwrap_or((0, 0));
1181            let built = env!("MEMRA_BUILT_CUDA_ARCH");
1182            let ok = matches!(
1183                (built, maj, min),
1184                ("120a", 12, 0) | ("120a", 12, 1) | ("100a", 10, 0) | ("90a", 9, 0) | ("89", 8, 9)
1185            );
1186            if !ok {
1187                return Err(format!(
1188                    "memra was built for sm_{built} but device {ordinal} reports compute \
1189                     capability {maj}.{min}. Rebuild on this machine (MEMRA_CUDA_ARCH \
1190                     auto-detects the GPU) or set MEMRA_ARCH_CHECK=0 to bypass."
1191                )
1192                .into());
1193            }
1194        }
1195        // Default async-pool RELEASE_THRESHOLD is 0: freed blocks return to the OS at every
1196        // sync, so cuMemAllocAsync NODES inside captured graphs re-map memory on EVERY
1197        // cuGraphLaunch (measured 226us/launch on the gemma graph door, 2026-07-23 osrt).
1198        // Pinning the threshold keeps the pool cached -> alloc nodes become pointer bumps.
1199        unsafe {
1200            use cudarc::driver::sys;
1201            let dev: sys::CUdevice = ordinal as sys::CUdevice;
1202            let mut pool: sys::CUmemoryPool = std::ptr::null_mut();
1203            if sys::cuDeviceGetDefaultMemPool(&mut pool, dev) == sys::CUresult::CUDA_SUCCESS {
1204                let mut thresh: u64 = u64::MAX;
1205                let _ = sys::cuMemPoolSetAttribute(
1206                    pool,
1207                    sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_RELEASE_THRESHOLD,
1208                    &mut thresh as *mut u64 as *mut core::ffi::c_void,
1209                );
1210            }
1211        }
1212        let module = gpu.ctx.load_module(Ptx::from_binary(FATBIN.to_vec()))?;
1213        let hybrid = gpu
1214            .ctx
1215            .load_module(Ptx::from_binary(HYBRID_FATBIN.to_vec()))?;
1216        let qmatvec = gpu
1217            .ctx
1218            .load_module(Ptx::from_binary(QMATVEC_FATBIN.to_vec()))?;
1219        let flash = gpu
1220            .ctx
1221            .load_module(Ptx::from_binary(flash_fatbin_bytes().to_vec()))?;
1222        let gemm = gpu
1223            .ctx
1224            .load_module(Ptx::from_binary(gemm_fatbin_bytes().into_owned()))?;
1225        let router = gpu
1226            .ctx
1227            .load_module(Ptx::from_binary(ROUTER_FATBIN.to_vec()))?;
1228        let sample = gpu
1229            .ctx
1230            .load_module(Ptx::from_binary(SAMPLE_FATBIN.to_vec()))?;
1231        let copy_stream = gpu.ctx.new_stream()?;
1232        // DECODE EVENT-TRACKING ELISION — DEFAULT ON (2026-07-05; MEMRA_EVT=1 = escape hatch).
1233        // cudarc is in multi-stream mode (main stream +
1234        // copy_stream are both created streams), so with tracking on EVERY launch arg records a
1235        // read/write CudaEvent and inserts cuStreamWaitEvent on prior events. On the 35B MoE decode
1236        // that is ~19k cuStreamWaitEvent + ~9k cuEventRecord + ~6k event create/destroy per token
1237        // (~7 ms/tok host time, measured nsys 2026-07-04 rtx6000), and +4.6% measured on 27B decode —
1238        // protecting NOTHING: every hot-path kernel/memcpy runs on the ONE gpu.stream.
1239        // CROSS-STREAM HAZARD AUDIT: MoeSlotCache in-memory prefetch uses copy_stream. Every
1240        // overwrite explicitly records the prior compute point and makes copy_stream wait; every
1241        // consumer explicitly waits for the copy completion event. The opt-in positioned-read
1242        // proof stays on gpu.stream and retains an explicit event solely to guard pinned-source
1243        // reuse. Graph-capture sites use only gpu.stream, so these handoffs never rely on cudarc's
1244        // implicit event tracking.
1245        // SAFETY: single-stream ordering is total; the runtime mem-pool is configured with
1246        // internal-dependency reuse (memra-runtime), so alloc reuse is stream-ordered too.
1247        if std::env::var("MEMRA_EVT")
1248            .map(|v| v == "1")
1249            .unwrap_or(false)
1250        {
1251            // escape hatch: keep cudarc's implicit cross-stream event tracking.
1252        } else {
1253            unsafe {
1254                gpu.ctx.disable_event_tracking();
1255            }
1256        }
1257        Ok(Self {
1258            gpu,
1259            module,
1260            hybrid,
1261            qmatvec,
1262            flash,
1263            flash_g: std::sync::OnceLock::new(),
1264            gemm,
1265            router,
1266            sample,
1267            moe_cache: Mutex::new(None),
1268            moe_cache_layout: Mutex::new(None),
1269            copy_stream,
1270            capture_keep_on: std::sync::atomic::AtomicBool::new(false),
1271            verify_exact: std::sync::atomic::AtomicBool::new(false),
1272            capture_keep: Mutex::new(Vec::new()),
1273            argmax_partials: Mutex::new(None),
1274            prime_deqw_ws: Mutex::new(None),
1275            router_stage: Mutex::new(None),
1276            fp8_scratch: Mutex::new(None),
1277            fa_vf16_scratch: Mutex::new(None),
1278            fa_part_pool: Mutex::new(None),
1279            fa_part_retired: Mutex::new(Vec::new()),
1280            fn_cache: Mutex::new(Default::default()),
1281            f16_scratch: Mutex::new(None),
1282            #[cfg(memra_cutlass)]
1283            cutlass_scratch: Mutex::new(None),
1284        })
1285    }
1286
1287    pub fn ctx(&self) -> &Arc<CudaContext> {
1288        &self.gpu.ctx
1289    }
1290
1291    /// Bytes the async pool holds MAPPED but NOT LIVE (reserved - used), i.e. freed blocks
1292    /// parked in the pool because `Engine::new` pins RELEASE_THRESHOLD to u64::MAX above.
1293    ///
1294    /// Why this is a public engine surface: `mem_get_info`'s `free` DOES NOT SEE these bytes —
1295    /// they are mapped to this process, so `free` counts them as gone, yet the very next
1296    /// `alloc_u8` is satisfied from them without touching `free` at all. Any admission or
1297    /// budget decision that reads `free` alone therefore under-counts real headroom by exactly
1298    /// this amount. Effective allocatable headroom is `free + pool_cached_bytes()`.
1299    ///
1300    /// MEASURED SIZE (c=64 serve burst, 9B NVFP4 + draft, 24GB card, 2026-08-06): 34-89 MB
1301    /// during the burst — SMALL. The admission gate adds it because a term that can only ever
1302    /// under-count headroom does not belong in a gate that queues real work, but the honest
1303    /// reading of this number is that pool caching is NOT where a long-running server's VRAM
1304    /// hides on this path: reserved ~= used throughout, so the memory the driver reports as
1305    /// gone is genuinely LIVE (see `pool_reserved_used` for the diagnostic pair).
1306    ///
1307    /// Returns 0 if the pool cannot be queried (never a false-positive headroom claim).
1308    pub fn pool_cached_bytes(&self) -> usize {
1309        let (reserved, used) = self.pool_reserved_used();
1310        reserved.saturating_sub(used)
1311    }
1312
1313    /// Raw async-pool occupancy: (RESERVED_MEM_CURRENT, USED_MEM_CURRENT) in bytes. Reserved is
1314    /// what the pool has mapped from the driver; used is what is live inside it. Exposed for
1315    /// admission/VRAM diagnostics — the pair distinguishes "memory is parked in the pool and
1316    /// `free` cannot see it" (reserved >> used) from "memory is genuinely held live by some
1317    /// owner" (reserved ~= used), which are opposite bugs with opposite fixes.
1318    /// (0, 0) if the pool cannot be queried.
1319    pub fn pool_reserved_used(&self) -> (usize, usize) {
1320        use cudarc::driver::sys;
1321        unsafe {
1322            let mut pool: sys::CUmemoryPool = std::ptr::null_mut();
1323            if sys::cuDeviceGetDefaultMemPool(&mut pool, self.gpu.ctx.ordinal() as sys::CUdevice)
1324                != sys::CUresult::CUDA_SUCCESS
1325            {
1326                return (0, 0);
1327            }
1328            let (mut reserved, mut used) = (0u64, 0u64);
1329            if sys::cuMemPoolGetAttribute(
1330                pool,
1331                sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_RESERVED_MEM_CURRENT,
1332                &mut reserved as *mut u64 as *mut core::ffi::c_void,
1333            ) != sys::CUresult::CUDA_SUCCESS
1334            {
1335                return (0, 0);
1336            }
1337            if sys::cuMemPoolGetAttribute(
1338                pool,
1339                sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_USED_MEM_CURRENT,
1340                &mut used as *mut u64 as *mut core::ffi::c_void,
1341            ) != sys::CUresult::CUDA_SUCCESS
1342            {
1343                return (0, 0);
1344            }
1345            (reserved as usize, used as usize)
1346        }
1347    }
1348
1349    /// Ambient stream (by value since M1-PP2 increment 2): the thread's pp2 stage stream
1350    /// when a stage scope is active, else the main compute stream — see `Gpu::stream`.
1351    pub fn stream(&self) -> Arc<CudaStream> {
1352        self.gpu.stream()
1353    }
1354    /// FP8-GLOBALS switch (MEMRA_GEMMA_GKV, default ON): gemma global (hd512) layers keep
1355    /// their KV in e4m3 — the dequant-latency arc (HANDOVER). Windowed layers stay q8_0/q5_1.
1356    pub fn gkv_on() -> bool {
1357        memra_kv::gkv_on()
1358    }
1359
1360    /// FP8-WINDOWED switch (MEMRA_GEMMA_WKV — measured 2026-07-12 in a validity-gated
1361    /// window: 1.7k 174.1-174.4 vs 168.6-169.4 default (+3%), 4.9k 158.7-160.4; vs llama
1362    /// same-window 159.5-160.2 / 140.6 = 1.09x / 1.13x): gemma windowed (hd256 SWA)
1363    /// layers hold e4m3 KV and ride the format-aware v4 lane from the kf8vf8 module.
1364    /// SERVING-MODE DEFAULT (2026-07-12, the 31B spec unlock): fp8-windowed KV GUTS the
1365    /// MTP drafter's acceptance — its single swa attention reads the windowed cache and
1366    /// e4m3 noise flips its argmaxes (31B short accept .758 -> 1.000 with q8/q5, spec 88
1367    /// -> 122.7 vs llama-mtp 112; depth .59 -> .78; 26B depth .57 -> .89). So the default
1368    /// keys on serving intent: SPEC serving (MEMRA_DRAFT set) -> OFF, plain -> ON (its
1369    /// depth-plain +3% stands). Explicit MEMRA_GEMMA_WKV always wins. GKV (globals) stays
1370    /// ON for both — no acceptance cost measured.
1371    pub fn wkv_on() -> bool {
1372        memra_kv::wkv_on()
1373    }
1374
1375    /// QWEN FP8-KV switch (MEMRA_KV_FP8 explicit; else the per-model KV_FP8_FORCE door set
1376    /// at model load; else OFF). Non-gemma full-attn layers hold e4m3 K/V via the kf8vf8
1377    /// module. Per-model verdict 2026-07-12: 9B +0.7-4% scaling with depth, 27B flat,
1378    /// 35B −2% (fp8 format-gates its v3 dp4a lane) — so the 9B class defaults ON
1379    /// (adopted 2026-07-28 with the deferred acceptance battery), others stay OFF.
1380    pub fn kv_fp8_on() -> bool {
1381        memra_kv::kv_fp8_on()
1382    }
1383
1384    /// fa kernel routed by head_dim: hd512 (gemma globals) resolves from the kf8vf8 module
1385    /// when the fp8-globals arm is on; everything else from the default flash module.
1386    fn fa_func(&self, name: &str, head_dim: usize) -> CudaFunction {
1387        if head_dim == 512 && Self::gkv_on() {
1388            self.func_g(name)
1389        } else {
1390            self.func(name)
1391        }
1392    }
1393
1394    /// Kernel from the FP8-GLOBALS (kf8vf8) flash module — gemma global-layer arm only.
1395    /// Format-AGNOSTIC kernels (e.g. fa_decode_combine_f32) are not compiled into the
1396    /// per-format fatbins; fall back to the base modules for those.
1397    fn func_g(&self, name: &str) -> CudaFunction {
1398        let m = self.flash_g.get_or_init(|| {
1399            self.gpu
1400                .ctx
1401                .load_module(cudarc::nvrtc::Ptx::from_binary(
1402                    FLASH_FATBIN_KF8VF8.to_vec(),
1403                ))
1404                .expect("load kf8vf8 flash fatbin (fp8-globals arm)")
1405        });
1406        let key = format!("g:{name}");
1407        if let Some(f) = self.fn_cache.lock().unwrap().get(&key) {
1408            return f.clone();
1409        }
1410        let f = match m.load_function(name) {
1411            Ok(f) => f,
1412            Err(_) => self.func(name),
1413        };
1414        self.fn_cache.lock().unwrap().insert(key, f.clone());
1415        f
1416    }
1417
1418    fn func(&self, name: &str) -> CudaFunction {
1419        // Resolution cache: cuModuleGetFunction fails inside a CUDA-graph capture region,
1420        // so capture-time lookups MUST be host-memory hits (warmups populate the cache).
1421        if let Some(f) = self.fn_cache.lock().unwrap().get(name) {
1422            return f.clone();
1423        }
1424        let f = self
1425            .module
1426            .load_function(name)
1427            .or_else(|_| self.hybrid.load_function(name))
1428            .or_else(|_| self.qmatvec.load_function(name))
1429            .or_else(|_| self.flash.load_function(name))
1430            .or_else(|_| self.gemm.load_function(name))
1431            .or_else(|_| self.router.load_function(name))
1432            .or_else(|_| self.sample.load_function(name))
1433            .unwrap_or_else(|_| panic!("kernel {name} not in any fatbin"));
1434        self.fn_cache
1435            .lock()
1436            .unwrap()
1437            .insert(name.to_string(), f.clone());
1438        f
1439    }
1440
1441    /// Scatter trimmed draft logits into full-vocab space: dst = -inf everywhere, then
1442    /// dst[d2t[i]] = src[i]. Two launches (fill, scatter) — no grid-wide sync needed.
1443    pub fn scatter_trim_logits(
1444        &self,
1445        src: &CudaSlice<f32>,
1446        d2t: &CudaSlice<u32>,
1447        dst: &mut CudaSlice<f32>,
1448        d_vocab: usize,
1449        n_vocab: usize,
1450    ) -> Result<(), Box<dyn std::error::Error>> {
1451        let f1 = self.func("scatter_trim_logits_f32");
1452        let f2 = self.func("scatter_trim_logits_pass2_f32");
1453        let (dv, nv) = (d_vocab as i32, n_vocab as i32);
1454        let cfg1 = LaunchConfig {
1455            grid_dim: (256, 1, 1),
1456            block_dim: (256, 1, 1),
1457            shared_mem_bytes: 0,
1458        };
1459        let __s_b1 = self.gpu.stream();
1460        let mut b1 = __s_b1.launch_builder(&f1);
1461        b1.arg(src).arg(d2t).arg(&mut *dst).arg(&dv).arg(&nv);
1462        unsafe {
1463            b1.launch(cfg1)?;
1464        }
1465        let cfg2 = LaunchConfig {
1466            grid_dim: (d_vocab.div_ceil(256) as u32, 1, 1),
1467            block_dim: (256, 1, 1),
1468            shared_mem_bytes: 0,
1469        };
1470        let __s_b2 = self.gpu.stream();
1471        let mut b2 = __s_b2.launch_builder(&f2);
1472        b2.arg(src).arg(d2t).arg(&mut *dst).arg(&dv);
1473        unsafe {
1474            b2.launch(cfg2)?;
1475        }
1476        Ok(())
1477    }
1478
1479    // ---- FILTERED-SPEC (feat/filtered-spec): top-k/p/min-p transforms applied symmetrically
1480    // to p and q — rejection sampling stays distribution-exact for the filtered target. ----
1481
1482    /// Per-row filtered-softmax stats: out[r] = (threshold_e, renorm_mass_e, row_max) for the
1483    /// filter (top_k, top_p, min_p) at `temp`. Rows index into x with row_stride f32s.
1484    #[allow(clippy::too_many_arguments)]
1485    pub fn filter_stats(
1486        &self,
1487        x: &CudaSlice<f32>,
1488        row_stride: usize,
1489        rows: &CudaSlice<i32>,
1490        out_th: &mut CudaSlice<f32>,
1491        out_z: &mut CudaSlice<f32>,
1492        out_max: &mut CudaSlice<f32>,
1493        n: usize,
1494        nrow: usize,
1495        temp: f32,
1496        top_k: i32,
1497        top_p: f32,
1498        min_p: f32,
1499    ) -> Result<(), Box<dyn std::error::Error>> {
1500        // A top-K-selection form of this kernel (3 vocab passes vs the search's ~51) was
1501        // implemented and REFUTED on 2026-08-21 (lane/moebatch-q35moe): the 248k-vocab row is
1502        // L2-resident, so the extra passes are near-free while the per-thread selection list
1503        // spills to local memory — B=8 tick 12.8/11.2 ms (cap 64/32) vs 10.4 ms for this
1504        // kernel, and serve c8 agg ~648 vs ~666. The receipts row is the record; the real
1505        // filtered-sampling win was batching the per-row launches (decode_batch.rs).
1506        //
1507        // COOPERATIVE MULTI-BLOCK FORM (lane/samplat, 2026-08-21): the surviving inefficiency
1508        // was WIDTH, not passes — one block per row left ~94% of the device idle for ~620us
1509        // per B=8 serve tick (5.9% of the tick, box4 nsys receipt). filter_stats_coop_f32
1510        // splits each row across 16 blocks with grid-synced bisection totals — same algorithm,
1511        // slice-partial f32 sums (accepted device-sampling class; sample-check arbitrates).
1512        // Admission: cooperative grid must co-reside (16*nrow blocks vs SM count).
1513        // MEMRA_FILTER_COOP=0 is the rollback seam to the single-block form.
1514        static COOP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1515        let coop_on =
1516            *COOP_ON.get_or_init(|| std::env::var("MEMRA_FILTER_COOP").as_deref() != Ok("0"));
1517        let (ni, nr, rs) = (n as i32, nrow as i32, row_stride as i64);
1518        if coop_on && 16 * nrow <= self.sm_count() as usize {
1519            let f = self.func("filter_stats_coop_f32");
1520            let mut ws = self.alloc_uninit::<f32>(nrow * (2 * 16 + 2))?;
1521            let cfg = LaunchConfig {
1522                grid_dim: (16, nrow as u32, 1),
1523                block_dim: (512, 1, 1),
1524                shared_mem_bytes: 0,
1525            };
1526            let __s_b = self.gpu.stream();
1527            let mut b = __s_b.launch_builder(&f);
1528            b.arg(x)
1529                .arg(&rs)
1530                .arg(rows)
1531                .arg(&mut *out_th)
1532                .arg(&mut *out_z)
1533                .arg(&mut *out_max)
1534                .arg(&mut ws)
1535                .arg(&ni)
1536                .arg(&nr)
1537                .arg(&temp)
1538                .arg(&top_k)
1539                .arg(&top_p)
1540                .arg(&min_p);
1541            unsafe {
1542                b.launch_cooperative(cfg)?;
1543            }
1544            return Ok(());
1545        }
1546        let f = self.func("filter_stats_f32");
1547        let cfg = LaunchConfig {
1548            grid_dim: (nrow as u32, 1, 1),
1549            block_dim: (1024, 1, 1),
1550            shared_mem_bytes: 0,
1551        };
1552        let __s_b = self.gpu.stream();
1553        let mut b = __s_b.launch_builder(&f);
1554        b.arg(x)
1555            .arg(&rs)
1556            .arg(rows)
1557            .arg(&mut *out_th)
1558            .arg(&mut *out_z)
1559            .arg(&mut *out_max)
1560            .arg(&ni)
1561            .arg(&nr)
1562            .arg(&temp)
1563            .arg(&top_k)
1564            .arg(&top_p)
1565            .arg(&min_p);
1566        unsafe {
1567            b.launch(cfg)?;
1568        }
1569        Ok(())
1570    }
1571
1572    /// out[pair] = filtered-softmax prob of ids[pair] in row rows[pair] (th/z per PAIR).
1573    #[allow(clippy::too_many_arguments)]
1574    pub fn softmax_gather_filtered(
1575        &self,
1576        x: &CudaSlice<f32>,
1577        row_stride: usize,
1578        ids: &CudaSlice<u32>,
1579        rows: &CudaSlice<i32>,
1580        th: &CudaSlice<f32>,
1581        z: &CudaSlice<f32>,
1582        out: &mut CudaSlice<f32>,
1583        n: usize,
1584        npair: usize,
1585        temp: f32,
1586    ) -> Result<(), Box<dyn std::error::Error>> {
1587        let f = self.func("softmax_gather_filtered_f32");
1588        let (ni, np, rs) = (n as i32, npair as i32, row_stride as i64);
1589        let cfg = LaunchConfig {
1590            grid_dim: (npair as u32, 1, 1),
1591            block_dim: (256, 1, 1),
1592            shared_mem_bytes: 0,
1593        };
1594        let __s_b = self.gpu.stream();
1595        let mut b = __s_b.launch_builder(&f);
1596        b.arg(x)
1597            .arg(&rs)
1598            .arg(ids)
1599            .arg(rows)
1600            .arg(th)
1601            .arg(z)
1602            .arg(&mut *out)
1603            .arg(&ni)
1604            .arg(&np)
1605            .arg(&temp);
1606        unsafe {
1607            b.launch(cfg)?;
1608        }
1609        Ok(())
1610    }
1611
1612    /// Filtered residual sample: token ~ norm(max(0, fp - fq)) with fp/fq the filtered softmaxes.
1613    #[allow(clippy::too_many_arguments)]
1614    pub fn residual_sample_filtered(
1615        &self,
1616        p: &CudaSlice<f32>,
1617        q: Option<&CudaSlice<f32>>,
1618        n: usize,
1619        temp: f32,
1620        seed: u64,
1621        stream_pos: u32,
1622        p_stats: (f32, f32, f32),
1623        q_stats: (f32, f32, f32),
1624        out_tok: &mut CudaSlice<u32>,
1625    ) -> Result<(), Box<dyn std::error::Error>> {
1626        let f = self.func("residual_sample_filtered_f32");
1627        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
1628        let has_q: i32 = q.is_some() as i32;
1629        let qbuf = q.unwrap_or(p);
1630        let (pm, pth, pz) = p_stats;
1631        let (qm, qth, qz) = q_stats;
1632        let cfg = LaunchConfig {
1633            grid_dim: (1, 1, 1),
1634            block_dim: (1024, 1, 1),
1635            shared_mem_bytes: 0,
1636        };
1637        let __s_b = self.gpu.stream();
1638        let mut b = __s_b.launch_builder(&f);
1639        b.arg(p)
1640            .arg(qbuf)
1641            .arg(&has_q)
1642            .arg(&ni)
1643            .arg(&temp)
1644            .arg(&slo)
1645            .arg(&shi)
1646            .arg(&stream_pos)
1647            .arg(&pm)
1648            .arg(&pth)
1649            .arg(&pz)
1650            .arg(&qm)
1651            .arg(&qth)
1652            .arg(&qz)
1653            .arg(&mut *out_tok);
1654        unsafe {
1655            b.launch(cfg)?;
1656        }
1657        Ok(())
1658    }
1659
1660    /// Sparse-q filtered residual sample: token ~ norm(max(0, fp - q)) where fp is the
1661    /// FILTERED softmax of `p` (stats from `filter_stats`) and q is a PROBABILITY vector
1662    /// supported on `cand_ids` (<=32 ids — the DFlash2 selector's candidate-set proposal;
1663    /// lane/dspark-sampled-admission-20260820). Same event semantics/Philox tag as
1664    /// `residual_sample_filtered` — one uniform per (seed, stream_pos).
1665    #[allow(clippy::too_many_arguments)]
1666    pub fn residual_sample_sparse_q(
1667        &self,
1668        p: &CudaSlice<f32>,
1669        cand_ids: &CudaSlice<u32>,
1670        q_probs: &CudaSlice<f32>,
1671        n_cand: usize,
1672        n: usize,
1673        temp: f32,
1674        seed: u64,
1675        stream_pos: u32,
1676        p_stats: (f32, f32, f32),
1677        out_tok: &mut CudaSlice<u32>,
1678    ) -> Result<(), Box<dyn std::error::Error>> {
1679        assert!(
1680            n_cand >= 1 && n_cand <= 32,
1681            "residual_sample_sparse_q supports 1..=32 candidates, got {n_cand}"
1682        );
1683        let f = self.func("residual_sample_sparse_q_f32");
1684        let (ni, nc) = (n as i32, n_cand as i32);
1685        let (slo, shi) = ((seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
1686        let (pm, pth, pz) = p_stats;
1687        let cfg = LaunchConfig {
1688            grid_dim: (1, 1, 1),
1689            block_dim: (1024, 1, 1),
1690            shared_mem_bytes: 0,
1691        };
1692        let __s_b = self.gpu.stream();
1693        let mut b = __s_b.launch_builder(&f);
1694        b.arg(p)
1695            .arg(cand_ids)
1696            .arg(q_probs)
1697            .arg(&nc)
1698            .arg(&ni)
1699            .arg(&temp)
1700            .arg(&slo)
1701            .arg(&shi)
1702            .arg(&stream_pos)
1703            .arg(&pm)
1704            .arg(&pth)
1705            .arg(&pz)
1706            .arg(&mut *out_tok);
1707        unsafe {
1708            b.launch(cfg)?;
1709        }
1710        Ok(())
1711    }
1712
1713    /// Gumbel-max draw from the FILTERED distribution (masked perturb; argmax after).
1714    #[allow(clippy::too_many_arguments)]
1715    pub fn gumbel_perturb_filtered(
1716        &self,
1717        x: &CudaSlice<f32>,
1718        y: &mut CudaSlice<f32>,
1719        n: usize,
1720        seed: u64,
1721        stream_pos: u32,
1722        temp: f32,
1723        row_max: f32,
1724        th: f32,
1725    ) -> Result<(), Box<dyn std::error::Error>> {
1726        let f = self.func("gumbel_perturb_filtered_f32");
1727        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
1728        let cfg = LaunchConfig {
1729            grid_dim: (n.div_ceil(256) as u32, 1, 1),
1730            block_dim: (256, 1, 1),
1731            shared_mem_bytes: 0,
1732        };
1733        let __s_b = self.gpu.stream();
1734        let mut b = __s_b.launch_builder(&f);
1735        b.arg(x)
1736            .arg(&mut *y)
1737            .arg(&ni)
1738            .arg(&slo)
1739            .arg(&shi)
1740            .arg(&stream_pos)
1741            .arg(&temp)
1742            .arg(&row_max)
1743            .arg(&th);
1744        unsafe {
1745            b.launch(cfg)?;
1746        }
1747        Ok(())
1748    }
1749
1750    /// Keskar penalties applied IN PLACE to a logits buffer: history token ids get
1751    /// rep-divided/multiplied + freq*count + presence subtracted. Symmetric p/q usage keeps
1752    /// filtered rejection sampling exact for the penalized target.
1753    #[allow(clippy::too_many_arguments)]
1754    pub fn penalize_logits(
1755        &self,
1756        x: &mut CudaSlice<f32>,
1757        hist: &CudaSlice<u32>,
1758        n_hist: usize,
1759        rep: f32,
1760        freq: f32,
1761        present: f32,
1762        n: usize,
1763    ) -> Result<(), Box<dyn std::error::Error>> {
1764        if n_hist == 0 {
1765            return Ok(());
1766        }
1767        let f = self.func("penalize_logits_f32");
1768        let (nh, ni) = (n_hist as i32, n as i32);
1769        let cfg = LaunchConfig {
1770            grid_dim: (n_hist.div_ceil(128) as u32, 1, 1),
1771            block_dim: (128, 1, 1),
1772            shared_mem_bytes: 0,
1773        };
1774        let __s_b = self.gpu.stream();
1775        let mut b = __s_b.launch_builder(&f);
1776        b.arg(&mut *x)
1777            .arg(hist)
1778            .arg(&nh)
1779            .arg(&rep)
1780            .arg(&freq)
1781            .arg(&present)
1782            .arg(&ni);
1783        unsafe {
1784            b.launch(cfg)?;
1785        }
1786        Ok(())
1787    }
1788
1789    /// Rows variant: penalize `nrow` contiguous rows of length n in one launch.
1790    #[allow(clippy::too_many_arguments)]
1791    pub fn penalize_logits_rows(
1792        &self,
1793        x: &mut CudaSlice<f32>,
1794        hist: &CudaSlice<u32>,
1795        n_hist: usize,
1796        rep: f32,
1797        freq: f32,
1798        present: f32,
1799        n: usize,
1800        nrow: usize,
1801    ) -> Result<(), Box<dyn std::error::Error>> {
1802        if n_hist == 0 || nrow == 0 {
1803            return Ok(());
1804        }
1805        let f = self.func("penalize_logits_rows_f32");
1806        let (nh, ni, nr) = (n_hist as i32, n as i32, nrow as i32);
1807        let cfg = LaunchConfig {
1808            grid_dim: (n_hist.div_ceil(128) as u32, nrow as u32, 1),
1809            block_dim: (128, 1, 1),
1810            shared_mem_bytes: 0,
1811        };
1812        let __s_b = self.gpu.stream();
1813        let mut b = __s_b.launch_builder(&f);
1814        b.arg(&mut *x)
1815            .arg(hist)
1816            .arg(&nh)
1817            .arg(&rep)
1818            .arg(&freq)
1819            .arg(&present)
1820            .arg(&ni)
1821            .arg(&nr);
1822        unsafe {
1823            b.launch(cfg)?;
1824        }
1825        Ok(())
1826    }
1827
1828    /// ROW-INCREMENTAL penalties (dspark penalized-sampled admission): row r of `x`
1829    /// penalizes over the last `min(win, n_hist0 + r)` entries of `hist[..n_hist0 + r]`,
1830    /// where `hist` = [session window (n_hist0) ++ per-row drafted tokens (nrow-1)]. This
1831    /// is the within-round evolving penalty state block drafting needs: verify row r's
1832    /// target is penalized by every token committed before it INCLUDING same-round
1833    /// accepts — `penalize_logits_rows` (one shared window) is the frozen-window
1834    /// approximation this exists to replace on the dspark route.
1835    #[allow(clippy::too_many_arguments)]
1836    pub fn penalize_logits_rows_inc(
1837        &self,
1838        x: &mut CudaSlice<f32>,
1839        hist: &CudaSlice<u32>,
1840        n_hist0: usize,
1841        rep: f32,
1842        freq: f32,
1843        present: f32,
1844        n: usize,
1845        nrow: usize,
1846        win: usize,
1847    ) -> Result<(), Box<dyn std::error::Error>> {
1848        if nrow == 0 || win == 0 || (n_hist0 == 0 && nrow == 1) {
1849            return Ok(());
1850        }
1851        debug_assert!(
1852            hist.len() >= n_hist0 + nrow - 1,
1853            "rows-inc hist must carry n_hist0 + nrow - 1 ids"
1854        );
1855        let f = self.func("penalize_logits_rows_inc_f32");
1856        let max_len = win.min(n_hist0 + nrow - 1).max(1);
1857        let (nh, ni, nr, wi) = (n_hist0 as i32, n as i32, nrow as i32, win as i32);
1858        let cfg = LaunchConfig {
1859            grid_dim: (max_len.div_ceil(128) as u32, nrow as u32, 1),
1860            block_dim: (128, 1, 1),
1861            shared_mem_bytes: 0,
1862        };
1863        let __s_b = self.gpu.stream();
1864        let mut b = __s_b.launch_builder(&f);
1865        b.arg(&mut *x)
1866            .arg(hist)
1867            .arg(&nh)
1868            .arg(&rep)
1869            .arg(&freq)
1870            .arg(&present)
1871            .arg(&ni)
1872            .arg(&nr)
1873            .arg(&wi);
1874        unsafe {
1875            b.launch(cfg)?;
1876        }
1877        Ok(())
1878    }
1879
1880    /// WEIGHT PREFETCH (SOTA item 3, 2026-07-13, DEFAULT ON): during a bandwidth-idle
1881    /// window (the fa launch reads KV, not weights) prefetch the NEXT matvec's
1882    /// decode-plane bytes into L2 so it reads L2-warm. Value-free scheduling op — same
1883    /// class as prefetch_l2 (numerics untouched by construction). Wired only where it
1884    /// measured positive: the E4B dc attn arm (+0.65%). 26B (flat — MoE ffn dominates),
1885    /// 31B (−0.2% — decode at the DRAM wall) and the ffn gate/up cascade (−1% — 29MB/layer
1886    /// floods the fill path) all probed and NOT wired. MEMRA_WPF=0 rollback seam.
1887    pub fn wpf_level() -> u32 {
1888        static ON: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
1889        *ON.get_or_init(|| {
1890            std::env::var("MEMRA_WPF")
1891                .ok()
1892                .and_then(|v| v.parse().ok())
1893                .unwrap_or(1)
1894        })
1895    }
1896
1897    /// PDL launch arm (SOTA item 2, 2026-07-13, DEFAULT ON): the six MEMRA_PDL_ENTRY glue
1898    /// kernels launch through cuLaunchKernelEx with PROGRAMMATIC_STREAM_SERIALIZATION — the
1899    /// grid launches while the predecessor drains (~120ns/kernel back, pdl_probe), the
1900    /// kernels' entry grid-dep sync restores read order (SASS-audited: ACQBULK precedes
1901    /// every LDG in all six). Valid windows: E4B +1.0-1.2% (128 AND 384-tok gens);
1902    /// 26B/31B/qwen flat no-harm. Battery: kernel-check GREEN, run-gen tokens IDENTICAL x3
1903    /// gemma, spec 64/64 E4B K=1/4/8 + 26B/31B K=4 + qwen PASS. Works eager AND under
1904    /// capture (capture encodes native programmatic edges — the post-capture edge-REWRITE
1905    /// arm died: engine graphs hold cuMemAllocAsync alloc nodes, edge edits on those return
1906    /// CUDA_ERROR_NOT_SUPPORTED). MEMRA_PDL=0 rollback seam.
1907    /// See the `verify_exact` field. Scoped by the dflash round around its t=16 verify.
1908    pub fn set_verify_exact(&self, on: bool) {
1909        self.verify_exact
1910            .store(on, std::sync::atomic::Ordering::Relaxed);
1911    }
1912    pub(crate) fn verify_exact_on(&self) -> bool {
1913        self.verify_exact.load(std::sync::atomic::Ordering::Relaxed)
1914    }
1915
1916    /// m=1 norm+rope+append fold seam (2026-07-23): MEMRA_QKV_APPEND=0 reverts to the
1917    /// fused-norm-rope + standalone-append pair (the exact-oracle bisect arm).
1918    pub fn qkv_append_on() -> bool {
1919        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1920        *ON.get_or_init(|| {
1921            std::env::var("MEMRA_QKV_APPEND")
1922                .map(|v| v != "0")
1923                .unwrap_or(true)
1924        })
1925    }
1926
1927    /// PDL wave-B1a seam: the four dense-glue kernels (rms_norm_f32, add_rms_norm_f32,
1928    /// add_scale_rms_norm_q8_1, quantize_q8_1). MEMRA_PDL_WB=0 reverts alone.
1929    pub fn pdl_wb_on() -> bool {
1930        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1931        *ON.get_or_init(|| {
1932            std::env::var("MEMRA_PDL_WB")
1933                .map(|v| v != "0")
1934                .unwrap_or(true)
1935        })
1936    }
1937
1938    /// Trunk-kernels norm ILP seam (lane/dspark-trunk-kernels-20260820): the T-row verify
1939    /// norms (rms_norm_f32 / add_rms_norm_f32 at grid=T, block=256) are serial-latency
1940    /// chains — 20 strided scalar load->fma rounds measured 11.8-12.2us/inst (nsys-B verify
1941    /// scope: 130 inst/rd = 1.51 ms/rd). The `_v2` twins unroll the element loop 4-deep
1942    /// (independent loads in flight; SAME per-thread element order into ONE accumulator,
1943    /// reduce VERBATIM) — BIT-IDENTICAL per row at every (ncols, blockDim).
1944    /// MEMRA_NORM_ILP=0 reverts to the v1 kernels alone.
1945    pub fn norm_ilp_on() -> bool {
1946        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1947        *ON.get_or_init(|| {
1948            std::env::var("MEMRA_NORM_ILP")
1949                .map(|v| v != "0")
1950                .unwrap_or(true)
1951        })
1952    }
1953
1954    /// Trunk-kernels FFN dual seam (lane/dspark-trunk-kernels-20260820): the qwen35
1955    /// t-parallel verify FFN pair rides the PROVEN dual gate+up doors
1956    /// (`matmul_decode_exact_dual_pre` + `silu_mul_scaled_q8_1`, the q27 verify shape —
1957    /// bit-identical per (tensor,token,row), kernel-check-pinned, MEMRA_SPEC_DUAL_T
1958    /// receipts) instead of two singles + silu_mul + a standalone quantize. The doors
1959    /// existed but the qwen35 body never called them (nsys-B verify scope: gate+up singles
1960    /// = 107 launches/rd at grid 4352). MEMRA_TK_FFN_DUAL=0 reverts to the singles chain.
1961    pub fn tk_ffn_dual_on() -> bool {
1962        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1963        *ON.get_or_init(|| {
1964            std::env::var("MEMRA_TK_FFN_DUAL")
1965                .map(|v| v != "0")
1966                .unwrap_or(true)
1967        })
1968    }
1969
1970    /// PDL wave-A seam: the mmvq matvec PDL launches only (the six glue kernels keep
1971    /// their own MEMRA_PDL master seam). MEMRA_PDL_MMVQ=0 reverts wave-A alone — the
1972    /// per-model no-harm bisect knob.
1973    pub fn pdl_mmvq_on() -> bool {
1974        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1975        *ON.get_or_init(|| {
1976            std::env::var("MEMRA_PDL_MMVQ")
1977                .map(|v| v != "0")
1978                .unwrap_or(true)
1979        })
1980    }
1981
1982    pub fn pdl_on() -> bool {
1983        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1984        *ON.get_or_init(|| std::env::var("MEMRA_PDL").map(|v| v != "0").unwrap_or(true))
1985    }
1986
1987    /// PDL wave-B seam (gap-diagnosis arc, GAP-DIAGNOSIS.md verdict 8): the gemma
1988    /// NVFP4mix decode chain's hot kernels — nvfp4 fused2/mr2 and the q8_0 `_rp`
1989    /// singles — join the wave-A launch class. Scheduling-only (the entry macro waits
1990    /// on the producer before any read), bit-identical by construction.
1991    /// MEMRA_PDL_NVFP4=0 reverts wave-B alone.
1992    pub fn pdl_nvfp4q8_on() -> bool {
1993        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1994        *ON.get_or_init(|| {
1995            std::env::var("MEMRA_PDL_NVFP4")
1996                .map(|v| v != "0")
1997                .unwrap_or(true)
1998        })
1999    }
2000
2001    /// Raw CUfunction for a PDL-attributed launch: the SAME kernels.fatbin loaded once more
2002    /// through the raw driver API (cudarc hides its CUfunction handles; a duplicate module
2003    /// of tiny glue kernels is free). Resolved lazily per name, cached process-wide.
2004    /// Fused t=1 q4_0 mr policy: env MEMRA_Q40_MR wins (1/2); else the per-model
2005    /// FUSED_MR1_DEFAULT (dense gemma = mr1, MoE = mr2 — see the static's doc).
2006    fn q40_mr1_on() -> bool {
2007        static Q40MR: std::sync::OnceLock<Option<u32>> = std::sync::OnceLock::new();
2008        match *Q40MR.get_or_init(|| {
2009            std::env::var("MEMRA_Q40_MR")
2010                .ok()
2011                .and_then(|v| v.parse().ok())
2012        }) {
2013            Some(v) => v == 1,
2014            None => crate::FUSED_MR1_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
2015        }
2016    }
2017
2018    /// PDL wave-B2: flash-module PDL functions. `g` selects the kf8vf8 flavor — the
2019    /// caller MUST pass the SAME flavor its builder launch would resolve (fa_func/func_g
2020    /// mirror); the flavors differ semantically (KV byte formats), a wrong-module launch
2021    /// writes wrong bytes silently.
2022    fn pdl_func_flash(
2023        &self,
2024        g: bool,
2025        name: &'static str,
2026    ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
2027        use cudarc::driver::sys as cu;
2028        // PER-CONTEXT caches (M1-PP2 cross-device fix, 8x box 2026-08-02): CUmodule and
2029        // CUfunction handles are CONTEXT-scoped, and a remote-stage Engine
2030        // (MEMRA_PP_DEVICES=a,b) lives in the other device's primary context. The old
2031        // process-wide OnceLock cache handed stage 1 the dev-a handles, so every stage-1
2032        // launch_pdl* died CUDA_ERROR_INVALID_HANDLE. Key module + function caches by
2033        // this engine's CUcontext; single-context runs behave exactly as before.
2034        static MODS: std::sync::Mutex<Option<std::collections::HashMap<(usize, bool), usize>>> =
2035            std::sync::Mutex::new(None);
2036        static FNS: std::sync::Mutex<
2037            Option<std::collections::HashMap<(usize, bool, &'static str), usize>>,
2038        > = std::sync::Mutex::new(None);
2039        let ctx_key = self.ctx().cu_ctx() as usize;
2040        if let Some(&f) = FNS
2041            .lock()
2042            .unwrap()
2043            .get_or_insert_with(Default::default)
2044            .get(&(ctx_key, g, name))
2045        {
2046            return Ok(f as cu::CUfunction);
2047        }
2048        let module = {
2049            let mut mods = MODS.lock().unwrap();
2050            let map = mods.get_or_insert_with(Default::default);
2051            match map.get(&(ctx_key, g)) {
2052                Some(&m) => m,
2053                None => {
2054                    let m = self.pdl_load_module_in_ctx(if g {
2055                        FLASH_FATBIN_KF8VF8
2056                    } else {
2057                        FLASH_FATBIN
2058                    })?;
2059                    map.insert((ctx_key, g), m);
2060                    m
2061                }
2062            }
2063        };
2064        let cname = std::ffi::CString::new(name)?;
2065        let mut f: cu::CUfunction = std::ptr::null_mut();
2066        let r = unsafe { cu::cuModuleGetFunction(&mut f, module as cu::CUmodule, cname.as_ptr()) };
2067        if r != cu::CUresult::CUDA_SUCCESS {
2068            return Err(format!("pdl_func_flash {name} (g={g}): {r:?}").into());
2069        }
2070        FNS.lock()
2071            .unwrap()
2072            .get_or_insert_with(Default::default)
2073            .insert((ctx_key, g, name), f as usize);
2074        Ok(f)
2075    }
2076
2077    /// Load a fatbin as a raw CUmodule IN THIS ENGINE'S CONTEXT. `cuModuleLoadData` binds
2078    /// the module to the thread's CURRENT context — a remote-stage engine must not
2079    /// inherit the primary's (the INVALID_HANDLE class above). Restores the caller's
2080    /// current context before returning.
2081    fn pdl_load_module_in_ctx(&self, bytes: &[u8]) -> Result<usize, Box<dyn std::error::Error>> {
2082        use cudarc::driver::sys as cu;
2083        let mut prev: cu::CUcontext = std::ptr::null_mut();
2084        unsafe {
2085            cu::cuCtxGetCurrent(&mut prev).result()?;
2086        }
2087        self.ctx().bind_to_thread()?;
2088        let mut m: cu::CUmodule = std::ptr::null_mut();
2089        let r = unsafe { cu::cuModuleLoadData(&mut m, bytes.as_ptr() as *const std::ffi::c_void) };
2090        let restore = if prev.is_null() {
2091            cu::CUresult::CUDA_SUCCESS
2092        } else {
2093            unsafe { cu::cuCtxSetCurrent(prev) }
2094        };
2095        if r != cu::CUresult::CUDA_SUCCESS {
2096            return Err(format!("pdl module load: {r:?}").into());
2097        }
2098        if restore != cu::CUresult::CUDA_SUCCESS {
2099            return Err(format!("pdl module load: ctx restore {restore:?}").into());
2100        }
2101        Ok(m as usize)
2102    }
2103
2104    /// Raw CUfunction for prebuilt-args dispatch experiments (M4 probe): same duplicate
2105    /// raw-module loading as the PDL path, WITHOUT the PDL launch attribute.
2106    pub fn raw_kernel_function(
2107        &self,
2108        name: &'static str,
2109    ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
2110        self.pdl_func(name)
2111    }
2112
2113    fn pdl_func(
2114        &self,
2115        name: &'static str,
2116    ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
2117        use cudarc::driver::sys as cu;
2118        // PER-CONTEXT caches — same M1-PP2 cross-device fix as pdl_func_flash (handles
2119        // are context-scoped; key everything by this engine's CUcontext).
2120        static MODULES: std::sync::Mutex<Option<std::collections::HashMap<usize, usize>>> =
2121            std::sync::Mutex::new(None);
2122        // PDL wave-A: the mmvq kernels live in the qmatvec fatbin, not kernels.cu — second
2123        // duplicate module, loaded lazily on the first kernels-module miss.
2124        static QMODULES: std::sync::Mutex<Option<std::collections::HashMap<usize, usize>>> =
2125            std::sync::Mutex::new(None);
2126        static FNS: std::sync::Mutex<
2127            Option<std::collections::HashMap<(usize, &'static str), usize>>,
2128        > = std::sync::Mutex::new(None);
2129        let ctx_key = self.ctx().cu_ctx() as usize;
2130        if let Some(&f) = FNS
2131            .lock()
2132            .unwrap()
2133            .get_or_insert_with(Default::default)
2134            .get(&(ctx_key, name))
2135        {
2136            return Ok(f as cu::CUfunction);
2137        }
2138        let module = {
2139            let mut mods = MODULES.lock().unwrap();
2140            let map = mods.get_or_insert_with(Default::default);
2141            match map.get(&ctx_key) {
2142                Some(&m) => m,
2143                None => {
2144                    let m = self.pdl_load_module_in_ctx(FATBIN)?;
2145                    map.insert(ctx_key, m);
2146                    m
2147                }
2148            }
2149        };
2150        let cname = std::ffi::CString::new(name)?;
2151        let mut f: cu::CUfunction = std::ptr::null_mut();
2152        let mut r =
2153            unsafe { cu::cuModuleGetFunction(&mut f, module as cu::CUmodule, cname.as_ptr()) };
2154        if r == cu::CUresult::CUDA_ERROR_NOT_FOUND {
2155            let qmodule = {
2156                let mut mods = QMODULES.lock().unwrap();
2157                let map = mods.get_or_insert_with(Default::default);
2158                match map.get(&ctx_key) {
2159                    Some(&m) => m,
2160                    None => {
2161                        let m = self.pdl_load_module_in_ctx(QMATVEC_FATBIN)?;
2162                        map.insert(ctx_key, m);
2163                        m
2164                    }
2165                }
2166            };
2167            r = unsafe { cu::cuModuleGetFunction(&mut f, qmodule as cu::CUmodule, cname.as_ptr()) };
2168        }
2169        if r != cu::CUresult::CUDA_SUCCESS {
2170            return Err(format!("pdl_func {name}: {r:?}").into());
2171        }
2172        FNS.lock()
2173            .unwrap()
2174            .get_or_insert_with(Default::default)
2175            .insert((ctx_key, name), f as usize);
2176        Ok(f)
2177    }
2178
2179    /// cuLaunchKernelEx with CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION on the
2180    /// compute stream. ONLY legal for kernels whose entry carries MEMRA_PDL_ENTRY.
2181    ///
2182    /// # Safety
2183    /// `params` must match the kernel's exact parameter list (order, types, count) —
2184    /// a mismatch corrupts the launch silently.
2185    /// Flash-module twin of `launch_pdl` — `g` picks the kf8vf8 flavor (must mirror the
2186    /// builder path's fa_func/func_g choice exactly).
2187    ///
2188    /// # Safety
2189    /// Same contract as `launch_pdl`.
2190    unsafe fn launch_pdl_flash(
2191        &self,
2192        g: bool,
2193        name: &'static str,
2194        grid: (u32, u32, u32),
2195        block: (u32, u32, u32),
2196        smem: u32,
2197        params: &mut [*mut std::ffi::c_void],
2198    ) -> Result<(), Box<dyn std::error::Error>> {
2199        use cudarc::driver::sys as cu;
2200        let f = self.pdl_func_flash(g, name)?;
2201        if smem > 0 {
2202            // mirror the builder path's opt-in ceiling (idempotent host-side set).
2203            let r =
2204                unsafe {
2205                    cu::cuFuncSetAttribute(f,
2206                cu::CUfunction_attribute_enum::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
2207                smem as i32)
2208                };
2209            if r != cu::CUresult::CUDA_SUCCESS {
2210                return Err(format!("pdl smem attr {name}: {r:?}").into());
2211            }
2212        }
2213        let mut attr = cu::CUlaunchAttribute {
2214            id: cu::CUlaunchAttributeID::CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION,
2215            pad: [0; 4],
2216            value: cu::CUlaunchAttributeValue {
2217                programmaticStreamSerializationAllowed: 1,
2218            },
2219        };
2220        let cfg = cu::CUlaunchConfig {
2221            gridDimX: grid.0,
2222            gridDimY: grid.1,
2223            gridDimZ: grid.2,
2224            blockDimX: block.0,
2225            blockDimY: block.1,
2226            blockDimZ: block.2,
2227            sharedMemBytes: smem,
2228            hStream: self.gpu.stream().cu_stream(),
2229            attrs: &mut attr,
2230            numAttrs: 1,
2231        };
2232        let r = unsafe { cu::cuLaunchKernelEx(&cfg, f, params.as_mut_ptr(), std::ptr::null_mut()) };
2233        if r != cu::CUresult::CUDA_SUCCESS {
2234            return Err(format!("launch_pdl_flash {name}: {r:?}").into());
2235        }
2236        Ok(())
2237    }
2238
2239    unsafe fn launch_pdl(
2240        &self,
2241        name: &'static str,
2242        grid: (u32, u32, u32),
2243        block: (u32, u32, u32),
2244        params: &mut [*mut std::ffi::c_void],
2245    ) -> Result<(), Box<dyn std::error::Error>> {
2246        use cudarc::driver::sys as cu;
2247        let f = self.pdl_func(name)?;
2248        let mut attr = cu::CUlaunchAttribute {
2249            id: cu::CUlaunchAttributeID::CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION,
2250            pad: [0; 4],
2251            value: cu::CUlaunchAttributeValue {
2252                programmaticStreamSerializationAllowed: 1,
2253            },
2254        };
2255        let cfg = cu::CUlaunchConfig {
2256            gridDimX: grid.0,
2257            gridDimY: grid.1,
2258            gridDimZ: grid.2,
2259            blockDimX: block.0,
2260            blockDimY: block.1,
2261            blockDimZ: block.2,
2262            sharedMemBytes: 0,
2263            hStream: self.gpu.stream().cu_stream(),
2264            attrs: &mut attr,
2265            numAttrs: 1,
2266        };
2267        let r = unsafe { cu::cuLaunchKernelEx(&cfg, f, params.as_mut_ptr(), std::ptr::null_mut()) };
2268        if r != cu::CUresult::CUDA_SUCCESS {
2269            return Err(format!("launch_pdl {name}: {r:?}").into());
2270        }
2271        Ok(())
2272    }
2273
2274    /// L2-prefetch a quant weight's DECODE plane (the rp4 split-plane mirror when present —
2275    /// that is what the m<=8 dispatch reads — else the raw block bytes). No-op on float arms.
2276    pub fn prefetch_weight_l2(
2277        &self,
2278        w: &crate::model::GpuTensor,
2279    ) -> Result<(), Box<dyn std::error::Error>> {
2280        if let crate::model::GpuTensor::Quant { bytes, rp4, .. } = w {
2281            let p = rp4.as_ref().unwrap_or(bytes);
2282            self.prefetch_l2(p, p.len())?;
2283        }
2284        Ok(())
2285    }
2286
2287    /// DSpark markov chain ops (dflash lane): gather one bf16 row of a [V, rank] table
2288    /// by the DEVICE token id at tok[idx] into f32.
2289    pub fn gather_row_bf16(
2290        &self,
2291        table: &CudaSlice<u8>,
2292        tok: &CudaSlice<u32>,
2293        idx: usize,
2294        dst: &mut CudaSlice<f32>,
2295        ncols: usize,
2296    ) -> Result<(), Box<dyn std::error::Error>> {
2297        let f = self.func("gather_row_bf16_f32");
2298        let cfg = LaunchConfig {
2299            grid_dim: (ncols.div_ceil(256) as u32, 1, 1),
2300            block_dim: (256, 1, 1),
2301            shared_mem_bytes: 0,
2302        };
2303        let (nc, ix) = (ncols as i32, idx as i32);
2304        let __s_b = self.gpu.stream();
2305        let mut b = __s_b.launch_builder(&f);
2306        b.arg(table).arg(tok).arg(&ix).arg(dst).arg(&nc);
2307        unsafe {
2308            b.launch(cfg)?;
2309        }
2310        Ok(())
2311    }
2312
2313    /// DFlash2 grouped dynamic causal conv (dflash lane, DFLASH2-EVAL-20260820.md):
2314    /// out[p,c] = sum_{o<ksize, o<=p} (base[half][o][c] + dyn[p][half][o][group(c)])
2315    /// * x[p-o][c]. `dyn_` is the kernel_projection GEMM output [rows, 2*ksize*groups];
2316    /// `base` is base_kernel [2, ksize, hidden] flattened; `half` picks prepare(0) /
2317    /// finish(1).
2318    #[allow(clippy::too_many_arguments)]
2319    pub fn dflash2_dynconv(
2320        &self,
2321        x: &CudaSlice<f32>,
2322        dyn_: &CudaSlice<f32>,
2323        base: &CudaSlice<f32>,
2324        out: &mut CudaSlice<f32>,
2325        rows: usize,
2326        hidden: usize,
2327        group_size: usize,
2328        ksize: usize,
2329        half: usize,
2330    ) -> Result<(), Box<dyn std::error::Error>> {
2331        assert_eq!(hidden % group_size, 0, "hidden % group_size != 0");
2332        let f = self.func("dflash2_dynconv_f32");
2333        let n = rows * hidden;
2334        let cfg = LaunchConfig {
2335            grid_dim: (n.div_ceil(256) as u32, 1, 1),
2336            block_dim: (256, 1, 1),
2337            shared_mem_bytes: 0,
2338        };
2339        let (ri, hi, gi, ki, hf) = (
2340            rows as i32,
2341            hidden as i32,
2342            group_size as i32,
2343            ksize as i32,
2344            half as i32,
2345        );
2346        let __s_b = self.gpu.stream();
2347        let mut b = __s_b.launch_builder(&f);
2348        b.arg(x)
2349            .arg(dyn_)
2350            .arg(base)
2351            .arg(out)
2352            .arg(&ri)
2353            .arg(&hi)
2354            .arg(&gi)
2355            .arg(&ki)
2356            .arg(&hf);
2357        unsafe {
2358            b.launch(cfg)?;
2359        }
2360        Ok(())
2361    }
2362
2363    /// Per-row top-k (k <= 32) over a [n_rows, n_cols] logits matrix (DFlash2
2364    /// candidate selector). Returns (values [n_rows, k], column indices [n_rows, k]),
2365    /// value-descending, ties to the lower index.
2366    pub fn topk_rows(
2367        &self,
2368        logits: &CudaSlice<f32>,
2369        n_rows: usize,
2370        n_cols: usize,
2371        k: usize,
2372    ) -> Result<(CudaSlice<f32>, CudaSlice<u32>), Box<dyn std::error::Error>> {
2373        assert!(k <= 32 && k >= 1, "topk_rows supports 1..=32, got {k}");
2374        assert!(k <= n_cols, "topk_rows: k {k} > n_cols {n_cols}");
2375        let f = self.func("topk_rows_f32");
2376        let nth = 256usize;
2377        let mut vals = self.uninit(n_rows * k)?;
2378        let mut idxs = self.gpu.stream().alloc_zeros::<u32>(n_rows * k)?;
2379        let cfg = LaunchConfig {
2380            grid_dim: (n_rows as u32, 1, 1),
2381            block_dim: (nth as u32, 1, 1),
2382            shared_mem_bytes: (nth * k * 8) as u32,
2383        };
2384        let (nr, nc, ki) = (n_rows as i32, n_cols as i32, k as i32);
2385        let __s_b = self.gpu.stream();
2386        let mut b = __s_b.launch_builder(&f);
2387        b.arg(logits)
2388            .arg(&nr)
2389            .arg(&nc)
2390            .arg(&ki)
2391            .arg(&mut vals)
2392            .arg(&mut idxs);
2393        unsafe {
2394            b.launch(cfg)?;
2395        }
2396        Ok((vals, idxs))
2397    }
2398
2399    /// logits[row_off .. row_off+n] += bias[0..n] (in place, one row).
2400    pub fn add_row_inplace(
2401        &self,
2402        logits: &mut CudaSlice<f32>,
2403        bias: &CudaSlice<f32>,
2404        n: usize,
2405        row_off: usize,
2406    ) -> Result<(), Box<dyn std::error::Error>> {
2407        let f = self.func("add_row_inplace_f32");
2408        let cfg = LaunchConfig {
2409            grid_dim: (n.div_ceil(256) as u32, 1, 1),
2410            block_dim: (256, 1, 1),
2411            shared_mem_bytes: 0,
2412        };
2413        let (ni, off) = (n as i32, row_off as i64);
2414        let __s_b = self.gpu.stream();
2415        let mut b = __s_b.launch_builder(&f);
2416        b.arg(logits).arg(bias).arg(&ni).arg(&off);
2417        unsafe {
2418            b.launch(cfg)?;
2419        }
2420        Ok(())
2421    }
2422
2423    /// L2 prefetch of a device byte range (latency-hiding arc; value-free scheduling op).
2424    pub fn prefetch_l2(
2425        &self,
2426        p: &CudaSlice<u8>,
2427        n: usize,
2428    ) -> Result<(), Box<dyn std::error::Error>> {
2429        let f = self.func("prefetch_l2_bytes");
2430        let lines = n.div_ceil(128);
2431        let ni = n as i64;
2432        let cfg = LaunchConfig {
2433            grid_dim: (lines.div_ceil(256) as u32, 1, 1),
2434            block_dim: (256, 1, 1),
2435            shared_mem_bytes: 0,
2436        };
2437        let __s_b = self.gpu.stream();
2438        let mut b = __s_b.launch_builder(&f);
2439        b.arg(p).arg(&ni);
2440        unsafe {
2441            b.launch(cfg)?;
2442        }
2443        Ok(())
2444    }
2445
2446    /// MoE router GEMV (MEMRA_ROUTER_KERNEL): deterministic warp-per-(expert,token) f32 dot.
2447    /// Different FP order than the cuBLAS path it replaces — battery-gated numeric config.
2448    pub fn router_gemv(
2449        &self,
2450        w: &CudaSlice<f32>,
2451        x: &CudaSlice<f32>,
2452        n_embd: usize,
2453        n_experts: usize,
2454        t: usize,
2455    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2456        // float4 v2 probed 2026-07-14: +0.25% but flips near-tie routing (new FP order,
2457        // stream differs) — too small to justify a numeric config change; deleted.
2458        // w8 twin (2026-07-31): on the 132-SM H100 the lone-warp form is 14.8% of the q35
2459        // decode step (latency-bound) — the calculus flipped. MEMRA_ROUTER_V2=0 reverts to
2460        // the warp form (rollback seam; new FP order, battery-arbitrated per model).
2461        let w8 = match std::env::var("MEMRA_ROUTER_V2").as_deref() {
2462            Ok("0") => false,
2463            Ok(_) => true,
2464            Err(_) => ROUTER_W8_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
2465        };
2466        // FAST-ROUTER batch twin (lane/fast-router, 2026-08-02): at prefill m the per-(e,tok)
2467        // w8 form re-streams both operand rows per output (GEMV program at GEMM shape — the
2468        // concat-prime exactness fix paid -10% q35 board-2048 prefill through it). The batch
2469        // twin (8x8 expert-x-token register tile) is BIT-IDENTICAL per row (same k order,
2470        // same tree, same fold — kernel-check sweeps m=1..2048 on real router weights), so
2471        // the crossover is pure perf, not a numeric config. MIN_T from the on-box sweep
2472        // (research/fast-router-20260802/crossover-router*.jsonl); decode t=1 and small-t
2473        // spec verify keep the plain w8 form. MEMRA_ROUTER_BATCH=0: rollback seam
2474        // (perf-only, bits equal).
2475        let batch = w8 && t >= ROUTER_BATCH_MIN_T && router_batch_on();
2476        self.router_gemv_form(w, x, n_embd, n_experts, t, w8, batch)
2477    }
2478
2479    /// Form-explicit router GEMV launch (kernel-check bit-identity gate + crossover bench
2480    /// force both forms; `batch` requires `w8`).
2481    pub fn router_gemv_form(
2482        &self,
2483        w: &CudaSlice<f32>,
2484        x: &CudaSlice<f32>,
2485        n_embd: usize,
2486        n_experts: usize,
2487        t: usize,
2488        w8: bool,
2489        batch: bool,
2490    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2491        debug_assert!(!batch || w8, "batch twin exists for the w8 form only");
2492        let mut y = self.alloc_uninit::<f32>(t * n_experts)?;
2493        let f = if batch {
2494            self.func("router_gemv_f32_w8_batch")
2495        } else if w8 {
2496            self.func("router_gemv_f32_w8")
2497        } else {
2498            self.func("router_gemv_f32")
2499        };
2500        let (ne, nx, ti) = (n_embd as i32, n_experts as i32, t as i32);
2501        let cfg = if batch {
2502            LaunchConfig {
2503                grid_dim: (n_experts.div_ceil(8) as u32, t.div_ceil(8) as u32, 1),
2504                block_dim: (32, 8, 1),
2505                shared_mem_bytes: 0,
2506            }
2507        } else {
2508            LaunchConfig {
2509                grid_dim: (n_experts as u32, t as u32, 1),
2510                block_dim: (32, if w8 { 8 } else { 1 }, 1),
2511                shared_mem_bytes: 0,
2512            }
2513        };
2514        let __s_b = self.gpu.stream();
2515        let mut b = __s_b.launch_builder(&f);
2516        b.arg(w).arg(x).arg(&mut y).arg(&ne).arg(&nx).arg(&ti);
2517        unsafe {
2518            b.launch(cfg)?;
2519        }
2520        Ok(y)
2521    }
2522
2523    /// `router_gemv` (decode form selection) writing into a caller-owned [t*n_experts]
2524    /// buffer — token-graph alloc-free.
2525    pub fn router_gemv_into(
2526        &self,
2527        w: &CudaSlice<f32>,
2528        x: &CudaSlice<f32>,
2529        y: &mut CudaSlice<f32>,
2530        n_embd: usize,
2531        n_experts: usize,
2532        t: usize,
2533    ) -> Result<(), Box<dyn std::error::Error>> {
2534        if y.len() < t * n_experts {
2535            return Err("router_gemv_into output too small".into());
2536        }
2537        let w8 = match std::env::var("MEMRA_ROUTER_V2").as_deref() {
2538            Ok("0") => false,
2539            Ok(_) => true,
2540            Err(_) => ROUTER_W8_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
2541        };
2542        let f = if w8 {
2543            self.func("router_gemv_f32_w8")
2544        } else {
2545            self.func("router_gemv_f32")
2546        };
2547        let (ne, nx, ti) = (n_embd as i32, n_experts as i32, t as i32);
2548        let cfg = LaunchConfig {
2549            grid_dim: (n_experts as u32, t as u32, 1),
2550            block_dim: (32, if w8 { 8 } else { 1 }, 1),
2551            shared_mem_bytes: 0,
2552        };
2553        let __s_b = self.gpu.stream();
2554        let mut b = __s_b.launch_builder(&f);
2555        b.arg(w).arg(x).arg(&mut *y).arg(&ne).arg(&nx).arg(&ti);
2556        unsafe {
2557            b.launch(cfg)?;
2558        }
2559        Ok(())
2560    }
2561
2562    /// f32 row permute: dst[idx[i], :] = src[i, :] (grouped-GEMM CSR -> pair-id reorder).
2563    pub fn rows_permute(
2564        &self,
2565        src: &CudaSlice<f32>,
2566        idx: &CudaSlice<i32>,
2567        nrows: usize,
2568        ncols: usize,
2569    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2570        let mut dst = self.alloc_uninit::<f32>(nrows * ncols)?;
2571        let f = self.func("rows_permute_f32");
2572        let (nc, nr) = (ncols as i32, nrows as i32);
2573        let cfg = LaunchConfig {
2574            grid_dim: (nrows as u32, 1, 1),
2575            block_dim: (256, 1, 1),
2576            shared_mem_bytes: 0,
2577        };
2578        let __s_b = self.gpu.stream();
2579        let mut b = __s_b.launch_builder(&f);
2580        b.arg(src).arg(idx).arg(&mut dst).arg(&nc).arg(&nr);
2581        unsafe {
2582            b.launch(cfg)?;
2583        }
2584        Ok(dst)
2585    }
2586
2587    /// shexp gate fused dot: g[tok] = sigmoid(dot(x[tok,:], w)) — replaces the per-layer
2588    /// cuBLASLt m=1 GEMM + separate sigmoid launch on the qwen35moe decode path (the
2589    /// splitKreduce x40/step dig, 2026-07-31). One fold order for every t, so the t=1
2590    /// decode chain and the small-t spec-verify chain match per row by construction.
2591    pub fn sigmoid_dot_rows(
2592        &self,
2593        x: &CudaSlice<f32>,
2594        w: &CudaSlice<f32>,
2595        n_embd: usize,
2596        t: usize,
2597    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2598        // MEMRA_SHEXP_DOT=0: rollback seam to the cuBLASLt linear + sigmoid pair (numeric
2599        // config; same class as MEMRA_ROUTER_V2).
2600        static OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2601        if *OFF.get_or_init(|| std::env::var("MEMRA_SHEXP_DOT").as_deref() == Ok("0")) {
2602            let gs = self.linear(x, w, t, n_embd, 1)?;
2603            let mut g = self.uninit(t)?;
2604            self.sigmoid(&gs, &mut g, t)?;
2605            return Ok(g);
2606        }
2607        // FAST-ROUTER lane note (2026-08-02): a register-tiled 8-token batch twin of this
2608        // kernel was built, proven bit-identical, and measured SLOWER at every prefill t on
2609        // the 5090 (0.62-0.89x — launch-latency-bound op, ~7us/layer at m=2048;
2610        // research/fast-router-20260802/crossover-router.jsonl). Dispatch arm killed per
2611        // flags doctrine; this per-token form serves every t.
2612        let mut g = self.alloc_uninit::<f32>(t)?;
2613        let f = self.func("sigmoid_dot_rows_f32");
2614        let (ne, ti) = (n_embd as i32, t as i32);
2615        let cfg = LaunchConfig {
2616            grid_dim: (t as u32, 1, 1),
2617            block_dim: (32, 8, 1),
2618            shared_mem_bytes: 0,
2619        };
2620        let __s_b = self.gpu.stream();
2621        let mut b = __s_b.launch_builder(&f);
2622        b.arg(x).arg(w).arg(&mut g).arg(&ne).arg(&ti);
2623        unsafe {
2624            b.launch(cfg)?;
2625        }
2626        Ok(g)
2627    }
2628
2629    /// `sigmoid_dot_rows` writing into a caller-owned [t] buffer (token-graph alloc-free).
2630    pub fn sigmoid_dot_rows_into(
2631        &self,
2632        x: &CudaSlice<f32>,
2633        w: &CudaSlice<f32>,
2634        g: &mut CudaSlice<f32>,
2635        n_embd: usize,
2636        t: usize,
2637    ) -> Result<(), Box<dyn std::error::Error>> {
2638        if g.len() < t {
2639            return Err("sigmoid_dot_rows_into output too small".into());
2640        }
2641        let f = self.func("sigmoid_dot_rows_f32");
2642        let (ne, ti) = (n_embd as i32, t as i32);
2643        let cfg = LaunchConfig {
2644            grid_dim: (t as u32, 1, 1),
2645            block_dim: (32, 8, 1),
2646            shared_mem_bytes: 0,
2647        };
2648        let __s_b = self.gpu.stream();
2649        let mut b = __s_b.launch_builder(&f);
2650        b.arg(x).arg(w).arg(&mut *g).arg(&ne).arg(&ti);
2651        unsafe {
2652            b.launch(cfg)?;
2653        }
2654        Ok(())
2655    }
2656
2657    /// ROUND-STREAM stream rollback: all counters <- pos_start + base + n_acc.
2658    pub fn spec_rollback_stream(
2659        &self,
2660        len_ptrs: &CudaSlice<u64>,
2661        pos_start: &CudaSlice<i32>,
2662        acc: &CudaSlice<u32>,
2663        base: usize,
2664        n_rows: usize,
2665    ) -> Result<(), Box<dyn std::error::Error>> {
2666        let f = self.func("spec_rollback_stream");
2667        let (b, nr) = (base as i32, n_rows as i32);
2668        let cfg = LaunchConfig {
2669            grid_dim: (n_rows.div_ceil(64) as u32, 1, 1),
2670            block_dim: (64, 1, 1),
2671            shared_mem_bytes: 0,
2672        };
2673        let __s_bl = self.gpu.stream();
2674        let mut bl = __s_bl.launch_builder(&f);
2675        bl.arg(len_ptrs).arg(pos_start).arg(acc).arg(&b).arg(&nr);
2676        unsafe {
2677            bl.launch(cfg)?;
2678        }
2679        Ok(())
2680    }
2681
2682    /// PLAIN-DECODE GRAPH ring store: ring[(pos_start - base) % cap] = vam[0].
2683    pub fn plain_tok_ring(
2684        &self,
2685        vam: &CudaSlice<u32>,
2686        pos_start: &CudaSlice<i32>,
2687        base: usize,
2688        ring: &mut CudaSlice<u32>,
2689    ) -> Result<(), Box<dyn std::error::Error>> {
2690        let f = self.func("plain_tok_ring");
2691        let (b, cap) = (base as i32, ring.len() as i32);
2692        let cfg = LaunchConfig {
2693            grid_dim: (1, 1, 1),
2694            block_dim: (32, 1, 1),
2695            shared_mem_bytes: 0,
2696        };
2697        let __s_bl = self.gpu.stream();
2698        let mut bl = __s_bl.launch_builder(&f);
2699        bl.arg(vam).arg(pos_start).arg(&b).arg(&mut *ring).arg(&cap);
2700        unsafe {
2701            bl.launch(cfg)?;
2702        }
2703        Ok(())
2704    }
2705
2706    /// ROUND-STREAM stage (c) 4 epilogue: ring commit + tiny counter copies.
2707    pub fn spec_ring_commit(
2708        &self,
2709        vtok: &CudaSlice<u32>,
2710        acc: &CudaSlice<u32>,
2711        brk: &CudaSlice<u32>,
2712        ring: &mut CudaSlice<u32>,
2713        pend: &mut CudaSlice<u32>,
2714    ) -> Result<(), Box<dyn std::error::Error>> {
2715        let f = self.func("spec_ring_commit");
2716        let cfg = LaunchConfig {
2717            grid_dim: (1, 1, 1),
2718            block_dim: (32, 1, 1),
2719            shared_mem_bytes: 0,
2720        };
2721        let __s_b = self.gpu.stream();
2722        let mut b = __s_b.launch_builder(&f);
2723        b.arg(vtok).arg(acc).arg(brk).arg(ring).arg(pend);
2724        unsafe {
2725            b.launch(cfg)?;
2726        }
2727        Ok(())
2728    }
2729    pub fn i32_copy_add(
2730        &self,
2731        src: &CudaSlice<i32>,
2732        dst: &mut CudaSlice<i32>,
2733        delta: i32,
2734    ) -> Result<(), Box<dyn std::error::Error>> {
2735        let f = self.func("i32_copy_add");
2736        let cfg = LaunchConfig {
2737            grid_dim: (1, 1, 1),
2738            block_dim: (32, 1, 1),
2739            shared_mem_bytes: 0,
2740        };
2741        let __s_b = self.gpu.stream();
2742        let mut b = __s_b.launch_builder(&f);
2743        b.arg(src).arg(dst).arg(&delta);
2744        unsafe {
2745            b.launch(cfg)?;
2746        }
2747        Ok(())
2748    }
2749    pub fn u32_copy(
2750        &self,
2751        src: &CudaSlice<u32>,
2752        dst: &mut CudaSlice<u32>,
2753    ) -> Result<(), Box<dyn std::error::Error>> {
2754        let f = self.func("u32_copy");
2755        let cfg = LaunchConfig {
2756            grid_dim: (1, 1, 1),
2757            block_dim: (32, 1, 1),
2758            shared_mem_bytes: 0,
2759        };
2760        let __s_b = self.gpu.stream();
2761        let mut b = __s_b.launch_builder(&f);
2762        b.arg(src).arg(dst);
2763        unsafe {
2764            b.launch(cfg)?;
2765        }
2766        Ok(())
2767    }
2768
2769    /// ROUND-GRAPH adaptive depth: brk[0] <- clamp(acc[0] + 1, floor, cap) — the host
2770    /// adaptive policy as a captured device op (policy-identical: the accept walk depth
2771    /// caps acceptance exactly like drafting fewer tokens).
2772    pub fn spec_adapt_k(
2773        &self,
2774        acc: &CudaSlice<u32>,
2775        brk: &mut CudaSlice<u32>,
2776        floor: usize,
2777        cap: usize,
2778    ) -> Result<(), Box<dyn std::error::Error>> {
2779        let f = self.func("spec_adapt_k");
2780        let (fl, cp) = (floor as i32, cap as i32);
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(acc).arg(brk).arg(&fl).arg(&cp);
2789        unsafe {
2790            b.launch(cfg)?;
2791        }
2792        Ok(())
2793    }
2794
2795    /// ROUND-STREAM stage (c) 3: accept walk fully device-driven (brk + assembled vtok).
2796    pub fn spec_accept_greedy_dc(
2797        &self,
2798        preds: &CudaSlice<u32>,
2799        vtok: &CudaSlice<u32>,
2800        last_pred: &CudaSlice<u32>,
2801        brk: &CudaSlice<u32>,
2802        out: &mut CudaSlice<u32>,
2803    ) -> Result<(), Box<dyn std::error::Error>> {
2804        let f = self.func("spec_accept_greedy_dc");
2805        let cfg = LaunchConfig {
2806            grid_dim: (1, 1, 1),
2807            block_dim: (32, 1, 1),
2808            shared_mem_bytes: 0,
2809        };
2810        let __s_b = self.gpu.stream();
2811        let mut b = __s_b.launch_builder(&f);
2812        b.arg(preds).arg(vtok).arg(last_pred).arg(brk).arg(out);
2813        unsafe {
2814            b.launch(cfg)?;
2815        }
2816        Ok(())
2817    }
2818
2819    /// ROUND-STREAM stage (c) 2: verify-chain device-pos entries.
2820    pub fn pos_iota(
2821        &self,
2822        pos0: &CudaSlice<i32>,
2823        out: &mut CudaSlice<i32>,
2824        t: usize,
2825    ) -> Result<(), Box<dyn std::error::Error>> {
2826        let f = self.func("pos_iota_i32");
2827        let ti = t as i32;
2828        let cfg = LaunchConfig {
2829            grid_dim: (1, 1, 1),
2830            block_dim: (t.max(1) as u32, 1, 1),
2831            shared_mem_bytes: 0,
2832        };
2833        let __s_b = self.gpu.stream();
2834        let mut b = __s_b.launch_builder(&f);
2835        b.arg(pos0).arg(out).arg(&ti);
2836        unsafe {
2837            b.launch(cfg)?;
2838        }
2839        Ok(())
2840    }
2841    #[allow(clippy::too_many_arguments)]
2842    pub fn append_kv_quantized_rows_dc(
2843        &self,
2844        k_rows: &CudaSlice<f32>,
2845        v_rows: &CudaSlice<f32>,
2846        kc: &mut CudaSlice<u8>,
2847        vc: &mut CudaSlice<u8>,
2848        t0_dev: &CudaSlice<i32>,
2849        t: usize,
2850        kv_dim_k: usize,
2851        kv_dim_v: usize,
2852        k_tok_bytes: usize,
2853        v_tok_bytes: usize,
2854        g: bool,
2855    ) -> Result<(), Box<dyn std::error::Error>> {
2856        let f = if g {
2857            self.func_g("append_quantize_kv_q8_0_q5_1_rows_dc")
2858        } else {
2859            self.func("append_quantize_kv_q8_0_q5_1_rows_dc")
2860        };
2861        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
2862        let cfg = LaunchConfig {
2863            grid_dim: (nblk, t as u32, 1),
2864            block_dim: (32, 1, 1),
2865            shared_mem_bytes: 0,
2866        };
2867        let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
2868        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
2869        let __s_b = self.gpu.stream();
2870        let mut b = __s_b.launch_builder(&f);
2871        b.arg(k_rows)
2872            .arg(v_rows)
2873            .arg(kc)
2874            .arg(vc)
2875            .arg(t0_dev)
2876            .arg(&kdk)
2877            .arg(&kdv)
2878            .arg(&ktb)
2879            .arg(&vtb);
2880        unsafe {
2881            b.launch(cfg)?;
2882        }
2883        Ok(())
2884    }
2885
2886    /// t=1 dc append with a FUSED len_d increment (wave 5c) — one launch replaces
2887    /// append_rows_dc + inc_seqlen. Single block (read-before-inc ordering).
2888    #[allow(clippy::too_many_arguments)]
2889    pub fn append_kv_quantized_row_dc_inc(
2890        &self,
2891        k_row: &CudaSlice<f32>,
2892        v_row: &CudaSlice<f32>,
2893        kc: &mut CudaSlice<u8>,
2894        vc: &mut CudaSlice<u8>,
2895        t0_dev: &mut CudaSlice<i32>,
2896        kv_dim_k: usize,
2897        kv_dim_v: usize,
2898        k_tok_bytes: usize,
2899        v_tok_bytes: usize,
2900        g: bool,
2901    ) -> Result<(), Box<dyn std::error::Error>> {
2902        let f = if g {
2903            self.func_g("append_quantize_kv_q8_0_q5_1_dc_inc")
2904        } else {
2905            self.func("append_quantize_kv_q8_0_q5_1_dc_inc")
2906        };
2907        let nthreads = ((kv_dim_k.max(kv_dim_v) / 32) * 32).min(1024) as u32;
2908        let cfg = LaunchConfig {
2909            grid_dim: (1, 1, 1),
2910            block_dim: (nthreads, 1, 1),
2911            shared_mem_bytes: 0,
2912        };
2913        let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
2914        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
2915        let __s_b = self.gpu.stream();
2916        let mut b = __s_b.launch_builder(&f);
2917        b.arg(k_row)
2918            .arg(v_row)
2919            .arg(kc)
2920            .arg(vc)
2921            .arg(t0_dev)
2922            .arg(&kdk)
2923            .arg(&kdv)
2924            .arg(&ktb)
2925            .arg(&vtb);
2926        unsafe {
2927            b.launch(cfg)?;
2928        }
2929        Ok(())
2930    }
2931
2932    /// ROUND-STREAM: draft-chain pack + in-graph d2t remap (see kernels.cu headers).
2933    pub fn pack_tok_p(
2934        &self,
2935        tok: &CudaSlice<u32>,
2936        p: &CudaSlice<f32>,
2937        out: &mut CudaSlice<u32>,
2938        slot: usize,
2939    ) -> Result<(), Box<dyn std::error::Error>> {
2940        let f = self.func("pack_tok_p");
2941        let sl = slot as i32;
2942        let cfg = LaunchConfig {
2943            grid_dim: (1, 1, 1),
2944            block_dim: (32, 1, 1),
2945            shared_mem_bytes: 0,
2946        };
2947        let __s_b = self.gpu.stream();
2948        let mut b = __s_b.launch_builder(&f);
2949        b.arg(tok).arg(p).arg(out).arg(&sl);
2950        unsafe {
2951            b.launch(cfg)?;
2952        }
2953        Ok(())
2954    }
2955    pub fn tok_map_u32(
2956        &self,
2957        tok: &mut CudaSlice<u32>,
2958        map: &CudaSlice<u32>,
2959    ) -> Result<(), Box<dyn std::error::Error>> {
2960        let f = self.func("tok_map_u32");
2961        let cfg = LaunchConfig {
2962            grid_dim: (1, 1, 1),
2963            block_dim: (32, 1, 1),
2964            shared_mem_bytes: 0,
2965        };
2966        let __s_b = self.gpu.stream();
2967        let mut b = __s_b.launch_builder(&f);
2968        b.arg(tok).arg(map);
2969        unsafe {
2970            b.launch(cfg)?;
2971        }
2972        Ok(())
2973    }
2974
2975    /// ROUND-STREAM stage (c) 1: device verify-token assembly + p-min break derivation.
2976    #[allow(clippy::too_many_arguments)]
2977    pub fn spec_assemble_verify(
2978        &self,
2979        tokp: &CudaSlice<u32>,
2980        pend: &CudaSlice<u32>,
2981        d2t: Option<&CudaSlice<u32>>,
2982        vtok: &mut CudaSlice<u32>,
2983        brk: &mut CudaSlice<u32>,
2984        p_min: f32,
2985        k: usize,
2986        pmin0: bool,
2987    ) -> Result<(), Box<dyn std::error::Error>> {
2988        let f = self.func("spec_assemble_verify");
2989        let (ki, pm) = (k as i32, if pmin0 { 1i32 } else { 0i32 });
2990        let cfg = LaunchConfig {
2991            grid_dim: (1, 1, 1),
2992            block_dim: (32, 1, 1),
2993            shared_mem_bytes: 0,
2994        };
2995        let __s_b = self.gpu.stream();
2996        let mut b = __s_b.launch_builder(&f);
2997        match d2t {
2998            Some(m) => {
2999                b.arg(tokp)
3000                    .arg(pend)
3001                    .arg(m)
3002                    .arg(vtok)
3003                    .arg(brk)
3004                    .arg(&p_min)
3005                    .arg(&ki)
3006                    .arg(&pm);
3007                unsafe {
3008                    b.launch(cfg)?;
3009                }
3010            }
3011            None => {
3012                let null: u64 = 0;
3013                b.arg(tokp)
3014                    .arg(pend)
3015                    .arg(&null)
3016                    .arg(vtok)
3017                    .arg(brk)
3018                    .arg(&p_min)
3019                    .arg(&ki)
3020                    .arg(&pm);
3021                unsafe {
3022                    b.launch(cfg)?;
3023                }
3024            }
3025        }
3026        Ok(())
3027    }
3028
3029    /// ROUND-STREAM stage (b) 3b: recur-restore twins with device-j (see hybrid.cu headers).
3030    #[allow(clippy::too_many_arguments)]
3031    pub fn ssm_conv_ring_rebuild_dc(
3032        &self,
3033        qkv_tm: &CudaSlice<f32>,
3034        ring_old: &CudaSlice<f32>,
3035        conv_state: &mut CudaSlice<f32>,
3036        conv_dim: usize,
3037        acc: &CudaSlice<u32>,
3038        base: usize,
3039        t_v: usize,
3040        d_conv: usize,
3041    ) -> Result<(), Box<dyn std::error::Error>> {
3042        let f = self.func("ssm_conv_ring_rebuild_f32_dc");
3043        let n = conv_dim * (d_conv - 1);
3044        let cfg = LaunchConfig::for_num_elems(n as u32);
3045        let (cd, b0, tv, dc) = (conv_dim as i32, base as i32, t_v as i32, d_conv as i32);
3046        let __s_b = self.gpu.stream();
3047        let mut b = __s_b.launch_builder(&f);
3048        b.arg(qkv_tm)
3049            .arg(ring_old)
3050            .arg(conv_state)
3051            .arg(&cd)
3052            .arg(acc)
3053            .arg(&b0)
3054            .arg(&tv)
3055            .arg(&dc);
3056        unsafe {
3057            b.launch(cfg)?;
3058        }
3059        Ok(())
3060    }
3061    #[allow(clippy::too_many_arguments)]
3062    pub fn gdn_scan_s128_dc(
3063        &self,
3064        q: &CudaSlice<f32>,
3065        k: &CudaSlice<f32>,
3066        v: &CudaSlice<f32>,
3067        g: &CudaSlice<f32>,
3068        beta: &CudaSlice<f32>,
3069        state_in: &CudaSlice<f32>,
3070        state_out: &mut CudaSlice<f32>,
3071        o: &mut CudaSlice<f32>,
3072        n_head: usize,
3073        acc: &CudaSlice<u32>,
3074        base: usize,
3075        t_v: usize,
3076        scale: f32,
3077    ) -> Result<(), Box<dyn std::error::Error>> {
3078        let f = self.func("gdn_scan_s128_dc");
3079        const S_V: u32 = 128;
3080        const WARP: u32 = 32;
3081        const COLS_PER_BLOCK: u32 = 4;
3082        let cfg = LaunchConfig {
3083            grid_dim: (n_head as u32, 1, S_V / COLS_PER_BLOCK),
3084            block_dim: (WARP, COLS_PER_BLOCK, 1),
3085            shared_mem_bytes: 0,
3086        };
3087        let (h, b0, tv) = (n_head as i32, base as i32, t_v as i32);
3088        let __s_b = self.gpu.stream();
3089        let mut b = __s_b.launch_builder(&f);
3090        b.arg(q)
3091            .arg(k)
3092            .arg(v)
3093            .arg(g)
3094            .arg(beta)
3095            .arg(state_in)
3096            .arg(state_out)
3097            .arg(o)
3098            .arg(&h)
3099            .arg(acc)
3100            .arg(&b0)
3101            .arg(&tv)
3102            .arg(&scale);
3103        unsafe {
3104            b.launch(cfg)?;
3105        }
3106        Ok(())
3107    }
3108
3109    /// ROUND-STREAM stage (b) 3a: device per-layer KV-len rollback (see spec_rollback_kv).
3110    pub fn spec_rollback_kv(
3111        &self,
3112        len_ptrs: &CudaSlice<u64>,
3113        saved: &CudaSlice<i32>,
3114        acc: &CudaSlice<u32>,
3115        base: usize,
3116        n_layer: usize,
3117    ) -> Result<(), Box<dyn std::error::Error>> {
3118        let f = self.func("spec_rollback_kv");
3119        let (b, nl) = (base as i32, n_layer as i32);
3120        let cfg = LaunchConfig {
3121            grid_dim: (n_layer.div_ceil(64) as u32, 1, 1),
3122            block_dim: (64, 1, 1),
3123            shared_mem_bytes: 0,
3124        };
3125        let __s_bl = self.gpu.stream();
3126        let mut bl = __s_bl.launch_builder(&f);
3127        bl.arg(len_ptrs).arg(saved).arg(acc).arg(&b).arg(&nl);
3128        unsafe {
3129            bl.launch(cfg)?;
3130        }
3131        Ok(())
3132    }
3133
3134    /// OPTIPIPE increment 1: derive the K=1 successor-valid bit on device.
3135    pub fn spec_fork_valid(
3136        &self,
3137        acc: &CudaSlice<u32>,
3138        optimistic_pending: u32,
3139        valid: &mut CudaSlice<u32>,
3140    ) -> Result<(), Box<dyn std::error::Error>> {
3141        let f = self.func("spec_fork_valid");
3142        let cfg = LaunchConfig {
3143            grid_dim: (1, 1, 1),
3144            block_dim: (1, 1, 1),
3145            shared_mem_bytes: 0,
3146        };
3147        let __s_bl = self.gpu.stream();
3148        let mut bl = __s_bl.launch_builder(&f);
3149        bl.arg(acc).arg(&optimistic_pending).arg(valid);
3150        unsafe {
3151            bl.launch(cfg)?;
3152        }
3153        Ok(())
3154    }
3155
3156    /// OPTIPIPE increment 1: leave stage-local KV lengths on hit, restore them on miss.
3157    pub fn spec_fork_reconcile_kv(
3158        &self,
3159        len_ptrs: &CudaSlice<u64>,
3160        saved: &CudaSlice<i32>,
3161        acc: &CudaSlice<u32>,
3162        valid: &CudaSlice<u32>,
3163        base: usize,
3164        n_layer: usize,
3165    ) -> Result<(), Box<dyn std::error::Error>> {
3166        let f = self.func("spec_fork_reconcile_kv");
3167        let (b, nl) = (base as i32, n_layer as i32);
3168        let cfg = LaunchConfig {
3169            grid_dim: (n_layer.div_ceil(64) as u32, 1, 1),
3170            block_dim: (64, 1, 1),
3171            shared_mem_bytes: 0,
3172        };
3173        let __s_bl = self.gpu.stream();
3174        let mut bl = __s_bl.launch_builder(&f);
3175        bl.arg(len_ptrs)
3176            .arg(saved)
3177            .arg(acc)
3178            .arg(valid)
3179            .arg(&b)
3180            .arg(&nl);
3181        unsafe {
3182            bl.launch(cfg)?;
3183        }
3184        Ok(())
3185    }
3186
3187    /// OPTIPIPE increment 1: conditionally restore one stage-owned recurrent-state buffer.
3188    pub fn spec_fork_restore_f32(
3189        &self,
3190        snapshot: &CudaSlice<f32>,
3191        state: &mut CudaSlice<f32>,
3192        valid: &CudaSlice<u32>,
3193    ) -> Result<(), Box<dyn std::error::Error>> {
3194        assert_eq!(
3195            snapshot.len(),
3196            state.len(),
3197            "fork recurrent snapshot shape mismatch"
3198        );
3199        let f = self.func("spec_fork_restore_f32");
3200        let n = state.len() as i32;
3201        let blocks = state.len().div_ceil(256).min(65535).max(1) as u32;
3202        let cfg = LaunchConfig {
3203            grid_dim: (blocks, 1, 1),
3204            block_dim: (256, 1, 1),
3205            shared_mem_bytes: 0,
3206        };
3207        let __s_bl = self.gpu.stream();
3208        let mut bl = __s_bl.launch_builder(&f);
3209        bl.arg(snapshot).arg(state).arg(valid).arg(&n);
3210        unsafe {
3211            bl.launch(cfg)?;
3212        }
3213        Ok(())
3214    }
3215
3216    /// ROUND-STREAM stage (b): device next-round seed gather (see spec_seed_gather header).
3217    /// Caller D2Ds h_seed into fill_prev after (both slots carry the same value in every arm).
3218    pub fn spec_seed_gather(
3219        &self,
3220        vx: &CudaSlice<f32>,
3221        fill_prev: &CudaSlice<f32>,
3222        acc: &CudaSlice<u32>,
3223        h_seed: &mut CudaSlice<f32>,
3224        base: usize,
3225        n_embd: usize,
3226    ) -> Result<(), Box<dyn std::error::Error>> {
3227        let f = self.func("spec_seed_gather");
3228        let (b, ne) = (base as i32, n_embd as i32);
3229        let cfg = LaunchConfig {
3230            grid_dim: (n_embd.div_ceil(256) as u32, 1, 1),
3231            block_dim: (256, 1, 1),
3232            shared_mem_bytes: 0,
3233        };
3234        let __s_bl = self.gpu.stream();
3235        let mut bl = __s_bl.launch_builder(&f);
3236        bl.arg(vx)
3237            .arg(fill_prev)
3238            .arg(acc)
3239            .arg(h_seed)
3240            .arg(&b)
3241            .arg(&ne);
3242        unsafe {
3243            bl.launch(cfg)?;
3244        }
3245        Ok(())
3246    }
3247
3248    /// ROUND-STREAM stage (a): device greedy accept walk (see spec_accept_greedy header).
3249    pub fn spec_accept_greedy(
3250        &self,
3251        preds: &CudaSlice<u32>,
3252        draft: &CudaSlice<u32>,
3253        last_pred: u32,
3254        base: usize,
3255        k_round: usize,
3256        out: &mut CudaSlice<u32>,
3257    ) -> Result<(), Box<dyn std::error::Error>> {
3258        let f = self.func("spec_accept_greedy");
3259        let (b, k) = (base as i32, k_round as i32);
3260        let cfg = LaunchConfig {
3261            grid_dim: (1, 1, 1),
3262            block_dim: (32, 1, 1),
3263            shared_mem_bytes: 0,
3264        };
3265        let __s_bl = self.gpu.stream();
3266        let mut bl = __s_bl.launch_builder(&f);
3267        bl.arg(preds)
3268            .arg(draft)
3269            .arg(&last_pred)
3270            .arg(&b)
3271            .arg(&k)
3272            .arg(out);
3273        unsafe {
3274            bl.launch(cfg)?;
3275        }
3276        Ok(())
3277    }
3278
3279    // ================= SAMPLED-SPEC PRIMITIVES (spec_sample.cu, piece A) =================
3280    // Counter-based randomness: every call takes (seed, stream_pos) — the caller owns the
3281    // event counter (one per sampled token). temp <= 0 arms are exact greedy limits.
3282
3283    /// y = x/temp + Gumbel(Philox(seed, stream_pos)) over n logits (then run device argmax on y
3284    /// = one categorical sample at temperature `temp`). temp<=0: y = x (pure copy).
3285    pub fn gumbel_perturb(
3286        &self,
3287        x: &CudaSlice<f32>,
3288        y: &mut CudaSlice<f32>,
3289        n: usize,
3290        seed: u64,
3291        stream_pos: u32,
3292        temp: f32,
3293    ) -> Result<(), Box<dyn std::error::Error>> {
3294        let f = self.func("gumbel_perturb_f32");
3295        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3296        let cfg = LaunchConfig {
3297            grid_dim: (n.div_ceil(256) as u32, 1, 1),
3298            block_dim: (256, 1, 1),
3299            shared_mem_bytes: 0,
3300        };
3301        let __s_b = self.gpu.stream();
3302        let mut b = __s_b.launch_builder(&f);
3303        b.arg(x)
3304            .arg(&mut *y)
3305            .arg(&ni)
3306            .arg(&slo)
3307            .arg(&shi)
3308            .arg(&stream_pos)
3309            .arg(&temp);
3310        unsafe {
3311            b.launch(cfg)?;
3312        }
3313        Ok(())
3314    }
3315
3316    /// GRAMMAR TOKEN MASK (constrained decoding, lane/constrained-full): ban every vocab id
3317    /// whose bit is unset in the packed llguidance bitset, IN PLACE on row `col` of a stacked
3318    /// [B, n_vocab] logits buffer. `mask` = the SimpleVob u32 words H2D'd verbatim
3319    /// (~n_vocab/8 bytes/step — trivial on PCIe); ids >= 32*mask_words (padded lm_head tail)
3320    /// are banned too, the device twin of constrained::apply_mask. Banned value -FLT_MAX ==
3321    /// the argmax/gumbel kernels' init sentinel, so a fully-banned tail can never win and
3322    /// ordering matches the host -inf mask bit-for-bit for every finite logit.
3323    pub fn mask_logits_col(
3324        &self,
3325        logits: &mut CudaSlice<f32>,
3326        mask: &CudaSlice<u32>,
3327        col: usize,
3328        n: usize,
3329        mask_words: usize,
3330    ) -> Result<(), Box<dyn std::error::Error>> {
3331        let f = self.func("mask_logits_f32");
3332        let (ci, ni, mw) = (col as i32, n as i32, mask_words as i32);
3333        let cfg = LaunchConfig {
3334            grid_dim: (n.div_ceil(256).min(1024) as u32, 1, 1),
3335            block_dim: (256, 1, 1),
3336            shared_mem_bytes: 0,
3337        };
3338        let __s_b = self.gpu.stream();
3339        let mut b = __s_b.launch_builder(&f);
3340        b.arg(&mut *logits).arg(mask).arg(&ci).arg(&ni).arg(&mw);
3341        unsafe {
3342            b.launch(cfg)?;
3343        }
3344        Ok(())
3345    }
3346
3347    /// Column-`col` twin of `gumbel_perturb` over stacked logits [B, n_vocab] (the batched
3348    /// serving tick's device sampler): y = x[col]/temp + gumbel(seed, stream_pos, lane).
3349    /// SAME kernel/Philox mapping as `gumbel_perturb` — bit-identical perturbation for the
3350    /// same (seed, stream_pos, temp) regardless of which batch column the row sits in
3351    /// (the lane index is the in-row position; `col` only moves the input pointer). That
3352    /// pointer-invariance IS the serving isolation contract for sampled rows.
3353    pub fn gumbel_perturb_col(
3354        &self,
3355        x: &CudaSlice<f32>,
3356        col: usize,
3357        y: &mut CudaSlice<f32>,
3358        n: usize,
3359        seed: u64,
3360        stream_pos: u32,
3361        temp: f32,
3362    ) -> Result<(), Box<dyn std::error::Error>> {
3363        let f = self.func("gumbel_perturb_f32");
3364        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3365        let col_view = x.slice(col * n..(col + 1) * n);
3366        let cfg = LaunchConfig {
3367            grid_dim: (n.div_ceil(256) as u32, 1, 1),
3368            block_dim: (256, 1, 1),
3369            shared_mem_bytes: 0,
3370        };
3371        let __s_b = self.gpu.stream();
3372        let mut b = __s_b.launch_builder(&f);
3373        b.arg(&col_view)
3374            .arg(&mut *y)
3375            .arg(&ni)
3376            .arg(&slo)
3377            .arg(&shi)
3378            .arg(&stream_pos)
3379            .arg(&temp);
3380        unsafe {
3381            b.launch(cfg)?;
3382        }
3383        Ok(())
3384    }
3385
3386    /// Filtered twin of `gumbel_perturb_col`: the per-row (row_max, th) floor comes from
3387    /// DEVICE buffers (`filter_stats` output slots at `stat_idx`) — one filtered draw from
3388    /// the top-k/top-p/min-p-truncated softmax with no stat D2H and no row copy. Same
3389    /// Philox mapping as every gumbel kernel (pointer-invariant across batch columns —
3390    /// the serving isolation contract for sampled rows).
3391    #[allow(clippy::too_many_arguments)]
3392    pub fn gumbel_perturb_filtered_col(
3393        &self,
3394        x: &CudaSlice<f32>,
3395        col: usize,
3396        y: &mut CudaSlice<f32>,
3397        n: usize,
3398        seed: u64,
3399        stream_pos: u32,
3400        temp: f32,
3401        stat_max: &CudaSlice<f32>,
3402        stat_th: &CudaSlice<f32>,
3403        stat_idx: usize,
3404    ) -> Result<(), Box<dyn std::error::Error>> {
3405        let f = self.func("gumbel_perturb_filtered_col_f32");
3406        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3407        let (ci, si) = (col as i32, stat_idx as i32);
3408        let cfg = LaunchConfig {
3409            grid_dim: (n.div_ceil(256) as u32, 1, 1),
3410            block_dim: (256, 1, 1),
3411            shared_mem_bytes: 0,
3412        };
3413        let __s_b = self.gpu.stream();
3414        let mut b = __s_b.launch_builder(&f);
3415        b.arg(x)
3416            .arg(&ci)
3417            .arg(&mut *y)
3418            .arg(&ni)
3419            .arg(&slo)
3420            .arg(&shi)
3421            .arg(&stream_pos)
3422            .arg(&temp)
3423            .arg(stat_max)
3424            .arg(stat_th)
3425            .arg(&si);
3426        unsafe {
3427            b.launch(cfg)?;
3428        }
3429        Ok(())
3430    }
3431
3432    /// In-graph sampling-event counter bump (spec_sample.cu kernel 5): ctr[0] += 1. The sampled
3433    /// graph-draft chain replays with FIXED kernel args, so the Philox event counter must be
3434    /// DEVICE data — the host seeds it once per round; every replay bumps it before the perturb
3435    /// reads it (counter is data, not state — graph-replay-safe).
3436    pub fn sctr_inc(&self, ctr: &mut CudaSlice<u32>) -> Result<(), Box<dyn std::error::Error>> {
3437        let f = self.func("memra_sctr_inc");
3438        let cfg = LaunchConfig {
3439            grid_dim: (1, 1, 1),
3440            block_dim: (1, 1, 1),
3441            shared_mem_bytes: 0,
3442        };
3443        let __s_b = self.gpu.stream();
3444        let mut b = __s_b.launch_builder(&f);
3445        b.arg(&mut *ctr);
3446        unsafe {
3447            b.launch(cfg)?;
3448        }
3449        Ok(())
3450    }
3451
3452    /// Graph-capturable `gumbel_perturb`: the sampling-event counter comes from DEVICE memory
3453    /// (`ctr[0]`) instead of a host scalar. Identical math to `gumbel_perturb` at
3454    /// stream_pos == ctr[0] (same Philox call, same lane mapping) — the eager and graph sampled
3455    /// chains produce bit-identical perturbations for the same (seed, counter, temp).
3456    pub fn gumbel_perturb_ctr(
3457        &self,
3458        x: &CudaSlice<f32>,
3459        y: &mut CudaSlice<f32>,
3460        n: usize,
3461        seed: u64,
3462        ctr: &CudaSlice<u32>,
3463        temp: f32,
3464    ) -> Result<(), Box<dyn std::error::Error>> {
3465        let f = self.func("gumbel_perturb_ctr_f32");
3466        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3467        let cfg = LaunchConfig {
3468            grid_dim: (n.div_ceil(256) as u32, 1, 1),
3469            block_dim: (256, 1, 1),
3470            shared_mem_bytes: 0,
3471        };
3472        let __s_b = self.gpu.stream();
3473        let mut b = __s_b.launch_builder(&f);
3474        b.arg(x)
3475            .arg(&mut *y)
3476            .arg(&ni)
3477            .arg(&slo)
3478            .arg(&shi)
3479            .arg(ctr)
3480            .arg(&temp);
3481        unsafe {
3482            b.launch(cfg)?;
3483        }
3484        Ok(())
3485    }
3486
3487    /// out[pair] = softmax_temp(x[rows[pair]])[ids[pair]] for npair (row, id) pairs; rows index
3488    /// into x with `row_stride` f32s per row. temp<=0: out = 1.0 iff id is the row argmax
3489    /// (smallest-index tie-break — matches the argmax-gate contract).
3490    pub fn softmax_gather(
3491        &self,
3492        x: &CudaSlice<f32>,
3493        row_stride: usize,
3494        ids: &CudaSlice<u32>,
3495        rows: &CudaSlice<i32>,
3496        out: &mut CudaSlice<f32>,
3497        n: usize,
3498        npair: usize,
3499        temp: f32,
3500    ) -> Result<(), Box<dyn std::error::Error>> {
3501        let f = self.func("softmax_gather_f32");
3502        let (ni, rs) = (n as i32, row_stride as i64);
3503        let np = npair as i32;
3504        let cfg = LaunchConfig {
3505            grid_dim: (npair as u32, 1, 1),
3506            block_dim: (256, 1, 1),
3507            shared_mem_bytes: 0,
3508        };
3509        let __s_b = self.gpu.stream();
3510        let mut b = __s_b.launch_builder(&f);
3511        b.arg(x)
3512            .arg(&rs)
3513            .arg(ids)
3514            .arg(rows)
3515            .arg(&mut *out)
3516            .arg(&ni)
3517            .arg(&np)
3518            .arg(&temp);
3519        unsafe {
3520            b.launch(cfg)?;
3521        }
3522        Ok(())
3523    }
3524
3525    /// Sample token from norm(max(0, softmax_temp(p) - softmax_temp(q))) (q = None -> plain
3526    /// categorical from softmax_temp(p)). Row stats (max, sumexp at temp) must be precomputed
3527    /// (softmax_gather's pass-1 values; see spec.rs caller). Deterministic fixed-order CDF walk.
3528    pub fn residual_sample(
3529        &self,
3530        p: &CudaSlice<f32>,
3531        q: Option<&CudaSlice<f32>>,
3532        n: usize,
3533        temp: f32,
3534        seed: u64,
3535        stream_pos: u32,
3536        out_tok: &mut CudaSlice<u32>,
3537    ) -> Result<(), Box<dyn std::error::Error>> {
3538        let f = self.func("residual_sample_f32");
3539        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3540        let nth = 1024u32;
3541        let cfg = LaunchConfig {
3542            grid_dim: (1, 1, 1),
3543            block_dim: (nth, 1, 1),
3544            shared_mem_bytes: 0,
3545        };
3546        let has_q: i32 = q.is_some() as i32;
3547        let qbuf = q.unwrap_or(p); // dummy when absent; kernel gates on has_q
3548        let __s_b = self.gpu.stream();
3549        let mut b = __s_b.launch_builder(&f);
3550        b.arg(p)
3551            .arg(qbuf)
3552            .arg(&has_q)
3553            .arg(&ni)
3554            .arg(&temp)
3555            .arg(&slo)
3556            .arg(&shi)
3557            .arg(&stream_pos)
3558            .arg(&mut *out_tok);
3559        unsafe {
3560            b.launch(cfg)?;
3561        }
3562        Ok(())
3563    }
3564
3565    /// Access the shared MoE residency cache (EDGE-1 §B), building it on first use under
3566    /// MEMRA_MOE_CACHE. The closure runs while the lock is held — keep it to lookup/admit/issue, not
3567    /// the GEMM. `max_block_bytes` sizes the slots (largest of gate/up/down). Returns the closure's
3568    /// result. If MEMRA_MOE_CACHE is unset this is never called (the caller checks the env first).
3569    pub fn with_moe_cache<R>(
3570        &self,
3571        max_block_bytes: usize,
3572        f: impl FnOnce(
3573            &mut crate::moe_cache::MoeSlotCache,
3574            &Engine,
3575        ) -> Result<R, Box<dyn std::error::Error>>,
3576    ) -> Result<R, Box<dyn std::error::Error>> {
3577        let mut guard = self.moe_cache.lock().unwrap();
3578        if guard.is_none() {
3579            *guard = Some(crate::moe_cache::MoeSlotCache::new(self, max_block_bytes)?);
3580        }
3581        let cache = guard.as_mut().unwrap();
3582        f(cache, self)
3583    }
3584
3585    /// Freeze the already-built MoE residency set. This never constructs a cache: callers use it
3586    /// only after a real prefill has populated the machine-specific CPU/GPU working set.
3587    pub fn freeze_moe_cache(&self) {
3588        if let Some(cache) = self.moe_cache.lock().unwrap().as_mut() {
3589            cache.freeze();
3590        }
3591    }
3592
3593    /// The current residency set as (layer, proj, ex) triples, or None if no cache was built.
3594    /// Never constructs a cache.
3595    pub fn export_moe_residency(&self) -> Option<Vec<(u16, u8, u16)>> {
3596        self.moe_cache
3597            .lock()
3598            .unwrap()
3599            .as_ref()
3600            .map(crate::moe_cache::MoeSlotCache::export_residency)
3601    }
3602
3603    pub(crate) fn moe_cache_frozen(&self) -> bool {
3604        self.moe_cache
3605            .lock()
3606            .unwrap()
3607            .as_ref()
3608            .is_some_and(crate::moe_cache::MoeSlotCache::is_frozen)
3609    }
3610
3611    /// A frozen heterogeneous CPU/GPU expert split cannot use Hy3's ordinary batched prefill
3612    /// efficiently: T>=PRIME_MIN_T bypasses the CPU backend and transiently rereads every missing
3613    /// expert through the GPU spill path. Replay the short prompt through decode after freezing,
3614    /// while leaving the profiling warmup's established batched behavior untouched.
3615    /// (`pub`: run-gen's #46 batched-prime gate skips itself when generation will take the
3616    /// tokenwise arm anyway.)
3617    pub fn frozen_cpu_experts_prefer_tokenwise_prime(&self) -> bool {
3618        crate::cpu_experts::configured()
3619            && self.moe_cache_frozen()
3620            && std::env::var("MEMRA_CPU_EXPERT_BATCHED_PRIME").as_deref() != Ok("1")
3621    }
3622
3623    /// Install the loaded model's exact retained expert-block inventory before lazy cache build.
3624    pub(crate) fn configure_moe_cache_layout(&self, block_bytes: Vec<usize>) {
3625        assert!(
3626            self.moe_cache.lock().unwrap().is_none(),
3627            "MoE cache layout configured after cache construction"
3628        );
3629        *self.moe_cache_layout.lock().unwrap() = Some(block_bytes);
3630    }
3631
3632    pub(crate) fn moe_cache_layout(&self) -> Option<Vec<usize>> {
3633        self.moe_cache_layout.lock().unwrap().clone()
3634    }
3635
3636    /// True if the MoE residency cache is enabled (MEMRA_MOE_CACHE set).
3637    pub fn moe_cache_enabled() -> bool {
3638        std::env::var("MEMRA_MOE_CACHE").as_deref() != Ok("0")
3639    }
3640
3641    /// Snapshot the MoE cache counters (hits, misses, staged_bytes, n_slots) for the §D.4 PCIe gate.
3642    /// Returns None if the cache was never built (disabled or no MoE forward ran).
3643    pub fn moe_cache_stats(&self) -> Option<(u64, u64, u64, usize)> {
3644        let guard = self.moe_cache.lock().unwrap();
3645        guard
3646            .as_ref()
3647            .map(|c| (c.hits, c.misses, c.staged_bytes, c.n_slots()))
3648    }
3649
3650    /// Experimental CPU expert backend counters: completed layer calls, experts served, and the
3651    /// sum of backend wall nanoseconds. The timer includes explicit disk->RAM fills on cache misses;
3652    /// callers compare a before/after snapshot around a decode window.
3653    pub fn cpu_expert_stats(
3654        &self,
3655    ) -> Option<(u64, u64, u64, u64, u64, u64, u64, u64, u64, u64, u64)> {
3656        crate::cpu_experts::configured().then(crate::cpu_experts::stats)
3657    }
3658
3659    /// Caller-blocked nanoseconds at CPU expert joins. Compare before/after snapshots to measure
3660    /// the backend tail that resident-GPU expert work did not hide.
3661    pub fn cpu_expert_predictor_stats(&self) -> (u64, u64) {
3662        crate::cpu_experts::predictor_stats()
3663    }
3664
3665    pub fn cpu_expert_exposed_wait_ns(&self) -> Option<u64> {
3666        crate::cpu_experts::configured().then(crate::cpu_experts::exposed_wait_ns)
3667    }
3668
3669    /// CPU-routed expert selections grouped by how many of their three projections were already
3670    /// resident in HBM. This makes otherwise-stranded partial residency visible to tuning runs.
3671    pub fn cpu_expert_gpu_residency_stats(&self) -> Option<(u64, u64, u64)> {
3672        crate::cpu_experts::configured().then(crate::cpu_experts::incomplete_gpu_residency_stats)
3673    }
3674
3675    /// Positioned-read proof-backend counters:
3676    /// `(reads, bytes, read_errors, short_reads, mmap_fallbacks, buffer_waits, ring_full)`.
3677    pub fn moe_pread_stats(&self) -> Option<(u64, u64, u64, u64, u64, u64, u64)> {
3678        let guard = self.moe_cache.lock().unwrap();
3679        guard
3680            .as_ref()
3681            .and_then(|cache| cache.pread_stats())
3682            .map(|stats| {
3683                (
3684                    stats.reads,
3685                    stats.bytes,
3686                    stats.read_errors,
3687                    stats.short_reads,
3688                    stats.fallbacks,
3689                    stats.buffer_waits,
3690                    stats.ring_full,
3691                )
3692            })
3693    }
3694
3695    /// Spill configuration values that warned and substituted their documented defaults.
3696    pub fn spill_config_fallbacks(&self) -> u64 {
3697        crate::spill_pread::config_fallbacks()
3698    }
3699
3700    /// Reset the MoE cache perf counters (to separate warmup from steady-state windows).
3701    pub fn moe_cache_reset_counters(&self) {
3702        if let Some(c) = self.moe_cache.lock().unwrap().as_mut() {
3703            c.reset_counters();
3704        }
3705    }
3706
3707    pub fn htod_bytes(&self, v: &[u8]) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
3708        Ok(self.gpu.stream().clone_htod(v)?)
3709    }
3710
3711    /// `htod_bytes` with a mapped (uninit) tail pad: the wide-load expert dots read up to 6B
3712    /// past the final q4_0 block through their aligned window — the bytes never reach a
3713    /// result (funnelshift discards them) but must be mapped memory.
3714    pub fn htod_bytes_padded(
3715        &self,
3716        v: &[u8],
3717        pad: usize,
3718    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
3719        let mut d = self.alloc_u8_uninit(v.len() + pad)?;
3720        {
3721            let mut view = d.slice_mut(0..v.len());
3722            self.gpu.stream().memcpy_htod(v, &mut view)?;
3723        }
3724        Ok(d)
3725    }
3726
3727    /// Device-to-device copy of `src` into `dst[off..off+len]` (f32). For in-place KV append.
3728    pub fn copy_into(
3729        &self,
3730        dst: &mut CudaSlice<f32>,
3731        off: usize,
3732        src: &CudaSlice<f32>,
3733        len: usize,
3734    ) -> Result<(), Box<dyn std::error::Error>> {
3735        let mut view = dst.slice_mut(off..off + len);
3736        self.gpu
3737            .stream()
3738            .memcpy_dtod(&src.slice(0..len), &mut view)?;
3739        Ok(())
3740    }
3741
3742    /// View a sub-range of a device buffer (for attending over [0..len) of a KV cache).
3743    /// u8 twin of copy_into (D2D byte-range copy at an offset).
3744    pub fn copy_u8_into(
3745        &self,
3746        dst: &mut CudaSlice<u8>,
3747        off: usize,
3748        src: &CudaSlice<u8>,
3749        len: usize,
3750    ) -> Result<(), Box<dyn std::error::Error>> {
3751        let mut view = dst.slice_mut(off..off + len);
3752        self.gpu
3753            .stream()
3754            .memcpy_dtod(&src.slice(0..len), &mut view)?;
3755        Ok(())
3756    }
3757
3758    /// D2D byte-range copy with explicit source and destination offsets.
3759    pub fn copy_u8_range_into(
3760        &self,
3761        dst: &mut CudaSlice<u8>,
3762        dst_off: usize,
3763        src: &CudaSlice<u8>,
3764        src_off: usize,
3765        len: usize,
3766    ) -> Result<(), Box<dyn std::error::Error>> {
3767        let mut dst_view = dst.slice_mut(dst_off..dst_off + len);
3768        self.gpu
3769            .stream()
3770            .memcpy_dtod(&src.slice(src_off..src_off + len), &mut dst_view)?;
3771        Ok(())
3772    }
3773
3774    /// Resolve an absolute append slot to the Step35 SWA layer's physical rows. At wrap, copy
3775    /// only the aligned live prefix through temporary device storage and rebase it at row zero,
3776    /// keeping the audited attention range contiguous without changing its absolute start.
3777    pub fn prepare_kv_append(
3778        &self,
3779        kv: &mut crate::cache::KvLayer,
3780        retain_from: usize,
3781        append_rows: usize,
3782    ) -> Result<usize, Box<dyn std::error::Error>> {
3783        let Some(plan) = kv
3784            .ring
3785            .as_ref()
3786            .map(|ring| ring.append_plan(kv.len, retain_from, append_rows))
3787            .transpose()?
3788        else {
3789            return Ok(kv.len);
3790        };
3791        match plan {
3792            crate::cache::KvRingAppend::Contiguous { write_row } => Ok(write_row),
3793            crate::cache::KvRingAppend::Rebase {
3794                src_row,
3795                keep_rows,
3796                new_base,
3797                write_row,
3798            } => {
3799                if keep_rows > 0 {
3800                    let k_len = keep_rows * kv.k_tok_bytes;
3801                    let v_len = keep_rows * kv.v_tok_bytes;
3802                    let mut k_tmp = self.alloc_u8_uninit(k_len)?;
3803                    let mut v_tmp = self.alloc_u8_uninit(v_len)?;
3804                    self.copy_u8_range_into(&mut k_tmp, 0, &kv.k, src_row * kv.k_tok_bytes, k_len)?;
3805                    self.copy_u8_range_into(&mut v_tmp, 0, &kv.v, src_row * kv.v_tok_bytes, v_len)?;
3806                    self.copy_u8_into(&mut kv.k, 0, &k_tmp, k_len)?;
3807                    self.copy_u8_into(&mut kv.v, 0, &v_tmp, v_len)?;
3808                }
3809                kv.ring.as_mut().unwrap().apply_rebase(new_base);
3810                Ok(write_row)
3811            }
3812        }
3813    }
3814
3815    /// H2D write of `src` into `dst[off..off+src.len()]` (u8). In-place row updates for the
3816    /// adaptive trim head: no realloc, so captured graphs keep their baked addresses.
3817    pub fn htod_u8_into(
3818        &self,
3819        dst: &mut CudaSlice<u8>,
3820        off: usize,
3821        src: &[u8],
3822    ) -> Result<(), Box<dyn std::error::Error>> {
3823        let mut view = dst.slice_mut(off..off + src.len());
3824        self.gpu.stream().memcpy_htod(src, &mut view)?;
3825        Ok(())
3826    }
3827
3828    pub fn view<'a>(&self, b: &'a CudaSlice<f32>, len: usize) -> cudarc::driver::CudaView<'a, f32> {
3829        b.slice(0..len)
3830    }
3831
3832    /// View the first `len` BYTES of a u8 device buffer (quantized KV cache: [0..t_kv*tok_bytes)).
3833    /// Byte-range view (gemma4 R6 window offset into the quantized KV stream).
3834    pub fn view_u8_range<'a>(
3835        &self,
3836        b: &'a CudaSlice<u8>,
3837        start: usize,
3838        end: usize,
3839    ) -> cudarc::driver::CudaView<'a, u8> {
3840        b.slice(start..end)
3841    }
3842    pub fn view_u8<'a>(
3843        &self,
3844        b: &'a CudaSlice<u8>,
3845        len: usize,
3846    ) -> cudarc::driver::CudaView<'a, u8> {
3847        b.slice(0..len)
3848    }
3849
3850    /// Append-quantize ONE token's post-RoPE K (q8_0) and V (q5_1) into the resident byte caches at
3851    /// token index `t` (KVQUANT-PLAN §C). One CTA (one warp) per 32-element block; the kernel writes
3852    /// the f16 scale(s) + packed quants for K and V. k_row/v_row are f32 [kv_dim_k]/[kv_dim_v].
3853    pub fn append_kv_quantized(
3854        &self,
3855        k_row: &CudaSlice<f32>,
3856        v_row: &CudaSlice<f32>,
3857        kc: &mut CudaSlice<u8>,
3858        vc: &mut CudaSlice<u8>,
3859        t: usize,
3860        kv_dim_k: usize,
3861        kv_dim_v: usize,
3862        k_tok_bytes: usize,
3863        v_tok_bytes: usize,
3864        g: bool,
3865    ) -> Result<(), Box<dyn std::error::Error>> {
3866        let f = if g {
3867            self.func_g("append_quantize_kv_q8_0_q5_1")
3868        } else {
3869            self.func("append_quantize_kv_q8_0_q5_1")
3870        };
3871        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
3872        let cfg = LaunchConfig {
3873            grid_dim: (nblk, 1, 1),
3874            block_dim: (32, 1, 1),
3875            shared_mem_bytes: 0,
3876        };
3877        let (ti, kdk, kdv) = (t as i32, kv_dim_k as i32, kv_dim_v as i32);
3878        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3879        let __s_b = self.gpu.stream();
3880        let mut b = __s_b.launch_builder(&f);
3881        b.arg(k_row)
3882            .arg(v_row)
3883            .arg(kc)
3884            .arg(vc)
3885            .arg(&ti)
3886            .arg(&kdk)
3887            .arg(&kdv)
3888            .arg(&ktb)
3889            .arg(&vtb);
3890        unsafe {
3891            b.launch(cfg)?;
3892        }
3893        Ok(())
3894    }
3895
3896    /// Device-counter variant of `append_kv_quantized` (CUDA-GRAPH-PLAN Phase 2): the write slot
3897    /// `t` is read from `t_dev[0]` (a resident device i32[1]) instead of a host int arg, so the
3898    /// launch args are FIXED across decode steps (graph-capturable). Identical quant math.
3899    pub fn append_kv_quantized_dc(
3900        &self,
3901        k_row: &CudaSlice<f32>,
3902        v_row: &CudaSlice<f32>,
3903        kc: &mut CudaSlice<u8>,
3904        vc: &mut CudaSlice<u8>,
3905        t_dev: &CudaSlice<i32>,
3906        kv_dim_k: usize,
3907        kv_dim_v: usize,
3908        k_tok_bytes: usize,
3909        v_tok_bytes: usize,
3910        g: bool,
3911    ) -> Result<(), Box<dyn std::error::Error>> {
3912        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
3913        let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
3914        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3915        // PDL wave-B2: flash-module flavor mirrors the builder path's g flag exactly.
3916        if Self::pdl_on() && Self::pdl_wb_on() {
3917            use cudarc::driver::{DevicePtr, DevicePtrMut};
3918            let s = &self.gpu.stream();
3919            let (pk, _g0) = k_row.device_ptr(s);
3920            let (pv, _g1) = v_row.device_ptr(s);
3921            let (pkc, _g2) = kc.device_ptr_mut(s);
3922            let (pvc, _g3) = vc.device_ptr_mut(s);
3923            let (pt, _g4) = t_dev.device_ptr(s);
3924            let mut ps = [
3925                &pk as *const _ as *mut std::ffi::c_void,
3926                &pv as *const _ as *mut _,
3927                &pkc as *const _ as *mut _,
3928                &pvc as *const _ as *mut _,
3929                &pt as *const _ as *mut _,
3930                &kdk as *const _ as *mut _,
3931                &kdv as *const _ as *mut _,
3932                &ktb as *const _ as *mut _,
3933                &vtb as *const _ as *mut _,
3934            ];
3935            unsafe {
3936                self.launch_pdl_flash(
3937                    g,
3938                    "append_quantize_kv_q8_0_q5_1_dc",
3939                    (nblk, 1, 1),
3940                    (32, 1, 1),
3941                    0,
3942                    &mut ps,
3943                )?;
3944            }
3945            return Ok(());
3946        }
3947        let f = if g {
3948            self.func_g("append_quantize_kv_q8_0_q5_1_dc")
3949        } else {
3950            self.func("append_quantize_kv_q8_0_q5_1_dc")
3951        };
3952        let cfg = LaunchConfig {
3953            grid_dim: (nblk, 1, 1),
3954            block_dim: (32, 1, 1),
3955            shared_mem_bytes: 0,
3956        };
3957        let __s_b = self.gpu.stream();
3958        let mut b = __s_b.launch_builder(&f);
3959        b.arg(k_row)
3960            .arg(v_row)
3961            .arg(kc)
3962            .arg(vc)
3963            .arg(t_dev)
3964            .arg(&kdk)
3965            .arg(&kdv)
3966            .arg(&ktb)
3967            .arg(&vtb);
3968        unsafe {
3969            b.launch(cfg)?;
3970        }
3971        Ok(())
3972    }
3973
3974    /// Append-quantize T token rows in one shot (BATCHED PROMPT PRIME). k_rows/v_rows are
3975    /// token-major [T, kv_dim] post-RoPE f32; rows land at cache slots t0..t0+T. Default = the
3976    /// batched `_rows` kernel: one (nblk, T) launch whose per-(block,token) warp program is the
3977    /// per-token append kernel verbatim -> every written row is BIT-IDENTICAL to T sequential
3978    /// `append_kv_quantized_view` calls (kernel_check pins the bytes). MEMRA_PRIME_APPEND_LOOP=1
3979    /// forces the T-launch per-row loop (the A/B seam that measured the launch overhead).
3980    #[allow(clippy::too_many_arguments)]
3981    pub fn append_kv_quantized_rows(
3982        &self,
3983        k_rows: &CudaSlice<f32>,
3984        v_rows: &CudaSlice<f32>,
3985        kc: &mut CudaSlice<u8>,
3986        vc: &mut CudaSlice<u8>,
3987        t0: usize,
3988        t: usize,
3989        kv_dim_k: usize,
3990        kv_dim_v: usize,
3991        k_tok_bytes: usize,
3992        v_tok_bytes: usize,
3993        g: bool,
3994    ) -> Result<(), Box<dyn std::error::Error>> {
3995        if std::env::var("MEMRA_PRIME_APPEND_LOOP").is_ok() {
3996            for i in 0..t {
3997                let k_row = k_rows.slice(i * kv_dim_k..(i + 1) * kv_dim_k);
3998                let v_row = v_rows.slice(i * kv_dim_v..(i + 1) * kv_dim_v);
3999                self.append_kv_quantized_view(
4000                    &k_row,
4001                    &v_row,
4002                    kc,
4003                    vc,
4004                    t0 + i,
4005                    kv_dim_k,
4006                    kv_dim_v,
4007                    k_tok_bytes,
4008                    v_tok_bytes,
4009                    g,
4010                )?;
4011            }
4012            return Ok(());
4013        }
4014        let f = if g {
4015            self.func_g("append_quantize_kv_q8_0_q5_1_rows")
4016        } else {
4017            self.func("append_quantize_kv_q8_0_q5_1_rows")
4018        };
4019        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
4020        let cfg = LaunchConfig {
4021            grid_dim: (nblk, t as u32, 1),
4022            block_dim: (32, 1, 1),
4023            shared_mem_bytes: 0,
4024        };
4025        let (t0i, kdk, kdv) = (t0 as i32, kv_dim_k as i32, kv_dim_v as i32);
4026        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
4027        let __s_b = self.gpu.stream();
4028        let mut b = __s_b.launch_builder(&f);
4029        b.arg(k_rows)
4030            .arg(v_rows)
4031            .arg(kc)
4032            .arg(vc)
4033            .arg(&t0i)
4034            .arg(&kdk)
4035            .arg(&kdv)
4036            .arg(&ktb)
4037            .arg(&vtb);
4038        unsafe {
4039            b.launch(cfg)?;
4040        }
4041        Ok(())
4042    }
4043
4044    /// Increment a device i32[1] counter in place (p[0] += 1) via the resident `inc_i32` kernel.
4045    /// Used to advance the device-resident seqlen/pos counters inside the decode-dc path (and,
4046    /// later, inside a captured graph) without a host round-trip.
4047    pub fn inc_seqlen(&self, p: &mut CudaSlice<i32>) -> Result<(), Box<dyn std::error::Error>> {
4048        let f = self.func("inc_i32");
4049        let cfg = LaunchConfig {
4050            grid_dim: (1, 1, 1),
4051            block_dim: (1, 1, 1),
4052            shared_mem_bytes: 0,
4053        };
4054        let __s_b = self.gpu.stream();
4055        let mut b = __s_b.launch_builder(&f);
4056        b.arg(p);
4057        unsafe {
4058            b.launch(cfg)?;
4059        }
4060        Ok(())
4061    }
4062
4063    /// Like `append_kv_quantized` but k_row/v_row are CudaViews (one token's row sliced out of a
4064    /// token-major [T, kv_dim] activation buffer — the MTP verify path appends T tokens).
4065    pub fn append_kv_quantized_view(
4066        &self,
4067        k_row: &cudarc::driver::CudaView<f32>,
4068        v_row: &cudarc::driver::CudaView<f32>,
4069        kc: &mut CudaSlice<u8>,
4070        vc: &mut CudaSlice<u8>,
4071        t: usize,
4072        kv_dim_k: usize,
4073        kv_dim_v: usize,
4074        k_tok_bytes: usize,
4075        v_tok_bytes: usize,
4076        g: bool,
4077    ) -> Result<(), Box<dyn std::error::Error>> {
4078        let stream = self.gpu.stream();
4079        ensure_tensor_stream_device(k_row, &stream, "append_kv_quantized_view.k_row")?;
4080        ensure_tensor_stream_device(v_row, &stream, "append_kv_quantized_view.v_row")?;
4081        ensure_tensor_stream_device(kc, &stream, "append_kv_quantized_view.k_cache")?;
4082        ensure_tensor_stream_device(vc, &stream, "append_kv_quantized_view.v_cache")?;
4083        let f = if g {
4084            self.func_g("append_quantize_kv_q8_0_q5_1")
4085        } else {
4086            self.func("append_quantize_kv_q8_0_q5_1")
4087        };
4088        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
4089        let cfg = LaunchConfig {
4090            grid_dim: (nblk, 1, 1),
4091            block_dim: (32, 1, 1),
4092            shared_mem_bytes: 0,
4093        };
4094        let (ti, kdk, kdv) = (t as i32, kv_dim_k as i32, kv_dim_v as i32);
4095        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
4096        let mut b = stream.launch_builder(&f);
4097        b.arg(k_row)
4098            .arg(v_row)
4099            .arg(kc)
4100            .arg(vc)
4101            .arg(&ti)
4102            .arg(&kdk)
4103            .arg(&kdv)
4104            .arg(&ktb)
4105            .arg(&vtb);
4106        unsafe {
4107            b.launch(cfg)?;
4108        }
4109        Ok(())
4110    }
4111
4112    /// Device-to-device copy of a CudaView `src` into `dst[off..off+len]` (f32). Like `copy_into`
4113    /// but the source is a sub-view (e.g. one column of a token-major activation buffer).
4114    pub fn copy_view_into(
4115        &self,
4116        dst: &mut CudaSlice<f32>,
4117        off: usize,
4118        src: &cudarc::driver::CudaView<f32>,
4119        len: usize,
4120    ) -> Result<(), Box<dyn std::error::Error>> {
4121        let mut view = dst.slice_mut(off..off + len);
4122        self.gpu
4123            .stream()
4124            .memcpy_dtod(&src.slice(0..len), &mut view)?;
4125        Ok(())
4126    }
4127
4128    /// Real device-to-device COPY of `src` into a freshly allocated buffer (NOT an Arc clone).
4129    /// Used for cache snapshots (MTP-PLAN §D.4): `CudaSlice::clone()` only bumps a refcount and
4130    /// would alias the live buffer; this allocs new device memory and memcpy_dtod's the contents.
4131    pub fn clone_dtod(
4132        &self,
4133        src: &CudaSlice<f32>,
4134    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4135        let mut dst = self.gpu.stream().alloc_zeros::<f32>(src.len())?;
4136        self.gpu.stream().memcpy_dtod(src, &mut dst)?;
4137        Ok(dst)
4138    }
4139
4140    /// D2D row extraction: copy a view (e.g. one row of a [B, n] batch buffer) into `dst`.
4141    /// Stream-ordered, async — decode_batch's per-sequence row plumbing.
4142    pub fn dtod_copy_view(
4143        &self,
4144        src: &cudarc::driver::CudaView<f32>,
4145        dst: &mut CudaSlice<f32>,
4146    ) -> Result<(), Box<dyn std::error::Error>> {
4147        self.gpu.stream().memcpy_dtod(src, dst)?;
4148        Ok(())
4149    }
4150
4151    /// D2D i8 twin of `dtod_copy_view` (q8_1 activation rows).
4152    pub fn dtod_copy_view_i8(
4153        &self,
4154        src: &cudarc::driver::CudaView<i8>,
4155        dst: &mut CudaSlice<i8>,
4156    ) -> Result<(), Box<dyn std::error::Error>> {
4157        self.gpu.stream().memcpy_dtod(src, dst)?;
4158        Ok(())
4159    }
4160
4161    /// D2D row placement: copy `src` into `dst[offset .. offset+src.len()]`.
4162    pub fn dtod_copy_into(
4163        &self,
4164        src: &CudaSlice<f32>,
4165        dst: &mut CudaSlice<f32>,
4166        offset: usize,
4167    ) -> Result<(), Box<dyn std::error::Error>> {
4168        let n = src.len();
4169        let mut dv = dst.slice_mut(offset..offset + n);
4170        self.gpu.stream().memcpy_dtod(src, &mut dv)?;
4171        Ok(())
4172    }
4173
4174    /// Batched uniform-size D2D copy (engine-bundle slice 1, DSF-ROUNDCOST-20260820 §1.1):
4175    /// `n` disjoint regions of `words` f32 each; `table` = [src_0..src_{n-1}, dst_0..dst_{n-1}]
4176    /// raw device pointers. ONE kernel launch replaces `n` memcpy_dtod dispatches — the dspark
4177    /// round's snap/commit copy dribble (~0.9 ms/round of dispatch serialization measured).
4178    /// Bytes and stream order are identical to the memcpy sequence it replaces.
4179    pub fn copy_batch_uniform_f32(
4180        &self,
4181        table: &CudaSlice<u64>,
4182        n: usize,
4183        words: usize,
4184    ) -> Result<(), Box<dyn std::error::Error>> {
4185        if n == 0 || words == 0 {
4186            return Ok(());
4187        }
4188        debug_assert!(
4189            table.len() >= 2 * n,
4190            "pointer table must hold n srcs + n dsts"
4191        );
4192        let f = self.func("copy_batch_uniform_f32");
4193        // Enough blocks to stream a multi-MB region, few enough that (chunks x n) stays a
4194        // sane grid: 512K-word ssm regions get 48 grid-striding blocks each.
4195        let chunks = (words / 4).max(1).div_ceil(256).min(48) as u32;
4196        let (ni, wi) = (n as i32, words as i32);
4197        let cfg = LaunchConfig {
4198            grid_dim: (chunks, n as u32, 1),
4199            block_dim: (256, 1, 1),
4200            shared_mem_bytes: 0,
4201        };
4202        let __s = self.gpu.stream();
4203        let mut b = __s.launch_builder(&f);
4204        b.arg(table).arg(&ni).arg(&wi);
4205        unsafe {
4206            b.launch(cfg)?;
4207        }
4208        Ok(())
4209    }
4210
4211    /// H2D refresh of an EXISTING u64 pointer table IN PLACE (stable pointer — the batched
4212    /// state-copy tables are refreshed per round because the GDN ssm handles ping-pong).
4213    pub fn htod_u64_into(
4214        &self,
4215        v: &[u64],
4216        dst: &mut CudaSlice<u64>,
4217    ) -> Result<(), Box<dyn std::error::Error>> {
4218        let mut view = dst.slice_mut(0..v.len());
4219        self.gpu.stream().memcpy_htod(v, &mut view)?;
4220        Ok(())
4221    }
4222
4223    /// Indirect-source copy (engine-bundle slice 3): the src ADDRESS is loaded from a
4224    /// device pointer-table entry at run time, so a captured graph follows the gdn
4225    /// ping-pong through the same table its scan kernels read — a baked memcpy node
4226    /// would keep the capture-time physical buffer. `dst_off`/`words` in f32 elements.
4227    pub fn copy_indirect_src_f32(
4228        &self,
4229        src_entry: &cudarc::driver::CudaView<u64>,
4230        dst: &mut CudaSlice<f32>,
4231        dst_off: usize,
4232        words: usize,
4233    ) -> Result<(), Box<dyn std::error::Error>> {
4234        let f = self.func("copy_indirect_src_f32");
4235        let chunks = (words / 4).max(1).div_ceil(256).min(48) as u32;
4236        let wi = words as i32;
4237        let cfg = LaunchConfig {
4238            grid_dim: (chunks, 1, 1),
4239            block_dim: (256, 1, 1),
4240            shared_mem_bytes: 0,
4241        };
4242        let mut dv = dst.slice_mut(dst_off..dst_off + words);
4243        let __s = self.gpu.stream();
4244        let mut b = __s.launch_builder(&f);
4245        b.arg(src_entry).arg(&mut dv).arg(&wi);
4246        unsafe {
4247            b.launch(cfg)?;
4248        }
4249        Ok(())
4250    }
4251
4252    /// Uninitialized i8 device buffer (decode_batch q8_1 row scratch).
4253    pub fn uninit_i8(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
4254        self.alloc_uninit::<i8>(n)
4255    }
4256
4257    /// Resident-quantized linear (Stage-A: f32 dequant-in-kernel). y[m,out]=x[m,in]@W[out,in]^T.
4258    pub fn qmatvec(
4259        &self,
4260        w: &CudaSlice<u8>,
4261        x: &CudaSlice<f32>,
4262        m: usize,
4263        in_f: usize,
4264        out_f: usize,
4265        qtype: i32,
4266        row_bytes: usize,
4267    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4268        let f = self.func("qmatvec_f32");
4269        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
4270        let cfg = LaunchConfig {
4271            grid_dim: (out_f as u32, m as u32, 1),
4272            block_dim: (256, 1, 1),
4273            shared_mem_bytes: 0,
4274        };
4275        let (inf, outf, mi, qt, rb) =
4276            (in_f as i32, out_f as i32, m as i32, qtype, row_bytes as i64);
4277        let __s_b = self.gpu.stream();
4278        let mut b = __s_b.launch_builder(&f);
4279        b.arg(w)
4280            .arg(x)
4281            .arg(&mut y)
4282            .arg(&inf)
4283            .arg(&outf)
4284            .arg(&mi)
4285            .arg(&qt)
4286            .arg(&rb);
4287        unsafe {
4288            b.launch(cfg)?;
4289        }
4290        Ok(y)
4291    }
4292
4293    /// Allocate a reusable u8 GPU scratch buffer (for staged expert weights).
4294    pub fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
4295        let s = self.gpu.stream().alloc_zeros::<u8>(n)?;
4296        self.keep_if_capturing(&s);
4297        Ok(s)
4298    }
4299
4300    /// Uninitialized u8 scratch — skips alloc_zeros' memset. ONLY for staging buffers whose read
4301    /// range is fully overwritten by a stage_expert H2D before any kernel reads it (LAUNCH-STRUCTURE
4302    /// STAGE 2: the per-layer MoE scratch trio was 3 dead ~1MB memsets per layer per decode token).
4303    pub fn alloc_u8_uninit(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
4304        let s = unsafe { self.gpu.stream().alloc::<u8>(n)? };
4305        self.keep_if_capturing(&s);
4306        Ok(s)
4307    }
4308
4309    /// Zero a SUB-RANGE of an f32 buffer (CudaViewMut) — the row-sized memset the moe_out
4310    /// memset-elision uses for tokens that fall off the gdec fast path (LAUNCH-STRUCTURE STAGE 2).
4311    pub fn memset_zeros_view(
4312        &self,
4313        dst: &mut cudarc::driver::CudaViewMut<f32>,
4314    ) -> Result<(), Box<dyn std::error::Error>> {
4315        self.gpu.stream().memset_zeros(dst)?;
4316        Ok(())
4317    }
4318
4319    /// EDGE-1 staging: copy `host_bytes` (a sub-slice of a HostExps buffer) into `scratch`
4320    /// at byte offset `off` (async H2D on the default stream). Length is host_bytes.len().
4321    /// The qmatvec_view that reads `scratch[off..]` is enqueued on the SAME stream after this,
4322    /// so ordering is guaranteed without an explicit sync (Stage-1; Stage-2 prefetch on a 2nd
4323    /// stream would require an event).
4324    pub fn stage_expert(
4325        &self,
4326        host_bytes: &[u8],
4327        scratch: &mut CudaSlice<u8>,
4328        off: usize,
4329    ) -> Result<(), Box<dyn std::error::Error>> {
4330        let mut dst = scratch.slice_mut(off..off + host_bytes.len()); // CudaViewMut<u8>
4331        self.gpu.stream().memcpy_htod(host_bytes, &mut dst)?; // accepts &[u8] HostSlice src
4332        Ok(())
4333    }
4334
4335    /// EDGE-1 §A: fused MoE router. `logits` is the router output [t, n_expert] (device, f32, the
4336    /// `gate_inp @ z` result). Returns (sel_idx [t, n_used] i32, sel_w [t, n_used] f32): the top-k
4337    /// expert ids (DESC by prob, ascending-index tiebreak) and renormalized weights. Replaces the
4338    /// host dtoh + softmax-256 + stable DESC top-8 sort + renorm (hybrid_forward.rs ~281-298).
4339    /// One CTA per token row, 256 threads (one per expert).
4340    pub fn moe_router_topk(
4341        &self,
4342        logits: &CudaSlice<f32>,
4343        t: usize,
4344        n_expert: usize,
4345        n_used: usize,
4346    ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4347        let f = self.func("moe_router_topk_f32");
4348        let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?; // kernel fully overwrites
4349        let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?; // kernel fully overwrites
4350        let cfg = LaunchConfig {
4351            grid_dim: (t as u32, 1, 1),
4352            block_dim: (n_expert as u32, 1, 1),
4353            shared_mem_bytes: 0,
4354        };
4355        let (ne, nu) = (n_expert as i32, n_used as i32);
4356        let __s_b = self.gpu.stream();
4357        let mut b = __s_b.launch_builder(&f);
4358        b.arg(logits)
4359            .arg(&mut sel_idx)
4360            .arg(&mut sel_w)
4361            .arg(&ne)
4362            .arg(&nu);
4363        unsafe {
4364            b.launch(cfg)?;
4365        }
4366        Ok((sel_idx, sel_w))
4367    }
4368
4369    /// gemma4 twin: per-expert output scale folded into the topk renorm write (replaces the
4370    /// separate moe_w_exscale launch; value chain identical: (w/ws) * s[sel]).
4371    pub fn moe_router_topk_scaled(
4372        &self,
4373        logits: &CudaSlice<f32>,
4374        t: usize,
4375        n_expert: usize,
4376        n_used: usize,
4377        ex_scale: &CudaSlice<f32>,
4378    ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4379        // barrier-lean v2 twin (per-warp top-k + one-warp merge) FALSIFIED 2026-07-14:
4380        // bit-identical streams but −1.4% (26B plain N=3 interleaved) — at t=1 the grid is
4381        // ONE block, so the 6.6us is launch/dependency overhead, not the barrier chain;
4382        // fewer barriers bought nothing and the merge structure cost. jsonl is the record.
4383        let f = self.func("moe_router_topk_scaled_f32");
4384        let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?;
4385        let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?;
4386        let cfg = LaunchConfig {
4387            grid_dim: (t as u32, 1, 1),
4388            block_dim: (n_expert as u32, 1, 1),
4389            shared_mem_bytes: 0,
4390        };
4391        let (ne, nu) = (n_expert as i32, n_used as i32);
4392        let __s_b = self.gpu.stream();
4393        let mut b = __s_b.launch_builder(&f);
4394        b.arg(logits)
4395            .arg(&mut sel_idx)
4396            .arg(&mut sel_w)
4397            .arg(&ne)
4398            .arg(&nu)
4399            .arg(ex_scale);
4400        unsafe {
4401            b.launch(cfg)?;
4402        }
4403        Ok((sel_idx, sel_w))
4404    }
4405
4406    /// LAUNCH-STRUCTURE STAGE 1 (2026-07-05): fused router + SINGLE-SYNC host readback. The old
4407    /// MEMRA_FUSED_ROUTER path lost 2% at t=1 because it paid TWO full stream syncs (dtoh_i32 then
4408    /// dtoh, each = clone_dtoh + synchronize) + two alloc_zeros memsets per MoE layer, where the
4409    /// host route pays ONE sync on the 1KB logits dtoh. This variant: uninit outputs (kernel fully
4410    /// overwrites), both DtoH copies issued ASYNC into a persistent PINNED host staging buffer
4411    /// (flags=0 — cacheable, NOT cudarc's WRITECOMBINED default, so the host-side reads of sel/w
4412    /// stay cached), then ONE synchronize. Numerics identical to `moe_router_topk` (same kernel).
4413    pub fn moe_router_topk_host(
4414        &self,
4415        logits: &CudaSlice<f32>,
4416        t: usize,
4417        n_expert: usize,
4418        n_used: usize,
4419    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
4420        let f = self.func("moe_router_topk_f32");
4421        let n = t * n_used;
4422        let mut sel_idx = self.alloc_uninit::<i32>(n)?;
4423        let mut sel_w = self.alloc_uninit::<f32>(n)?;
4424        let cfg = LaunchConfig {
4425            grid_dim: (t as u32, 1, 1),
4426            block_dim: (n_expert as u32, 1, 1),
4427            shared_mem_bytes: 0,
4428        };
4429        let (ne, nu) = (n_expert as i32, n_used as i32);
4430        let __s_b = self.gpu.stream();
4431        let mut b = __s_b.launch_builder(&f);
4432        b.arg(logits)
4433            .arg(&mut sel_idx)
4434            .arg(&mut sel_w)
4435            .arg(&ne)
4436            .arg(&nu);
4437        unsafe {
4438            b.launch(cfg)?;
4439        }
4440        // single-sync readback: sel (i32) at offset 0, w (f32) at offset n*4 of the pinned stage.
4441        let bytes = n * 8;
4442        let mut guard = self.router_stage.lock().unwrap();
4443        if guard.as_ref().map(|p| p.cap < bytes).unwrap_or(true) {
4444            *guard = Some(PinnedStage::new(bytes.max(4096))?);
4445        }
4446        let stage = guard.as_mut().unwrap();
4447        let (si, sw) = unsafe {
4448            (
4449                std::slice::from_raw_parts_mut(stage.ptr as *mut i32, n),
4450                std::slice::from_raw_parts_mut(stage.ptr.add(n * 4) as *mut f32, n),
4451            )
4452        };
4453        self.gpu.stream().memcpy_dtoh(&sel_idx, si)?; // async (pinned dst)
4454        self.gpu.stream().memcpy_dtoh(&sel_w, sw)?; // async (pinned dst)
4455        self.gpu.stream().synchronize()?; // ONE sync for both
4456        Ok((si.iter().map(|&i| i as u32).collect(), sw.to_vec()))
4457    }
4458
4459    /// Device sigmoid router for Step-3.7 / DeepSeek-V3-class MoEs. `correction_bias` is added
4460    /// only to the top-k key; returned weights use the un-biased sigmoid score. `active` masks
4461    /// original expert ids before top-k. Exact key ties choose the smaller original id.
4462    #[allow(clippy::too_many_arguments)]
4463    pub fn moe_router_sigmoid_topk(
4464        &self,
4465        logits: &CudaSlice<f32>,
4466        t: usize,
4467        n_expert: usize,
4468        n_used: usize,
4469        active_count: usize,
4470        correction_bias: &CudaSlice<f32>,
4471        active: &CudaSlice<u8>,
4472        scaling_factor: f32,
4473        route_norm: bool,
4474    ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4475        crate::sigrouter_contract::validate_active_count(n_used, active_count)?;
4476        if n_expert == 0 || n_expert > 1024 || n_used == 0 || n_used > n_expert {
4477            return Err(format!(
4478                "sigmoid router shape unsupported: n_expert={n_expert}, n_used={n_used}",
4479            )
4480            .into());
4481        }
4482        if logits.len() < t * n_expert
4483            || correction_bias.len() != n_expert
4484            || active.len() != n_expert
4485        {
4486            return Err(format!(
4487                "sigmoid router buffer mismatch: logits={} bias={} active={} expected logits>={} row={}",
4488                logits.len(), correction_bias.len(), active.len(), t * n_expert, n_expert,
4489            ).into());
4490        }
4491        let f = if crate::sig_expf_dev_on() && crate::topk_fast_on() {
4492            // Latency twin of the dexp arm (identical outputs): barrier-lean top-k
4493            // over the dexp scoring class. Composes the two doors it rides.
4494            self.func("moe_router_sigmoid_topk_f32_dexp_fast")
4495        } else if crate::sig_expf_dev_on() {
4496            self.func("moe_router_sigmoid_topk_f32_dexp")
4497        } else if crate::topk_fast_on() {
4498            self.func("moe_router_sigmoid_topk_f32_fast")
4499        } else {
4500            self.func("moe_router_sigmoid_topk_f32")
4501        };
4502        let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?;
4503        let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?;
4504        let threads = n_expert.div_ceil(32) * 32;
4505        let cfg = LaunchConfig {
4506            grid_dim: (t as u32, 1, 1),
4507            block_dim: (threads as u32, 1, 1),
4508            shared_mem_bytes: 0,
4509        };
4510        let (ne, nu, rn) = (n_expert as i32, n_used as i32, i32::from(route_norm));
4511        let __s_b = self.gpu.stream();
4512        let mut b = __s_b.launch_builder(&f);
4513        b.arg(logits)
4514            .arg(correction_bias)
4515            .arg(active)
4516            .arg(&mut sel_idx)
4517            .arg(&mut sel_w)
4518            .arg(&ne)
4519            .arg(&nu)
4520            .arg(&scaling_factor)
4521            .arg(&rn);
4522        unsafe {
4523            b.launch(cfg)?;
4524        }
4525        Ok((sel_idx, sel_w))
4526    }
4527
4528    /// `moe_router_sigmoid_topk` writing into caller-owned buffers (alloc-free: child graphs
4529    /// cannot contain mem nodes, so the token-graph e-sections pre-own every output).
4530    #[allow(clippy::too_many_arguments)]
4531    pub fn moe_router_sigmoid_topk_into(
4532        &self,
4533        logits: &CudaSlice<f32>,
4534        t: usize,
4535        n_expert: usize,
4536        n_used: usize,
4537        active_count: usize,
4538        correction_bias: &CudaSlice<f32>,
4539        active: &CudaSlice<u8>,
4540        scaling_factor: f32,
4541        route_norm: bool,
4542        sel_idx: &mut CudaSlice<i32>,
4543        sel_w: &mut CudaSlice<f32>,
4544    ) -> Result<(), Box<dyn std::error::Error>> {
4545        crate::sigrouter_contract::validate_active_count(n_used, active_count)?;
4546        if n_expert == 0
4547            || n_expert > 1024
4548            || n_used == 0
4549            || n_used > n_expert
4550            || logits.len() < t * n_expert
4551            || correction_bias.len() != n_expert
4552            || active.len() != n_expert
4553            || sel_idx.len() < t * n_used
4554            || sel_w.len() < t * n_used
4555        {
4556            return Err("sigmoid router _into geometry mismatch".into());
4557        }
4558        let f = if crate::sig_expf_dev_on() {
4559            self.func("moe_router_sigmoid_topk_f32_dexp")
4560        } else if crate::topk_fast_on() {
4561            self.func("moe_router_sigmoid_topk_f32_fast")
4562        } else {
4563            self.func("moe_router_sigmoid_topk_f32")
4564        };
4565        let threads = n_expert.div_ceil(32) * 32;
4566        let cfg = LaunchConfig {
4567            grid_dim: (t as u32, 1, 1),
4568            block_dim: (threads as u32, 1, 1),
4569            shared_mem_bytes: 0,
4570        };
4571        let (ne, nu, rn) = (n_expert as i32, n_used as i32, i32::from(route_norm));
4572        let __s_b = self.gpu.stream();
4573        let mut b = __s_b.launch_builder(&f);
4574        b.arg(logits)
4575            .arg(correction_bias)
4576            .arg(active)
4577            .arg(&mut *sel_idx)
4578            .arg(&mut *sel_w)
4579            .arg(&ne)
4580            .arg(&nu)
4581            .arg(&scaling_factor)
4582            .arg(&rn);
4583        unsafe {
4584            b.launch(cfg)?;
4585        }
4586        Ok(())
4587    }
4588
4589    /// Single-sync pinned readback twin of `moe_router_sigmoid_topk`. This preserves the existing
4590    /// grouped/staged dispatch contract while replacing the full-logit DtoH plus host sigmoid/sort.
4591    #[allow(clippy::too_many_arguments)]
4592    pub fn moe_router_sigmoid_topk_host(
4593        &self,
4594        logits: &CudaSlice<f32>,
4595        t: usize,
4596        n_expert: usize,
4597        n_used: usize,
4598        active_count: usize,
4599        correction_bias: &CudaSlice<f32>,
4600        active: &CudaSlice<u8>,
4601        scaling_factor: f32,
4602        route_norm: bool,
4603    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
4604        let (sel_idx, sel_w) = self.moe_router_sigmoid_topk(
4605            logits,
4606            t,
4607            n_expert,
4608            n_used,
4609            active_count,
4610            correction_bias,
4611            active,
4612            scaling_factor,
4613            route_norm,
4614        )?;
4615        let n = t * n_used;
4616        let bytes = n * 8;
4617        let mut guard = self.router_stage.lock().unwrap();
4618        if guard.as_ref().map(|p| p.cap < bytes).unwrap_or(true) {
4619            *guard = Some(PinnedStage::new(bytes.max(4096))?);
4620        }
4621        let stage = guard.as_mut().unwrap();
4622        let (si, sw) = unsafe {
4623            (
4624                std::slice::from_raw_parts_mut(stage.ptr as *mut i32, n),
4625                std::slice::from_raw_parts_mut(stage.ptr.add(n * 4) as *mut f32, n),
4626            )
4627        };
4628        self.gpu.stream().memcpy_dtoh(&sel_idx, si)?;
4629        self.gpu.stream().memcpy_dtoh(&sel_w, sw)?;
4630        self.gpu.stream().synchronize()?;
4631        Ok((si.iter().map(|&i| i as u32).collect(), sw.to_vec()))
4632    }
4633
4634    /// EDGE-1 §C.2: async H2D of `host_bytes` into `scratch[off..]` on the COPY stream, returning a
4635    /// recorded event the compute stream can `wait` on before the dependent GEMM. Used for in-token
4636    /// expert prefetch (pipeline by one). `host_bytes` should be pinned for a true DMA (§C.1).
4637    pub fn stage_expert_async(
4638        &self,
4639        host_bytes: &[u8],
4640        scratch: &mut CudaSlice<u8>,
4641        off: usize,
4642    ) -> Result<cudarc::driver::CudaEvent, Box<dyn std::error::Error>> {
4643        let mut dst = scratch.slice_mut(off..off + host_bytes.len());
4644        self.copy_stream.memcpy_htod(host_bytes, &mut dst)?;
4645        Ok(self.copy_stream.record_event(None)?)
4646    }
4647
4648    /// Make the compute stream wait for an async copy event (the consumer side of `stage_expert_async`).
4649    pub fn compute_wait(
4650        &self,
4651        ev: &cudarc::driver::CudaEvent,
4652    ) -> Result<(), Box<dyn std::error::Error>> {
4653        self.gpu.stream().wait(ev)?;
4654        Ok(())
4655    }
4656
4657    /// qmatvec over a byte sub-range of a (resident/scratch) CudaSlice<u8> holding ONE expert
4658    /// matrix. x is a CudaView<f32> (a sliced row of z, or a sliced activation). Reuses the
4659    /// validated qmatvec_f32 dequant path (NOT a fast path — the correctness gate). The
4660    /// CudaView base+offset pointer is honored by the launch arg.
4661    pub fn qmatvec_view(
4662        &self,
4663        w: &CudaSlice<u8>,
4664        range: std::ops::Range<usize>,
4665        x: &cudarc::driver::CudaView<f32>,
4666        m: usize,
4667        in_f: usize,
4668        out_f: usize,
4669        qtype: i32,
4670        row_bytes: usize,
4671    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4672        let f = self.func("qmatvec_f32");
4673        let wv = w.slice(range); // CudaView<u8>, offset honored
4674        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
4675        let cfg = LaunchConfig {
4676            grid_dim: (out_f as u32, m as u32, 1),
4677            block_dim: (256, 1, 1),
4678            shared_mem_bytes: 0,
4679        };
4680        let (inf, outf, mi, qt, rb) =
4681            (in_f as i32, out_f as i32, m as i32, qtype, row_bytes as i64);
4682        let __s_b = self.gpu.stream();
4683        let mut b = __s_b.launch_builder(&f);
4684        b.arg(&wv)
4685            .arg(x)
4686            .arg(&mut y)
4687            .arg(&inf)
4688            .arg(&outf)
4689            .arg(&mi)
4690            .arg(&qt)
4691            .arg(&rb);
4692        unsafe {
4693            b.launch(cfg)?;
4694        }
4695        Ok(y)
4696    }
4697
4698    /// STAGE-2 GROUPED DECODE (2026-07-04): one MoE layer's gate+up+SiLU for all `n_used` routed
4699    /// experts of ONE token in ONE launch (replaces 8x qmatvec(gate) + 8x qmatvec(up) + 8x
4700    /// silu_mul = 24 launches). `gp`/`up` are the 8 expert weight-block device pointers (SLRU
4701    /// cache slots — fixed-address, stable for the launch). Returns act [n_used, n_ff].
4702    /// BIT-IDENTICAL to the sequential chain: each dot reproduces qmatvec_f32's exact 256-thread
4703    /// reduction; the SiLU epilogue is silu_mul_f32's exact expression (see kernel header).
4704    #[allow(clippy::too_many_arguments)]
4705    /// dp4a q8 twins (MoE expert dp4a arc, 2026-07-06): same contract as the _f32 versions but
4706    /// consume a PRE-QUANTIZED q8_1 activation. FP-order differs from _f32 (int dot + warp tree)
4707    /// — the argmax/stream-identity battery arbitrates; MEMRA_MOE_Q8=0 restores f32.
4708    pub fn moe_gate_up_silu8_q8(
4709        &self,
4710        gp: WPtr8,
4711        up: WPtr8,
4712        aq: &CudaSlice<i8>,
4713        ad: &CudaSlice<f32>,
4714        in_f: usize,
4715        n_ff: usize,
4716        n_used: usize,
4717        qt_g: i32,
4718        qt_u: i32,
4719        rb_g: usize,
4720        rb_u: usize,
4721    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4722        let f = self.func("moe_gate_up_silu8_q8");
4723        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
4724        let cfg = LaunchConfig {
4725            grid_dim: (n_ff as u32, n_used as u32, 1),
4726            block_dim: (32, 1, 1),
4727            shared_mem_bytes: 0,
4728        };
4729        let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
4730        let __s_b = self.gpu.stream();
4731        let mut b = __s_b.launch_builder(&f);
4732        b.arg(&gp)
4733            .arg(&up)
4734            .arg(aq)
4735            .arg(ad)
4736            .arg(&mut act)
4737            .arg(&inf)
4738            .arg(&nff)
4739            .arg(&qt_g)
4740            .arg(&qt_u)
4741            .arg(&rbg)
4742            .arg(&rbu);
4743        unsafe {
4744            b.launch(cfg)?;
4745        }
4746        Ok(act)
4747    }
4748
4749    #[allow(clippy::too_many_arguments)]
4750    pub fn moe_down8_fma_q8(
4751        &self,
4752        dp: WPtr8,
4753        w: F32x8,
4754        aq2: &CudaSlice<i8>,
4755        ad2: &CudaSlice<f32>,
4756        dst: &mut cudarc::driver::CudaViewMut<f32>,
4757        in_f: usize,
4758        out_f: usize,
4759        n_used: usize,
4760        qt: i32,
4761        rb: usize,
4762    ) -> Result<(), Box<dyn std::error::Error>> {
4763        let f = self.func("moe_down8_fma_q8");
4764        let cfg = LaunchConfig {
4765            grid_dim: (out_f as u32, 1, 1),
4766            block_dim: (32, 1, 1),
4767            shared_mem_bytes: 0,
4768        };
4769        let (inf, outf, nu, rbi) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
4770        let __s_b = self.gpu.stream();
4771        let mut b = __s_b.launch_builder(&f);
4772        b.arg(&dp)
4773            .arg(&w)
4774            .arg(aq2)
4775            .arg(ad2)
4776            .arg(dst)
4777            .arg(&inf)
4778            .arg(&outf)
4779            .arg(&nu)
4780            .arg(&qt)
4781            .arg(&rbi);
4782        unsafe {
4783            b.launch(cfg)?;
4784        }
4785        Ok(())
4786    }
4787
4788    /// q8 sequential expert matvec (staged path twin of qmatvec_view for IQ3_S/IQ4_XS).
4789    pub fn qmatvec_expert_q8(
4790        &self,
4791        w: &CudaSlice<u8>,
4792        range: std::ops::Range<usize>,
4793        aq: &CudaSlice<i8>,
4794        ad: &CudaSlice<f32>,
4795        m: usize,
4796        in_f: usize,
4797        out_f: usize,
4798        qtype: i32,
4799        row_bytes: usize,
4800    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4801        let f = self.func("qmatvec_expert_q8");
4802        let wv = w.slice(range);
4803        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
4804        const ROWS: u32 = 4; // MEMRA_MMVQ_ROWS
4805        let cfg = LaunchConfig {
4806            grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, m as u32, 1),
4807            block_dim: (32, ROWS, 1),
4808            shared_mem_bytes: 0,
4809        };
4810        let (inf, outf, mi, rbi) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
4811        let __s_b = self.gpu.stream();
4812        let mut b = __s_b.launch_builder(&f);
4813        b.arg(&wv)
4814            .arg(aq)
4815            .arg(ad)
4816            .arg(&mut y)
4817            .arg(&inf)
4818            .arg(&outf)
4819            .arg(&mi)
4820            .arg(&qtype)
4821            .arg(&rbi);
4822        unsafe {
4823            b.launch(cfg)?;
4824        }
4825        Ok(y)
4826    }
4827
4828    pub fn moe_gate_up_silu8(
4829        &self,
4830        gp: WPtr8,
4831        up: WPtr8,
4832        x: &cudarc::driver::CudaView<f32>,
4833        in_f: usize,
4834        n_ff: usize,
4835        n_used: usize,
4836        qt_g: i32,
4837        qt_u: i32,
4838        rb_g: usize,
4839        rb_u: usize,
4840    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4841        let f = self.func("moe_gate_up_silu8_f32");
4842        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?; // fully overwritten
4843        let cfg = LaunchConfig {
4844            grid_dim: (n_ff as u32, n_used as u32, 1),
4845            block_dim: (256, 1, 1),
4846            shared_mem_bytes: 0,
4847        };
4848        let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
4849        let __s_b = self.gpu.stream();
4850        let mut b = __s_b.launch_builder(&f);
4851        b.arg(&gp)
4852            .arg(&up)
4853            .arg(x)
4854            .arg(&mut act)
4855            .arg(&inf)
4856            .arg(&nff)
4857            .arg(&qt_g)
4858            .arg(&qt_u)
4859            .arg(&rbg)
4860            .arg(&rbu);
4861        unsafe {
4862            b.launch(cfg)?;
4863        }
4864        Ok(act)
4865    }
4866
4867    /// STAGE-2 GROUPED DECODE: one MoE layer's down-proj + weighted accumulation for all `n_used`
4868    /// routed experts in ONE launch (replaces 8x qmatvec(down) + 8x axpy = 16 launches), writing
4869    /// the token's moe_out row DIRECTLY (`dst` is the zeroed row; the in-kernel slot-ordered
4870    /// __fmaf_rn chain starting at 0.0f reproduces the sequential axpy_f32 accumulation into the
4871    /// zeroed row bit-for-bit — the A2 byte-identity scheme at m=1).
4872    #[allow(clippy::too_many_arguments)]
4873    pub fn moe_down8_fma_into(
4874        &self,
4875        dp: WPtr8,
4876        w: F32x8,
4877        act: &CudaSlice<f32>,
4878        dst: &mut cudarc::driver::CudaViewMut<f32>,
4879        in_f: usize,
4880        out_f: usize,
4881        n_used: usize,
4882        qt: i32,
4883        rb: usize,
4884    ) -> Result<(), Box<dyn std::error::Error>> {
4885        let f = self.func("moe_down8_fma_f32");
4886        let cfg = LaunchConfig {
4887            grid_dim: (out_f as u32, 1, 1),
4888            block_dim: (256, 1, 1),
4889            shared_mem_bytes: 0,
4890        };
4891        let (inf, outf, nu, rbv) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
4892        let __s_b = self.gpu.stream();
4893        let mut b = __s_b.launch_builder(&f);
4894        b.arg(&dp)
4895            .arg(&w)
4896            .arg(act)
4897            .arg(dst)
4898            .arg(&inf)
4899            .arg(&outf)
4900            .arg(&nu)
4901            .arg(&qt)
4902            .arg(&rbv);
4903        unsafe {
4904            b.launch(cfg)?;
4905        }
4906        Ok(())
4907    }
4908
4909    /// LAUNCH-STRUCTURE STAGE 3: device-dispatch twin of `moe_gate_up_silu8` for FULLY-RESIDENT
4910    /// layers. The expert ids come from the router kernel's DEVICE `sel` output (no DtoH) and the
4911    /// weight pointers from the per-layer device table `[3, n_expert]` of slot base addresses.
4912    /// BIT-IDENTICAL math (same grid/block/reduction; only the pointer/id source differs).
4913    #[allow(clippy::too_many_arguments)]
4914    /// dp4a q8 twin of the _dev pair (resident-experts arc).
4915    ///
4916    /// GEOMETRY VARIANTS (multirow/occupancy arc 2026-07-05): all outputs are BIT-IDENTICAL to
4917    /// the base one-warp-per-(row,slot) kernel (same expert_dot_g g-order + warp tree per row;
4918    /// down's FMA chain stays slot-ordered serial). Seams:
4919    ///   MEMRA_MOE_DEVQ8_GU   = 0(base) | 1 | 2 | 4 -> _r{1,2,4} multirow twin (RPW rows/warp)
4920    ///                       | s2 (gate/up warp split) | s2z (s2 + WPB rows packed per block)
4921    ///                       | gs4 (gate/up x low/high-group 4-warp split, nsb==64 only)
4922    ///                       | u64 (nsb==64 unrolled ILP twin, geometry unchanged)
4923    ///   MEMRA_MOE_DEVQ8_WPB  = warps per block for _r twins / z-rows for s2z (default 4)
4924    ///   MEMRA_MOE_DEVQ8_DOWN = auto(default: w8h2 when in_f==512 & n_used<=8 — measured +3.8%
4925    ///                       decode on 35B/rtx6000) | 0 (base one-warp serial-slot) | 1 | 2 | 4 ->
4926    ///                       _w8r{1,2,4} slot-parallel twin | h2 (half-warp dual-row, nsb==16
4927    ///                       only) | w8h2 (h2 x slot-parallel)
4928    #[allow(clippy::too_many_arguments)]
4929    /// MoE PREFILL pair-batch matvec: one launch covers all (token,expert) pairs for one proj.
4930    #[allow(clippy::too_many_arguments)]
4931    pub fn moe_pairs_matvec_q8(
4932        &self,
4933        table: &CudaSlice<u64>,
4934        proj: i32,
4935        pair_tok: &CudaSlice<i32>,
4936        pair_ex: &CudaSlice<i32>,
4937        aq: &CudaSlice<i8>,
4938        ad: &CudaSlice<f32>,
4939        in_f: usize,
4940        out_f: usize,
4941        n_expert: usize,
4942        n_pairs: usize,
4943        qtype: i32,
4944        row_bytes: usize,
4945    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4946        let f = self.func("moe_pairs_matvec_q8");
4947        let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
4948        const ROWS: u32 = 4;
4949        let cfg = LaunchConfig {
4950            grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_pairs as u32, 1),
4951            block_dim: (32, ROWS, 1),
4952            shared_mem_bytes: 0,
4953        };
4954        let (inf, outf, ne, np, rbi) = (
4955            in_f as i32,
4956            out_f as i32,
4957            n_expert as i32,
4958            n_pairs as i32,
4959            row_bytes as i64,
4960        );
4961        let __s_b = self.gpu.stream();
4962        let mut b = __s_b.launch_builder(&f);
4963        b.arg(table)
4964            .arg(&proj)
4965            .arg(pair_tok)
4966            .arg(pair_ex)
4967            .arg(aq)
4968            .arg(ad)
4969            .arg(&mut y)
4970            .arg(&inf)
4971            .arg(&outf)
4972            .arg(&ne)
4973            .arg(&np)
4974            .arg(&qtype)
4975            .arg(&rbi);
4976        unsafe {
4977            b.launch(cfg)?;
4978        }
4979        Ok(y)
4980    }
4981
4982    /// Expert-major pair matvec (weight-reuse across each expert's token group).
4983    #[allow(clippy::too_many_arguments)]
4984    pub fn moe_pairs_matvec_q8_em(
4985        &self,
4986        table: &CudaSlice<u64>,
4987        proj: i32,
4988        ex_ids: &CudaSlice<i32>,
4989        ex_off: &CudaSlice<i32>,
4990        ex_pairs: &CudaSlice<i32>,
4991        pair_tok: &CudaSlice<i32>,
4992        aq: &CudaSlice<i8>,
4993        ad: &CudaSlice<f32>,
4994        in_f: usize,
4995        out_f: usize,
4996        n_expert: usize,
4997        n_active: usize,
4998        n_pairs: usize,
4999        qtype: i32,
5000        row_bytes: usize,
5001    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5002        let f = self.func("moe_pairs_matvec_q8_em");
5003        let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
5004        const ROWS: u32 = 4;
5005        let cfg = LaunchConfig {
5006            grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_active as u32, 1),
5007            block_dim: (32, ROWS, 1),
5008            shared_mem_bytes: 0,
5009        };
5010        let (inf, outf, ne, na, rbi) = (
5011            in_f as i32,
5012            out_f as i32,
5013            n_expert as i32,
5014            n_active as i32,
5015            row_bytes as i64,
5016        );
5017        let __s_b = self.gpu.stream();
5018        let mut b = __s_b.launch_builder(&f);
5019        b.arg(table)
5020            .arg(&proj)
5021            .arg(ex_ids)
5022            .arg(ex_off)
5023            .arg(ex_pairs)
5024            .arg(pair_tok)
5025            .arg(aq)
5026            .arg(ad)
5027            .arg(&mut y)
5028            .arg(&inf)
5029            .arg(&outf)
5030            .arg(&ne)
5031            .arg(&na)
5032            .arg(&qtype)
5033            .arg(&rbi);
5034        unsafe {
5035            b.launch(cfg)?;
5036        }
5037        Ok(y)
5038    }
5039
5040    // Decode-once expert-major MMQ (rung 3). Same CSR inputs/geometry as _em; kernel dequants each
5041    // weight group once per (row,group) then dp4a's across the expert's token group.
5042    #[allow(clippy::too_many_arguments)]
5043    pub fn moe_pairs_matvec_q8_dec(
5044        &self,
5045        table: &CudaSlice<u64>,
5046        proj: i32,
5047        ex_ids: &CudaSlice<i32>,
5048        ex_off: &CudaSlice<i32>,
5049        ex_pairs: &CudaSlice<i32>,
5050        pair_tok: &CudaSlice<i32>,
5051        aq: &CudaSlice<i8>,
5052        ad: &CudaSlice<f32>,
5053        in_f: usize,
5054        out_f: usize,
5055        n_expert: usize,
5056        n_active: usize,
5057        n_pairs: usize,
5058        qtype: i32,
5059        row_bytes: usize,
5060    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5061        let f = self.func("moe_pairs_matvec_q8_dec");
5062        let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
5063        const ROWS: u32 = 4;
5064        let cfg = LaunchConfig {
5065            grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_active as u32, 1),
5066            block_dim: (32, ROWS, 1),
5067            shared_mem_bytes: 0,
5068        };
5069        let (inf, outf, ne, na, rbi) = (
5070            in_f as i32,
5071            out_f as i32,
5072            n_expert as i32,
5073            n_active as i32,
5074            row_bytes as i64,
5075        );
5076        let __s_b = self.gpu.stream();
5077        let mut b = __s_b.launch_builder(&f);
5078        b.arg(table)
5079            .arg(&proj)
5080            .arg(ex_ids)
5081            .arg(ex_off)
5082            .arg(ex_pairs)
5083            .arg(pair_tok)
5084            .arg(aq)
5085            .arg(ad)
5086            .arg(&mut y)
5087            .arg(&inf)
5088            .arg(&outf)
5089            .arg(&ne)
5090            .arg(&na)
5091            .arg(&qtype)
5092            .arg(&rbi);
5093        unsafe {
5094            b.launch(cfg)?;
5095        }
5096        Ok(y)
5097    }
5098
5099    pub fn moe_pairs_gelu_mul(
5100        &self,
5101        gate: &CudaSlice<f32>,
5102        up: &CudaSlice<f32>,
5103        n: usize,
5104    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5105        let f = self.func("moe_pairs_gelu_mul");
5106        let mut act = self.alloc_uninit::<f32>(n)?;
5107        let cfg = LaunchConfig::for_num_elems(n as u32);
5108        let nl = n as i64;
5109        let __s_b = self.gpu.stream();
5110        let mut b = __s_b.launch_builder(&f);
5111        b.arg(gate).arg(up).arg(&mut act).arg(&nl);
5112        unsafe {
5113            b.launch(cfg)?;
5114        }
5115        Ok(act)
5116    }
5117
5118    pub fn moe_pairs_silu_mul(
5119        &self,
5120        gate: &CudaSlice<f32>,
5121        up: &CudaSlice<f32>,
5122        n: usize,
5123    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5124        let f = self.func("moe_pairs_silu_mul");
5125        let mut act = self.alloc_uninit::<f32>(n)?;
5126        let cfg = LaunchConfig::for_num_elems(n as u32);
5127        let nl = n as i64;
5128        let __s_b = self.gpu.stream();
5129        let mut b = __s_b.launch_builder(&f);
5130        b.arg(gate).arg(up).arg(&mut act).arg(&nl);
5131        unsafe {
5132            b.launch(cfg)?;
5133        }
5134        Ok(act)
5135    }
5136
5137    #[allow(clippy::too_many_arguments)]
5138    pub fn moe_pairs_scatter(
5139        &self,
5140        y_down: &CudaSlice<f32>,
5141        pair_w: &CudaSlice<f32>,
5142        tok_pair_off: &CudaSlice<i32>,
5143        tok_pair_ids: &CudaSlice<i32>,
5144        moe_out: &mut CudaSlice<f32>,
5145        t: usize,
5146        n_embd: usize,
5147    ) -> Result<(), Box<dyn std::error::Error>> {
5148        let f = self.func("moe_pairs_scatter");
5149        let cfg = LaunchConfig {
5150            grid_dim: (((n_embd + 255) / 256) as u32, t as u32, 1),
5151            block_dim: (256, 1, 1),
5152            shared_mem_bytes: 0,
5153        };
5154        let ne = n_embd as i32;
5155        let __s_b = self.gpu.stream();
5156        let mut b = __s_b.launch_builder(&f);
5157        b.arg(y_down)
5158            .arg(pair_w)
5159            .arg(tok_pair_off)
5160            .arg(tok_pair_ids)
5161            .arg(moe_out)
5162            .arg(&ne);
5163        unsafe {
5164            b.launch(cfg)?;
5165        }
5166        Ok(())
5167    }
5168
5169    /// gemma4 GELU twin of moe_gate_up_silu8_dev_q8 (base geometry — slot-packed j8/j8r2
5170    /// twins probed 2026-08-01 g26 decode dig: bit-identical rows, -2.5%/-2.9% whole-model
5171    /// decode x3 interleaved -> refuted and killed; research/g26-decode-20260801/receipts.md).
5172    #[allow(clippy::too_many_arguments)]
5173    pub fn moe_gate_up_gelu8_dev_q8(
5174        &self,
5175        table: &CudaSlice<u64>,
5176        sel: &cudarc::driver::CudaView<i32>,
5177        aq: &CudaSlice<i8>,
5178        ad: &CudaSlice<f32>,
5179        in_f: usize,
5180        n_ff: usize,
5181        n_used: usize,
5182        n_expert: usize,
5183        qt_g: i32,
5184        qt_u: i32,
5185        rb_g: usize,
5186        rb_u: usize,
5187    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5188        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
5189        let (inf, nff, ne, rbg, rbu) = (
5190            in_f as i32,
5191            n_ff as i32,
5192            n_expert as i32,
5193            rb_g as i64,
5194            rb_u as i64,
5195        );
5196        let f = self.func("moe_gate_up_gelu8_dev_q8");
5197        let cfg = LaunchConfig {
5198            grid_dim: (n_ff as u32, n_used as u32, 1),
5199            block_dim: (32, 1, 1),
5200            shared_mem_bytes: 0,
5201        };
5202        let __s_b = self.gpu.stream();
5203        let mut b = __s_b.launch_builder(&f);
5204        b.arg(table)
5205            .arg(sel)
5206            .arg(aq)
5207            .arg(ad)
5208            .arg(&mut act)
5209            .arg(&inf)
5210            .arg(&nff)
5211            .arg(&ne)
5212            .arg(&qt_g)
5213            .arg(&qt_u)
5214            .arg(&rbg)
5215            .arg(&rbu);
5216        unsafe {
5217            b.launch(cfg)?;
5218        }
5219        Ok(act)
5220    }
5221
5222    /// gemma4 GELU rows twin (verify): one launch over (n_ff, n_used, t).
5223    #[allow(clippy::too_many_arguments)]
5224    pub fn moe_gate_up_gelu8_dev_q8_rows(
5225        &self,
5226        table: &CudaSlice<u64>,
5227        sel: &CudaSlice<i32>,
5228        aq: &CudaSlice<i8>,
5229        ad: &CudaSlice<f32>,
5230        t: usize,
5231        in_f: usize,
5232        n_ff: usize,
5233        n_used: usize,
5234        n_expert: usize,
5235        qt_g: i32,
5236        qt_u: i32,
5237        rb_g: usize,
5238        rb_u: usize,
5239    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5240        let mut act = self.alloc_uninit::<f32>(t * n_used * n_ff)?;
5241        let (inf, nff, ne, rbg, rbu, nu) = (
5242            in_f as i32,
5243            n_ff as i32,
5244            n_expert as i32,
5245            rb_g as i64,
5246            rb_u as i64,
5247            n_used as i32,
5248        );
5249        let f = self.func("moe_gate_up_gelu8_dev_q8_rows");
5250        let cfg = LaunchConfig {
5251            grid_dim: (n_ff as u32, n_used as u32, t as u32),
5252            block_dim: (32, 1, 1),
5253            shared_mem_bytes: 0,
5254        };
5255        let __s_b = self.gpu.stream();
5256        let mut b = __s_b.launch_builder(&f);
5257        b.arg(table)
5258            .arg(sel)
5259            .arg(aq)
5260            .arg(ad)
5261            .arg(&mut act)
5262            .arg(&inf)
5263            .arg(&nff)
5264            .arg(&ne)
5265            .arg(&qt_g)
5266            .arg(&qt_u)
5267            .arg(&rbg)
5268            .arg(&rbu)
5269            .arg(&nu);
5270        unsafe {
5271            b.launch(cfg)?;
5272        }
5273        Ok(act)
5274    }
5275
5276    /// gemma4 GELU CSR twin (verify dedup: owner block serves every pair of its expert).
5277    #[allow(clippy::too_many_arguments)]
5278    pub fn moe_gate_up_gelu8_dev_q8_csr(
5279        &self,
5280        table: &CudaSlice<u64>,
5281        sel: &CudaSlice<i32>,
5282        aq: &CudaSlice<i8>,
5283        ad: &CudaSlice<f32>,
5284        n_pairs: usize,
5285        in_f: usize,
5286        n_ff: usize,
5287        n_used: usize,
5288        n_expert: usize,
5289        qt_g: i32,
5290        qt_u: i32,
5291        rb_g: usize,
5292        rb_u: usize,
5293    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5294        let mut act = self.alloc_uninit::<f32>(n_pairs * n_ff)?;
5295        let (inf, nff, ne, rbg, rbu, nu, npi) = (
5296            in_f as i32,
5297            n_ff as i32,
5298            n_expert as i32,
5299            rb_g as i64,
5300            rb_u as i64,
5301            n_used as i32,
5302            n_pairs as i32,
5303        );
5304        let f = self.func("moe_gate_up_gelu8_dev_q8_csr");
5305        let cfg = LaunchConfig {
5306            grid_dim: (n_ff as u32, n_pairs as u32, 1),
5307            block_dim: (32, 1, 1),
5308            shared_mem_bytes: 0,
5309        };
5310        let __s_b = self.gpu.stream();
5311        let mut b = __s_b.launch_builder(&f);
5312        b.arg(table)
5313            .arg(sel)
5314            .arg(aq)
5315            .arg(ad)
5316            .arg(&mut act)
5317            .arg(&inf)
5318            .arg(&nff)
5319            .arg(&ne)
5320            .arg(&qt_g)
5321            .arg(&qt_u)
5322            .arg(&rbg)
5323            .arg(&rbu)
5324            .arg(&nu)
5325            .arg(&npi);
5326        unsafe {
5327            b.launch(cfg)?;
5328        }
5329        Ok(act)
5330    }
5331
5332    /// gemma4 generic down rows twin (verify): one launch over (out_f, 1, t).
5333    #[allow(clippy::too_many_arguments)]
5334    pub fn moe_down8_fma_dev_q8_rows_g(
5335        &self,
5336        table: &CudaSlice<u64>,
5337        sel: &CudaSlice<i32>,
5338        w: &CudaSlice<f32>,
5339        aq2: &CudaSlice<i8>,
5340        ad2: &CudaSlice<f32>,
5341        dst: &mut CudaSlice<f32>,
5342        t: usize,
5343        in_f: usize,
5344        out_f: usize,
5345        n_used: usize,
5346        n_expert: usize,
5347        qt: i32,
5348        rb: usize,
5349    ) -> Result<(), Box<dyn std::error::Error>> {
5350        let (inf, outf, nu, ne, rbi) = (
5351            in_f as i32,
5352            out_f as i32,
5353            n_used as i32,
5354            n_expert as i32,
5355            rb as i64,
5356        );
5357        // Exact Step-3.7 B=1 shape: expose the eight independent slot dots as
5358        // eight warps, then replay the original slot-ordered FMA chain. Every
5359        // other shape retains the generic one-warp rows kernel.
5360        let step_b1_w8 = t == 1 && in_f == 1280 && out_f == 4096 && n_used == 8 && qt == QT_IQ4_XS;
5361        let f = self.func(if step_b1_w8 {
5362            "moe_down8_fma_dev_q8_rows_w8"
5363        } else {
5364            "moe_down8_fma_dev_q8_rows_g"
5365        });
5366        let cfg = LaunchConfig {
5367            grid_dim: (out_f as u32, 1, t as u32),
5368            block_dim: (32, if step_b1_w8 { 8 } else { 1 }, 1),
5369            shared_mem_bytes: 0,
5370        };
5371        let __s_b = self.gpu.stream();
5372        let mut b = __s_b.launch_builder(&f);
5373        b.arg(table)
5374            .arg(sel)
5375            .arg(w)
5376            .arg(aq2)
5377            .arg(ad2)
5378            .arg(dst)
5379            .arg(&inf)
5380            .arg(&outf)
5381            .arg(&nu)
5382            .arg(&ne)
5383            .arg(&qt)
5384            .arg(&rbi);
5385        unsafe {
5386            b.launch(cfg)?;
5387        }
5388        Ok(())
5389    }
5390
5391    /// rp_q4 microprobe (2026-07-10 verify-trunk lever): b4 GGUF-block layout vs the Q4_0
5392    /// split-plane twin on the wq-class shape. Returns (blk_us, rp_us) after asserting bitwise
5393    /// identity. Bench-only surface (rp_q4_probe bin); no production dispatch reads this.
5394    pub fn rp_probe_q4(&self, m: usize) -> Result<(f64, f64), Box<dyn std::error::Error>> {
5395        let (out_f, in_f) = (2048usize, 2816usize);
5396        let nblk = in_f / 32;
5397        let mut seed = 0x9E3779B97F4A7C15u64;
5398        let mut rng = move || {
5399            seed = seed
5400                .wrapping_mul(6364136223846793005)
5401                .wrapping_add(1442695040888963407);
5402            (seed >> 33) as u8
5403        };
5404        let mut w = vec![0u8; out_f * nblk * 18];
5405        for b in w.iter_mut() {
5406            *b = rng();
5407        }
5408        for r in 0..out_f {
5409            for g in 0..nblk {
5410                let off = (r * nblk + g) * 18;
5411                w[off] = 0x00;
5412                w[off + 1] = 0x2C; // sane half d
5413            }
5414        }
5415        let qplane = out_f * nblk * 16;
5416        let mut wrp = vec![0u8; w.len()];
5417        for r in 0..out_f {
5418            for g in 0..nblk {
5419                let src = &w[(r * nblk + g) * 18..(r * nblk + g) * 18 + 18];
5420                wrp[qplane + (r * nblk + g) * 2..qplane + (r * nblk + g) * 2 + 2]
5421                    .copy_from_slice(&src[0..2]);
5422                wrp[(r * nblk + g) * 16..(r * nblk + g) * 16 + 16].copy_from_slice(&src[2..18]);
5423            }
5424        }
5425        let w_d = self.htod_bytes(&w)?;
5426        let wrp_d = self.htod_bytes(&wrp)?;
5427        let mut aq = vec![0i8; m * in_f];
5428        for v in aq.iter_mut() {
5429            *v = rng() as i8;
5430        }
5431        let aq_d = self.htod_i8(&aq)?;
5432        let ad_d = self.htod(&vec![0.03125f32; m * nblk])?;
5433        let mut y0 = self.alloc_uninit::<f32>(m * out_f)?;
5434        let mut y1 = self.alloc_uninit::<f32>(m * out_f)?;
5435        const RPB: u32 = 4;
5436        let cfg = LaunchConfig {
5437            grid_dim: ((out_f as u32).div_ceil(RPB), 1, 1),
5438            block_dim: (32, RPB, 1),
5439            shared_mem_bytes: 0,
5440        };
5441        let (inf, outf, mi) = (in_f as i32, out_f as i32, m as i32);
5442        let (rb, qp) = ((nblk * 18) as i64, qplane as i64);
5443        let fb = self.func("qmatvec_q4_0_mmvq_b4");
5444        let fr = self.func("qmatvec_q4_0_mmvq_b4_rp");
5445        {
5446            let __s_b = self.gpu.stream();
5447            let mut b = __s_b.launch_builder(&fb);
5448            b.arg(&w_d)
5449                .arg(&aq_d)
5450                .arg(&ad_d)
5451                .arg(&mut y0)
5452                .arg(&inf)
5453                .arg(&outf)
5454                .arg(&mi)
5455                .arg(&rb);
5456            unsafe {
5457                b.launch(cfg)?;
5458            }
5459            let __s_b = self.gpu.stream();
5460            let mut b = __s_b.launch_builder(&fr);
5461            b.arg(&wrp_d)
5462                .arg(&aq_d)
5463                .arg(&ad_d)
5464                .arg(&mut y1)
5465                .arg(&inf)
5466                .arg(&outf)
5467                .arg(&mi)
5468                .arg(&qp);
5469            unsafe {
5470                b.launch(cfg)?;
5471            }
5472        }
5473        self.gpu.stream().synchronize()?;
5474        let (h0, h1) = (self.dtoh(&y0)?, self.dtoh(&y1)?);
5475        let nd = h0
5476            .iter()
5477            .zip(&h1)
5478            .filter(|(a, b)| a.to_bits() != b.to_bits())
5479            .count();
5480        if nd != 0 {
5481            return Err(format!("rp twin not bitwise: {nd}/{} diffs", h0.len()).into());
5482        }
5483        let mut time = |rp: bool| -> Result<f64, Box<dyn std::error::Error>> {
5484            self.gpu.stream().synchronize()?;
5485            let t0 = std::time::Instant::now();
5486            for _ in 0..500 {
5487                if rp {
5488                    let __s_b = self.gpu.stream();
5489                    let mut b = __s_b.launch_builder(&fr);
5490                    b.arg(&wrp_d)
5491                        .arg(&aq_d)
5492                        .arg(&ad_d)
5493                        .arg(&mut y1)
5494                        .arg(&inf)
5495                        .arg(&outf)
5496                        .arg(&mi)
5497                        .arg(&qp);
5498                    unsafe {
5499                        b.launch(cfg)?;
5500                    }
5501                } else {
5502                    let __s_b = self.gpu.stream();
5503                    let mut b = __s_b.launch_builder(&fb);
5504                    b.arg(&w_d)
5505                        .arg(&aq_d)
5506                        .arg(&ad_d)
5507                        .arg(&mut y0)
5508                        .arg(&inf)
5509                        .arg(&outf)
5510                        .arg(&mi)
5511                        .arg(&rb);
5512                    unsafe {
5513                        b.launch(cfg)?;
5514                    }
5515                }
5516            }
5517            self.gpu.stream().synchronize()?;
5518            Ok(t0.elapsed().as_secs_f64() * 1e6 / 500.0)
5519        };
5520        let _ = time(false)?;
5521        let _ = time(true)?; // warm
5522        Ok((time(false)?, time(true)?))
5523    }
5524
5525    /// Build the Q4_0 split-plane decode mirror for a 2D Quant tensor (device-side permutation,
5526    /// q4_0_split_rp_build). Raw bytes stay resident (prefill/gemm/Stage-A); the m<=8 decode
5527    /// dispatch prefers the mirror (_rp twins). No-op unless (Q4_0, 2D, mirror absent).
5528    /// VRAM cost == the tensor's weight size. MEMRA_Q4RP=0 disables at the call sites.
5529    pub fn build_q4_rp4(
5530        &self,
5531        t: &mut crate::model::GpuTensor,
5532    ) -> Result<(), Box<dyn std::error::Error>> {
5533        use crate::model::GpuTensor;
5534        let GpuTensor::Quant {
5535            bytes,
5536            qtype,
5537            row_bytes,
5538            ne,
5539            rp4,
5540            ..
5541        } = t
5542        else {
5543            return Ok(());
5544        };
5545        if *qtype != QT_Q4_0 || rp4.is_some() || ne.len() != 2 {
5546            return Ok(());
5547        }
5548        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5549        if in_f % 32 != 0 || *row_bytes != (in_f / 32) * 18 {
5550            return Ok(());
5551        }
5552        let nblk = in_f / 32;
5553        let mut dst = self.alloc_uninit::<u8>(out_f * nblk * 18)?;
5554        let f = self.func("q4_0_split_rp_build");
5555        let n = (out_f * nblk) as i32;
5556        let cfg = LaunchConfig {
5557            grid_dim: (((out_f * nblk) as u32).div_ceil(256), 1, 1),
5558            block_dim: (256, 1, 1),
5559            shared_mem_bytes: 0,
5560        };
5561        let (of, nb) = (out_f as i32, nblk as i32);
5562        let _ = n;
5563        let __s_b = self.gpu.stream();
5564        let mut b = __s_b.launch_builder(&f);
5565        b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
5566        unsafe {
5567            b.launch(cfg)?;
5568        }
5569        *rp4 = Some(dst);
5570        Ok(())
5571    }
5572
5573    /// Q8_0 twin of `build_q4_rp4` (H100 coalescing fix, 2026-07-26 ncu: GGUF 34B-stride
5574    /// weight loads hold Max Bandwidth at 41-46%; the split mirror makes them aligned 16B
5575    /// ldcs). Raw bytes stay resident (prefill GEMM/MMQ/fused m=1 launches read GGUF layout);
5576    /// the mmvq/batched decode arms prefer the mirror via `rp4`. Bit-identical outputs.
5577    pub fn build_q8_rp4(
5578        &self,
5579        t: &mut crate::model::GpuTensor,
5580    ) -> Result<(), Box<dyn std::error::Error>> {
5581        use crate::model::GpuTensor;
5582        let GpuTensor::Quant {
5583            bytes,
5584            qtype,
5585            row_bytes,
5586            ne,
5587            rp4,
5588            ..
5589        } = t
5590        else {
5591            return Ok(());
5592        };
5593        if *qtype != QT_Q8_0 || rp4.is_some() || ne.len() != 2 {
5594            return Ok(());
5595        }
5596        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5597        if in_f % 32 != 0 || *row_bytes != (in_f / 32) * 34 {
5598            return Ok(());
5599        }
5600        *rp4 = Some(self.build_q8_rp4_raw(bytes, in_f, out_f)?);
5601        Ok(())
5602    }
5603
5604    /// Raw rp-mirror build for gates/benches: split GGUF Q8_0 bytes into the qplane+dplane
5605    /// mirror without a GpuTensor (same kernel the loader path above uses).
5606    pub fn build_q8_rp4_raw(
5607        &self,
5608        bytes: &CudaSlice<u8>,
5609        in_f: usize,
5610        out_f: usize,
5611    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
5612        assert!(in_f % 32 == 0);
5613        let nblk = in_f / 32;
5614        let mut dst = self.alloc_uninit::<u8>(out_f * nblk * 34)?;
5615        let f = self.func("q8_0_split_rp_build");
5616        let cfg = LaunchConfig {
5617            grid_dim: (((out_f * nblk) as u32).div_ceil(256), 1, 1),
5618            block_dim: (256, 1, 1),
5619            shared_mem_bytes: 0,
5620        };
5621        let (of, nb) = (out_f as i32, nblk as i32);
5622        let __s_b = self.gpu.stream();
5623        let mut b = __s_b.launch_builder(&f);
5624        b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
5625        unsafe {
5626            b.launch(cfg)?;
5627        }
5628        Ok(dst)
5629    }
5630
5631    /// K-quant twins of `build_q8_rp4` (H100 K-quant coalescing fix, 2026-08-01 ncu on the
5632    /// q27 Q4_K_M decode: q4_K mmvq DRAM 41-54% with 65% excessive sectors, q6_K 40% with
5633    /// 78% — the 144B/210B superblock strides land every 4B weight load off-sector). The
5634    /// mirror re-packs each tensor into planes (q4_K: qs ++ 16B meta; q6_K: ql ++ qh ++
5635    /// scales ++ d — same total bytes) so every quant fetch is an aligned 16B ldcs. Raw
5636    /// bytes stay resident (prefill GEMM/dequant/Stage-A read GGUF layout); the mmvq/batched
5637    /// decode arms prefer the mirror via `rp4`. Bit-identical outputs.
5638    pub fn build_q4k_rp4(
5639        &self,
5640        t: &mut crate::model::GpuTensor,
5641    ) -> Result<(), Box<dyn std::error::Error>> {
5642        use crate::model::GpuTensor;
5643        let GpuTensor::Quant {
5644            bytes,
5645            qtype,
5646            row_bytes,
5647            ne,
5648            rp4,
5649            ..
5650        } = t
5651        else {
5652            return Ok(());
5653        };
5654        if *qtype != QT_Q4_K || rp4.is_some() || ne.len() != 2 {
5655            return Ok(());
5656        }
5657        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5658        if in_f % 256 != 0 || *row_bytes != (in_f / 256) * 144 {
5659            return Ok(());
5660        }
5661        *rp4 = Some(self.build_kq_rp4_raw(bytes, in_f, out_f, QT_Q4_K)?);
5662        Ok(())
5663    }
5664
5665    pub fn build_q6k_rp4(
5666        &self,
5667        t: &mut crate::model::GpuTensor,
5668    ) -> Result<(), Box<dyn std::error::Error>> {
5669        use crate::model::GpuTensor;
5670        let GpuTensor::Quant {
5671            bytes,
5672            qtype,
5673            row_bytes,
5674            ne,
5675            rp4,
5676            ..
5677        } = t
5678        else {
5679            return Ok(());
5680        };
5681        if *qtype != QT_Q6_K || rp4.is_some() || ne.len() != 2 {
5682            return Ok(());
5683        }
5684        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5685        if in_f % 256 != 0 || *row_bytes != (in_f / 256) * 210 {
5686            return Ok(());
5687        }
5688        *rp4 = Some(self.build_kq_rp4_raw(bytes, in_f, out_f, QT_Q6_K)?);
5689        Ok(())
5690    }
5691
5692    /// Raw K-quant rp-mirror build for gates/benches (same kernels the loader path uses).
5693    pub fn build_kq_rp4_raw(
5694        &self,
5695        bytes: &CudaSlice<u8>,
5696        in_f: usize,
5697        out_f: usize,
5698        qtype: i32,
5699    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
5700        assert!(in_f % 256 == 0);
5701        let nsbk = in_f / 256;
5702        let (sb_bytes, kname) = match qtype {
5703            QT_Q4_K => (144usize, "q4_K_split_rp_build"),
5704            QT_Q6_K => (210usize, "q6_K_split_rp_build"),
5705            _ => return Err(format!("build_kq_rp4_raw: qtype {qtype} has no rp mirror").into()),
5706        };
5707        let mut dst = self.alloc_uninit::<u8>(out_f * nsbk * sb_bytes)?;
5708        let f = self.func(kname);
5709        let cfg = LaunchConfig {
5710            grid_dim: (((out_f * nsbk) as u32).div_ceil(256), 1, 1),
5711            block_dim: (256, 1, 1),
5712            shared_mem_bytes: 0,
5713        };
5714        let (of, nb) = (out_f as i32, nsbk as i32);
5715        let __s_b = self.gpu.stream();
5716        let mut b = __s_b.launch_builder(&f);
5717        b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
5718        unsafe {
5719            b.launch(cfg)?;
5720        }
5721        Ok(dst)
5722    }
5723
5724    /// MEMRA_KQRP seam: the K-quant (q4_K/q6_K) split-plane decode mirrors at model load.
5725    /// Default follows the Q8RP convention — ON on the Hopper lane (80GB pays the mirror
5726    /// VRAM), OFF elsewhere (a 24GB card cannot hold model + mirror + KV for the big trunks).
5727    pub fn kqrp_enabled() -> bool {
5728        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5729        *ON.get_or_init(|| match std::env::var("MEMRA_KQRP").as_deref() {
5730            Ok("0") => false,
5731            Ok(_) => true,
5732            Err(_) => cfg!(memra_hopper_mma),
5733        })
5734    }
5735
5736    /// IN-PLACE split-plane swap (the 31B dense arc): build the split layout and REPLACE the
5737    /// GGUF bytes (zero extra steady-state VRAM — the transient peak is one tensor's size).
5738    /// The tensor's `rp` flag then routes every consumer (mmvq/batched `_rp` twins, the
5739    /// `qmatvec_gemm_q4_0_rp` prefill kernel). Callers gate on the fast path being active —
5740    /// the Stage-A f32 oracle (`MEMRA_FAST=0`) reads GGUF layout and must never see a swap.
5741    pub fn build_q4_rp_swap(
5742        &self,
5743        t: &mut crate::model::GpuTensor,
5744    ) -> Result<bool, Box<dyn std::error::Error>> {
5745        use crate::model::GpuTensor;
5746        // QTYPE GUARD IN THE SWAP ITSELF (gemma4 NVFP4mix prefill-NaN, 2026-08-17):
5747        // `rp4` is a SHARED field — the Q8RP walk parks Q8_0 MIRRORS there, and this
5748        // fn used to `take()` whatever it found. On a Q8_0-carrying gemma4-dense trunk
5749        // the swap hijacked those mirrors: `bytes` became split-plane in place, the
5750        // m<=16 `_rp` dispatch read them correctly (masking the corruption from every
5751        // decode pin), and every GGUF-layout prefill consumer (MMQ + GEMM) read the
5752        // fp16 d-plane as weights -> layer-0 NaN, <pad>-spam serving. Only a tensor
5753        // this fn's OWN builder serves may ever be swapped; everything else refuses
5754        // here, regardless of walk ordering.
5755        if !matches!(t, GpuTensor::Quant { qtype, .. } if *qtype == QT_Q4_0) {
5756            return Ok(false);
5757        }
5758        self.build_q4_rp4(t)?;
5759        self.gpu.stream().synchronize()?; // build kernel reads the GGUF bytes — drain BEFORE dropping them
5760        let GpuTensor::Quant { bytes, rp4, rp, .. } = t else {
5761            return Ok(false);
5762        };
5763        match rp4.take() {
5764            Some(split) => {
5765                *bytes = split; // the GGUF-layout buffer drops here
5766                *rp = true;
5767                Ok(true)
5768            }
5769            None => Ok(false),
5770        }
5771    }
5772
5773    /// MEMRA_Q4RP seam (default ON): the Q4_0 split-plane decode mirror at model load.
5774    pub fn q4rp_enabled() -> bool {
5775        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5776        *ON.get_or_init(|| {
5777            std::env::var("MEMRA_Q4RP")
5778                .map(|v| v != "0")
5779                .unwrap_or(true)
5780        })
5781    }
5782
5783    /// gemma4-E4B: dense [t][row_elems] gather of layer il's rows from the strided prologue
5784    /// buffer ([t][n_layer][n_epl]; off = il*n_epl, stride = n_layer*n_epl).
5785    pub fn copy_rows_strided(
5786        &self,
5787        src: &CudaSlice<f32>,
5788        dst: &mut CudaSlice<f32>,
5789        row_elems: usize,
5790        n_rows: usize,
5791        src_stride: usize,
5792        src_off: usize,
5793    ) -> Result<(), Box<dyn std::error::Error>> {
5794        let f = self.func("copy_rows_strided_f32");
5795        let cfg = LaunchConfig {
5796            grid_dim: (((row_elems as u32 + 255) / 256).max(1), n_rows as u32, 1),
5797            block_dim: (256, 1, 1),
5798            shared_mem_bytes: 0,
5799        };
5800        let (re, nr) = (row_elems as i32, n_rows as i32);
5801        let (st, off) = (src_stride as i64, src_off as i64);
5802        let __s_b = self.gpu.stream();
5803        let mut b = __s_b.launch_builder(&f);
5804        b.arg(src)
5805            .arg(&mut *dst)
5806            .arg(&re)
5807            .arg(&nr)
5808            .arg(&st)
5809            .arg(&off);
5810        unsafe {
5811            b.launch(cfg)?;
5812        }
5813        Ok(())
5814    }
5815
5816    /// Place dense `[row][row_elems]` source rows into one column range of a strided destination.
5817    ///
5818    /// This is a byte-preserving layout operation. It exists so multi-GPU collectives can move
5819    /// one dense shard per rank and reconstruct the canonical token-major matrix without issuing
5820    /// one peer copy per token.
5821    pub fn place_rows_strided(
5822        &self,
5823        src: &CudaSlice<f32>,
5824        dst: &mut CudaSlice<f32>,
5825        row_elems: usize,
5826        n_rows: usize,
5827        dst_stride: usize,
5828        dst_off: usize,
5829    ) -> Result<(), Box<dyn std::error::Error>> {
5830        if row_elems == 0 || n_rows == 0 {
5831            return Err("strided row placement requires nonzero rows and row width".into());
5832        }
5833        let src_len = n_rows
5834            .checked_mul(row_elems)
5835            .ok_or("strided row placement source size overflow")?;
5836        let dst_len = n_rows
5837            .checked_sub(1)
5838            .and_then(|rows| rows.checked_mul(dst_stride))
5839            .and_then(|base| base.checked_add(dst_off))
5840            .and_then(|base| base.checked_add(row_elems))
5841            .ok_or("strided row placement destination size overflow")?;
5842        let row_end = dst_off
5843            .checked_add(row_elems)
5844            .ok_or("strided row placement row size overflow")?;
5845        if src.len() < src_len || dst.len() < dst_len || row_end > dst_stride {
5846            return Err(format!(
5847                "strided row placement geometry mismatch: src={} need_src={src_len} \
5848                 dst={} need_dst={dst_len} row_elems={row_elems} rows={n_rows} \
5849                 dst_stride={dst_stride} dst_off={dst_off}",
5850                src.len(),
5851                dst.len(),
5852            )
5853            .into());
5854        }
5855        if row_elems > i32::MAX as usize || n_rows > i32::MAX as usize {
5856            return Err("strided row placement exceeds CUDA kernel geometry".into());
5857        }
5858        let f = self.func("place_rows_strided_f32");
5859        let cfg = LaunchConfig {
5860            grid_dim: (((row_elems as u32 + 255) / 256).max(1), n_rows as u32, 1),
5861            block_dim: (256, 1, 1),
5862            shared_mem_bytes: 0,
5863        };
5864        let (re, nr) = (row_elems as i32, n_rows as i32);
5865        let (st, off) = (dst_stride as i64, dst_off as i64);
5866        let __s_b = self.gpu.stream();
5867        let mut b = __s_b.launch_builder(&f);
5868        b.arg(src)
5869            .arg(&mut *dst)
5870            .arg(&re)
5871            .arg(&nr)
5872            .arg(&st)
5873            .arg(&off);
5874        unsafe {
5875            b.launch(cfg)?;
5876        }
5877        Ok(())
5878    }
5879
5880    /// Async device u32 store (value rides the kernel ARG — no host-memory transfer/sync).
5881    pub fn u32_set_k(
5882        &self,
5883        dst: &mut CudaSlice<u32>,
5884        v: u32,
5885        idx: usize,
5886    ) -> Result<(), Box<dyn std::error::Error>> {
5887        let f = self.func("u32_set_k");
5888        let cfg = LaunchConfig {
5889            grid_dim: (1, 1, 1),
5890            block_dim: (1, 1, 1),
5891            shared_mem_bytes: 0,
5892        };
5893        let ii = idx as i32;
5894        let __s_b = self.gpu.stream();
5895        let mut b = __s_b.launch_builder(&f);
5896        b.arg(dst).arg(&v).arg(&ii);
5897        unsafe {
5898            b.launch(cfg)?;
5899        }
5900        Ok(())
5901    }
5902
5903    /// counter += v (device-slot append advance; the +1 twin is `inc_seqlen`).
5904    pub fn i32_add_k(
5905        &self,
5906        d: &mut CudaSlice<i32>,
5907        v: i32,
5908    ) -> Result<(), Box<dyn std::error::Error>> {
5909        let f = self.func("i32_add_k");
5910        let cfg = LaunchConfig {
5911            grid_dim: (1, 1, 1),
5912            block_dim: (32, 1, 1),
5913            shared_mem_bytes: 0,
5914        };
5915        let __s_b = self.gpu.stream();
5916        let mut b = __s_b.launch_builder(&f);
5917        b.arg(d).arg(&v);
5918        unsafe {
5919            b.launch(cfg)?;
5920        }
5921        Ok(())
5922    }
5923
5924    /// pos rows from a device counter: dst[i] = ctr[0] + i (verify-stream rope positions).
5925    pub fn i32_iota_from(
5926        &self,
5927        ctr: &CudaSlice<i32>,
5928        dst: &mut CudaSlice<i32>,
5929        n: usize,
5930    ) -> Result<(), Box<dyn std::error::Error>> {
5931        let f = self.func("i32_iota_from");
5932        let cfg = LaunchConfig::for_num_elems(n as u32);
5933        let ni = n as i32;
5934        let __s_b = self.gpu.stream();
5935        let mut b = __s_b.launch_builder(&f);
5936        b.arg(ctr).arg(dst).arg(&ni);
5937        unsafe {
5938            b.launch(cfg)?;
5939        }
5940        Ok(())
5941    }
5942
5943    /// In-place trim-id translate: buf[idx] = map[buf[idx]] (FR-Spec d2t, async single-slot).
5944    pub fn u32_map_k(
5945        &self,
5946        buf: &mut CudaSlice<u32>,
5947        map: &CudaSlice<u32>,
5948        idx: usize,
5949    ) -> Result<(), Box<dyn std::error::Error>> {
5950        let f = self.func("u32_map_k");
5951        let cfg = LaunchConfig {
5952            grid_dim: (1, 1, 1),
5953            block_dim: (1, 1, 1),
5954            shared_mem_bytes: 0,
5955        };
5956        let ii = idx as i32;
5957        let __s_b = self.gpu.stream();
5958        let mut b = __s_b.launch_builder(&f);
5959        b.arg(buf).arg(map).arg(&ii);
5960        unsafe {
5961            b.launch(cfg)?;
5962        }
5963        Ok(())
5964    }
5965
5966    /// Pack a[off..off+n1] ++ b[0..n2] into one buffer (single dtoh follows).
5967    #[allow(clippy::too_many_arguments)]
5968    pub fn u32_pack2(
5969        &self,
5970        a: &CudaSlice<u32>,
5971        off_a: usize,
5972        n1: usize,
5973        b_in: &CudaSlice<u32>,
5974        n2: usize,
5975        out: &mut CudaSlice<u32>,
5976    ) -> Result<(), Box<dyn std::error::Error>> {
5977        let f = self.func("u32_pack2");
5978        let cfg = LaunchConfig::for_num_elems((n1 + n2) as u32);
5979        let (oa, i1, i2) = (off_a as i32, n1 as i32, n2 as i32);
5980        let __s_b = self.gpu.stream();
5981        let mut b = __s_b.launch_builder(&f);
5982        b.arg(a).arg(&oa).arg(&i1).arg(b_in).arg(&i2).arg(out);
5983        unsafe {
5984            b.launch(cfg)?;
5985        }
5986        Ok(())
5987    }
5988
5989    /// gemma4 R3 device fold: w[i] *= s[sel[i]] over the router's [n] (sel, w) pair.
5990    pub fn moe_w_exscale(
5991        &self,
5992        w: &mut CudaSlice<f32>,
5993        sel: &CudaSlice<i32>,
5994        s: &CudaSlice<f32>,
5995        n: usize,
5996    ) -> Result<(), Box<dyn std::error::Error>> {
5997        let f = self.func("moe_w_exscale");
5998        let cfg = LaunchConfig::for_num_elems(n as u32);
5999        let ni = n as i32;
6000        let __s_b = self.gpu.stream();
6001        let mut b = __s_b.launch_builder(&f);
6002        b.arg(w).arg(sel).arg(s).arg(&ni);
6003        unsafe {
6004            b.launch(cfg)?;
6005        }
6006        Ok(())
6007    }
6008
6009    /// Down-projection macro fold: w[i] *= macros[2*n_expert + sel[i]] on the device router
6010    /// weights (one launch per MoE layer, only for macro-carrying artifacts — see MoeWeights).
6011    pub fn moe_w_scale_by_expert(
6012        &self,
6013        w: &mut CudaSlice<f32>,
6014        sel: &CudaSlice<i32>,
6015        macros: &CudaSlice<f32>,
6016        n_expert: usize,
6017        n: usize,
6018    ) -> Result<(), Box<dyn std::error::Error>> {
6019        let f = self.func("moe_w_scale_by_expert");
6020        let cfg = LaunchConfig {
6021            grid_dim: (n.div_ceil(64) as u32, 1, 1),
6022            block_dim: (64, 1, 1),
6023            shared_mem_bytes: 0,
6024        };
6025        let (ne, nn) = (n_expert as i32, n as i32);
6026        let __s_b = self.gpu.stream();
6027        let mut b = __s_b.launch_builder(&f);
6028        b.arg(w).arg(sel).arg(macros).arg(&ne).arg(&nn);
6029        unsafe {
6030            b.launch(cfg)?;
6031        }
6032        Ok(())
6033    }
6034
6035    pub fn moe_gate_up_silu8_dev_q8(
6036        &self,
6037        table: &CudaSlice<u64>,
6038        sel: &cudarc::driver::CudaView<i32>,
6039        aq: &CudaSlice<i8>,
6040        ad: &CudaSlice<f32>,
6041        in_f: usize,
6042        n_ff: usize,
6043        n_used: usize,
6044        n_expert: usize,
6045        qt_g: i32,
6046        qt_u: i32,
6047        rb_g: usize,
6048        rb_u: usize,
6049        macros: &CudaSlice<f32>,
6050    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6051        static GU: std::sync::OnceLock<(String, u32)> = std::sync::OnceLock::new();
6052        let (mode, wpb) = GU.get_or_init(|| {
6053            let mode = std::env::var("MEMRA_MOE_DEVQ8_GU").unwrap_or_default();
6054            let wpb = std::env::var("MEMRA_MOE_DEVQ8_WPB")
6055                .ok()
6056                .and_then(|v| v.parse().ok())
6057                .unwrap_or(4u32)
6058                .clamp(1, 16);
6059            (mode, wpb)
6060        });
6061        let (mode, wpb) = (mode.as_str(), *wpb);
6062        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
6063        let (inf, nff, ne, rbg, rbu) = (
6064            in_f as i32,
6065            n_ff as i32,
6066            n_expert as i32,
6067            rb_g as i64,
6068            rb_u as i64,
6069        );
6070        let (f, cfg) = match mode {
6071            "1" | "2" | "4" => {
6072                let rpw: u32 = mode.parse().unwrap();
6073                let f = self.func(match rpw {
6074                    1 => "moe_gate_up_silu8_dev_q8_r1",
6075                    2 => "moe_gate_up_silu8_dev_q8_r2",
6076                    _ => "moe_gate_up_silu8_dev_q8_r4",
6077                });
6078                let rows_per_block = (rpw * wpb) as usize;
6079                let gx = n_ff.div_ceil(rows_per_block) as u32;
6080                (
6081                    f,
6082                    LaunchConfig {
6083                        grid_dim: (gx, n_used as u32, 1),
6084                        block_dim: (32, wpb, 1),
6085                        shared_mem_bytes: 0,
6086                    },
6087                )
6088            }
6089            "j8" if n_used <= 32 => (
6090                self.func("moe_gate_up_silu8_dev_q8_j8"),
6091                LaunchConfig {
6092                    grid_dim: (n_ff as u32, 1, 1),
6093                    block_dim: (32, n_used as u32, 1),
6094                    shared_mem_bytes: 0,
6095                },
6096            ),
6097            // SMEM-GRID twins (IQ3_S 2KB grid copied to shared, static smem — bit-identical dots)
6098            "vsm2" => {
6099                let f = self.func("moe_gate_up_silu8_dev_q8_vsm2");
6100                let sh = (rb_g + rb_u) as u32;
6101                use cudarc::driver::sys::CUfunction_attribute_enum as A;
6102                f.set_attribute(
6103                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
6104                    sh as i32,
6105                )?;
6106                (
6107                    f,
6108                    LaunchConfig {
6109                        grid_dim: (n_ff as u32, n_used as u32, 1),
6110                        block_dim: (32, 1, 1),
6111                        shared_mem_bytes: sh,
6112                    },
6113                )
6114            }
6115            "vsm" => {
6116                let f = self.func("moe_gate_up_silu8_dev_q8_vsm");
6117                let sh = (rb_g + rb_u) as u32;
6118                use cudarc::driver::sys::CUfunction_attribute_enum as A;
6119                f.set_attribute(
6120                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
6121                    sh as i32,
6122                )?;
6123                (
6124                    f,
6125                    LaunchConfig {
6126                        grid_dim: (n_ff as u32, n_used as u32, 1),
6127                        block_dim: (32, 1, 1),
6128                        shared_mem_bytes: sh,
6129                    },
6130                )
6131            }
6132            "sg" => (
6133                self.func("moe_gate_up_silu8_dev_q8_sg"),
6134                LaunchConfig {
6135                    grid_dim: (n_ff as u32, n_used as u32, 1),
6136                    block_dim: (32, 1, 1),
6137                    shared_mem_bytes: 0,
6138                },
6139            ),
6140            "j8sg" if n_used <= 32 => (
6141                self.func("moe_gate_up_silu8_dev_q8_j8sg"),
6142                LaunchConfig {
6143                    grid_dim: (n_ff as u32, 1, 1),
6144                    block_dim: (32, n_used as u32, 1),
6145                    shared_mem_bytes: 0,
6146                },
6147            ),
6148            "u64" if in_f == 2048 => (
6149                self.func("moe_gate_up_silu8_dev_q8_u64"),
6150                LaunchConfig {
6151                    grid_dim: (n_ff as u32, n_used as u32, 1),
6152                    block_dim: (32, 1, 1),
6153                    shared_mem_bytes: 0,
6154                },
6155            ),
6156            "gs4" if in_f == 2048 => (
6157                self.func("moe_gate_up_silu8_dev_q8_gs4"),
6158                LaunchConfig {
6159                    grid_dim: (n_ff as u32, n_used as u32, 1),
6160                    block_dim: (32, 4, 1),
6161                    shared_mem_bytes: 0,
6162                },
6163            ),
6164            // _v twin (down8 lane 2026-07-08): wide-load IQ4_XS dot, base geometry, bit-identical.
6165            "v" | "" => (
6166                self.func("moe_gate_up_silu8_dev_q8_v"),
6167                LaunchConfig {
6168                    grid_dim: (n_ff as u32, n_used as u32, 1),
6169                    block_dim: (32, 1, 1),
6170                    shared_mem_bytes: 0,
6171                },
6172            ),
6173            "s2" => (
6174                self.func("moe_gate_up_silu8_dev_q8_s2"),
6175                LaunchConfig {
6176                    grid_dim: (n_ff as u32, n_used as u32, 1),
6177                    block_dim: (32, 2, 1),
6178                    shared_mem_bytes: 0,
6179                },
6180            ),
6181            "s2z" => {
6182                let rz = wpb.min(16); // s2z smem tile is [16][2]
6183                (
6184                    self.func("moe_gate_up_silu8_dev_q8_s2z"),
6185                    LaunchConfig {
6186                        grid_dim: (n_ff.div_ceil(rz as usize) as u32, n_used as u32, 1),
6187                        block_dim: (32, 2, rz),
6188                        shared_mem_bytes: 0,
6189                    },
6190                )
6191            }
6192            _ => (
6193                self.func("moe_gate_up_silu8_dev_q8"),
6194                LaunchConfig {
6195                    grid_dim: (n_ff as u32, n_used as u32, 1),
6196                    block_dim: (32, 1, 1),
6197                    shared_mem_bytes: 0,
6198                },
6199            ),
6200        };
6201        let __s_b = self.gpu.stream();
6202        let mut b = __s_b.launch_builder(&f);
6203        b.arg(table)
6204            .arg(sel)
6205            .arg(aq)
6206            .arg(ad)
6207            .arg(&mut act)
6208            .arg(&inf)
6209            .arg(&nff)
6210            .arg(&ne)
6211            .arg(&qt_g)
6212            .arg(&qt_u)
6213            .arg(&rbg)
6214            .arg(&rbu)
6215            .arg(macros);
6216        unsafe {
6217            b.launch(cfg)?;
6218        }
6219        Ok(act)
6220    }
6221
6222    #[allow(clippy::too_many_arguments)]
6223    pub fn moe_down8_fma_dev_q8(
6224        &self,
6225        table: &CudaSlice<u64>,
6226        sel: &cudarc::driver::CudaView<i32>,
6227        w: &cudarc::driver::CudaView<f32>,
6228        aq2: &CudaSlice<i8>,
6229        ad2: &CudaSlice<f32>,
6230        dst: &mut cudarc::driver::CudaViewMut<f32>,
6231        in_f: usize,
6232        out_f: usize,
6233        n_used: usize,
6234        n_expert: usize,
6235        qt: i32,
6236        rb: usize,
6237    ) -> Result<(), Box<dyn std::error::Error>> {
6238        static DOWN: std::sync::OnceLock<String> = std::sync::OnceLock::new();
6239        let mode = DOWN.get_or_init(|| std::env::var("MEMRA_MOE_DEVQ8_DOWN").unwrap_or_default());
6240        let (inf, outf, nu, ne, rbi) = (
6241            in_f as i32,
6242            out_f as i32,
6243            n_used as i32,
6244            n_expert as i32,
6245            rb as i64,
6246        );
6247        // the w8 twins' smem tile is [RPW][8] — n_used must fit the 8-slot tile;
6248        // the h2 twins are nsb==16 (in_f==512) shape-gated.
6249        let (f, cfg) = match mode.as_str() {
6250            m @ ("1" | "2" | "4") if n_used <= 8 => {
6251                let rpw: usize = m.parse().unwrap();
6252                let f = self.func(match rpw {
6253                    1 => "moe_down8_fma_dev_q8_w8r1",
6254                    2 => "moe_down8_fma_dev_q8_w8r2",
6255                    _ => "moe_down8_fma_dev_q8_w8r4",
6256                });
6257                (
6258                    f,
6259                    LaunchConfig {
6260                        grid_dim: (out_f.div_ceil(rpw) as u32, 1, 1),
6261                        block_dim: (32, n_used as u32, 1),
6262                        shared_mem_bytes: 0,
6263                    },
6264                )
6265            }
6266            "h2" if in_f == 512 => (
6267                self.func("moe_down8_fma_dev_q8_h2"),
6268                LaunchConfig {
6269                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6270                    block_dim: (32, 1, 1),
6271                    shared_mem_bytes: 0,
6272                },
6273            ),
6274            // "" = AUTO gemma shape (in_f==704): w8r2 measured +1 tok/s vs base (sweep
6275            // 1/2/4 -> 133.6/134.2/133.6, 2026-07-10); slot-ordered chain preserved.
6276            "" if in_f == 704 && n_used <= 8 => (
6277                self.func("moe_down8_fma_dev_q8_w8r2"),
6278                LaunchConfig {
6279                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6280                    block_dim: (32, n_used as u32, 1),
6281                    shared_mem_bytes: 0,
6282                },
6283            ),
6284            // "" = AUTO: the measured winner for the 35B expert shape (arc 2026-07-05, +3.8%);
6285            // any shape the h2 kernels can't take (nsb!=16 / n_used>8) falls to base via `_`.
6286            // _v twins (down8 lane 2026-07-08): wide-load IQ4_XS dot, bit-identical outputs.
6287            "w8h2v" | "" if in_f == 512 && n_used <= 8 => (
6288                self.func("moe_down8_fma_dev_q8_w8h2v"),
6289                LaunchConfig {
6290                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6291                    block_dim: (32, n_used as u32, 1),
6292                    shared_mem_bytes: 0,
6293                },
6294            ),
6295            "w8h2r2v" if in_f == 512 && n_used <= 8 => (
6296                self.func("moe_down8_fma_dev_q8_w8h2r2v"),
6297                LaunchConfig {
6298                    grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
6299                    block_dim: (32, n_used as u32, 1),
6300                    shared_mem_bytes: 0,
6301                },
6302            ),
6303            "w8h2r2" if in_f == 512 && n_used <= 8 => (
6304                self.func("moe_down8_fma_dev_q8_w8h2r2"),
6305                LaunchConfig {
6306                    grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
6307                    block_dim: (32, n_used as u32, 1),
6308                    shared_mem_bytes: 0,
6309                },
6310            ),
6311            "w8h2" if in_f == 512 && n_used <= 8 => (
6312                self.func("moe_down8_fma_dev_q8_w8h2"),
6313                LaunchConfig {
6314                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6315                    block_dim: (32, n_used as u32, 1),
6316                    shared_mem_bytes: 0,
6317                },
6318            ),
6319            _ => (
6320                self.func("moe_down8_fma_dev_q8"),
6321                LaunchConfig {
6322                    grid_dim: (out_f as u32, 1, 1),
6323                    block_dim: (32, 1, 1),
6324                    shared_mem_bytes: 0,
6325                },
6326            ),
6327        };
6328        let __s_b = self.gpu.stream();
6329        let mut b = __s_b.launch_builder(&f);
6330        b.arg(table)
6331            .arg(sel)
6332            .arg(w)
6333            .arg(aq2)
6334            .arg(ad2)
6335            .arg(dst)
6336            .arg(&inf)
6337            .arg(&outf)
6338            .arg(&nu)
6339            .arg(&ne)
6340            .arg(&qt)
6341            .arg(&rbi);
6342        unsafe {
6343            b.launch(cfg)?;
6344        }
6345        Ok(())
6346    }
6347
6348    /// SMALL-M VERIFY rows twin (MEMRA_SPEC_M2, lane/spec-m2): ONE launch covers all `t` tokens
6349    /// of the spec verify's MoE dev gate/up (grid.z = token) — the _v geometry per token, with
6350    /// tok-offset sel/aq/ad/act pointers matching the serial loop's slices. BIT-IDENTICAL per
6351    /// token (see the kernel header). aq/ad are the BATCHED z-quantize ([t, in_f] rows —
6352    /// quantize_q8_1's per-32-block program is row-independent, so batched rows == the serial
6353    /// loop's per-token quantize_q8_1_view bytes). Returns act [t, n_used, n_ff].
6354    #[allow(clippy::too_many_arguments)]
6355    pub fn moe_gate_up_silu8_dev_q8_rows(
6356        &self,
6357        table: &CudaSlice<u64>,
6358        sel: &CudaSlice<i32>,
6359        aq: &CudaSlice<i8>,
6360        ad: &CudaSlice<f32>,
6361        t: usize,
6362        in_f: usize,
6363        n_ff: usize,
6364        n_used: usize,
6365        n_expert: usize,
6366        qt_g: i32,
6367        qt_u: i32,
6368        rb_g: usize,
6369        rb_u: usize,
6370        macros: &CudaSlice<f32>,
6371    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6372        let f = self.func("moe_gate_up_silu8_dev_q8_v_rows");
6373        let mut act = self.alloc_uninit::<f32>(t * n_used * n_ff)?;
6374        let cfg = LaunchConfig {
6375            grid_dim: (n_ff as u32, n_used as u32, t as u32),
6376            block_dim: (32, 1, 1),
6377            shared_mem_bytes: 0,
6378        };
6379        let (inf, nff, ne, nu, rbg, rbu) = (
6380            in_f as i32,
6381            n_ff as i32,
6382            n_expert as i32,
6383            n_used as i32,
6384            rb_g as i64,
6385            rb_u as i64,
6386        );
6387        let __s_b = self.gpu.stream();
6388        let mut b = __s_b.launch_builder(&f);
6389        b.arg(table)
6390            .arg(sel)
6391            .arg(aq)
6392            .arg(ad)
6393            .arg(&mut act)
6394            .arg(&inf)
6395            .arg(&nff)
6396            .arg(&ne)
6397            .arg(&qt_g)
6398            .arg(&qt_u)
6399            .arg(&rbg)
6400            .arg(&rbu)
6401            .arg(&nu)
6402            .arg(macros);
6403        unsafe {
6404            b.launch(cfg)?;
6405        }
6406        Ok(act)
6407    }
6408
6409    /// SMALL-M VERIFY rows twin of the down proj: w8h2v geometry per token on a grid.z token
6410    /// axis. Caller gates the w8h2v shape contract (in_f == 512, n_used <= 8) — same gate as
6411    /// the AUTO dispatch in `moe_down8_fma_dev_q8`. aq2/ad2 = batched act quantize
6412    /// ([t*n_used, in_f] rows). dst rows are FULLY overwritten per token.
6413    #[allow(clippy::too_many_arguments)]
6414    pub fn moe_down8_fma_dev_q8_rows(
6415        &self,
6416        table: &CudaSlice<u64>,
6417        sel: &CudaSlice<i32>,
6418        w: &CudaSlice<f32>,
6419        aq2: &CudaSlice<i8>,
6420        ad2: &CudaSlice<f32>,
6421        dst: &mut CudaSlice<f32>,
6422        t: usize,
6423        in_f: usize,
6424        out_f: usize,
6425        n_used: usize,
6426        n_expert: usize,
6427        qt: i32,
6428        rb: usize,
6429    ) -> Result<(), Box<dyn std::error::Error>> {
6430        assert!(
6431            in_f == 512 && n_used <= 8,
6432            "down rows twin is w8h2v shape-gated"
6433        );
6434        let f = self.func("moe_down8_fma_dev_q8_w8h2v_rows");
6435        let cfg = LaunchConfig {
6436            grid_dim: (out_f.div_ceil(2) as u32, 1, t as u32),
6437            block_dim: (32, n_used as u32, 1),
6438            shared_mem_bytes: 0,
6439        };
6440        let (inf, outf, nu, ne, rbi) = (
6441            in_f as i32,
6442            out_f as i32,
6443            n_used as i32,
6444            n_expert as i32,
6445            rb as i64,
6446        );
6447        let __s_b = self.gpu.stream();
6448        let mut b = __s_b.launch_builder(&f);
6449        b.arg(table)
6450            .arg(sel)
6451            .arg(w)
6452            .arg(aq2)
6453            .arg(ad2)
6454            .arg(dst)
6455            .arg(&inf)
6456            .arg(&outf)
6457            .arg(&nu)
6458            .arg(&ne)
6459            .arg(&qt)
6460            .arg(&rbi);
6461        unsafe {
6462            b.launch(cfg)?;
6463        }
6464        Ok(())
6465    }
6466
6467    /// CSR gate/up v3 (owner-scan dedup, no build kernel): qtypes {IQ4_XS, IQ3_S} (caller
6468    /// gates), grid.y = pair index; the first pair of each expert serves all its pairs.
6469    /// Bit-identical to moe_gate_up_silu8_dev_q8_v_rows (explicit-intrinsic accumulate).
6470    #[allow(clippy::too_many_arguments)]
6471    pub fn moe_gate_up_silu8_dev_q8_csr(
6472        &self,
6473        table: &CudaSlice<u64>,
6474        sel: &CudaSlice<i32>,
6475        aq: &CudaSlice<i8>,
6476        ad: &CudaSlice<f32>,
6477        n_pairs: usize,
6478        in_f: usize,
6479        n_ff: usize,
6480        n_used: usize,
6481        n_expert: usize,
6482        qt_g: i32,
6483        qt_u: i32,
6484        rb_g: usize,
6485        rb_u: usize,
6486    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6487        // NVFP4 experts take the NVFP4-specialized owner-scan twin (lane/moebatch-q35moe);
6488        // host gate guarantees qt_g == qt_u within a supported class.
6489        let f = if qt_g == crate::QT_NVFP4 {
6490            self.func("moe_gate_up_silu8_dev_q8_csr_nvfp4")
6491        } else {
6492            self.func("moe_gate_up_silu8_dev_q8_csr_iq4")
6493        };
6494        let mut act = self.alloc_uninit::<f32>(n_pairs * n_ff)?;
6495        let cfg = LaunchConfig {
6496            grid_dim: (n_ff as u32, n_pairs as u32, 1),
6497            block_dim: (32, 1, 1),
6498            shared_mem_bytes: 0,
6499        };
6500        let (inf, nff, ne, nu, npi, rbg, rbu) = (
6501            in_f as i32,
6502            n_ff as i32,
6503            n_expert as i32,
6504            n_used as i32,
6505            n_pairs as i32,
6506            rb_g as i64,
6507            rb_u as i64,
6508        );
6509        let __s_b = self.gpu.stream();
6510        let mut b = __s_b.launch_builder(&f);
6511        b.arg(table)
6512            .arg(sel)
6513            .arg(aq)
6514            .arg(ad)
6515            .arg(&mut act)
6516            .arg(&inf)
6517            .arg(&nff)
6518            .arg(&ne)
6519            .arg(&qt_g)
6520            .arg(&qt_u)
6521            .arg(&rbg)
6522            .arg(&rbu)
6523            .arg(&nu)
6524            .arg(&npi);
6525        unsafe {
6526            b.launch(cfg)?;
6527        }
6528        Ok(act)
6529    }
6530
6531    /// TEST SEAM (down8 lane 2026-07-08): launch a down dev_q8 variant BY NAME with its
6532    /// canonical geometry, bypassing the env-cached dispatch so moe-devq8-check can byte-
6533    /// compare variants in one process. Variants: "base", "w8h2", "w8h2r2", "w8h2v", "w8h2r2v".
6534    #[allow(clippy::too_many_arguments)]
6535    pub fn moe_down8_fma_dev_q8_variant(
6536        &self,
6537        variant: &str,
6538        table: &CudaSlice<u64>,
6539        sel: &cudarc::driver::CudaView<i32>,
6540        w: &cudarc::driver::CudaView<f32>,
6541        aq2: &CudaSlice<i8>,
6542        ad2: &CudaSlice<f32>,
6543        dst: &mut cudarc::driver::CudaViewMut<f32>,
6544        in_f: usize,
6545        out_f: usize,
6546        n_used: usize,
6547        n_expert: usize,
6548        qt: i32,
6549        rb: usize,
6550    ) -> Result<(), Box<dyn std::error::Error>> {
6551        let (inf, outf, nu, ne, rbi) = (
6552            in_f as i32,
6553            out_f as i32,
6554            n_used as i32,
6555            n_expert as i32,
6556            rb as i64,
6557        );
6558        let (f, cfg) = match variant {
6559            "w8h2" | "w8h2v" => (
6560                self.func(if variant == "w8h2" {
6561                    "moe_down8_fma_dev_q8_w8h2"
6562                } else {
6563                    "moe_down8_fma_dev_q8_w8h2v"
6564                }),
6565                LaunchConfig {
6566                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6567                    block_dim: (32, n_used as u32, 1),
6568                    shared_mem_bytes: 0,
6569                },
6570            ),
6571            "w8h2r2" | "w8h2r2v" => (
6572                self.func(if variant == "w8h2r2" {
6573                    "moe_down8_fma_dev_q8_w8h2r2"
6574                } else {
6575                    "moe_down8_fma_dev_q8_w8h2r2v"
6576                }),
6577                LaunchConfig {
6578                    grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
6579                    block_dim: (32, n_used as u32, 1),
6580                    shared_mem_bytes: 0,
6581                },
6582            ),
6583            _ => (
6584                self.func("moe_down8_fma_dev_q8"),
6585                LaunchConfig {
6586                    grid_dim: (out_f as u32, 1, 1),
6587                    block_dim: (32, 1, 1),
6588                    shared_mem_bytes: 0,
6589                },
6590            ),
6591        };
6592        let __s_b = self.gpu.stream();
6593        let mut b = __s_b.launch_builder(&f);
6594        b.arg(table)
6595            .arg(sel)
6596            .arg(w)
6597            .arg(aq2)
6598            .arg(ad2)
6599            .arg(dst)
6600            .arg(&inf)
6601            .arg(&outf)
6602            .arg(&nu)
6603            .arg(&ne)
6604            .arg(&qt)
6605            .arg(&rbi);
6606        unsafe {
6607            b.launch(cfg)?;
6608        }
6609        Ok(())
6610    }
6611
6612    /// TEST SEAM (down8 lane): gate_up twin of the above. Variants: "base", "v".
6613    #[allow(clippy::too_many_arguments)]
6614    pub fn moe_gate_up_silu8_dev_q8_variant(
6615        &self,
6616        variant: &str,
6617        table: &CudaSlice<u64>,
6618        sel: &cudarc::driver::CudaView<i32>,
6619        aq: &CudaSlice<i8>,
6620        ad: &CudaSlice<f32>,
6621        in_f: usize,
6622        n_ff: usize,
6623        n_used: usize,
6624        n_expert: usize,
6625        qt_g: i32,
6626        qt_u: i32,
6627        rb_g: usize,
6628        rb_u: usize,
6629    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6630        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
6631        let (inf, nff, ne, rbg, rbu) = (
6632            in_f as i32,
6633            n_ff as i32,
6634            n_expert as i32,
6635            rb_g as i64,
6636            rb_u as i64,
6637        );
6638        let f = self.func(if variant == "v" {
6639            "moe_gate_up_silu8_dev_q8_v"
6640        } else {
6641            "moe_gate_up_silu8_dev_q8"
6642        });
6643        let cfg = LaunchConfig {
6644            grid_dim: (n_ff as u32, n_used as u32, 1),
6645            block_dim: (32, 1, 1),
6646            shared_mem_bytes: 0,
6647        };
6648        let __s_b = self.gpu.stream();
6649        let mut b = __s_b.launch_builder(&f);
6650        b.arg(table)
6651            .arg(sel)
6652            .arg(aq)
6653            .arg(ad)
6654            .arg(&mut act)
6655            .arg(&inf)
6656            .arg(&nff)
6657            .arg(&ne)
6658            .arg(&qt_g)
6659            .arg(&qt_u)
6660            .arg(&rbg)
6661            .arg(&rbu);
6662        unsafe {
6663            b.launch(cfg)?;
6664        }
6665        Ok(act)
6666    }
6667
6668    pub fn moe_gate_up_silu8_dev(
6669        &self,
6670        table: &CudaSlice<u64>,
6671        sel: &cudarc::driver::CudaView<i32>,
6672        x: &cudarc::driver::CudaView<f32>,
6673        in_f: usize,
6674        n_ff: usize,
6675        n_used: usize,
6676        n_expert: usize,
6677        qt_g: i32,
6678        qt_u: i32,
6679        rb_g: usize,
6680        rb_u: usize,
6681        macros: &CudaSlice<f32>,
6682    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6683        let f = self.func("moe_gate_up_silu8_dev");
6684        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?; // fully overwritten
6685        let cfg = LaunchConfig {
6686            grid_dim: (n_ff as u32, n_used as u32, 1),
6687            block_dim: (256, 1, 1),
6688            shared_mem_bytes: 0,
6689        };
6690        let (inf, nff, ne, rbg, rbu) = (
6691            in_f as i32,
6692            n_ff as i32,
6693            n_expert as i32,
6694            rb_g as i64,
6695            rb_u as i64,
6696        );
6697        let __s_b = self.gpu.stream();
6698        let mut b = __s_b.launch_builder(&f);
6699        b.arg(table)
6700            .arg(sel)
6701            .arg(x)
6702            .arg(&mut act)
6703            .arg(&inf)
6704            .arg(&nff)
6705            .arg(&ne)
6706            .arg(&qt_g)
6707            .arg(&qt_u)
6708            .arg(&rbg)
6709            .arg(&rbu)
6710            .arg(macros);
6711        unsafe {
6712            b.launch(cfg)?;
6713        }
6714        Ok(act)
6715    }
6716
6717    /// LAUNCH-STRUCTURE STAGE 3: device-dispatch twin of `moe_down8_fma_into` — expert ids AND
6718    /// renormalized weights read from the router kernel's device output. BIT-IDENTICAL chain.
6719    #[allow(clippy::too_many_arguments)]
6720    pub fn moe_down8_fma_dev(
6721        &self,
6722        table: &CudaSlice<u64>,
6723        sel: &cudarc::driver::CudaView<i32>,
6724        w: &cudarc::driver::CudaView<f32>,
6725        act: &CudaSlice<f32>,
6726        dst: &mut cudarc::driver::CudaViewMut<f32>,
6727        in_f: usize,
6728        out_f: usize,
6729        n_used: usize,
6730        n_expert: usize,
6731        qt: i32,
6732        rb: usize,
6733    ) -> Result<(), Box<dyn std::error::Error>> {
6734        let f = self.func("moe_down8_fma_dev");
6735        let cfg = LaunchConfig {
6736            grid_dim: (out_f as u32, 1, 1),
6737            block_dim: (256, 1, 1),
6738            shared_mem_bytes: 0,
6739        };
6740        let (inf, outf, nu, ne, rbv) = (
6741            in_f as i32,
6742            out_f as i32,
6743            n_used as i32,
6744            n_expert as i32,
6745            rb as i64,
6746        );
6747        let __s_b = self.gpu.stream();
6748        let mut b = __s_b.launch_builder(&f);
6749        b.arg(table)
6750            .arg(sel)
6751            .arg(w)
6752            .arg(act)
6753            .arg(dst)
6754            .arg(&inf)
6755            .arg(&outf)
6756            .arg(&nu)
6757            .arg(&ne)
6758            .arg(&qt)
6759            .arg(&rbv);
6760        unsafe {
6761            b.launch(cfg)?;
6762        }
6763        Ok(())
6764    }
6765
6766    /// dst[i] += alpha * src[i], i in 0..n. dst is a CudaViewMut (a row of moe_out).
6767    pub fn axpy_into(
6768        &self,
6769        src: &CudaSlice<f32>,
6770        alpha: f32,
6771        dst: &mut cudarc::driver::CudaViewMut<f32>,
6772        n: usize,
6773    ) -> Result<(), Box<dyn std::error::Error>> {
6774        let f = self.func("axpy_f32");
6775        let cfg = LaunchConfig::for_num_elems(n as u32);
6776        let (a, ni) = (alpha, n as i32);
6777        let __s_b = self.gpu.stream();
6778        let mut b = __s_b.launch_builder(&f);
6779        b.arg(src).arg(dst).arg(&a).arg(&ni);
6780        unsafe {
6781            b.launch(cfg)?;
6782        }
6783        Ok(())
6784    }
6785
6786    /// Host-oracle twin of `axpy_into` with separate RN multiply and add.
6787    pub fn axpy_host_into(
6788        &self,
6789        src: &cudarc::driver::CudaView<'_, f32>,
6790        alpha: f32,
6791        dst: &mut cudarc::driver::CudaViewMut<f32>,
6792        n: usize,
6793    ) -> Result<(), Box<dyn std::error::Error>> {
6794        let f = self.func("axpy_host_f32");
6795        let cfg = LaunchConfig::for_num_elems(n as u32);
6796        let (a, ni) = (alpha, n as i32);
6797        let __s_b = self.gpu.stream();
6798        let mut b = __s_b.launch_builder(&f);
6799        b.arg(src).arg(dst).arg(&a).arg(&ni);
6800        unsafe {
6801            b.launch(cfg)?;
6802        }
6803        Ok(())
6804    }
6805
6806    /// dst[r*ncols + c] += src[r*ncols + c] * scale[r]. Per-row scalar accumulate (shared expert).
6807    pub fn add_scaled_rows(
6808        &self,
6809        src: &CudaSlice<f32>,
6810        scale: &CudaSlice<f32>,
6811        dst: &mut CudaSlice<f32>,
6812        ncols: usize,
6813        nrows: usize,
6814    ) -> Result<(), Box<dyn std::error::Error>> {
6815        let f = self.func("add_scaled_rows_f32");
6816        let cfg = LaunchConfig::for_num_elems((ncols * nrows) as u32);
6817        let (nc, nr) = (ncols as i32, nrows as i32);
6818        let __s_b = self.gpu.stream();
6819        let mut b = __s_b.launch_builder(&f);
6820        b.arg(src).arg(scale).arg(dst).arg(&nc).arg(&nr);
6821        unsafe {
6822            b.launch(cfg)?;
6823        }
6824        Ok(())
6825    }
6826
6827    // ======== A2 GROUPED MoE PREFILL KERNELS ========
6828
6829    /// Gather m_e rows from src[T, ncols] into dst[m_e, ncols] using index array idx[m_e].
6830    pub fn gather_rows(
6831        &self,
6832        src: &CudaSlice<f32>,
6833        idx: &CudaSlice<i32>,
6834        dst: &mut CudaSlice<f32>,
6835        ncols: usize,
6836        m_e: usize,
6837    ) -> Result<(), Box<dyn std::error::Error>> {
6838        let f = self.func("gather_rows_f32");
6839        let cfg = LaunchConfig::for_num_elems((m_e * ncols) as u32);
6840        let (nc, me) = (ncols as i32, m_e as i32);
6841        let __s_b = self.gpu.stream();
6842        let mut b = __s_b.launch_builder(&f);
6843        b.arg(src).arg(idx).arg(dst).arg(&nc).arg(&me);
6844        unsafe {
6845            b.launch(cfg)?;
6846        }
6847        Ok(())
6848    }
6849
6850    /// Scatter expert outputs into per-token slots: dst[tok_idx[r], slot_idx[r], :] = src[r, :] * weight[r].
6851    /// dst is [T, n_used, ncols], zero-initialized. Each (expert, token) pair maps to a unique slot.
6852    /// Scatter expert outputs into per-token slots (raw copy, no weight multiply).
6853    /// Weight stored into wbuf[tok*n_used + slot] for FMA in reduce step.
6854    pub fn scatter_slot(
6855        &self,
6856        src: &CudaSlice<f32>,
6857        tok_idx: &CudaSlice<i32>,
6858        slot_idx: &CudaSlice<i32>,
6859        weight: &CudaSlice<f32>,
6860        dst: &mut CudaSlice<f32>,
6861        wbuf: &mut CudaSlice<f32>,
6862        ncols: usize,
6863        n_used: usize,
6864        m_e: usize,
6865    ) -> Result<(), Box<dyn std::error::Error>> {
6866        let f = self.func("scatter_add_slot_f32");
6867        let cfg = LaunchConfig::for_num_elems((m_e * ncols) as u32);
6868        let (nc, nu, me) = (ncols as i32, n_used as i32, m_e as i32);
6869        let __s_b = self.gpu.stream();
6870        let mut b = __s_b.launch_builder(&f);
6871        b.arg(src)
6872            .arg(tok_idx)
6873            .arg(slot_idx)
6874            .arg(weight)
6875            .arg(dst)
6876            .arg(wbuf)
6877            .arg(&nc)
6878            .arg(&nu)
6879            .arg(&me);
6880        unsafe {
6881            b.launch(cfg)?;
6882        }
6883        Ok(())
6884    }
6885
6886    /// Reduce n_used slots per token: dst[t, col] = sum_s slots[t, s, col].
6887    /// Reduce n_used slots per token: dst[t, col] = sum_s FMA(wbuf[t,s], slots[t,s,col], acc).
6888    /// Uses FMA for bit-identity with the sequential axpy path.
6889    pub fn reduce_slots(
6890        &self,
6891        slots: &CudaSlice<f32>,
6892        wbuf: &CudaSlice<f32>,
6893        dst: &mut CudaSlice<f32>,
6894        ncols: usize,
6895        n_used: usize,
6896        t: usize,
6897    ) -> Result<(), Box<dyn std::error::Error>> {
6898        let f = self.func("reduce_slots_f32");
6899        let cfg = LaunchConfig::for_num_elems((t * ncols) as u32);
6900        let (nc, nu, ti) = (ncols as i32, n_used as i32, t as i32);
6901        let __s_b = self.gpu.stream();
6902        let mut b = __s_b.launch_builder(&f);
6903        b.arg(slots).arg(wbuf).arg(dst).arg(&nc).arg(&nu).arg(&ti);
6904        unsafe {
6905            b.launch(cfg)?;
6906        }
6907        Ok(())
6908    }
6909
6910    /// Canonical slot-order reduction with separately rounded multiply and add.
6911    ///
6912    /// This is the one-launch twin of repeated `axpy_host_into` calls. It preserves the official
6913    /// Step host-oracle arithmetic while allowing owner outputs to remain device-resident.
6914    pub fn reduce_slots_host(
6915        &self,
6916        slots: &CudaSlice<f32>,
6917        wbuf: &CudaSlice<f32>,
6918        dst: &mut CudaSlice<f32>,
6919        ncols: usize,
6920        n_used: usize,
6921        t: usize,
6922    ) -> Result<(), Box<dyn std::error::Error>> {
6923        let f = self.func("reduce_slots_host_f32");
6924        let cfg = LaunchConfig::for_num_elems((t * ncols) as u32);
6925        let (nc, nu, ti) = (ncols as i32, n_used as i32, t as i32);
6926        let __s_b = self.gpu.stream();
6927        let mut b = __s_b.launch_builder(&f);
6928        b.arg(slots).arg(wbuf).arg(dst).arg(&nc).arg(&nu).arg(&ti);
6929        unsafe {
6930            b.launch(cfg)?;
6931        }
6932        Ok(())
6933    }
6934
6935    /// Stage-B: quantize activation [m,in] f32 -> q8_1 (int8 qs + per-block f32 scale).
6936    /// Quantize an activation [m, in_f] to q8_1 (int8 qs + per-32 f32 scale). Public so the
6937    /// forward can quantize a SHARED activation ONCE and feed it to several matmuls (gate+up
6938    /// share `z`; q/k/v and wqkv/gate/beta/alpha share `h`) — quantize_q8_1 was 13.5% of decode
6939    /// GPU time, ~half of it redundant re-quantization of the same row.
6940    /// quantize_q8_1 over a CudaView (a sliced z-row) — same kernel, offset-honoring arg.
6941    pub fn quantize_q8_1_view(
6942        &self,
6943        x: &cudarc::driver::CudaView<f32>,
6944        m: usize,
6945        in_f: usize,
6946    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6947        let f = self.func("quantize_q8_1");
6948        let nblk = in_f / 32;
6949        let mut q = self.alloc_uninit::<i8>(m * in_f)?;
6950        let mut d = self.alloc_uninit::<f32>(m * nblk)?;
6951        let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
6952        let (inf, mi) = (in_f as i32, m as i32);
6953        let __s_b = self.gpu.stream();
6954        let mut b = __s_b.launch_builder(&f);
6955        b.arg(x).arg(&mut q).arg(&mut d).arg(&inf).arg(&mi);
6956        unsafe {
6957            b.launch(cfg)?;
6958        }
6959        Ok((q, d))
6960    }
6961
6962    pub fn quantize_q8_1(
6963        &self,
6964        x: &CudaSlice<f32>,
6965        m: usize,
6966        in_f: usize,
6967    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6968        let nblk = in_f / 32;
6969        let mut q = self.alloc_uninit::<i8>(m * in_f)?; // full-overwrite output: skip memset
6970        let mut d = self.alloc_uninit::<f32>(m * nblk)?; // full-overwrite output: skip memset
6971        // WARP-PER-BLOCK kernel: one warp per 32-block -> m*in_f threads total.
6972        let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
6973        let (inf, mi) = (in_f as i32, m as i32);
6974        if Self::pdl_on() && Self::pdl_wb_on() {
6975            {
6976                use cudarc::driver::{DevicePtr, DevicePtrMut};
6977                let s = &self.gpu.stream();
6978                let (px, _g0) = x.device_ptr(s);
6979                let (pq, _g1) = q.device_ptr_mut(s);
6980                let (pd, _g2) = d.device_ptr_mut(s);
6981                let mut ps = [
6982                    &px as *const _ as *mut std::ffi::c_void,
6983                    &pq as *const _ as *mut _,
6984                    &pd as *const _ as *mut _,
6985                    &inf as *const _ as *mut _,
6986                    &mi as *const _ as *mut _,
6987                ];
6988                unsafe {
6989                    self.launch_pdl("quantize_q8_1", cfg.grid_dim, cfg.block_dim, &mut ps)?;
6990                }
6991            }
6992            return Ok((q, d));
6993        }
6994        let f = self.func("quantize_q8_1");
6995        let __s_b = self.gpu.stream();
6996        let mut b = __s_b.launch_builder(&f);
6997        b.arg(x).arg(&mut q).arg(&mut d).arg(&inf).arg(&mi);
6998        unsafe {
6999            b.launch(cfg)?;
7000        }
7001        Ok((q, d))
7002    }
7003
7004    /// Stage-C FP4: quantize activation [m,in] f32 -> e2m1 nibbles (aq4: u32 [m, in/8]) + per-16
7005    /// UE4M3 scale (ad4: u8 [m, in/16]), the layout the mxf4nvf4 block-scale GEMM B-operand wants.
7006    /// in_f must be a multiple of 64 (one NVFP4 K-block). One thread per (token, 16-block).
7007    pub fn quantize_fp4_act(
7008        &self,
7009        x: &CudaSlice<f32>,
7010        m: usize,
7011        in_f: usize,
7012    ) -> Result<(CudaSlice<u32>, CudaSlice<u8>), Box<dyn std::error::Error>> {
7013        let f = self.func("quantize_fp4_act");
7014        let nb16 = in_f / 16;
7015        let mut aq4 = self.alloc_uninit::<u32>(m * (in_f / 8))?; // full-overwrite output: skip memset
7016        let mut ad4 = self.alloc_uninit::<u8>(m * nb16)?; // full-overwrite output: skip memset
7017        let cfg = LaunchConfig::for_num_elems((m * nb16) as u32);
7018        let (inf, mi) = (in_f as i32, m as i32);
7019        let __s_b = self.gpu.stream();
7020        let mut b = __s_b.launch_builder(&f);
7021        b.arg(x).arg(&mut aq4).arg(&mut ad4).arg(&inf).arg(&mi);
7022        unsafe {
7023            b.launch(cfg)?;
7024        }
7025        Ok((aq4, ad4))
7026    }
7027
7028    /// Stage-C FP4 GEMM (NVFP4 weights): native mxf4nvf4 block-scale tensor-core matmul. Feeds raw
7029    /// e2m1 weight nibbles + raw UE4M3 micro-scales directly to mma.sync.m16n8k64 (762 TFLOP/s peak,
7030    /// 3.5x int8). Activation `x` is quantized to FP4 e2m1 here. NVFP4 per-tensor macro-scale applied
7031    /// post (scale==1.0 -> no-op). `bytes` = raw NVFP4 weight rows. Used by the MEMRA_FP4 prefill path.
7032    pub fn qmatvec_gemm_nvfp4_fp4(
7033        &self,
7034        bytes: &CudaSlice<u8>,
7035        x: &CudaSlice<f32>,
7036        m: usize,
7037        in_f: usize,
7038        out_f: usize,
7039        row_bytes: usize,
7040        scale: f32,
7041    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7042        assert!(
7043            in_f % 64 == 0,
7044            "FP4 GEMM requires in_f % 64 == 0, got {in_f}"
7045        );
7046        let (aq4, ad4) = self.quantize_fp4_act(x, m, in_f)?;
7047        let mut y = self.fp4_gemm_launch(bytes, &aq4, &ad4, m, in_f, out_f, row_bytes)?;
7048        if scale != 1.0 {
7049            self.scale_inplace(&mut y, scale, m * out_f)?;
7050        }
7051        Ok(y)
7052    }
7053
7054    /// Shared mxf4 GEMM launch (pre-quantized FP4 activation aq4/ad4). Same CTA tile as the int8 GEMM
7055    /// (BM=64 rows x BN=128 tokens, 4 warps). No macro-scale applied here.
7056    fn fp4_gemm_launch(
7057        &self,
7058        bytes: &CudaSlice<u8>,
7059        aq4: &CudaSlice<u32>,
7060        ad4: &CudaSlice<u8>,
7061        m: usize,
7062        in_f: usize,
7063        out_f: usize,
7064        row_bytes: usize,
7065    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7066        let f = self.func("qmatvec_gemm_nvfp4_fp4");
7067        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
7068        const BM: u32 = 64;
7069        const BN: u32 = 256;
7070        let cfg = LaunchConfig {
7071            grid_dim: ((out_f as u32 + BM - 1) / BM, (m as u32 + BN - 1) / BN, 1),
7072            block_dim: (32, 4, 1),
7073            shared_mem_bytes: 0,
7074        };
7075        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7076        let __s_b = self.gpu.stream();
7077        let mut b = __s_b.launch_builder(&f);
7078        b.arg(bytes)
7079            .arg(aq4)
7080            .arg(ad4)
7081            .arg(&mut y)
7082            .arg(&inf)
7083            .arg(&outf)
7084            .arg(&mi)
7085            .arg(&rb);
7086        unsafe {
7087            b.launch(cfg)?;
7088        }
7089        Ok(y)
7090    }
7091
7092    /// Test entry (kernel_check): run the FP4 GEMM from raw bytes; NO macro-scale (caller compares bare).
7093    pub fn qmatvec_gemm_nvfp4_fp4_raw(
7094        &self,
7095        bytes: &CudaSlice<u8>,
7096        x: &CudaSlice<f32>,
7097        m: usize,
7098        in_f: usize,
7099        out_f: usize,
7100        row_bytes: usize,
7101    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7102        assert!(
7103            in_f % 64 == 0,
7104            "FP4 GEMM requires in_f % 64 == 0, got {in_f}"
7105        );
7106        let (aq4, ad4) = self.quantize_fp4_act(x, m, in_f)?;
7107        self.fp4_gemm_launch(bytes, &aq4, &ad4, m, in_f, out_f, row_bytes)
7108    }
7109
7110    /// Stage-B: Q8_0 weight x q8_1 activation int8 dp4a matmul. y[m,out]=x@W^T.
7111    pub fn qmatvec_q8_0_fast(
7112        &self,
7113        w: &CudaSlice<u8>,
7114        x: &CudaSlice<f32>,
7115        m: usize,
7116        in_f: usize,
7117        out_f: usize,
7118        row_bytes: usize,
7119    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7120        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
7121        let f = self.func("qmatvec_q8_0_dp4a");
7122        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
7123        let cfg = LaunchConfig {
7124            grid_dim: (out_f as u32, m as u32, 1),
7125            block_dim: (128, 1, 1),
7126            shared_mem_bytes: 0,
7127        };
7128        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7129        let __s_b = self.gpu.stream();
7130        let mut b = __s_b.launch_builder(&f);
7131        b.arg(w)
7132            .arg(&aq)
7133            .arg(&ad)
7134            .arg(&mut y)
7135            .arg(&inf)
7136            .arg(&outf)
7137            .arg(&mi)
7138            .arg(&rb);
7139        unsafe {
7140            b.launch(cfg)?;
7141        }
7142        Ok(y)
7143    }
7144
7145    /// Stage-B: Q4_K weight x q8_1 activation int8 dp4a (decode). Min-offset via q8_1 sum term.
7146    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
7147    pub fn qmatvec_q4_K_fast(
7148        &self,
7149        w: &CudaSlice<u8>,
7150        x: &CudaSlice<f32>,
7151        m: usize,
7152        in_f: usize,
7153        out_f: usize,
7154        row_bytes: usize,
7155    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7156        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
7157        let f = self.func("qmatvec_q4_K_dp4a");
7158        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
7159        let cfg = LaunchConfig {
7160            grid_dim: (out_f as u32, m as u32, 1),
7161            block_dim: (128, 1, 1),
7162            shared_mem_bytes: 0,
7163        };
7164        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7165        let __s_b = self.gpu.stream();
7166        let mut b = __s_b.launch_builder(&f);
7167        b.arg(w)
7168            .arg(&aq)
7169            .arg(&ad)
7170            .arg(&mut y)
7171            .arg(&inf)
7172            .arg(&outf)
7173            .arg(&mi)
7174            .arg(&rb);
7175        unsafe {
7176            b.launch(cfg)?;
7177        }
7178        Ok(y)
7179    }
7180
7181    /// Stage-B: Q6_K weight x q8_1 activation int8 dp4a (decode, symmetric).
7182    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
7183    pub fn qmatvec_q6_K_fast(
7184        &self,
7185        w: &CudaSlice<u8>,
7186        x: &CudaSlice<f32>,
7187        m: usize,
7188        in_f: usize,
7189        out_f: usize,
7190        row_bytes: usize,
7191    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7192        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
7193        let f = self.func("qmatvec_q6_K_dp4a");
7194        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
7195        let cfg = LaunchConfig {
7196            grid_dim: (out_f as u32, m as u32, 1),
7197            block_dim: (128, 1, 1),
7198            shared_mem_bytes: 0,
7199        };
7200        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7201        let __s_b = self.gpu.stream();
7202        let mut b = __s_b.launch_builder(&f);
7203        b.arg(w)
7204            .arg(&aq)
7205            .arg(&ad)
7206            .arg(&mut y)
7207            .arg(&inf)
7208            .arg(&outf)
7209            .arg(&mi)
7210            .arg(&rb);
7211        unsafe {
7212            b.launch(cfg)?;
7213        }
7214        Ok(y)
7215    }
7216
7217    /// Stage-B: Q5_K weight x q8_1 activation int8 dp4a (decode). Min-offset via q8_1 sum term.
7218    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
7219    pub fn qmatvec_q5_K_fast(
7220        &self,
7221        w: &CudaSlice<u8>,
7222        x: &CudaSlice<f32>,
7223        m: usize,
7224        in_f: usize,
7225        out_f: usize,
7226        row_bytes: usize,
7227    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7228        self.qmatvec_dp4a_named(
7229            "qmatvec_q5_K_dp4a",
7230            &w.slice(0..w.len()),
7231            x,
7232            m,
7233            in_f,
7234            out_f,
7235            row_bytes,
7236        )
7237    }
7238    /// Stage-B: Q3_K weight x q8_1 activation int8 dp4a (decode, symmetric).
7239    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
7240    pub fn qmatvec_q3_K_fast(
7241        &self,
7242        w: &CudaSlice<u8>,
7243        x: &CudaSlice<f32>,
7244        m: usize,
7245        in_f: usize,
7246        out_f: usize,
7247        row_bytes: usize,
7248    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7249        self.qmatvec_dp4a_named(
7250            "qmatvec_q3_K_dp4a",
7251            &w.slice(0..w.len()),
7252            x,
7253            m,
7254            in_f,
7255            out_f,
7256            row_bytes,
7257        )
7258    }
7259    /// A6 split-plane twin of `qmatvec_nvfp4_fast` (weights repacked; used by the rp gates).
7260    pub fn qmatvec_nvfp4_fast_rp(
7261        &self,
7262        w: &CudaSlice<u8>,
7263        x: &CudaSlice<f32>,
7264        m: usize,
7265        in_f: usize,
7266        out_f: usize,
7267        row_bytes: usize,
7268    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7269        assert!(
7270            in_f % 64 == 0,
7271            "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
7272        );
7273        self.qmatvec_dp4a_named(
7274            "qmatvec_nvfp4_dp4a_rp",
7275            &w.slice(0..w.len()),
7276            x,
7277            m,
7278            in_f,
7279            out_f,
7280            row_bytes,
7281        )
7282    }
7283    /// Stage-B: NVFP4 weight x q8_1 activation int8 dp4a (decode, symmetric, codebook lookup).
7284    pub fn qmatvec_nvfp4_fast(
7285        &self,
7286        w: &cudarc::driver::CudaView<'_, u8>,
7287        x: &CudaSlice<f32>,
7288        m: usize,
7289        in_f: usize,
7290        out_f: usize,
7291        row_bytes: usize,
7292    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7293        // B1: the NVFP4 dp4a kernel maps two 32-elem q8_1 blocks onto one 64-elem block_nvfp4
7294        // (sblk = g >> 1). in_f must be a multiple of 64 or the last block reads a partial superblock.
7295        assert!(
7296            in_f % 64 == 0,
7297            "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
7298        );
7299        self.qmatvec_dp4a_named("qmatvec_nvfp4_dp4a", w, x, m, in_f, out_f, row_bytes)
7300    }
7301    /// v2-layout twin of `qmatvec_nvfp4_fast` for the slot-major expert banks
7302    /// (MEMRA_NVFP4_BANK_V2) — bit-identical per row, coalesced reads.
7303    pub fn qmatvec_nvfp4_fast_v2(
7304        &self,
7305        w: &cudarc::driver::CudaView<'_, u8>,
7306        x: &CudaSlice<f32>,
7307        m: usize,
7308        in_f: usize,
7309        out_f: usize,
7310        row_bytes: usize,
7311    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7312        assert!(
7313            in_f % 64 == 0,
7314            "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
7315        );
7316        self.qmatvec_dp4a_named("qmatvec_nvfp4_dp4a_v2", w, x, m, in_f, out_f, row_bytes)
7317    }
7318    /// Stage-B (optional perf): IQ4_XS codebook int8 dp4a.
7319    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
7320    pub fn qmatvec_iq4_XS_fast(
7321        &self,
7322        w: &CudaSlice<u8>,
7323        x: &CudaSlice<f32>,
7324        m: usize,
7325        in_f: usize,
7326        out_f: usize,
7327        row_bytes: usize,
7328    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7329        self.qmatvec_dp4a_named(
7330            "qmatvec_iq4_XS_dp4a",
7331            &w.slice(0..w.len()),
7332            x,
7333            m,
7334            in_f,
7335            out_f,
7336            row_bytes,
7337        )
7338    }
7339
7340    /// Shared dp4a launcher: quantize_q8_1 then call the named kernel (grid (out,m), block 64).
7341    fn qmatvec_dp4a_named(
7342        &self,
7343        name: &str,
7344        w: &cudarc::driver::CudaView<'_, u8>,
7345        x: &CudaSlice<f32>,
7346        m: usize,
7347        in_f: usize,
7348        out_f: usize,
7349        row_bytes: usize,
7350    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7351        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
7352        let f = self.func(name);
7353        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
7354        let cfg = LaunchConfig {
7355            grid_dim: (out_f as u32, m as u32, 1),
7356            block_dim: (128, 1, 1),
7357            shared_mem_bytes: 0,
7358        };
7359        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7360        let __s_b = self.gpu.stream();
7361        let mut b = __s_b.launch_builder(&f);
7362        b.arg(w)
7363            .arg(&aq)
7364            .arg(&ad)
7365            .arg(&mut y)
7366            .arg(&inf)
7367            .arg(&outf)
7368            .arg(&mi)
7369            .arg(&rb);
7370        unsafe {
7371            b.launch(cfg)?;
7372        }
7373        Ok(y)
7374    }
7375
7376    /// NVFP4 dp4a matvec over PRE-QUANTIZED q8_1 activations, writing a caller-provided output.
7377    /// Same kernel and math as `qmatvec_nvfp4_fast` (which quantizes internally and allocates
7378    /// its output); this entry exists so a routed-expert program can quantize one activation
7379    /// ONCE and reuse it across every expert's gate/up, feed `silu_mul_scaled_q8_1`'s q8_1
7380    /// straight into down, and keep persistent output workspaces — zero per-expert allocations.
7381    #[allow(clippy::too_many_arguments)]
7382    pub fn qmatvec_nvfp4_fast_prequant_into(
7383        &self,
7384        w: &CudaSlice<u8>,
7385        aq: &CudaSlice<i8>,
7386        ad: &CudaSlice<f32>,
7387        y: &mut CudaSlice<f32>,
7388        m: usize,
7389        in_f: usize,
7390        out_f: usize,
7391        row_bytes: usize,
7392    ) -> Result<(), Box<dyn std::error::Error>> {
7393        assert!(
7394            in_f % 64 == 0,
7395            "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
7396        );
7397        if y.len() < m * out_f {
7398            return Err(format!(
7399                "NVFP4 prequant output {} is shorter than {m}x{out_f}",
7400                y.len()
7401            )
7402            .into());
7403        }
7404        let f = self.func("qmatvec_nvfp4_dp4a");
7405        let cfg = LaunchConfig {
7406            grid_dim: (out_f as u32, m as u32, 1),
7407            block_dim: (128, 1, 1),
7408            shared_mem_bytes: 0,
7409        };
7410        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7411        let __s_b = self.gpu.stream();
7412        let mut b = __s_b.launch_builder(&f);
7413        b.arg(w)
7414            .arg(aq)
7415            .arg(ad)
7416            .arg(y)
7417            .arg(&inf)
7418            .arg(&outf)
7419            .arg(&mi)
7420            .arg(&rb);
7421        unsafe {
7422            b.launch(cfg)?;
7423        }
7424        Ok(())
7425    }
7426
7427    /// Fused QKV F32 matvec (one launch for all three rank-local projections; see the kernel
7428    /// doc for the numeric-class note). Requires `in_f % 4 == 0`.
7429    #[allow(clippy::too_many_arguments)]
7430    pub fn matvec_f32_qkv_into(
7431        &self,
7432        wq: &CudaSlice<f32>,
7433        wk: &CudaSlice<f32>,
7434        wv: &CudaSlice<f32>,
7435        wg: &CudaSlice<f32>,
7436        x: &CudaSlice<f32>,
7437        yq: &mut CudaSlice<f32>,
7438        yk: &mut CudaSlice<f32>,
7439        yv: &mut CudaSlice<f32>,
7440        yg: &mut CudaSlice<f32>,
7441        in_f: usize,
7442        out_q: usize,
7443        out_kv: usize,
7444        out_g: usize,
7445    ) -> Result<(), Box<dyn std::error::Error>> {
7446        if in_f % 4 != 0
7447            || wq.len() != out_q * in_f
7448            || wk.len() != out_kv * in_f
7449            || wv.len() != out_kv * in_f
7450            || wg.len() < out_g * in_f
7451            || x.len() < in_f
7452            || yq.len() < out_q
7453            || yk.len() < out_kv
7454            || yv.len() < out_kv
7455            || (out_g > 0 && yg.len() < out_g)
7456        {
7457            return Err(format!(
7458                "fused QKV geometry in={in_f} out_q={out_q} out_kv={out_kv} out_g={out_g} \
7459                 wq={} wk={} wv={} wg={}",
7460                wq.len(),
7461                wk.len(),
7462                wv.len(),
7463                wg.len()
7464            )
7465            .into());
7466        }
7467        let f = self.func("matvec_f32_qkv");
7468        let cfg = LaunchConfig {
7469            grid_dim: ((out_q + 2 * out_kv + out_g) as u32, 1, 1),
7470            block_dim: (128, 1, 1),
7471            shared_mem_bytes: 0,
7472        };
7473        let (inf, oq, okv, og) = (in_f as i32, out_q as i32, out_kv as i32, out_g as i32);
7474        let __s_b = self.gpu.stream();
7475        let mut b = __s_b.launch_builder(&f);
7476        b.arg(wq)
7477            .arg(wk)
7478            .arg(wv)
7479            .arg(wg)
7480            .arg(x)
7481            .arg(yq)
7482            .arg(yk)
7483            .arg(yv)
7484            .arg(yg)
7485            .arg(&inf)
7486            .arg(&oq)
7487            .arg(&okv)
7488            .arg(&og);
7489        unsafe {
7490            b.launch(cfg)?;
7491        }
7492        Ok(())
7493    }
7494
7495    /// Selected-experts batched twin of `qmatvec_nvfp4_fast_prequant_into`: one launch covers
7496    /// every selected expert, weights indexed `sel[t] * expert_stride` into a contiguous
7497    /// per-rank bank, activations advancing `act_row_stride`/`ad_row_stride` elements per
7498    /// selection (0 for a shared input). Per (expert, row) bit-identical to the per-expert
7499    /// kernel — the batching only removes host launch latency.
7500    #[allow(clippy::too_many_arguments)]
7501    /// FUSION #2a: gate+up sweeps in one launch (v2 banks only; identical geometry both
7502    /// banks, caller-guarded). Per-row bit-identical to two qmatvec_nvfp4_sel_into calls.
7503    #[allow(clippy::too_many_arguments)]
7504    pub fn qmatvec_nvfp4_sel_gu_into(
7505        &self,
7506        gate_bank: &CudaSlice<u8>,
7507        up_bank: &CudaSlice<u8>,
7508        sel: &CudaSlice<i32>,
7509        aq: &CudaSlice<i8>,
7510        ad: &CudaSlice<f32>,
7511        yg: &mut CudaSlice<f32>,
7512        yu: &mut CudaSlice<f32>,
7513        n_sel: usize,
7514        in_f: usize,
7515        out_f: usize,
7516        row_bytes: usize,
7517        expert_stride: usize,
7518    ) -> Result<(), Box<dyn std::error::Error>> {
7519        assert!(in_f % 64 == 0, "NVFP4 dp4a requires in_f % 64 == 0");
7520        if yg.len() < n_sel * out_f || yu.len() < n_sel * out_f || sel.len() < n_sel {
7521            return Err("NVFP4 gu sel geometry".into());
7522        }
7523        let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_gu");
7524        let cfg = LaunchConfig {
7525            grid_dim: ((2 * out_f) as u32, n_sel as u32, 1),
7526            block_dim: (128, 1, 1),
7527            shared_mem_bytes: 0,
7528        };
7529        let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
7530        let (rb, es) = (row_bytes as i64, expert_stride as i64);
7531        let (ars, adrs) = (0i64, 0i64);
7532        let __s_b = self.gpu.stream();
7533        let mut b = __s_b.launch_builder(&f);
7534        b.arg(gate_bank)
7535            .arg(up_bank)
7536            .arg(sel)
7537            .arg(aq)
7538            .arg(ad)
7539            .arg(yg)
7540            .arg(yu)
7541            .arg(&inf)
7542            .arg(&outf)
7543            .arg(&ns)
7544            .arg(&rb)
7545            .arg(&es)
7546            .arg(&ars)
7547            .arg(&adrs);
7548        unsafe {
7549            b.launch(cfg)?;
7550        }
7551        Ok(())
7552    }
7553
7554    pub fn qmatvec_nvfp4_sel_into(
7555        &self,
7556        bank: &CudaSlice<u8>,
7557        sel: &CudaSlice<i32>,
7558        aq: &CudaSlice<i8>,
7559        ad: &CudaSlice<f32>,
7560        y: &mut CudaSlice<f32>,
7561        n_sel: usize,
7562        in_f: usize,
7563        out_f: usize,
7564        row_bytes: usize,
7565        expert_stride: usize,
7566        act_row_stride: usize,
7567        ad_row_stride: usize,
7568    ) -> Result<(), Box<dyn std::error::Error>> {
7569        assert!(
7570            in_f % 64 == 0,
7571            "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
7572        );
7573        if y.len() < n_sel * out_f || sel.len() < n_sel {
7574            return Err(format!(
7575                "NVFP4 sel output {} / sel {} shorter than {n_sel}x{out_f}",
7576                y.len(),
7577                sel.len()
7578            )
7579            .into());
7580        }
7581        // MEMRA_SEL_MR=1: 4-concurrent-row-groups twin — per row bit-identical (same 128-thread
7582        // striding + reduction). MEASURED SLOWER on the 188-SM card (40.8 vs 42.9 tok/s e2e,
7583        // 2026-08-21: 512-thread blocks trade occupancy for launch-tail savings and lose; the
7584        // sequential-rows variant was flat). Default stays the single-row form.
7585        // MEMRA_SEL_STREAM=1: 16-rows-per-block streaming twin with next-row register
7586        // prefetch (bit-identical per row; one group per thread, so in_f <= 4096 only).
7587        static MR: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
7588        let mode = *MR.get_or_init(|| {
7589            if crate::tp::nvfp4_bank_v2_on() {
7590                3
7591            } else if std::env::var("MEMRA_SEL_STREAM").as_deref() == Ok("1") {
7592                2
7593            } else if std::env::var("MEMRA_SEL_MR").as_deref() == Ok("1") {
7594                1
7595            } else {
7596                0
7597            }
7598        });
7599        let mode = if mode == 2 && in_f > 4096 { 0 } else { mode };
7600        // v2s streaming twin (MEMRA_SEL_V2S=1 on top of the v2 bank): 8 contiguous rows per
7601        // block with next-row int4 prefetch; needs 16B-aligned rows (gate/up 2304B yes, down
7602        // 360B no -> single-row v2) and one slot per thread (in_f <= 4096).
7603        static V2S: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7604        let v2s = mode == 3
7605            && *V2S.get_or_init(|| std::env::var("MEMRA_SEL_V2S").as_deref() == Ok("1"))
7606            && row_bytes % 16 == 0
7607            && in_f <= 4096;
7608        let f = match (mode, v2s) {
7609            (3, true) => self.func("qmatvec_nvfp4_dp4a_sel_v2s"),
7610            (3, false) => self.func("qmatvec_nvfp4_dp4a_sel_v2"),
7611            (2, _) => self.func("qmatvec_nvfp4_dp4a_sel_stream"),
7612            (1, _) => self.func("qmatvec_nvfp4_dp4a_sel_mr4"),
7613            _ => self.func("qmatvec_nvfp4_dp4a_sel"),
7614        };
7615        // Thread-fit block for narrow rows (the DOWN sweep: in_f=640 -> nsb=20 slots left
7616        // 108 of 128 threads idle AND thread-capped resident blocks). blockDim >= nsb keeps
7617        // thread g on slot g; the dropped threads contributed exact 0.0 partials to the
7618        // reduce, so the result bits are unchanged. Applies to the single-row forms only.
7619        let nsb = in_f >> 5;
7620        let fit_block: u32 = if (mode == 0 || mode == 3) && !v2s && nsb <= 32 {
7621            32
7622        } else if mode == 1 {
7623            512
7624        } else {
7625            128
7626        };
7627        let cfg = LaunchConfig {
7628            grid_dim: (
7629                if v2s {
7630                    (out_f as u32).div_ceil(8)
7631                } else {
7632                    match mode {
7633                        2 => (out_f as u32).div_ceil(16),
7634                        1 => (out_f as u32).div_ceil(4),
7635                        _ => out_f as u32,
7636                    }
7637                },
7638                n_sel as u32,
7639                1,
7640            ),
7641            block_dim: (fit_block, 1, 1),
7642            shared_mem_bytes: 0,
7643        };
7644        let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
7645        let (rb, es, ars, adrs) = (
7646            row_bytes as i64,
7647            expert_stride as i64,
7648            act_row_stride as i64,
7649            ad_row_stride as i64,
7650        );
7651        let __s_b = self.gpu.stream();
7652        let mut b = __s_b.launch_builder(&f);
7653        b.arg(bank)
7654            .arg(sel)
7655            .arg(aq)
7656            .arg(ad)
7657            .arg(y)
7658            .arg(&inf)
7659            .arg(&outf)
7660            .arg(&ns)
7661            .arg(&rb)
7662            .arg(&es)
7663            .arg(&ars)
7664            .arg(&adrs);
7665        unsafe {
7666            b.launch(cfg)?;
7667        }
7668        Ok(())
7669    }
7670
7671    /// Selected-experts batched twin of `silu_mul_scaled_q8_1`: [n_sel, n_per] rows, macros
7672    /// from device arrays indexed via sel. Per expert row bit-identical to the scalar kernel.
7673    /// `limit` = the step35 routed SwiGLU clamp (min(silu, limit) * clamp(up, +-limit)); None
7674    /// takes the plain SiLU kernel.
7675    #[allow(clippy::too_many_arguments)]
7676    pub fn silu_mul_scaled_q8_1_sel_into(
7677        &self,
7678        gate: &CudaSlice<f32>,
7679        up: &CudaSlice<f32>,
7680        gmac: &CudaSlice<f32>,
7681        umac: &CudaSlice<f32>,
7682        sel: &CudaSlice<i32>,
7683        limit: Option<f32>,
7684        out_q: &mut CudaSlice<i8>,
7685        out_d: &mut CudaSlice<f32>,
7686        n_per: usize,
7687        n_sel: usize,
7688    ) -> Result<(), Box<dyn std::error::Error>> {
7689        let n = n_per * n_sel;
7690        if n_per % 32 != 0 || out_q.len() < n || out_d.len() < n / 32 {
7691            return Err(format!(
7692                "silu sel geometry n_per={n_per} n_sel={n_sel} q={} d={}",
7693                out_q.len(),
7694                out_d.len()
7695            )
7696            .into());
7697        }
7698        if let Some(limit) = limit {
7699            if limit <= 1e-6 {
7700                return Err(format!(
7701                    "silu sel clamp limit {limit} is at or below the 1e-6 eps gate"
7702                )
7703                .into());
7704            }
7705            let f = self.func("silu_mul_scaled_q8_1_sel_clamp");
7706            let cfg = LaunchConfig::for_num_elems(n as u32);
7707            let (np, ns) = (n_per as i32, n_sel as i32);
7708            let __s_b = self.gpu.stream();
7709            let mut b = __s_b.launch_builder(&f);
7710            b.arg(gate)
7711                .arg(up)
7712                .arg(gmac)
7713                .arg(umac)
7714                .arg(sel)
7715                .arg(&limit)
7716                .arg(out_q)
7717                .arg(out_d)
7718                .arg(&np)
7719                .arg(&ns);
7720            unsafe {
7721                b.launch(cfg)?;
7722            }
7723            return Ok(());
7724        }
7725        let f = self.func("silu_mul_scaled_q8_1_sel");
7726        let cfg = LaunchConfig::for_num_elems(n as u32);
7727        let (np, ns) = (n_per as i32, n_sel as i32);
7728        let __s_b = self.gpu.stream();
7729        let mut b = __s_b.launch_builder(&f);
7730        b.arg(gate)
7731            .arg(up)
7732            .arg(gmac)
7733            .arg(umac)
7734            .arg(sel)
7735            .arg(out_q)
7736            .arg(out_d)
7737            .arg(&np)
7738            .arg(&ns);
7739        unsafe {
7740            b.launch(cfg)?;
7741        }
7742        Ok(())
7743    }
7744
7745    pub fn htod(&self, v: &[f32]) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7746        Ok(self.gpu.stream().clone_htod(v)?)
7747    }
7748    pub fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
7749        Ok(self.gpu.stream().clone_htod(v)?)
7750    }
7751    /// i8 upload (moe-devq8-check: synthetic q8_1 activation bytes).
7752    pub fn htod_i8(&self, v: &[i8]) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
7753        Ok(self.gpu.stream().clone_htod(v)?)
7754    }
7755    pub fn htod_u64(&self, v: &[u64]) -> Result<CudaSlice<u64>, Box<dyn std::error::Error>> {
7756        Ok(self.gpu.stream().clone_htod(v)?)
7757    }
7758    /// View twin of `dtoh` (lean-logits component 3: D2H one row of a [B, n_vocab] stack).
7759    pub fn dtoh_view(
7760        &self,
7761        d: &cudarc::driver::CudaView<f32>,
7762    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
7763        let v = self.gpu.stream().clone_dtoh(d)?;
7764        self.gpu.stream().synchronize()?;
7765        Ok(v)
7766    }
7767    pub fn dtoh(&self, d: &CudaSlice<f32>) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
7768        let v = self.gpu.stream().clone_dtoh(d)?;
7769        self.gpu.stream().synchronize()?;
7770        Ok(v)
7771    }
7772    /// Queue two f32 device-to-host copies on the compute stream, then establish one host
7773    /// boundary for both. Hy3's CPU/GPU expert split needs the router logits and the MoE input;
7774    /// issuing them together avoids a second stream synchronization in every trunk layer.
7775    pub fn dtoh_pair(
7776        &self,
7777        a: &CudaSlice<f32>,
7778        b: &CudaSlice<f32>,
7779    ) -> Result<(Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
7780        let av = self.gpu.stream().clone_dtoh(a)?;
7781        let bv = self.gpu.stream().clone_dtoh(b)?;
7782        self.gpu.stream().synchronize()?;
7783        Ok((av, bv))
7784    }
7785    /// View-scoped twin of `dtoh_pair` for reusable capacity buffers whose inactive tail must not
7786    /// cross a shape-sensitive host boundary.
7787    pub fn dtoh_pair_views(
7788        &self,
7789        a: &cudarc::driver::CudaView<f32>,
7790        b: &cudarc::driver::CudaView<f32>,
7791    ) -> Result<(Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
7792        let av = self.gpu.stream().clone_dtoh(a)?;
7793        let bv = self.gpu.stream().clone_dtoh(b)?;
7794        self.gpu.stream().synchronize()?;
7795        Ok((av, bv))
7796    }
7797    /// Device-to-host copy of an i32 buffer (fused-router sel_idx readback).
7798    pub fn dtoh_i32(&self, d: &CudaSlice<i32>) -> Result<Vec<i32>, Box<dyn std::error::Error>> {
7799        let v = self.gpu.stream().clone_dtoh(d)?;
7800        self.gpu.stream().synchronize()?;
7801        Ok(v)
7802    }
7803    /// Device-to-host copy of a u8 buffer (used to read back the quantized KV cache for validation).
7804    pub fn dtoh_u8(&self, d: &CudaSlice<u8>) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
7805        let v = self.gpu.stream().clone_dtoh(d)?;
7806        self.gpu.stream().synchronize()?;
7807        Ok(v)
7808    }
7809    pub fn dtoh_u8_view(
7810        &self,
7811        d: &cudarc::driver::CudaView<u8>,
7812    ) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
7813        let v = self.gpu.stream().clone_dtoh(d)?;
7814        self.gpu.stream().synchronize()?;
7815        Ok(v)
7816    }
7817    pub fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7818        let s = self.gpu.stream().alloc_zeros::<f32>(n)?;
7819        self.keep_if_capturing(&s);
7820        Ok(s)
7821    }
7822
7823    /// GPU-resident greedy argmax (CUDA-GRAPH-PLAN Phase 1): logits[n_vocab] -> token id in a
7824    /// resident device u32 [1]. PARALLEL 2-pass (RANK1 LEVER): the old single-CTA scan (one 256-thread
7825    /// block on one SM over 248K logits) was memory-starved at ~426us/token. Now pass 1 fans NB=256
7826    /// blocks across the SMs to saturate HBM, pass 2 reduces the NB partials. Bit-identical to host
7827    /// `argmax` (smallest index on tie). The whole point is NOT to dtoh logits — only a [1] u32 is read
7828    /// back (or kept resident for graph replay). Returns the device token buffer.
7829    /// Softmax probability of the (already-argmaxed) token `tok` under `logits` — the spec-decode
7830    /// p-min confidence signal. 2-pass like the parallel argmax; returns a device [1] f32.
7831    pub fn prob_of_token_device(
7832        &self,
7833        logits: &CudaSlice<f32>,
7834        tok: &CudaSlice<u32>,
7835        n_vocab: usize,
7836    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7837        let nb = ARGMAX_NB;
7838        let mut part = self.alloc_uninit::<f32>(nb)?;
7839        let mut p = self.alloc_uninit::<f32>(1)?;
7840        let f1 = self.func("prob_of_token_partial_f32");
7841        let cfg1 = LaunchConfig {
7842            grid_dim: (nb as u32, 1, 1),
7843            block_dim: (256, 1, 1),
7844            shared_mem_bytes: 0,
7845        };
7846        let nv = n_vocab as i32;
7847        let __s_b1 = self.gpu.stream();
7848        let mut b1 = __s_b1.launch_builder(&f1);
7849        b1.arg(logits).arg(tok).arg(&mut part).arg(&nv);
7850        unsafe {
7851            b1.launch(cfg1)?;
7852        }
7853        let f2 = self.func("prob_of_token_final_f32");
7854        let cfg2 = LaunchConfig {
7855            grid_dim: (1, 1, 1),
7856            block_dim: (256, 1, 1),
7857            shared_mem_bytes: 0,
7858        };
7859        let nbi = nb as i32;
7860        let __s_b2 = self.gpu.stream();
7861        let mut b2 = __s_b2.launch_builder(&f2);
7862        b2.arg(&part).arg(&mut p).arg(&nbi);
7863        unsafe {
7864            b2.launch(cfg2)?;
7865        }
7866        Ok(p)
7867    }
7868
7869    /// Like `prob_of_token_device` but writes into a PERSISTENT `p_out` buffer (stable pointer).
7870    /// Required for CUDA-graph capture of the draft chain: the captured prob kernels must write
7871    /// where the host reads the p-min confidence between replays. Same kernels, same math.
7872    /// Slot-addressed twin of `prob_of_token_device_into`: token read from `tok_all[tok_idx]`
7873    /// (a view at the slot), probability written to `p_out[p_idx]` — same two kernels, the
7874    /// pointers just land mid-buffer. Zero-sync (gemma confidence-adaptive draft depth).
7875    pub fn prob_of_token_device_col(
7876        &self,
7877        logits: &CudaSlice<f32>,
7878        tok_all: &CudaSlice<u32>,
7879        tok_idx: usize,
7880        p_out: &mut CudaSlice<f32>,
7881        p_idx: usize,
7882        n_vocab: usize,
7883    ) -> Result<(), Box<dyn std::error::Error>> {
7884        let tok_v = tok_all.slice(tok_idx..tok_idx + 1);
7885        let mut p_v = p_out.slice_mut(p_idx..p_idx + 1);
7886        let nb = ARGMAX_NB;
7887        let mut part = self.alloc_uninit::<f32>(nb)?;
7888        let f1 = self.func("prob_of_token_partial_f32");
7889        let cfg1 = LaunchConfig {
7890            grid_dim: (nb as u32, 1, 1),
7891            block_dim: (256, 1, 1),
7892            shared_mem_bytes: 0,
7893        };
7894        let nv = n_vocab as i32;
7895        let __s_b1 = self.gpu.stream();
7896        let mut b1 = __s_b1.launch_builder(&f1);
7897        b1.arg(logits).arg(&tok_v).arg(&mut part).arg(&nv);
7898        unsafe {
7899            b1.launch(cfg1)?;
7900        }
7901        let f2 = self.func("prob_of_token_final_f32");
7902        let cfg2 = LaunchConfig {
7903            grid_dim: (1, 1, 1),
7904            block_dim: (256, 1, 1),
7905            shared_mem_bytes: 0,
7906        };
7907        let nbi = nb as i32;
7908        let __s_b2 = self.gpu.stream();
7909        let mut b2 = __s_b2.launch_builder(&f2);
7910        b2.arg(&part).arg(&mut p_v).arg(&nbi);
7911        unsafe {
7912            b2.launch(cfg2)?;
7913        }
7914        Ok(())
7915    }
7916
7917    pub fn prob_of_token_device_into(
7918        &self,
7919        logits: &CudaSlice<f32>,
7920        tok: &CudaSlice<u32>,
7921        p_out: &mut CudaSlice<f32>,
7922        n_vocab: usize,
7923    ) -> Result<(), Box<dyn std::error::Error>> {
7924        let nb = ARGMAX_NB;
7925        let mut part = self.alloc_uninit::<f32>(nb)?;
7926        let f1 = self.func("prob_of_token_partial_f32");
7927        let cfg1 = LaunchConfig {
7928            grid_dim: (nb as u32, 1, 1),
7929            block_dim: (256, 1, 1),
7930            shared_mem_bytes: 0,
7931        };
7932        let nv = n_vocab as i32;
7933        let __s_b1 = self.gpu.stream();
7934        let mut b1 = __s_b1.launch_builder(&f1);
7935        b1.arg(logits).arg(tok).arg(&mut part).arg(&nv);
7936        unsafe {
7937            b1.launch(cfg1)?;
7938        }
7939        let f2 = self.func("prob_of_token_final_f32");
7940        let cfg2 = LaunchConfig {
7941            grid_dim: (1, 1, 1),
7942            block_dim: (256, 1, 1),
7943            shared_mem_bytes: 0,
7944        };
7945        let nbi = nb as i32;
7946        let __s_b2 = self.gpu.stream();
7947        let mut b2 = __s_b2.launch_builder(&f2);
7948        b2.arg(&part).arg(p_out).arg(&nbi);
7949        unsafe {
7950            b2.launch(cfg2)?;
7951        }
7952        Ok(())
7953    }
7954
7955    /// Token-graph chunk loop: hist[idx] = *tok; idx += 1 — device-indexed history append
7956    /// (graph-constant params, device-varying index). Capture-safe.
7957    pub fn u32_hist_append(
7958        &self,
7959        tok: &CudaSlice<u32>,
7960        hist: &mut CudaSlice<u32>,
7961        idx: &mut CudaSlice<i32>,
7962    ) -> Result<(), Box<dyn std::error::Error>> {
7963        let f = self.func("u32_hist_append");
7964        let cfg = LaunchConfig {
7965            grid_dim: (1, 1, 1),
7966            block_dim: (32, 1, 1),
7967            shared_mem_bytes: 0,
7968        };
7969        let __s_b = self.gpu.stream();
7970        let mut b = __s_b.launch_builder(&f);
7971        b.arg(tok).arg(&mut *hist).arg(&mut *idx);
7972        unsafe {
7973            b.launch(cfg)?;
7974        }
7975        Ok(())
7976    }
7977
7978    pub fn argmax_token_device(
7979        &self,
7980        logits: &CudaSlice<f32>,
7981        n_vocab: usize,
7982    ) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
7983        let mut tok = unsafe { self.gpu.stream().alloc::<u32>(1)? };
7984        self.argmax_token_device_into(logits, &mut tok, n_vocab)?;
7985        Ok(tok)
7986    }
7987    /// Like `argmax_token_device` but writes into a PERSISTENT `tok` buffer (stable pointer) instead
7988    /// of allocating a fresh one. Required for CUDA-graph capture: the captured argmax must write the
7989    /// next token into the SAME device buffer the next replay's embed_gather reads, so the buffer
7990    /// pointer is baked once and the token id never round-trips to host inside steady state. The
7991    /// pass-1 partials scratch (`argmax_partials`) is also a resident stable-pointer buffer so both
7992    /// captured passes bake fixed addresses.
7993    pub fn argmax_token_device_into(
7994        &self,
7995        logits: &CudaSlice<f32>,
7996        tok: &mut CudaSlice<u32>,
7997        n_vocab: usize,
7998    ) -> Result<(), Box<dyn std::error::Error>> {
7999        let nb = ARGMAX_NB;
8000        let f1 = self.func("argmax_partial_f32");
8001        let f2 = self.func("argmax_final_f32");
8002        let mut guard = self.argmax_partials.lock().unwrap();
8003        if guard.is_none() {
8004            // allocate ONCE; under generate_graph this runs in the tracking-off prime window so the
8005            // buffers carry no cudarc events (illegal inside capture).
8006            let pv = self.gpu.stream().alloc_zeros::<f32>(nb)?;
8007            let pi = self.gpu.stream().alloc_zeros::<i32>(nb)?;
8008            *guard = Some((pv, pi));
8009        }
8010        let (part_v, part_i) = guard.as_mut().unwrap();
8011        let nv = n_vocab as i32;
8012        let nbi = nb as i32;
8013        // pass 1: NB blocks x 256 threads grid-stride scan -> per-block (val, idx) partials.
8014        let cfg1 = LaunchConfig {
8015            grid_dim: (nb as u32, 1, 1),
8016            block_dim: (256, 1, 1),
8017            shared_mem_bytes: 0,
8018        };
8019        let __s_b1 = self.gpu.stream();
8020        let mut b1 = __s_b1.launch_builder(&f1);
8021        b1.arg(logits).arg(&mut *part_v).arg(&mut *part_i).arg(&nv);
8022        unsafe {
8023            b1.launch(cfg1)?;
8024        }
8025        // pass 2: one block reduces NB partials -> token_out[0].
8026        let cfg2 = LaunchConfig {
8027            grid_dim: (1, 1, 1),
8028            block_dim: (256, 1, 1),
8029            shared_mem_bytes: 0,
8030        };
8031        let __s_b2 = self.gpu.stream();
8032        let mut b2 = __s_b2.launch_builder(&f2);
8033        b2.arg(&*part_v).arg(&*part_i).arg(tok).arg(&nbi);
8034        unsafe {
8035            b2.launch(cfg2)?;
8036        }
8037        Ok(())
8038    }
8039    /// Column-`col` device argmax over a stacked verify-logits buffer [t, n_vocab] (spec accept
8040    /// walk): toks[out_idx] = argmax(logits[col*n_vocab .. (col+1)*n_vocab]). SAME 2-pass kernels
8041    /// and tie-break contract as `argmax_token_device_into` (bit-identical to host argmax,
8042    /// argmax_gate-validated) — only the input pointer (a column view) and the output slot differ.
8043    /// Lets the accept walk read ONE [t] u32 instead of dtoh'ing the full [t, n_vocab] logits.
8044    pub fn argmax_token_device_col(
8045        &self,
8046        logits: &CudaSlice<f32>,
8047        col: usize,
8048        n_vocab: usize,
8049        toks: &mut CudaSlice<u32>,
8050        out_idx: usize,
8051    ) -> Result<(), Box<dyn std::error::Error>> {
8052        let nb = ARGMAX_NB;
8053        let f1 = self.func("argmax_partial_f32");
8054        let f2 = self.func("argmax_final_f32");
8055        let mut guard = self.argmax_partials.lock().unwrap();
8056        if guard.is_none() {
8057            let pv = self.gpu.stream().alloc_zeros::<f32>(nb)?;
8058            let pi = self.gpu.stream().alloc_zeros::<i32>(nb)?;
8059            *guard = Some((pv, pi));
8060        }
8061        let (part_v, part_i) = guard.as_mut().unwrap();
8062        let col_view = logits.slice(col * n_vocab..(col + 1) * n_vocab);
8063        let nv = n_vocab as i32;
8064        let nbi = nb as i32;
8065        let cfg1 = LaunchConfig {
8066            grid_dim: (nb as u32, 1, 1),
8067            block_dim: (256, 1, 1),
8068            shared_mem_bytes: 0,
8069        };
8070        let __s_b1 = self.gpu.stream();
8071        let mut b1 = __s_b1.launch_builder(&f1);
8072        b1.arg(&col_view)
8073            .arg(&mut *part_v)
8074            .arg(&mut *part_i)
8075            .arg(&nv);
8076        unsafe {
8077            b1.launch(cfg1)?;
8078        }
8079        let mut tok_view = toks.slice_mut(out_idx..out_idx + 1);
8080        let cfg2 = LaunchConfig {
8081            grid_dim: (1, 1, 1),
8082            block_dim: (256, 1, 1),
8083            shared_mem_bytes: 0,
8084        };
8085        let __s_b2 = self.gpu.stream();
8086        let mut b2 = __s_b2.launch_builder(&f2);
8087        b2.arg(&*part_v).arg(&*part_i).arg(&mut tok_view).arg(&nbi);
8088        unsafe {
8089            b2.launch(cfg2)?;
8090        }
8091        Ok(())
8092    }
8093    /// Read back a device u32 buffer (the spec accept walk's [t] per-column argmax tokens).
8094    pub fn htod_u32_v(&self, v: &[u32]) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
8095        Ok(self.gpu.stream().clone_htod(v)?)
8096    }
8097    pub fn dtoh_u32(&self, d: &CudaSlice<u32>) -> Result<Vec<u32>, Box<dyn std::error::Error>> {
8098        let v = self.gpu.stream().clone_dtoh(d)?;
8099        self.gpu.stream().synchronize()?;
8100        Ok(v)
8101    }
8102    /// Allocate a zeroed device u32 buffer (persistent spec-loop prediction slots).
8103    /// H2D into an EXISTING u32 buffer (stable pointer — the per-step grammar-mask upload:
8104    /// contents change every step, the address must not, so a captured graph can read it).
8105    pub fn htod_u32_into(
8106        &self,
8107        dst: &mut CudaSlice<u32>,
8108        src: &[u32],
8109    ) -> Result<(), Box<dyn std::error::Error>> {
8110        let mut view = dst.slice_mut(0..src.len());
8111        self.gpu.stream().memcpy_htod(src, &mut view)?;
8112        Ok(())
8113    }
8114
8115    /// H2D into an existing i32 buffer. OPTIPIPE uses this to refresh a stage-local saved-len
8116    /// table without changing the device address its reconcile kernel consumes.
8117    pub fn htod_i32_into(
8118        &self,
8119        dst: &mut CudaSlice<i32>,
8120        src: &[i32],
8121    ) -> Result<(), Box<dyn std::error::Error>> {
8122        let mut view = dst.slice_mut(0..src.len());
8123        self.gpu.stream().memcpy_htod(src, &mut view)?;
8124        Ok(())
8125    }
8126
8127    pub fn alloc_u32_zeroed(&self, n: usize) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
8128        let s = self.gpu.stream().alloc_zeros::<u32>(n)?;
8129        self.keep_if_capturing(&s);
8130        Ok(s)
8131    }
8132    /// embed_gather into a PERSISTENT `x_out` buffer (stable pointer) for CUDA-graph capture (the
8133    /// embed output starts the per-step kernel chain and must be at a fixed address across replays).
8134    pub fn embed_gather_device_into(
8135        &self,
8136        embd: &CudaSlice<u8>,
8137        token_d: &CudaSlice<u32>,
8138        x_out: &mut CudaSlice<f32>,
8139        n_embd: usize,
8140        qtype: i32,
8141        row_bytes: usize,
8142    ) -> Result<(), Box<dyn std::error::Error>> {
8143        let f = self.func("embed_gather_u32");
8144        let cfg = LaunchConfig {
8145            grid_dim: (((n_embd as u32 + 255) / 256).max(1), 1, 1),
8146            block_dim: (256, 1, 1),
8147            shared_mem_bytes: 0,
8148        };
8149        let (ne, qt, rb) = (n_embd as i32, qtype, row_bytes as i64);
8150        let __s_b = self.gpu.stream();
8151        let mut b = __s_b.launch_builder(&f);
8152        b.arg(embd)
8153            .arg(token_d)
8154            .arg(x_out)
8155            .arg(&ne)
8156            .arg(&qt)
8157            .arg(&rb);
8158        unsafe {
8159            b.launch(cfg)?;
8160        }
8161        Ok(())
8162    }
8163    /// Read a [1] i32 device counter (pos / seqlen) back to host. Tiny D2H + sync.
8164    pub fn dtoh_i32_one(&self, d: &CudaSlice<i32>) -> Result<i32, Box<dyn std::error::Error>> {
8165        let v = self.gpu.stream().clone_dtoh(d)?;
8166        self.gpu.stream().synchronize()?;
8167        Ok(v[0])
8168    }
8169    /// Set a [1] i32 device counter IN PLACE (keeps the buffer pointer stable — required for the
8170    /// graph-resident pos/seqlen counters whose addresses are baked into captured graphs). Restores
8171    /// the counter value after the throwaway capture warmups corrupt it.
8172    /// ASYNC i32 single-slot store (value rides the kernel arg — no host-memory transfer/sync).
8173    /// The graph-arc device-len counters use this; set_i32_one below is the SYNCING pageable
8174    /// copy (fine at stream-idle boundaries, poison mid-round).
8175    pub fn i32_set_k(
8176        &self,
8177        dst: &mut CudaSlice<i32>,
8178        v: i32,
8179    ) -> Result<(), Box<dyn std::error::Error>> {
8180        let f = self.func("i32_set_k");
8181        let cfg = LaunchConfig {
8182            grid_dim: (1, 1, 1),
8183            block_dim: (1, 1, 1),
8184            shared_mem_bytes: 0,
8185        };
8186        let idx = 0i32;
8187        let __s_b = self.gpu.stream();
8188        let mut b = __s_b.launch_builder(&f);
8189        b.arg(dst).arg(&v).arg(&idx);
8190        unsafe {
8191            b.launch(cfg)?;
8192        }
8193        Ok(())
8194    }
8195
8196    pub fn set_i32_one(
8197        &self,
8198        d: &mut CudaSlice<i32>,
8199        v: i32,
8200    ) -> Result<(), Box<dyn std::error::Error>> {
8201        self.gpu.stream().memcpy_htod(&[v], d)?;
8202        Ok(())
8203    }
8204    /// Set a [1] u32 device buffer IN PLACE (stable pointer) — for the resident `token_d` counter
8205    /// during priming / capture-state restore.
8206    pub fn set_u32_one(
8207        &self,
8208        d: &mut CudaSlice<u32>,
8209        v: u32,
8210    ) -> Result<(), Box<dyn std::error::Error>> {
8211        self.gpu.stream().memcpy_htod(&[v], d)?;
8212        Ok(())
8213    }
8214    /// Read back a [1] u32 device buffer (the argmax token). One tiny D2H + sync.
8215    pub fn dtoh_u32_one(&self, d: &CudaSlice<u32>) -> Result<u32, Box<dyn std::error::Error>> {
8216        let v = self.gpu.stream().clone_dtoh(d)?;
8217        self.gpu.stream().synchronize()?;
8218        Ok(v[0])
8219    }
8220    /// Upload raw bytes to a resident device u8 buffer (e.g. the embed table for device gather).
8221    pub fn upload_u8(&self, bytes: &[u8]) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
8222        Ok(self.gpu.stream().clone_htod(bytes)?)
8223    }
8224    /// Embed-from-device (CUDA-GRAPH-PLAN Phase 1): gather+dequant the row for the token id in
8225    /// `token_d[0]` from the resident embed table `embd` -> x_out[n_embd]. Bit-identical to host
8226    /// EmbedHost::gather (same per-dtype `deq`). No host round-trip of the token id.
8227    pub fn embed_gather_device(
8228        &self,
8229        embd: &CudaSlice<u8>,
8230        token_d: &CudaSlice<u32>,
8231        n_embd: usize,
8232        qtype: i32,
8233        row_bytes: usize,
8234    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8235        let f = self.func("embed_gather_u32");
8236        let mut x = self.alloc_uninit::<f32>(n_embd)?;
8237        let cfg = LaunchConfig {
8238            grid_dim: (((n_embd as u32 + 255) / 256).max(1), 1, 1),
8239            block_dim: (256, 1, 1),
8240            shared_mem_bytes: 0,
8241        };
8242        let (ne, qt, rb) = (n_embd as i32, qtype, row_bytes as i64);
8243        let __s_b = self.gpu.stream();
8244        let mut b = __s_b.launch_builder(&f);
8245        b.arg(embd)
8246            .arg(token_d)
8247            .arg(&mut x)
8248            .arg(&ne)
8249            .arg(&qt)
8250            .arg(&rb);
8251        unsafe {
8252            b.launch(cfg)?;
8253        }
8254        Ok(x)
8255    }
8256
8257    /// T-token device embed gather (spec verify/replay): tokens uploaded as a tiny [T] u32 htod,
8258    /// rows dequanted on-device -> x[T, n_embd]. Replaces host per-row dequant + T*n_embd*4B htod
8259    /// (nsys: 84% of spec API time was HtoD). Bit-identical rows (same per-dtype deq).
8260    pub fn embed_gather_device_t(
8261        &self,
8262        embd: &CudaSlice<u8>,
8263        tokens: &[u32],
8264        n_embd: usize,
8265        qtype: i32,
8266        row_bytes: usize,
8267    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8268        let t = tokens.len();
8269        let tok_d = self.gpu.stream().clone_htod(tokens)?;
8270        let f = self.func("embed_gather_u32_t");
8271        let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
8272        let cfg = LaunchConfig {
8273            grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
8274            block_dim: (256, 1, 1),
8275            shared_mem_bytes: 0,
8276        };
8277        let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
8278        let __s_b = self.gpu.stream();
8279        let mut b = __s_b.launch_builder(&f);
8280        b.arg(embd)
8281            .arg(&tok_d)
8282            .arg(&mut x)
8283            .arg(&ne)
8284            .arg(&qt)
8285            .arg(&rb)
8286            .arg(&ti);
8287        unsafe {
8288            b.launch(cfg)?;
8289        }
8290        Ok(x)
8291    }
8292
8293    /// T-token embed gather from a DEVICE token buffer (round-stream stage c: the verify tokens
8294    /// are assembled on-device from the draft-chain pack slots; no host round trip). Same kernel
8295    /// as embed_gather_device_t — bit-identical rows.
8296    /// embed_gather over a token VIEW (spec round: tokens live in the round's batch buffer).
8297    pub fn embed_gather_device_tv(
8298        &self,
8299        embd: &CudaSlice<u8>,
8300        tok_v: &cudarc::driver::CudaView<u32>,
8301        t: usize,
8302        n_embd: usize,
8303        qtype: i32,
8304        row_bytes: usize,
8305    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8306        let f = self.func("embed_gather_u32_t");
8307        let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
8308        let cfg = LaunchConfig {
8309            grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
8310            block_dim: (256, 1, 1),
8311            shared_mem_bytes: 0,
8312        };
8313        let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
8314        let __s_b = self.gpu.stream();
8315        let mut b = __s_b.launch_builder(&f);
8316        b.arg(embd)
8317            .arg(tok_v)
8318            .arg(&mut x)
8319            .arg(&ne)
8320            .arg(&qt)
8321            .arg(&rb)
8322            .arg(&ti);
8323        unsafe {
8324            b.launch(cfg)?;
8325        }
8326        Ok(x)
8327    }
8328
8329    pub fn embed_gather_device_td(
8330        &self,
8331        embd: &CudaSlice<u8>,
8332        tok_d: &CudaSlice<u32>,
8333        t: usize,
8334        n_embd: usize,
8335        qtype: i32,
8336        row_bytes: usize,
8337    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8338        let f = self.func("embed_gather_u32_t");
8339        let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
8340        let cfg = LaunchConfig {
8341            grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
8342            block_dim: (256, 1, 1),
8343            shared_mem_bytes: 0,
8344        };
8345        let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
8346        let __s_b = self.gpu.stream();
8347        let mut b = __s_b.launch_builder(&f);
8348        b.arg(embd)
8349            .arg(tok_d)
8350            .arg(&mut x)
8351            .arg(&ne)
8352            .arg(&qt)
8353            .arg(&rb)
8354            .arg(&ti);
8355        unsafe {
8356            b.launch(cfg)?;
8357        }
8358        Ok(x)
8359    }
8360
8361    /// Uninitialized device buffer — SKIPS the memset that `alloc_zeros` always issues. Decode
8362    /// profile (nsys): ~1050 memsets/token = 6.5% of decode GPU time + ~half the launch count, the
8363    /// dominant contributor to the 19% inter-kernel idle gap and a blocker for clean CUDA-graph
8364    /// capture. Use ONLY for buffers a kernel FULLY overwrites (every element written, no `+=`).
8365    /// SAFETY: caller guarantees the producing kernel writes every element before any read.
8366    #[inline]
8367    /// Keep an allocation alive for the current capture (no-op when retain mode is off).
8368    fn keep_if_capturing<T: cudarc::driver::DeviceRepr + Send + 'static>(&self, s: &CudaSlice<T>) {
8369        if self
8370            .capture_keep_on
8371            .load(std::sync::atomic::Ordering::Relaxed)
8372        {
8373            self.capture_keep.lock().unwrap().push(Box::new(s.clone()));
8374        }
8375    }
8376
8377    fn alloc_uninit<T: cudarc::driver::DeviceRepr + Send + 'static>(
8378        &self,
8379        n: usize,
8380    ) -> Result<CudaSlice<T>, Box<dyn std::error::Error>> {
8381        let mut s = unsafe { self.gpu.stream().alloc::<T>(n)? };
8382        // MEMRA_DEBUG_ZERO_ALLOCS=1 (task #14 defect hunt): memset EVERY engine allocation —
8383        // the global uninit-read discriminator (the prime-fn-scoped zeroing experiment could
8384        // not cover engine-internal buffers). Debug-only: massive launch overhead.
8385        {
8386            static Z: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8387            if *Z.get_or_init(|| std::env::var("MEMRA_DEBUG_ZERO_ALLOCS").as_deref() == Ok("1")) {
8388                // raw D8 memset (T lacks ValidAsZeroBits in the generic bound)
8389                use cudarc::driver::DevicePtrMut;
8390                let n_bytes = s.len() * std::mem::size_of::<T>();
8391                let stream = self.gpu.stream();
8392                let (p_, _g) = s.device_ptr_mut(&stream);
8393                unsafe {
8394                    cudarc::driver::sys::cuMemsetD8Async(p_, 0, n_bytes, stream.cu_stream())
8395                        .result()?;
8396                }
8397            }
8398        }
8399        self.keep_if_capturing(&s);
8400        Ok(s)
8401    }
8402
8403    /// Public f32 uninitialized scratch (see `alloc_uninit`). For decode/forward scratch a kernel
8404    /// fully overwrites. SAFETY: producing kernel must write every element before any read.
8405    /// Uninitialized q8_1 activation pair (int8 + per-32 scales) — the fa combine q8-emit
8406    /// consumers alloc through this (m=1 decode arms).
8407    pub fn uninit_q8_pair(
8408        &self,
8409        n: usize,
8410    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8411        Ok((
8412            self.alloc_uninit::<i8>(n)?,
8413            self.alloc_uninit::<f32>(n / 32)?,
8414        ))
8415    }
8416
8417    pub fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8418        self.alloc_uninit::<f32>(n)
8419    }
8420
8421    /// i8 uninitialized scratch (same contract as `uninit`).
8422    pub fn alloc_i8_uninit(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
8423        self.alloc_uninit::<i8>(n)
8424    }
8425
8426    /// RMSNorm: x[ncols,nrows] row-major, weight[ncols] -> dst. One block/row, 256 threads.
8427    /// gemma4: 3 rms_norms of the SAME input in one launch (one reduction, three weights).
8428    /// Per-output bit-identical to three rms_norm calls (verbatim reduction/scale chain).
8429    #[allow(clippy::too_many_arguments)]
8430    pub fn rms_norm3(
8431        &self,
8432        x: &CudaSlice<f32>,
8433        w0: &CudaSlice<f32>,
8434        w1: &CudaSlice<f32>,
8435        w2: &CudaSlice<f32>,
8436        d0: &mut CudaSlice<f32>,
8437        d1: &mut CudaSlice<f32>,
8438        d2: &mut CudaSlice<f32>,
8439        ncols: usize,
8440        nrows: usize,
8441        eps: f32,
8442    ) -> Result<(), Box<dyn std::error::Error>> {
8443        let f = self.func("rms_norm3_f32");
8444        let cfg = LaunchConfig {
8445            grid_dim: (nrows as u32, 1, 1),
8446            block_dim: (rms_block(), 1, 1),
8447            shared_mem_bytes: 0,
8448        };
8449        let (nc, e) = (ncols as i32, eps);
8450        let __s_b = self.gpu.stream();
8451        let mut b = __s_b.launch_builder(&f);
8452        b.arg(x)
8453            .arg(w0)
8454            .arg(w1)
8455            .arg(w2)
8456            .arg(d0)
8457            .arg(d1)
8458            .arg(d2)
8459            .arg(&nc)
8460            .arg(&e);
8461        unsafe {
8462            b.launch(cfg)?;
8463        }
8464        Ok(())
8465    }
8466
8467    /// gemma4 fused q/k/v head norms (one launch, per-row rms_norm_f32-verbatim).
8468    #[allow(clippy::too_many_arguments)]
8469    /// True when the warp-per-row qkv norm would engage for (rows, ncols) — the emit lane
8470    /// piggybacks on the same conditions.
8471    pub fn qkvnorm_w_on_prefill(rows: usize, ncols: usize) -> bool {
8472        static WARP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8473        *WARP_ON.get_or_init(|| {
8474            std::env::var("MEMRA_QKVNORM_W")
8475                .map(|v| v != "0")
8476                .unwrap_or(true)
8477        }) && ncols % 4 == 0
8478            && rows >= 64
8479    }
8480
8481    /// w4 norm with bf16 V EMIT (31B glue lane): the v segment also writes its normed rows as
8482    /// bf16 (the FA V operand — bit-identical to a post-hoc f32_to_bf16). Prefill-depth only.
8483    #[allow(clippy::too_many_arguments)]
8484    pub fn rms_norm_qkv_w4b(
8485        &self,
8486        q: &CudaSlice<f32>,
8487        k: &CudaSlice<f32>,
8488        v: &CudaSlice<f32>,
8489        wq: &CudaSlice<f32>,
8490        wk: &CudaSlice<f32>,
8491        wv: &CudaSlice<f32>,
8492        dq: &mut CudaSlice<f32>,
8493        dk: &mut CudaSlice<f32>,
8494        dv: &mut CudaSlice<f32>,
8495        dvb: &mut CudaSlice<u8>,
8496        ncols: usize,
8497        rq: usize,
8498        rk: usize,
8499        eps: f32,
8500        vf16: bool,
8501    ) -> Result<(), Box<dyn std::error::Error>> {
8502        assert!(ncols % 4 == 0 && rq + 2 * rk >= 64);
8503        let f = self.func("rms_norm_qkv_w4b_f32");
8504        let rows = (rq + 2 * rk) as u32;
8505        let cfg = LaunchConfig {
8506            grid_dim: (rows.div_ceil(8), 1, 1),
8507            block_dim: (256, 1, 1),
8508            shared_mem_bytes: 0,
8509        };
8510        let (nc, rqi, rki, rvi, e) = (ncols as i32, rq as i32, rk as i32, rk as i32, eps);
8511        let vf = vf16 as i32;
8512        let __s_b = self.gpu.stream();
8513        let mut b = __s_b.launch_builder(&f);
8514        b.arg(q)
8515            .arg(k)
8516            .arg(v)
8517            .arg(wq)
8518            .arg(wk)
8519            .arg(wv)
8520            .arg(dq)
8521            .arg(dk)
8522            .arg(dv)
8523            .arg(&mut *dvb)
8524            .arg(&nc)
8525            .arg(&rqi)
8526            .arg(&rki)
8527            .arg(&rvi)
8528            .arg(&e)
8529            .arg(&vf);
8530        unsafe {
8531            b.launch(cfg)?;
8532        }
8533        Ok(())
8534    }
8535
8536    pub fn rms_norm_qkv(
8537        &self,
8538        q: &CudaSlice<f32>,
8539        k: &CudaSlice<f32>,
8540        v: &CudaSlice<f32>,
8541        wq: &CudaSlice<f32>,
8542        wk: &CudaSlice<f32>,
8543        wv: &CudaSlice<f32>,
8544        dq: &mut CudaSlice<f32>,
8545        dk: &mut CudaSlice<f32>,
8546        dv: &mut CudaSlice<f32>,
8547        ncols: usize,
8548        rq: usize,
8549        rk: usize,
8550        eps: f32,
8551    ) -> Result<(), Box<dyn std::error::Error>> {
8552        // Warp-per-row float4 twin (default; MEMRA_QKVNORM_W=0 reverts): the block-per-row form
8553        // spends 767us/launch on 17k+ 2KB rows at prefill depth (launch/reduce latency-bound,
8554        // ~92GB/s). Own numeric config (reduce order differs) — battery-gated.
8555        static WARP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8556        let warp_on = *WARP_ON.get_or_init(|| {
8557            std::env::var("MEMRA_QKVNORM_W")
8558                .map(|v| v != "0")
8559                .unwrap_or(true)
8560        });
8561        // rows >= 64 keeps decode (nh + 2*nkv rows) on the block-tree kernel — decode/verify/
8562        // replay numerics are untouched on every model; only prefill depth takes the new config.
8563        if warp_on && ncols % 4 == 0 && rq + 2 * rk >= 64 {
8564            let f = self.func("rms_norm_qkv_w4_f32");
8565            let rows = (rq + 2 * rk) as u32;
8566            let cfg = LaunchConfig {
8567                grid_dim: (rows.div_ceil(8), 1, 1),
8568                block_dim: (256, 1, 1),
8569                shared_mem_bytes: 0,
8570            };
8571            let (nc, rqi, rki, rvi, e) = (ncols as i32, rq as i32, rk as i32, rk as i32, eps);
8572            let __s_b = self.gpu.stream();
8573            let mut b = __s_b.launch_builder(&f);
8574            b.arg(q)
8575                .arg(k)
8576                .arg(v)
8577                .arg(wq)
8578                .arg(wk)
8579                .arg(wv)
8580                .arg(dq)
8581                .arg(dk)
8582                .arg(dv)
8583                .arg(&nc)
8584                .arg(&rqi)
8585                .arg(&rki)
8586                .arg(&rvi)
8587                .arg(&e);
8588            unsafe {
8589                b.launch(cfg)?;
8590            }
8591            return Ok(());
8592        }
8593        let f = self.func("rms_norm_qkv_f32");
8594        let grid = (rq + 2 * rk) as u32;
8595        let cfg = LaunchConfig {
8596            grid_dim: (grid, 1, 1),
8597            block_dim: (rms_block(), 1, 1),
8598            shared_mem_bytes: 0,
8599        };
8600        let (nc, rqi, rki, e) = (ncols as i32, rq as i32, rk as i32, eps);
8601        let __s_b = self.gpu.stream();
8602        let mut b = __s_b.launch_builder(&f);
8603        b.arg(q)
8604            .arg(k)
8605            .arg(v)
8606            .arg(wq)
8607            .arg(wk)
8608            .arg(wv)
8609            .arg(dq)
8610            .arg(dk)
8611            .arg(dv)
8612            .arg(&nc)
8613            .arg(&rqi)
8614            .arg(&rki)
8615            .arg(&e);
8616        unsafe {
8617            b.launch(cfg)?;
8618        }
8619        Ok(())
8620    }
8621
8622    /// gemma4 fused pair of rms_norms over two different inputs (same width).
8623    #[allow(clippy::too_many_arguments)]
8624    pub fn rms_norm2x(
8625        &self,
8626        a: &CudaSlice<f32>,
8627        bb: &CudaSlice<f32>,
8628        wa: &CudaSlice<f32>,
8629        wb: &CudaSlice<f32>,
8630        da: &mut CudaSlice<f32>,
8631        db: &mut CudaSlice<f32>,
8632        ncols: usize,
8633        nrows: usize,
8634        eps: f32,
8635    ) -> Result<(), Box<dyn std::error::Error>> {
8636        let f = self.func("rms_norm2x_f32");
8637        let cfg = LaunchConfig {
8638            grid_dim: (2 * nrows as u32, 1, 1),
8639            block_dim: (rms_block(), 1, 1),
8640            shared_mem_bytes: 0,
8641        };
8642        let (nc, nr, e) = (ncols as i32, nrows as i32, eps);
8643        let __s_b = self.gpu.stream();
8644        let mut b = __s_b.launch_builder(&f);
8645        b.arg(a)
8646            .arg(bb)
8647            .arg(wa)
8648            .arg(wb)
8649            .arg(da)
8650            .arg(db)
8651            .arg(&nc)
8652            .arg(&nr)
8653            .arg(&e);
8654        unsafe {
8655            b.launch(cfg)?;
8656        }
8657        Ok(())
8658    }
8659
8660    /// gemma4 R4: in-place final-logit softcap y = cap*tanh(y/cap).
8661    pub fn softcap(
8662        &self,
8663        y: &mut CudaSlice<f32>,
8664        cap: f32,
8665        n: usize,
8666    ) -> Result<(), Box<dyn std::error::Error>> {
8667        let f = self.func("softcap_f32");
8668        let cfg = LaunchConfig::for_num_elems(n as u32);
8669        let ni = n as i32;
8670        let __s_b = self.gpu.stream();
8671        let mut b = __s_b.launch_builder(&f);
8672        b.arg(y).arg(&cap).arg(&ni);
8673        unsafe {
8674            b.launch(cfg)?;
8675        }
8676        Ok(())
8677    }
8678
8679    /// gemma4 suppress-token mask: y[row][ids[j]] = -inf over t logits rows (fixed-arg launch —
8680    /// graph-capture safe; NOT monotonic like softcap, so it must run before any argmax).
8681    pub fn mask_ids_rows(
8682        &self,
8683        y: &mut CudaSlice<f32>,
8684        ids: &CudaSlice<i32>,
8685        n_ids: usize,
8686        n_vocab: usize,
8687        t: usize,
8688    ) -> Result<(), Box<dyn std::error::Error>> {
8689        let f = self.func("mask_ids_rows_f32");
8690        let cfg = LaunchConfig::for_num_elems((n_ids * t) as u32);
8691        let (ni, nv, ti) = (n_ids as i32, n_vocab as i32, t as i32);
8692        let __s_b = self.gpu.stream();
8693        let mut b = __s_b.launch_builder(&f);
8694        b.arg(y).arg(ids).arg(&ni).arg(&nv).arg(&ti);
8695        unsafe {
8696            b.launch(cfg)?;
8697        }
8698        Ok(())
8699    }
8700
8701    /// gemma4: res = (a+b)*c AND dst = rms_norm(res, w) in one launch.
8702    #[allow(clippy::too_many_arguments)]
8703    pub fn add_scale_rms_norm(
8704        &self,
8705        a: &CudaSlice<f32>,
8706        b_in: &CudaSlice<f32>,
8707        c: f32,
8708        w: &CudaSlice<f32>,
8709        res: &mut CudaSlice<f32>,
8710        dst: &mut CudaSlice<f32>,
8711        ncols: usize,
8712        nrows: usize,
8713        eps: f32,
8714    ) -> Result<(), Box<dyn std::error::Error>> {
8715        let f = self.func("add_scale_rms_norm_f32");
8716        let cfg = LaunchConfig {
8717            grid_dim: (nrows as u32, 1, 1),
8718            block_dim: (rms_block(), 1, 1),
8719            shared_mem_bytes: 0,
8720        };
8721        let (nc, e2) = (ncols as i32, eps);
8722        let __s_b = self.gpu.stream();
8723        let mut b = __s_b.launch_builder(&f);
8724        b.arg(a)
8725            .arg(b_in)
8726            .arg(&c)
8727            .arg(w)
8728            .arg(res)
8729            .arg(dst)
8730            .arg(&nc)
8731            .arg(&e2);
8732        unsafe {
8733            b.launch(cfg)?;
8734        }
8735        Ok(())
8736    }
8737
8738    /// gemma4: res = (a+b)*c AND the next layer's attn_norm EMITTED q8_1 in one launch.
8739    /// Quantize epilogue bit-identical to quantize_q8_1 (the rms_norm_q8_1 form).
8740    #[allow(clippy::too_many_arguments)]
8741    pub fn add_scale_rms_norm_q8_1(
8742        &self,
8743        a: &CudaSlice<f32>,
8744        b_in: &CudaSlice<f32>,
8745        c: f32,
8746        w: &CudaSlice<f32>,
8747        res: &mut CudaSlice<f32>,
8748        ncols: usize,
8749        nrows: usize,
8750        eps: f32,
8751    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8752        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
8753        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
8754        let (nc, e2) = (ncols as i32, eps);
8755        if Self::pdl_on() && Self::pdl_wb_on() {
8756            {
8757                use cudarc::driver::{DevicePtr, DevicePtrMut};
8758                let s = &self.gpu.stream();
8759                let (pa, _g0) = a.device_ptr(s);
8760                let (pb, _g1) = b_in.device_ptr(s);
8761                let (pw, _g2) = w.device_ptr(s);
8762                let (pr, _g3) = res.device_ptr_mut(s);
8763                let (pq, _g4) = out_q.device_ptr_mut(s);
8764                let (pd, _g5) = out_d.device_ptr_mut(s);
8765                let mut ps = [
8766                    &pa as *const _ as *mut std::ffi::c_void,
8767                    &pb as *const _ as *mut _,
8768                    &c as *const _ as *mut _,
8769                    &pw as *const _ as *mut _,
8770                    &pr as *const _ as *mut _,
8771                    &pq as *const _ as *mut _,
8772                    &pd as *const _ as *mut _,
8773                    &nc as *const _ as *mut _,
8774                    &e2 as *const _ as *mut _,
8775                ];
8776                unsafe {
8777                    self.launch_pdl(
8778                        "add_scale_rms_norm_q8_1",
8779                        (nrows as u32, 1, 1),
8780                        (rms_block(), 1, 1),
8781                        &mut ps,
8782                    )?;
8783                }
8784            }
8785            return Ok((out_q, out_d));
8786        }
8787        let f = self.func("add_scale_rms_norm_q8_1");
8788        let cfg = LaunchConfig {
8789            grid_dim: (nrows as u32, 1, 1),
8790            block_dim: (rms_block(), 1, 1),
8791            shared_mem_bytes: 0,
8792        };
8793        let __s_b = self.gpu.stream();
8794        let mut b = __s_b.launch_builder(&f);
8795        b.arg(a)
8796            .arg(b_in)
8797            .arg(&c)
8798            .arg(w)
8799            .arg(res)
8800            .arg(&mut out_q)
8801            .arg(&mut out_d)
8802            .arg(&nc)
8803            .arg(&e2);
8804        unsafe {
8805            b.launch(cfg)?;
8806        }
8807        Ok((out_q, out_d))
8808    }
8809
8810    /// Slot-fed add_scale_rms_norm_q8_1 twin (alloc-free capture lane).
8811    #[allow(clippy::too_many_arguments)]
8812    pub fn add_scale_rms_norm_q8_1_into(
8813        &self,
8814        a: &CudaSlice<f32>,
8815        b_in: &CudaSlice<f32>,
8816        c: f32,
8817        w: &CudaSlice<f32>,
8818        res: &mut CudaSlice<f32>,
8819        ncols: usize,
8820        nrows: usize,
8821        eps: f32,
8822        out_q: &mut CudaSlice<i8>,
8823        out_d: &mut CudaSlice<f32>,
8824    ) -> Result<(), Box<dyn std::error::Error>> {
8825        debug_assert!(out_q.len() >= nrows * ncols && out_d.len() >= nrows * (ncols / 32));
8826        let (nc, e2) = (ncols as i32, eps);
8827        if Self::pdl_on() && Self::pdl_wb_on() {
8828            use cudarc::driver::{DevicePtr, DevicePtrMut};
8829            let s = &self.gpu.stream();
8830            let (pa, _g0) = a.device_ptr(s);
8831            let (pb, _g1) = b_in.device_ptr(s);
8832            let (pw, _g2) = w.device_ptr(s);
8833            let (pr, _g3) = res.device_ptr_mut(s);
8834            let (pq, _g4) = out_q.device_ptr_mut(s);
8835            let (pd, _g5) = out_d.device_ptr_mut(s);
8836            let mut ps = [
8837                &pa as *const _ as *mut std::ffi::c_void,
8838                &pb as *const _ as *mut _,
8839                &c as *const _ as *mut _,
8840                &pw as *const _ as *mut _,
8841                &pr as *const _ as *mut _,
8842                &pq as *const _ as *mut _,
8843                &pd as *const _ as *mut _,
8844                &nc as *const _ as *mut _,
8845                &e2 as *const _ as *mut _,
8846            ];
8847            unsafe {
8848                self.launch_pdl(
8849                    "add_scale_rms_norm_q8_1",
8850                    (nrows as u32, 1, 1),
8851                    (rms_block(), 1, 1),
8852                    &mut ps,
8853                )?;
8854            }
8855            return Ok(());
8856        }
8857        let f = self.func("add_scale_rms_norm_q8_1");
8858        let cfg = LaunchConfig {
8859            grid_dim: (nrows as u32, 1, 1),
8860            block_dim: (rms_block(), 1, 1),
8861            shared_mem_bytes: 0,
8862        };
8863        let __s_b = self.gpu.stream();
8864        let mut b = __s_b.launch_builder(&f);
8865        b.arg(a)
8866            .arg(b_in)
8867            .arg(&c)
8868            .arg(w)
8869            .arg(res)
8870            .arg(&mut *out_q)
8871            .arg(&mut *out_d)
8872            .arg(&nc)
8873            .arg(&e2);
8874        unsafe {
8875            b.launch(cfg)?;
8876        }
8877        Ok(())
8878    }
8879
8880    /// E4B glue fusion: rms(a, wa) prologue + the add_scale_rms_norm_q8_1 program — one launch
8881    /// replaces the per-layer rms_norm_f32(y) + emit pair in the PLE tail.
8882    #[allow(clippy::too_many_arguments)]
8883    pub fn rms_pre_add_scale_rms_norm_q8_1(
8884        &self,
8885        a: &CudaSlice<f32>,
8886        wa: &CudaSlice<f32>,
8887        b_in: &CudaSlice<f32>,
8888        c: f32,
8889        w: &CudaSlice<f32>,
8890        res: &mut CudaSlice<f32>,
8891        ncols: usize,
8892        nrows: usize,
8893        eps: f32,
8894    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8895        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
8896        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
8897        let (nc, e2) = (ncols as i32, eps);
8898        if Self::pdl_on() {
8899            {
8900                use cudarc::driver::{DevicePtr, DevicePtrMut};
8901                let s = &self.gpu.stream();
8902                let (pa, _g0) = a.device_ptr(s);
8903                let (pwa, _g1) = wa.device_ptr(s);
8904                let (pb, _g2) = b_in.device_ptr(s);
8905                let (pw, _g3) = w.device_ptr(s);
8906                let (pr, _g4) = res.device_ptr_mut(s);
8907                let (pq, _g5) = out_q.device_ptr_mut(s);
8908                let (pd, _g6) = out_d.device_ptr_mut(s);
8909                let mut ps = [
8910                    &pa as *const _ as *mut std::ffi::c_void,
8911                    &pwa as *const _ as *mut _,
8912                    &pb as *const _ as *mut _,
8913                    &c as *const _ as *mut _,
8914                    &pw as *const _ as *mut _,
8915                    &pr as *const _ as *mut _,
8916                    &pq as *const _ as *mut _,
8917                    &pd as *const _ as *mut _,
8918                    &nc as *const _ as *mut _,
8919                    &e2 as *const _ as *mut _,
8920                ];
8921                unsafe {
8922                    self.launch_pdl(
8923                        "rms_pre_add_scale_rms_norm_q8_1",
8924                        (nrows as u32, 1, 1),
8925                        (rms_block(), 1, 1),
8926                        &mut ps,
8927                    )?;
8928                }
8929            }
8930            return Ok((out_q, out_d));
8931        }
8932        let f = self.func("rms_pre_add_scale_rms_norm_q8_1");
8933        let cfg = LaunchConfig {
8934            grid_dim: (nrows as u32, 1, 1),
8935            block_dim: (rms_block(), 1, 1),
8936            shared_mem_bytes: 0,
8937        };
8938        let __s_b = self.gpu.stream();
8939        let mut b = __s_b.launch_builder(&f);
8940        b.arg(a)
8941            .arg(wa)
8942            .arg(b_in)
8943            .arg(&c)
8944            .arg(w)
8945            .arg(res)
8946            .arg(&mut out_q)
8947            .arg(&mut out_d)
8948            .arg(&nc)
8949            .arg(&e2);
8950        unsafe {
8951            b.launch(cfg)?;
8952        }
8953        Ok((out_q, out_d))
8954    }
8955
8956    /// GELU(tanh)*up with the activation emitted q8_1 alongside f32 (glue-fusion lane): the
8957    /// consumer matmul rides matmul_pre, killing its standalone quantize_q8_1 launch.
8958    pub fn gelu_tanh_mul_q8_1(
8959        &self,
8960        gate: &CudaSlice<f32>,
8961        up: &cudarc::driver::CudaView<f32>,
8962        act: &mut CudaSlice<f32>,
8963        ncols: usize,
8964        nrows: usize,
8965    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8966        debug_assert!(ncols % 128 == 0);
8967        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
8968        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
8969        let nc = ncols as i32;
8970        if Self::pdl_on() {
8971            {
8972                use cudarc::driver::{DevicePtr, DevicePtrMut};
8973                let s = &self.gpu.stream();
8974                let (pg, _g0) = gate.device_ptr(s);
8975                let (pu, _g1) = up.device_ptr(s);
8976                let (pact, _g2) = act.device_ptr_mut(s);
8977                let (pq, _g3) = out_q.device_ptr_mut(s);
8978                let (pd, _g4) = out_d.device_ptr_mut(s);
8979                let mut ps = [
8980                    &pg as *const _ as *mut std::ffi::c_void,
8981                    &pu as *const _ as *mut _,
8982                    &pact as *const _ as *mut _,
8983                    &pq as *const _ as *mut _,
8984                    &pd as *const _ as *mut _,
8985                    &nc as *const _ as *mut _,
8986                ];
8987                unsafe {
8988                    self.launch_pdl(
8989                        "gelu_tanh_mul_q8_1",
8990                        (nrows as u32, 1, 1),
8991                        (rms_block(), 1, 1),
8992                        &mut ps,
8993                    )?;
8994                }
8995            }
8996            return Ok((out_q, out_d));
8997        }
8998        let f = self.func("gelu_tanh_mul_q8_1");
8999        let cfg = LaunchConfig {
9000            grid_dim: (nrows as u32, 1, 1),
9001            block_dim: (rms_block(), 1, 1),
9002            shared_mem_bytes: 0,
9003        };
9004        let __s_b = self.gpu.stream();
9005        let mut b = __s_b.launch_builder(&f);
9006        b.arg(gate)
9007            .arg(up)
9008            .arg(act)
9009            .arg(&mut out_q)
9010            .arg(&mut out_d)
9011            .arg(&nc);
9012        unsafe {
9013            b.launch(cfg)?;
9014        }
9015        Ok((out_q, out_d))
9016    }
9017
9018    /// Slot-fed gelu_tanh_mul_q8_1 twin (alloc-free capture lane; incl. the PDL arm).
9019    #[allow(clippy::too_many_arguments)]
9020    pub fn gelu_tanh_mul_q8_1_into(
9021        &self,
9022        gate: &CudaSlice<f32>,
9023        up: &cudarc::driver::CudaView<f32>,
9024        act: &mut CudaSlice<f32>,
9025        ncols: usize,
9026        nrows: usize,
9027        out_q: &mut CudaSlice<i8>,
9028        out_d: &mut CudaSlice<f32>,
9029    ) -> Result<(), Box<dyn std::error::Error>> {
9030        debug_assert!(ncols % 128 == 0);
9031        debug_assert!(out_q.len() >= nrows * ncols && out_d.len() >= nrows * (ncols / 32));
9032        let nc = ncols as i32;
9033        if Self::pdl_on() {
9034            use cudarc::driver::{DevicePtr, DevicePtrMut};
9035            let s = &self.gpu.stream();
9036            let (pg, _g0) = gate.device_ptr(s);
9037            let (pu, _g1) = up.device_ptr(s);
9038            let (pact, _g2) = act.device_ptr_mut(s);
9039            let (pq, _g3) = out_q.device_ptr_mut(s);
9040            let (pd, _g4) = out_d.device_ptr_mut(s);
9041            let mut ps = [
9042                &pg as *const _ as *mut std::ffi::c_void,
9043                &pu as *const _ as *mut _,
9044                &pact as *const _ as *mut _,
9045                &pq as *const _ as *mut _,
9046                &pd as *const _ as *mut _,
9047                &nc as *const _ as *mut _,
9048            ];
9049            unsafe {
9050                self.launch_pdl(
9051                    "gelu_tanh_mul_q8_1",
9052                    (nrows as u32, 1, 1),
9053                    (rms_block(), 1, 1),
9054                    &mut ps,
9055                )?;
9056            }
9057            return Ok(());
9058        }
9059        let f = self.func("gelu_tanh_mul_q8_1");
9060        let cfg = LaunchConfig {
9061            grid_dim: (nrows as u32, 1, 1),
9062            block_dim: (rms_block(), 1, 1),
9063            shared_mem_bytes: 0,
9064        };
9065        let __s_b = self.gpu.stream();
9066        let mut b = __s_b.launch_builder(&f);
9067        b.arg(gate)
9068            .arg(up)
9069            .arg(&mut *act)
9070            .arg(&mut *out_q)
9071            .arg(&mut *out_d)
9072            .arg(&nc);
9073        unsafe {
9074            b.launch(cfg)?;
9075        }
9076        Ok(())
9077    }
9078
9079    /// gemma4: add + rms_norm3 with outputs 0/2 emitted q8_1 (zsh + moe_in) and 1 f32 (router).
9080    #[allow(clippy::too_many_arguments)]
9081    pub fn add_rms_norm3_q8z(
9082        &self,
9083        a: &CudaSlice<f32>,
9084        b_in: &CudaSlice<f32>,
9085        w0: &CudaSlice<f32>,
9086        w1: &CudaSlice<f32>,
9087        w2: &CudaSlice<f32>,
9088        res: &mut CudaSlice<f32>,
9089        out1: &mut CudaSlice<f32>,
9090        ncols: usize,
9091        nrows: usize,
9092        eps: f32,
9093    ) -> Result<
9094        (
9095            (CudaSlice<i8>, CudaSlice<f32>),
9096            (CudaSlice<i8>, CudaSlice<f32>),
9097        ),
9098        Box<dyn std::error::Error>,
9099    > {
9100        let mut q0 = self.alloc_uninit::<i8>(nrows * ncols)?;
9101        let mut d0 = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9102        let mut q2 = self.alloc_uninit::<i8>(nrows * ncols)?;
9103        let mut d2 = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9104        let f = self.func("add_rms_norm3_q8z_f32");
9105        let cfg = LaunchConfig {
9106            grid_dim: (nrows as u32, 1, 1),
9107            block_dim: (rms_block(), 1, 1),
9108            shared_mem_bytes: 0,
9109        };
9110        let (nc, e2) = (ncols as i32, eps);
9111        let __s_b = self.gpu.stream();
9112        let mut b = __s_b.launch_builder(&f);
9113        b.arg(a)
9114            .arg(b_in)
9115            .arg(w0)
9116            .arg(w1)
9117            .arg(w2)
9118            .arg(res)
9119            .arg(&mut q0)
9120            .arg(&mut d0)
9121            .arg(out1)
9122            .arg(&mut q2)
9123            .arg(&mut d2)
9124            .arg(&nc)
9125            .arg(&e2);
9126        unsafe {
9127            b.launch(cfg)?;
9128        }
9129        Ok(((q0, d0), (q2, d2)))
9130    }
9131
9132    /// gemma4: res = a+b AND the three rms_norms of res in one launch.
9133    #[allow(clippy::too_many_arguments)]
9134    pub fn add_rms_norm3(
9135        &self,
9136        a: &CudaSlice<f32>,
9137        b_in: &CudaSlice<f32>,
9138        w0: &CudaSlice<f32>,
9139        w1: &CudaSlice<f32>,
9140        w2: &CudaSlice<f32>,
9141        res: &mut CudaSlice<f32>,
9142        d0: &mut CudaSlice<f32>,
9143        d1: &mut CudaSlice<f32>,
9144        d2: &mut CudaSlice<f32>,
9145        ncols: usize,
9146        nrows: usize,
9147        eps: f32,
9148    ) -> Result<(), Box<dyn std::error::Error>> {
9149        let f = self.func("add_rms_norm3_f32");
9150        let cfg = LaunchConfig {
9151            grid_dim: (nrows as u32, 1, 1),
9152            block_dim: (rms_block(), 1, 1),
9153            shared_mem_bytes: 0,
9154        };
9155        let (nc, e2) = (ncols as i32, eps);
9156        let __s_b = self.gpu.stream();
9157        let mut b = __s_b.launch_builder(&f);
9158        b.arg(a)
9159            .arg(b_in)
9160            .arg(w0)
9161            .arg(w1)
9162            .arg(w2)
9163            .arg(res)
9164            .arg(d0)
9165            .arg(d1)
9166            .arg(d2)
9167            .arg(&nc)
9168            .arg(&e2);
9169        unsafe {
9170            b.launch(cfg)?;
9171        }
9172        Ok(())
9173    }
9174
9175    /// dst = (a + b) * c (residual add + layer scale, one launch).
9176    pub fn add_scale(
9177        &self,
9178        a: &CudaSlice<f32>,
9179        b_in: &CudaSlice<f32>,
9180        c: f32,
9181        dst: &mut CudaSlice<f32>,
9182        n: usize,
9183    ) -> Result<(), Box<dyn std::error::Error>> {
9184        let f = self.func("add_scale_f32");
9185        let cfg = LaunchConfig::for_num_elems(n as u32);
9186        let ni = n as i32;
9187        let __s_b = self.gpu.stream();
9188        let mut b = __s_b.launch_builder(&f);
9189        b.arg(a).arg(b_in).arg(&c).arg(dst).arg(&ni);
9190        unsafe {
9191            b.launch(cfg)?;
9192        }
9193        Ok(())
9194    }
9195
9196    /// Vision-tower LayerNorm (with bias) over [nrows, ncols] — lane/vision.
9197    pub fn layer_norm_bias(
9198        &self,
9199        x: &CudaSlice<f32>,
9200        w: &CudaSlice<f32>,
9201        b: &CudaSlice<f32>,
9202        dst: &mut CudaSlice<f32>,
9203        ncols: usize,
9204        nrows: usize,
9205        eps: f32,
9206    ) -> Result<(), Box<dyn std::error::Error>> {
9207        let f = self.func("layer_norm_bias_f32");
9208        let (nc, e) = (ncols as i32, eps);
9209        let cfg = LaunchConfig {
9210            grid_dim: (nrows as u32, 1, 1),
9211            block_dim: (256, 1, 1),
9212            shared_mem_bytes: 0,
9213        };
9214        let __s_b = self.gpu.stream();
9215        let mut lb = __s_b.launch_builder(&f);
9216        lb.arg(x).arg(w).arg(b).arg(&mut *dst).arg(&nc).arg(&e);
9217        unsafe {
9218            lb.launch(cfg)?;
9219        }
9220        Ok(())
9221    }
9222
9223    /// gelu_pytorch_tanh elementwise (vision tower MLP activation).
9224    pub fn gelu_tanh(
9225        &self,
9226        x: &CudaSlice<f32>,
9227        dst: &mut CudaSlice<f32>,
9228        n: usize,
9229    ) -> Result<(), Box<dyn std::error::Error>> {
9230        let f = self.func("gelu_tanh_f32");
9231        let ni = n as i64;
9232        let cfg = LaunchConfig {
9233            grid_dim: (n.div_ceil(256) as u32, 1, 1),
9234            block_dim: (256, 1, 1),
9235            shared_mem_bytes: 0,
9236        };
9237        let __s_b = self.gpu.stream();
9238        let mut lb = __s_b.launch_builder(&f);
9239        lb.arg(x).arg(&mut *dst).arg(&ni);
9240        unsafe {
9241            lb.launch(cfg)?;
9242        }
9243        Ok(())
9244    }
9245
9246    /// In-place row softmax over [nrows, ncols] (bidirectional vision attention).
9247    pub fn row_softmax(
9248        &self,
9249        x: &mut CudaSlice<f32>,
9250        ncols: usize,
9251        nrows: usize,
9252    ) -> Result<(), Box<dyn std::error::Error>> {
9253        let f = self.func("row_softmax_f32");
9254        let nc = ncols as i32;
9255        let cfg = LaunchConfig {
9256            grid_dim: (nrows as u32, 1, 1),
9257            block_dim: (256, 1, 1),
9258            shared_mem_bytes: 0,
9259        };
9260        let __s_b = self.gpu.stream();
9261        let mut lb = __s_b.launch_builder(&f);
9262        lb.arg(&mut *x).arg(&nc);
9263        unsafe {
9264            lb.launch(cfg)?;
9265        }
9266        Ok(())
9267    }
9268
9269    pub fn rms_norm(
9270        &self,
9271        x: &CudaSlice<f32>,
9272        w: &CudaSlice<f32>,
9273        dst: &mut CudaSlice<f32>,
9274        ncols: usize,
9275        nrows: usize,
9276        eps: f32,
9277    ) -> Result<(), Box<dyn std::error::Error>> {
9278        let (nc, e) = (ncols as i32, eps);
9279        let kname = if Self::norm_ilp_on() {
9280            "rms_norm_f32_v2"
9281        } else {
9282            "rms_norm_f32"
9283        };
9284        if Self::pdl_on() && Self::pdl_wb_on() {
9285            use cudarc::driver::{DevicePtr, DevicePtrMut};
9286            let s = &self.gpu.stream();
9287            let (px, _g0) = x.device_ptr(s);
9288            let (pw, _g1) = w.device_ptr(s);
9289            let (pd, _g2) = dst.device_ptr_mut(s);
9290            let mut ps = [
9291                &px as *const _ as *mut std::ffi::c_void,
9292                &pw as *const _ as *mut _,
9293                &pd as *const _ as *mut _,
9294                &nc as *const _ as *mut _,
9295                &e as *const _ as *mut _,
9296            ];
9297            unsafe {
9298                self.launch_pdl(kname, (nrows as u32, 1, 1), (rms_block(), 1, 1), &mut ps)?;
9299            }
9300            return Ok(());
9301        }
9302        let f = self.func(kname);
9303        let cfg = LaunchConfig {
9304            grid_dim: (nrows as u32, 1, 1),
9305            block_dim: (rms_block(), 1, 1),
9306            shared_mem_bytes: 0,
9307        };
9308        let __s_b = self.gpu.stream();
9309        let mut b = __s_b.launch_builder(&f);
9310        b.arg(x).arg(w).arg(dst).arg(&nc).arg(&e);
9311        unsafe {
9312            b.launch(cfg)?;
9313        }
9314        Ok(())
9315    }
9316
9317    /// RMS-norm with blockDim=1024 — BIT-IDENTICAL to the fused `rms_norm_q8_1` and
9318    /// `add_rms_norm_q8_1` kernels' sum-of-squares reduction. The spec verify path MUST use this
9319    /// to match decode's FP accumulation order: the standard `rms_norm` at blockDim=256 has a
9320    /// different per-thread stride (ncols/256 partials vs ncols/1024 partials) and therefore a
9321    /// different shfl-tree reduction that can shift `scale = rsqrt(sum/n + eps)` by ULPs, causing
9322    /// divergence through the GDN scan and argmax flips on the 9B text prompt. The underlying
9323    /// `rms_norm_f32` kernel supports any blockDim (generic reduce with shared[32]).
9324    pub fn rms_norm_decode(
9325        &self,
9326        x: &CudaSlice<f32>,
9327        w: &CudaSlice<f32>,
9328        dst: &mut CudaSlice<f32>,
9329        ncols: usize,
9330        nrows: usize,
9331        eps: f32,
9332    ) -> Result<(), Box<dyn std::error::Error>> {
9333        let f = self.func(if Self::norm_ilp_on() {
9334            "rms_norm_f32_v2"
9335        } else {
9336            "rms_norm_f32"
9337        });
9338        let cfg = LaunchConfig {
9339            grid_dim: (nrows as u32, 1, 1),
9340            block_dim: (1024, 1, 1),
9341            shared_mem_bytes: 0,
9342        };
9343        let (nc, e) = (ncols as i32, eps);
9344        let __s_b = self.gpu.stream();
9345        let mut b = __s_b.launch_builder(&f);
9346        b.arg(x).arg(w).arg(dst).arg(&nc).arg(&e);
9347        unsafe {
9348            b.launch(cfg)?;
9349        }
9350        Ok(())
9351    }
9352
9353    /// DECODE GLUE-FUSION LEVER: `z = rms_norm(x)*w` emitted DIRECTLY as q8_1 (no f32 `z` materialized,
9354    /// no standalone quantize_q8_1 launch). Returns (out_q [nrows*ncols i8], out_d [nrows*nblk f32])
9355    /// ready to feed matmul_pre. BIT-IDENTICAL to rms_norm + quantize_q8_1. ncols % 32 == 0.
9356    pub fn rms_norm_q8_1(
9357        &self,
9358        x: &CudaSlice<f32>,
9359        w: &CudaSlice<f32>,
9360        ncols: usize,
9361        nrows: usize,
9362        eps: f32,
9363    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9364        let nblk = ncols / 32;
9365        let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
9366        let mut d = self.alloc_uninit::<f32>(nrows * nblk)?;
9367        let (nc, e) = (ncols as i32, eps);
9368        if Self::pdl_on() {
9369            {
9370                use cudarc::driver::{DevicePtr, DevicePtrMut};
9371                let s = &self.gpu.stream();
9372                let (px, _g0) = x.device_ptr(s);
9373                let (pw, _g1) = w.device_ptr(s);
9374                let (pq, _g2) = q.device_ptr_mut(s);
9375                let (pd, _g3) = d.device_ptr_mut(s);
9376                let mut ps = [
9377                    &px as *const _ as *mut std::ffi::c_void,
9378                    &pw as *const _ as *mut _,
9379                    &pq as *const _ as *mut _,
9380                    &pd as *const _ as *mut _,
9381                    &nc as *const _ as *mut _,
9382                    &e as *const _ as *mut _,
9383                ];
9384                unsafe {
9385                    self.launch_pdl("rms_norm_q8_1", (nrows as u32, 1, 1), (1024, 1, 1), &mut ps)?;
9386                }
9387            }
9388            return Ok((q, d));
9389        }
9390        let f = self.func("rms_norm_q8_1");
9391        // 1024 threads: decode is nrows=1 -> ONE CTA; 32 warps hide the pass1->pass2 latency
9392        // (s[32] reduce already sized for 32 warps). Same shape math at any blockDim.
9393        let cfg = LaunchConfig {
9394            grid_dim: (nrows as u32, 1, 1),
9395            block_dim: (1024, 1, 1),
9396            shared_mem_bytes: 0,
9397        };
9398        let __s_b = self.gpu.stream();
9399        let mut b = __s_b.launch_builder(&f);
9400        b.arg(x).arg(w).arg(&mut q).arg(&mut d).arg(&nc).arg(&e);
9401        unsafe {
9402            b.launch(cfg)?;
9403        }
9404        Ok((q, d))
9405    }
9406
9407    /// Slot-fed rms_norm_q8_1 twin (alloc-free capture lane): identical launch (incl. the
9408    /// PDL arm), caller-owned outputs.
9409    pub fn rms_norm_q8_1_into(
9410        &self,
9411        x: &CudaSlice<f32>,
9412        w: &CudaSlice<f32>,
9413        ncols: usize,
9414        nrows: usize,
9415        eps: f32,
9416        q: &mut CudaSlice<i8>,
9417        d: &mut CudaSlice<f32>,
9418    ) -> Result<(), Box<dyn std::error::Error>> {
9419        let nblk = ncols / 32;
9420        debug_assert!(q.len() >= nrows * ncols && d.len() >= nrows * nblk);
9421        let (nc, e) = (ncols as i32, eps);
9422        if Self::pdl_on() {
9423            use cudarc::driver::{DevicePtr, DevicePtrMut};
9424            let s = &self.gpu.stream();
9425            let (px, _g0) = x.device_ptr(s);
9426            let (pw, _g1) = w.device_ptr(s);
9427            let (pq, _g2) = q.device_ptr_mut(s);
9428            let (pd, _g3) = d.device_ptr_mut(s);
9429            let mut ps = [
9430                &px as *const _ as *mut std::ffi::c_void,
9431                &pw as *const _ as *mut _,
9432                &pq as *const _ as *mut _,
9433                &pd as *const _ as *mut _,
9434                &nc as *const _ as *mut _,
9435                &e as *const _ as *mut _,
9436            ];
9437            unsafe {
9438                self.launch_pdl("rms_norm_q8_1", (nrows as u32, 1, 1), (1024, 1, 1), &mut ps)?;
9439            }
9440            return Ok(());
9441        }
9442        let f = self.func("rms_norm_q8_1");
9443        let cfg = LaunchConfig {
9444            grid_dim: (nrows as u32, 1, 1),
9445            block_dim: (1024, 1, 1),
9446            shared_mem_bytes: 0,
9447        };
9448        let __s_b = self.gpu.stream();
9449        let mut b = __s_b.launch_builder(&f);
9450        b.arg(x).arg(w).arg(&mut *q).arg(&mut *d).arg(&nc).arg(&e);
9451        unsafe {
9452            b.launch(cfg)?;
9453        }
9454        Ok(())
9455    }
9456
9457    /// Slot-fed quantize_q8_1 twin (alloc-free capture lane).
9458    pub fn quantize_q8_1_into(
9459        &self,
9460        x: &CudaSlice<f32>,
9461        m: usize,
9462        in_f: usize,
9463        q: &mut CudaSlice<i8>,
9464        d: &mut CudaSlice<f32>,
9465    ) -> Result<(), Box<dyn std::error::Error>> {
9466        let nblk = in_f / 32;
9467        debug_assert!(q.len() >= m * in_f && d.len() >= m * nblk);
9468        let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
9469        let (inf, mi) = (in_f as i32, m as i32);
9470        if Self::pdl_on() && Self::pdl_wb_on() {
9471            use cudarc::driver::{DevicePtr, DevicePtrMut};
9472            let s = &self.gpu.stream();
9473            let (px, _g0) = x.device_ptr(s);
9474            let (pq, _g1) = q.device_ptr_mut(s);
9475            let (pd, _g2) = d.device_ptr_mut(s);
9476            let mut ps = [
9477                &px as *const _ as *mut std::ffi::c_void,
9478                &pq as *const _ as *mut _,
9479                &pd as *const _ as *mut _,
9480                &inf as *const _ as *mut _,
9481                &mi as *const _ as *mut _,
9482            ];
9483            unsafe {
9484                self.launch_pdl("quantize_q8_1", cfg.grid_dim, cfg.block_dim, &mut ps)?;
9485            }
9486            return Ok(());
9487        }
9488        let f = self.func("quantize_q8_1");
9489        let __s_b = self.gpu.stream();
9490        let mut b = __s_b.launch_builder(&f);
9491        b.arg(x).arg(&mut *q).arg(&mut *d).arg(&inf).arg(&mi);
9492        unsafe {
9493            b.launch(cfg)?;
9494        }
9495        Ok(())
9496    }
9497
9498    /// DECODE GLUE-FUSION LEVER: `res = a+b; z = rms_norm(res)*w` with z emitted as q8_1. `res` is
9499    /// still written (the post-ffn residual add reads it). Fuses add_rms_norm + quantize_q8_1.
9500    /// Returns (out_q, out_d) for matmul_pre. BIT-IDENTICAL. ncols % 32 == 0.
9501    pub fn add_rms_norm_q8_1(
9502        &self,
9503        a: &CudaSlice<f32>,
9504        b_in: &CudaSlice<f32>,
9505        w: &CudaSlice<f32>,
9506        res: &mut CudaSlice<f32>,
9507        ncols: usize,
9508        nrows: usize,
9509        eps: f32,
9510    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9511        let nblk = ncols / 32;
9512        let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
9513        let mut d = self.alloc_uninit::<f32>(nrows * nblk)?;
9514        let f = self.func("add_rms_norm_q8_1");
9515        // 1024 threads: same single-CTA-at-decode reasoning as rms_norm_q8_1.
9516        let cfg = LaunchConfig {
9517            grid_dim: (nrows as u32, 1, 1),
9518            block_dim: (1024, 1, 1),
9519            shared_mem_bytes: 0,
9520        };
9521        let (nc, e) = (ncols as i32, eps);
9522        let __s_bld = self.gpu.stream();
9523        let mut bld = __s_bld.launch_builder(&f);
9524        bld.arg(a)
9525            .arg(b_in)
9526            .arg(w)
9527            .arg(res)
9528            .arg(&mut q)
9529            .arg(&mut d)
9530            .arg(&nc)
9531            .arg(&e);
9532        unsafe {
9533            bld.launch(cfg)?;
9534        }
9535        Ok((q, d))
9536    }
9537
9538    /// RANK3 LEVER (add+rmsnorm fuse): `res = a + b; dst = rms_norm(res) * w` in ONE launch. Fuses
9539    /// e.add(a,b,res) + e.rms_norm(res,w,dst), removing one launch + one HBM read of the residual per
9540    /// residual+norm pair. BIT-IDENTICAL to the two-kernel sequence (same IEEE add, same reduction).
9541    /// O-PROJ TAIL FUSION M2: mixed = a0+a1 composed in-register, then the VERBATIM
9542    /// add_rms_norm program. Raw UVA pointers for the join partials (persistent ws rows).
9543    #[allow(clippy::too_many_arguments)]
9544    pub fn join_add_rms_norm_raw(
9545        &self,
9546        a0_raw: u64,
9547        a1_raw: u64,
9548        x: &CudaSlice<f32>,
9549        w: &CudaSlice<f32>,
9550        res: &mut CudaSlice<f32>,
9551        dst: &mut CudaSlice<f32>,
9552        ncols: usize,
9553        eps: f32,
9554    ) -> Result<(), Box<dyn std::error::Error>> {
9555        if a0_raw == 0 || a1_raw == 0 || x.len() < ncols || res.len() < ncols || dst.len() < ncols {
9556            return Err("join_add_rms_norm geometry".into());
9557        }
9558        let f = self.func("join_add_rms_norm_f32");
9559        let cfg = LaunchConfig {
9560            grid_dim: (1, 1, 1),
9561            block_dim: (rms_block(), 1, 1),
9562            shared_mem_bytes: 0,
9563        };
9564        let (nc, e) = (ncols as i32, eps);
9565        let __s_b = self.gpu.stream();
9566        let mut b = __s_b.launch_builder(&f);
9567        b.arg(&a0_raw)
9568            .arg(&a1_raw)
9569            .arg(x)
9570            .arg(w)
9571            .arg(&mut *res)
9572            .arg(&mut *dst)
9573            .arg(&nc)
9574            .arg(&e);
9575        unsafe {
9576            b.launch(cfg)?;
9577        }
9578        Ok(())
9579    }
9580
9581    pub fn add_rms_norm(
9582        &self,
9583        a: &CudaSlice<f32>,
9584        b: &CudaSlice<f32>,
9585        w: &CudaSlice<f32>,
9586        res: &mut CudaSlice<f32>,
9587        dst: &mut CudaSlice<f32>,
9588        ncols: usize,
9589        nrows: usize,
9590        eps: f32,
9591    ) -> Result<(), Box<dyn std::error::Error>> {
9592        let (nc, e) = (ncols as i32, eps);
9593        let kname = if Self::norm_ilp_on() {
9594            "add_rms_norm_f32_v2"
9595        } else {
9596            "add_rms_norm_f32"
9597        };
9598        if Self::pdl_on() && Self::pdl_wb_on() {
9599            use cudarc::driver::{DevicePtr, DevicePtrMut};
9600            let s = &self.gpu.stream();
9601            let (pa, _g0) = a.device_ptr(s);
9602            let (pb, _g1) = b.device_ptr(s);
9603            let (pw, _g2) = w.device_ptr(s);
9604            let (pr, _g3) = res.device_ptr_mut(s);
9605            let (pd, _g4) = dst.device_ptr_mut(s);
9606            let mut ps = [
9607                &pa as *const _ as *mut std::ffi::c_void,
9608                &pb as *const _ as *mut _,
9609                &pw as *const _ as *mut _,
9610                &pr as *const _ as *mut _,
9611                &pd as *const _ as *mut _,
9612                &nc as *const _ as *mut _,
9613                &e as *const _ as *mut _,
9614            ];
9615            unsafe {
9616                self.launch_pdl(kname, (nrows as u32, 1, 1), (rms_block(), 1, 1), &mut ps)?;
9617            }
9618            return Ok(());
9619        }
9620        let f = self.func(kname);
9621        let cfg = LaunchConfig {
9622            grid_dim: (nrows as u32, 1, 1),
9623            block_dim: (rms_block(), 1, 1),
9624            shared_mem_bytes: 0,
9625        };
9626        let __s_b2 = self.gpu.stream();
9627        let mut b2 = __s_b2.launch_builder(&f);
9628        b2.arg(a)
9629            .arg(b)
9630            .arg(w)
9631            .arg(&mut *res)
9632            .arg(&mut *dst)
9633            .arg(&nc)
9634            .arg(&e);
9635        unsafe {
9636            b2.launch(cfg)?;
9637        }
9638        Ok(())
9639    }
9640
9641    /// E4B glue fusion: rms(a, wa) prologue + add_rms_norm — folds the post-attn norm into
9642    /// the tail entry (res = rms(a)*wa + b; dst = rms(res)*w).
9643    #[allow(clippy::too_many_arguments)]
9644    pub fn rms_pre_add_rms_norm(
9645        &self,
9646        a: &CudaSlice<f32>,
9647        wa: &CudaSlice<f32>,
9648        b: &CudaSlice<f32>,
9649        w: &CudaSlice<f32>,
9650        res: &mut CudaSlice<f32>,
9651        dst: &mut CudaSlice<f32>,
9652        ncols: usize,
9653        nrows: usize,
9654        eps: f32,
9655    ) -> Result<(), Box<dyn std::error::Error>> {
9656        let f = self.func("rms_pre_add_rms_norm_f32");
9657        let cfg = LaunchConfig {
9658            grid_dim: (nrows as u32, 1, 1),
9659            block_dim: (rms_block(), 1, 1),
9660            shared_mem_bytes: 0,
9661        };
9662        let (nc, e) = (ncols as i32, eps);
9663        let __s_b2 = self.gpu.stream();
9664        let mut b2 = __s_b2.launch_builder(&f);
9665        b2.arg(a)
9666            .arg(wa)
9667            .arg(b)
9668            .arg(w)
9669            .arg(&mut *res)
9670            .arg(&mut *dst)
9671            .arg(&nc)
9672            .arg(&e);
9673        unsafe {
9674            b2.launch(cfg)?;
9675        }
9676        Ok(())
9677    }
9678
9679    /// wave-2 fold: rms(a,wa) + add + ffn-norm with zsh EMITTED q8_1 (fused2 consumes it).
9680    #[allow(clippy::too_many_arguments)]
9681    pub fn rms_pre_add_rms_norm_q8z(
9682        &self,
9683        a: &CudaSlice<f32>,
9684        wa: &CudaSlice<f32>,
9685        b: &CudaSlice<f32>,
9686        w: &CudaSlice<f32>,
9687        res: &mut CudaSlice<f32>,
9688        dst: &mut CudaSlice<f32>,
9689        ncols: usize,
9690        nrows: usize,
9691        eps: f32,
9692    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9693        debug_assert!(ncols % 128 == 0);
9694        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
9695        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9696        let (nc, e) = (ncols as i32, eps);
9697        if Self::pdl_on() {
9698            {
9699                use cudarc::driver::{DevicePtr, DevicePtrMut};
9700                let s = &self.gpu.stream();
9701                let (pa, _g0) = a.device_ptr(s);
9702                let (pwa, _g1) = wa.device_ptr(s);
9703                let (pb, _g2) = b.device_ptr(s);
9704                let (pw, _g3) = w.device_ptr(s);
9705                let (pr, _g4) = res.device_ptr_mut(s);
9706                let (pdst, _g5) = dst.device_ptr_mut(s);
9707                let (pq, _g6) = out_q.device_ptr_mut(s);
9708                let (pd, _g7) = out_d.device_ptr_mut(s);
9709                let mut ps = [
9710                    &pa as *const _ as *mut std::ffi::c_void,
9711                    &pwa as *const _ as *mut _,
9712                    &pb as *const _ as *mut _,
9713                    &pw as *const _ as *mut _,
9714                    &pr as *const _ as *mut _,
9715                    &pdst as *const _ as *mut _,
9716                    &pq as *const _ as *mut _,
9717                    &pd as *const _ as *mut _,
9718                    &nc as *const _ as *mut _,
9719                    &e as *const _ as *mut _,
9720                ];
9721                unsafe {
9722                    self.launch_pdl(
9723                        "rms_pre_add_rms_norm_q8z_f32",
9724                        (nrows as u32, 1, 1),
9725                        (rms_block(), 1, 1),
9726                        &mut ps,
9727                    )?;
9728                }
9729            }
9730            return Ok((out_q, out_d));
9731        }
9732        let f = self.func("rms_pre_add_rms_norm_q8z_f32");
9733        let cfg = LaunchConfig {
9734            grid_dim: (nrows as u32, 1, 1),
9735            block_dim: (rms_block(), 1, 1),
9736            shared_mem_bytes: 0,
9737        };
9738        let __s_b2 = self.gpu.stream();
9739        let mut b2 = __s_b2.launch_builder(&f);
9740        b2.arg(a)
9741            .arg(wa)
9742            .arg(b)
9743            .arg(w)
9744            .arg(&mut *res)
9745            .arg(&mut *dst)
9746            .arg(&mut out_q)
9747            .arg(&mut out_d)
9748            .arg(&nc)
9749            .arg(&e);
9750        unsafe {
9751            b2.launch(cfg)?;
9752        }
9753        Ok((out_q, out_d))
9754    }
9755
9756    /// Slot-fed twin of `rms_pre_add_rms_norm_q8z` (gemma4 pn-fold, slotted/graph arm):
9757    /// identical kernel, caller-owned outputs, PLAIN launch only — the dc_slotted capture
9758    /// body must stay attribute-free (the fused2_into precedent).
9759    #[allow(clippy::too_many_arguments)]
9760    pub fn rms_pre_add_rms_norm_q8z_into(
9761        &self,
9762        a: &CudaSlice<f32>,
9763        wa: &CudaSlice<f32>,
9764        b: &CudaSlice<f32>,
9765        w: &CudaSlice<f32>,
9766        res: &mut CudaSlice<f32>,
9767        dst: &mut CudaSlice<f32>,
9768        ncols: usize,
9769        nrows: usize,
9770        eps: f32,
9771        out_q: &mut CudaSlice<i8>,
9772        out_d: &mut CudaSlice<f32>,
9773    ) -> Result<(), Box<dyn std::error::Error>> {
9774        debug_assert!(ncols % 128 == 0);
9775        let (nc, e) = (ncols as i32, eps);
9776        let f = self.func("rms_pre_add_rms_norm_q8z_f32");
9777        let cfg = LaunchConfig {
9778            grid_dim: (nrows as u32, 1, 1),
9779            block_dim: (rms_block(), 1, 1),
9780            shared_mem_bytes: 0,
9781        };
9782        let __s_b = self.gpu.stream();
9783        let mut b2 = __s_b.launch_builder(&f);
9784        b2.arg(a)
9785            .arg(wa)
9786            .arg(b)
9787            .arg(w)
9788            .arg(&mut *res)
9789            .arg(&mut *dst)
9790            .arg(&mut *out_q)
9791            .arg(&mut *out_d)
9792            .arg(&nc)
9793            .arg(&e);
9794        unsafe {
9795            b2.launch(cfg)?;
9796        }
9797        Ok(())
9798    }
9799
9800    /// Slot-fed twin of `rms_pre_add_scale_rms_norm_q8_1` (gemma4 pn-fold exit, slotted
9801    /// arm): identical kernel, caller-owned outputs, PLAIN launch (capture-safe).
9802    #[allow(clippy::too_many_arguments)]
9803    pub fn rms_pre_add_scale_rms_norm_q8_1_into(
9804        &self,
9805        a: &CudaSlice<f32>,
9806        wa: &CudaSlice<f32>,
9807        b_in: &CudaSlice<f32>,
9808        c: f32,
9809        w: &CudaSlice<f32>,
9810        res: &mut CudaSlice<f32>,
9811        ncols: usize,
9812        nrows: usize,
9813        eps: f32,
9814        out_q: &mut CudaSlice<i8>,
9815        out_d: &mut CudaSlice<f32>,
9816    ) -> Result<(), Box<dyn std::error::Error>> {
9817        debug_assert!(ncols % 128 == 0);
9818        let (nc, e2) = (ncols as i32, eps);
9819        let f = self.func("rms_pre_add_scale_rms_norm_q8_1");
9820        let cfg = LaunchConfig {
9821            grid_dim: (nrows as u32, 1, 1),
9822            block_dim: (rms_block(), 1, 1),
9823            shared_mem_bytes: 0,
9824        };
9825        let __s_b = self.gpu.stream();
9826        let mut b2 = __s_b.launch_builder(&f);
9827        b2.arg(a)
9828            .arg(wa)
9829            .arg(b_in)
9830            .arg(&c)
9831            .arg(w)
9832            .arg(&mut *res)
9833            .arg(&mut *out_q)
9834            .arg(&mut *out_d)
9835            .arg(&nc)
9836            .arg(&e2);
9837        unsafe {
9838            b2.launch(cfg)?;
9839        }
9840        Ok(())
9841    }
9842
9843    /// gemma4 pn-fold seam (GAP-DIAGNOSIS verdict 7, the E4B glue backport): the dense
9844    /// decode/verify/slotted trio folds post_attn_norm into the tail entry
9845    /// (rms_pre_add_rms_norm[_q8z]) and post_ffw_norm into the residual exit
9846    /// (rms_pre_add_scale_rms_norm_q8_1). BITS-CHANGING vs the two-launch chain (the
9847    /// single-phase reduction's expansion rounding — E4B receipts); every arm moves
9848    /// together so decode == verify == graph parity holds BY CONSTRUCTION within either
9849    /// seam value. MEMRA_G4_PNFOLD=0 restores the unfused chain everywhere.
9850    pub fn g4_pnfold_on() -> bool {
9851        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9852        *ON.get_or_init(|| {
9853            std::env::var("MEMRA_G4_PNFOLD")
9854                .map(|v| v != "0")
9855                .unwrap_or(true)
9856        })
9857    }
9858
9859    /// wave-4b: OUT-dim concat of three Q4_0 tensors (same in_features; rows are independent
9860    /// blocks, so the concat is a D2D byte concat of the GGUF-layout planes). Returns None
9861    /// off-class (non-Q4_0, mismatched widths, or any tensor already rp-swapped in place).
9862    pub fn build_q4_out_concat3(
9863        &self,
9864        w0: &crate::model::GpuTensor,
9865        w1: &crate::model::GpuTensor,
9866        w2: &crate::model::GpuTensor,
9867    ) -> Result<Option<crate::model::GpuTensor>, Box<dyn std::error::Error>> {
9868        use crate::model::GpuTensor;
9869        let part = |w: &GpuTensor| -> Option<(usize, usize)> {
9870            match w {
9871                GpuTensor::Quant {
9872                    qtype,
9873                    row_bytes,
9874                    rp,
9875                    ..
9876                } if *qtype == QT_Q4_0 && !*rp => Some((*row_bytes, w.out_features())),
9877                _ => None,
9878            }
9879        };
9880        let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (part(w0), part(w1), part(w2))
9881        else {
9882            return Ok(None);
9883        };
9884        if rb0 != rb1
9885            || rb0 != rb2
9886            || w0.in_features() != w1.in_features()
9887            || w0.in_features() != w2.in_features()
9888        {
9889            return Ok(None);
9890        }
9891        fn bytes_of(w: &crate::model::GpuTensor) -> &CudaSlice<u8> {
9892            match w {
9893                crate::model::GpuTensor::Quant { bytes, .. } => bytes,
9894                _ => unreachable!(),
9895            }
9896        }
9897        let (b0, b1, b2) = (bytes_of(w0), bytes_of(w1), bytes_of(w2));
9898        let total = rb0 * (o0 + o1 + o2);
9899        let mut cat = self.alloc_u8(total)?;
9900        self.copy_u8_into(&mut cat, 0, b0, rb0 * o0)?;
9901        self.copy_u8_into(&mut cat, rb0 * o0, b1, rb1 * o1)?;
9902        self.copy_u8_into(&mut cat, rb0 * (o0 + o1), b2, rb2 * o2)?;
9903        Ok(Some(GpuTensor::Quant {
9904            bytes: cat,
9905            qtype: QT_Q4_0,
9906            row_bytes: rb0,
9907            ne: vec![w0.in_features() as u64, (o0 + o1 + o2) as u64],
9908            scale: 1.0,
9909            rp: false,
9910            #[cfg(memra_cutlass)]
9911            cutlass: None,
9912            fp8: None,
9913            blk: None,
9914            rp4: None,
9915            f16: None,
9916        }))
9917    }
9918
9919    /// FULL-WIDTH-ROPE CONTRACT for the fused rms_norm+qkv+rope kernels
9920    /// (lane/graph-s-key-exactness-20260819, probe O-10 follow-up).
9921    ///
9922    /// `rms_norm_qkv_rope_f32` / `_cat_f32` (`cu/kernels.cu`) and
9923    /// `rms_norm_qkv_rope_append_body` (`cu/flash_attn.cu`) take NO `n_dims`/`n_rot` argument.
9924    /// They compute `int half = ncols / 2` and rotate the FULL head width by construction — the
9925    /// standalone `rope_neox*` kernels take `n_dims` and early-return above it, these do not.
9926    ///
9927    /// Every call site today is a gemma-4 arm, where `n_rot == head_dim`, so nothing is wrong
9928    /// now. But this is the n_rot bug ONE FUSION away: fuse a partial-rotary arch onto these
9929    /// kernels (qwen3.5 = 64 rotary dims of a 256-wide head; step35 full-attn = 64 of 128) and
9930    /// 192 dims that must pass through unrotated get rotated silently — no error, no NaN, just a
9931    /// wrong model. The n_rot lane already paid for that class once, in the config derivation.
9932    ///
9933    /// So the fusions now take the layer's DERIVED rope width and refuse anything but full
9934    /// width. A future partial-rotary caller fails at its first launch with the geometry named
9935    /// instead of serving quietly wrong logits.
9936    fn full_width_rope_only(
9937        kernel: &str,
9938        n_rot: usize,
9939        head_dim: usize,
9940    ) -> Result<(), Box<dyn std::error::Error>> {
9941        if n_rot == head_dim {
9942            return Ok(());
9943        }
9944        Err(format!(
9945            "{kernel}: PARTIAL ROTARY REFUSED — n_rot {n_rot} != head_dim {head_dim}. This fused \
9946             rms_norm+qkv+rope kernel carries no n_dims parameter and rotates the full head \
9947             width (half = ncols/2), so it would rotate dims {n_rot}..{head_dim} that must pass \
9948             through unrotated. Use the split path (rms_norm_qkv + rope_neox/rope_neox2 with \
9949             n_dims={n_rot}), or add an n_dims early-return to the kernel and widen this guard."
9950        )
9951        .into())
9952    }
9953
9954    /// wave-4b: the qkv-cat twin — one contiguous [rq+2*rk, hd] input from the concat matvec.
9955    /// `n_rot` is the layer's derived rotary width and MUST equal `head_dim`
9956    /// ([`Engine::full_width_rope_only`]).
9957    #[allow(clippy::too_many_arguments)]
9958    pub fn rms_norm_qkv_rope_cat(
9959        &self,
9960        qkv: &CudaSlice<f32>,
9961        wq: &CudaSlice<f32>,
9962        wk: &CudaSlice<f32>,
9963        wv: &CudaSlice<f32>,
9964        q: &mut CudaSlice<f32>,
9965        k: &mut CudaSlice<f32>,
9966        v: &mut CudaSlice<f32>,
9967        head_dim: usize,
9968        n_rot: usize,
9969        rq: usize,
9970        rk: usize,
9971        pos: &CudaSlice<i32>,
9972        nh_q: usize,
9973        nh_k: usize,
9974        base: f32,
9975        freq_scale: f32,
9976        ff: Option<&CudaSlice<f32>>,
9977        eps: f32,
9978    ) -> Result<(), Box<dyn std::error::Error>> {
9979        Self::full_width_rope_only("rms_norm_qkv_rope_cat", n_rot, head_dim)?;
9980        let rows = rq + rk + rk;
9981        let theta_scale = base.powf(-2.0 / head_dim as f32);
9982        let (nc, rqi, rki, nhq, nhk) = (
9983            head_dim as i32,
9984            rq as i32,
9985            rk as i32,
9986            nh_q as i32,
9987            nh_k as i32,
9988        );
9989        if Self::pdl_on() {
9990            use cudarc::driver::{DevicePtr, DevicePtrMut};
9991            let s = &self.gpu.stream();
9992            let (pqkv, _g0) = qkv.device_ptr(s);
9993            let (pwq, _g1) = wq.device_ptr(s);
9994            let (pwk, _g2) = wk.device_ptr(s);
9995            let (pwv, _g3) = wv.device_ptr(s);
9996            let (pq, _g4) = q.device_ptr_mut(s);
9997            let (pk, _g5) = k.device_ptr_mut(s);
9998            let (pv, _g6) = v.device_ptr_mut(s);
9999            let (ppos, _g7) = pos.device_ptr(s);
10000            let (pff, _g8) = match ff {
10001                Some(t) => {
10002                    let (p, g) = t.device_ptr(s);
10003                    (p, Some(g))
10004                }
10005                None => (0, None),
10006            };
10007            let mut ps = [
10008                &pqkv as *const _ as *mut std::ffi::c_void,
10009                &pwq as *const _ as *mut _,
10010                &pwk as *const _ as *mut _,
10011                &pwv as *const _ as *mut _,
10012                &pq as *const _ as *mut _,
10013                &pk as *const _ as *mut _,
10014                &pv as *const _ as *mut _,
10015                &nc as *const _ as *mut _,
10016                &rqi as *const _ as *mut _,
10017                &rki as *const _ as *mut _,
10018                &ppos as *const _ as *mut _,
10019                &nhq as *const _ as *mut _,
10020                &nhk as *const _ as *mut _,
10021                &theta_scale as *const _ as *mut _,
10022                &freq_scale as *const _ as *mut _,
10023                &pff as *const _ as *mut _,
10024                &eps as *const _ as *mut _,
10025            ];
10026            unsafe {
10027                self.launch_pdl(
10028                    "rms_norm_qkv_rope_cat_f32",
10029                    (rows as u32, 1, 1),
10030                    (rms_block(), 1, 1),
10031                    &mut ps,
10032                )?;
10033            }
10034            return Ok(());
10035        }
10036        let f = self.func("rms_norm_qkv_rope_cat_f32");
10037        let cfg = LaunchConfig {
10038            grid_dim: (rows as u32, 1, 1),
10039            block_dim: (rms_block(), 1, 1),
10040            shared_mem_bytes: 0,
10041        };
10042        let __s_b = self.gpu.stream();
10043        let mut b = __s_b.launch_builder(&f);
10044        match ff {
10045            Some(t) => {
10046                b.arg(qkv)
10047                    .arg(wq)
10048                    .arg(wk)
10049                    .arg(wv)
10050                    .arg(&mut *q)
10051                    .arg(&mut *k)
10052                    .arg(&mut *v)
10053                    .arg(&nc)
10054                    .arg(&rqi)
10055                    .arg(&rki)
10056                    .arg(pos)
10057                    .arg(&nhq)
10058                    .arg(&nhk)
10059                    .arg(&theta_scale)
10060                    .arg(&freq_scale)
10061                    .arg(t)
10062                    .arg(&eps);
10063                unsafe {
10064                    b.launch(cfg)?;
10065                }
10066            }
10067            None => {
10068                let null: u64 = 0;
10069                b.arg(qkv)
10070                    .arg(wq)
10071                    .arg(wk)
10072                    .arg(wv)
10073                    .arg(&mut *q)
10074                    .arg(&mut *k)
10075                    .arg(&mut *v)
10076                    .arg(&nc)
10077                    .arg(&rqi)
10078                    .arg(&rki)
10079                    .arg(pos)
10080                    .arg(&nhq)
10081                    .arg(&nhk)
10082                    .arg(&theta_scale)
10083                    .arg(&freq_scale)
10084                    .arg(&null)
10085                    .arg(&eps);
10086                unsafe {
10087                    b.launch(cfg)?;
10088                }
10089            }
10090        }
10091        Ok(())
10092    }
10093
10094    /// wave-3 fold: rms_norm_qkv + rope_neox2 in ONE launch (n_dims == head_dim; ff nullable).
10095    /// `n_rot` is the layer's derived rotary width and MUST equal `head_dim`
10096    /// ([`Engine::full_width_rope_only`]).
10097    #[allow(clippy::too_many_arguments)]
10098    pub fn rms_norm_qkv_rope(
10099        &self,
10100        q0: &CudaSlice<f32>,
10101        k0: &CudaSlice<f32>,
10102        v0: &CudaSlice<f32>,
10103        wq: &CudaSlice<f32>,
10104        wk: &CudaSlice<f32>,
10105        wv: &CudaSlice<f32>,
10106        q: &mut CudaSlice<f32>,
10107        k: &mut CudaSlice<f32>,
10108        v: &mut CudaSlice<f32>,
10109        head_dim: usize,
10110        n_rot: usize,
10111        rq: usize,
10112        rk: usize,
10113        pos: &CudaSlice<i32>,
10114        nh_q: usize,
10115        nh_k: usize,
10116        base: f32,
10117        freq_scale: f32,
10118        ff: Option<&CudaSlice<f32>>,
10119        eps: f32,
10120    ) -> Result<(), Box<dyn std::error::Error>> {
10121        Self::full_width_rope_only("rms_norm_qkv_rope", n_rot, head_dim)?;
10122        let f = self.func("rms_norm_qkv_rope_f32");
10123        let rows = rq + rk + rk; // q rows + k rows + v rows (rk == rv)
10124        let cfg = LaunchConfig {
10125            grid_dim: (rows as u32, 1, 1),
10126            block_dim: (rms_block(), 1, 1),
10127            shared_mem_bytes: 0,
10128        };
10129        let theta_scale = base.powf(-2.0 / head_dim as f32);
10130        let (nc, rqi, rki, nhq, nhk) = (
10131            head_dim as i32,
10132            rq as i32,
10133            rk as i32,
10134            nh_q as i32,
10135            nh_k as i32,
10136        );
10137        let __s_b = self.gpu.stream();
10138        let mut b = __s_b.launch_builder(&f);
10139        match ff {
10140            Some(t) => {
10141                b.arg(q0)
10142                    .arg(k0)
10143                    .arg(v0)
10144                    .arg(wq)
10145                    .arg(wk)
10146                    .arg(wv)
10147                    .arg(&mut *q)
10148                    .arg(&mut *k)
10149                    .arg(&mut *v)
10150                    .arg(&nc)
10151                    .arg(&rqi)
10152                    .arg(&rki)
10153                    .arg(pos)
10154                    .arg(&nhq)
10155                    .arg(&nhk)
10156                    .arg(&theta_scale)
10157                    .arg(&freq_scale)
10158                    .arg(t)
10159                    .arg(&eps);
10160                unsafe {
10161                    b.launch(cfg)?;
10162                }
10163            }
10164            None => {
10165                let null: u64 = 0;
10166                b.arg(q0)
10167                    .arg(k0)
10168                    .arg(v0)
10169                    .arg(wq)
10170                    .arg(wk)
10171                    .arg(wv)
10172                    .arg(&mut *q)
10173                    .arg(&mut *k)
10174                    .arg(&mut *v)
10175                    .arg(&nc)
10176                    .arg(&rqi)
10177                    .arg(&rki)
10178                    .arg(pos)
10179                    .arg(&nhq)
10180                    .arg(&nhk)
10181                    .arg(&theta_scale)
10182                    .arg(&freq_scale)
10183                    .arg(&null)
10184                    .arg(&eps);
10185                unsafe {
10186                    b.launch(cfg)?;
10187                }
10188            }
10189        }
10190        Ok(())
10191    }
10192
10193    /// FUSED norm+rope+APPEND (m=1 decode, 2026-07-23): one launch replaces the
10194    /// rms_norm_qkv_rope + append_kv_quantized_dc pair. Kernel lives in the flash fatbins
10195    /// (format-flavored quant tail) — `g` must mirror the append path's flavor exactly.
10196    /// `n_rot` is the layer's derived rotary width and MUST equal `head_dim`
10197    /// ([`Engine::full_width_rope_only`]).
10198    #[allow(clippy::too_many_arguments)]
10199    pub fn rms_norm_qkv_rope_append_dc(
10200        &self,
10201        q0: &CudaSlice<f32>,
10202        k0: &CudaSlice<f32>,
10203        v0: &CudaSlice<f32>,
10204        wq: &CudaSlice<f32>,
10205        wk: &CudaSlice<f32>,
10206        wv: &CudaSlice<f32>,
10207        q: &mut CudaSlice<f32>,
10208        k: &mut CudaSlice<f32>,
10209        v: &mut CudaSlice<f32>,
10210        head_dim: usize,
10211        n_rot: usize,
10212        rq: usize,
10213        rk: usize,
10214        pos: &CudaSlice<i32>,
10215        nh_q: usize,
10216        nh_k: usize,
10217        base: f32,
10218        freq_scale: f32,
10219        ff: Option<&CudaSlice<f32>>,
10220        eps: f32,
10221        kc: &mut CudaSlice<u8>,
10222        vc: &mut CudaSlice<u8>,
10223        t_dev: &CudaSlice<i32>,
10224        k_tok_bytes: usize,
10225        v_tok_bytes: usize,
10226        g: bool,
10227    ) -> Result<(), Box<dyn std::error::Error>> {
10228        Self::full_width_rope_only("rms_norm_qkv_rope_append_dc", n_rot, head_dim)?;
10229        let rows = rq + rk + rk;
10230        let theta_scale = base.powf(-2.0 / head_dim as f32);
10231        let (nc, rqi, rki, nhq, nhk) = (
10232            head_dim as i32,
10233            rq as i32,
10234            rk as i32,
10235            nh_q as i32,
10236            nh_k as i32,
10237        );
10238        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
10239        if Self::pdl_on() && Self::pdl_wb_on() {
10240            use cudarc::driver::{DevicePtr, DevicePtrMut};
10241            let s = &self.gpu.stream();
10242            let (p0, _a0) = q0.device_ptr(s);
10243            let (p1, _a1) = k0.device_ptr(s);
10244            let (p2, _a2) = v0.device_ptr(s);
10245            let (pwq, _a3) = wq.device_ptr(s);
10246            let (pwk, _a4) = wk.device_ptr(s);
10247            let (pwv, _a5) = wv.device_ptr(s);
10248            let (pq, _a6) = q.device_ptr_mut(s);
10249            let (pk, _a7) = k.device_ptr_mut(s);
10250            let (pv, _a8) = v.device_ptr_mut(s);
10251            let (pp, _a9) = pos.device_ptr(s);
10252            let pff: u64 = match ff {
10253                Some(t) => {
10254                    let (p, _gg) = t.device_ptr(s);
10255                    p as u64
10256                }
10257                None => 0,
10258            };
10259            let (pkc, _a10) = kc.device_ptr_mut(s);
10260            let (pvc, _a11) = vc.device_ptr_mut(s);
10261            let (pt, _a12) = t_dev.device_ptr(s);
10262            let mut ps = [
10263                &p0 as *const _ as *mut std::ffi::c_void,
10264                &p1 as *const _ as *mut _,
10265                &p2 as *const _ as *mut _,
10266                &pwq as *const _ as *mut _,
10267                &pwk as *const _ as *mut _,
10268                &pwv as *const _ as *mut _,
10269                &pq as *const _ as *mut _,
10270                &pk as *const _ as *mut _,
10271                &pv as *const _ as *mut _,
10272                &nc as *const _ as *mut _,
10273                &rqi as *const _ as *mut _,
10274                &rki as *const _ as *mut _,
10275                &pp as *const _ as *mut _,
10276                &nhq as *const _ as *mut _,
10277                &nhk as *const _ as *mut _,
10278                &theta_scale as *const _ as *mut _,
10279                &freq_scale as *const _ as *mut _,
10280                &pff as *const _ as *mut _,
10281                &eps as *const _ as *mut _,
10282                &pkc as *const _ as *mut _,
10283                &pvc as *const _ as *mut _,
10284                &pt as *const _ as *mut _,
10285                &ktb as *const _ as *mut _,
10286                &vtb as *const _ as *mut _,
10287            ];
10288            unsafe {
10289                self.launch_pdl_flash(
10290                    g,
10291                    "rms_norm_qkv_rope_append_dc_f32",
10292                    (rows as u32, 1, 1),
10293                    (rms_block(), 1, 1),
10294                    0,
10295                    &mut ps,
10296                )?;
10297            }
10298            return Ok(());
10299        }
10300        let f = if g {
10301            self.func_g("rms_norm_qkv_rope_append_dc_f32")
10302        } else {
10303            self.func("rms_norm_qkv_rope_append_dc_f32")
10304        };
10305        let cfg = LaunchConfig {
10306            grid_dim: (rows as u32, 1, 1),
10307            block_dim: (rms_block(), 1, 1),
10308            shared_mem_bytes: 0,
10309        };
10310        let __s_b = self.gpu.stream();
10311        let mut b = __s_b.launch_builder(&f);
10312        match ff {
10313            Some(t) => {
10314                b.arg(q0)
10315                    .arg(k0)
10316                    .arg(v0)
10317                    .arg(wq)
10318                    .arg(wk)
10319                    .arg(wv)
10320                    .arg(&mut *q)
10321                    .arg(&mut *k)
10322                    .arg(&mut *v)
10323                    .arg(&nc)
10324                    .arg(&rqi)
10325                    .arg(&rki)
10326                    .arg(pos)
10327                    .arg(&nhq)
10328                    .arg(&nhk)
10329                    .arg(&theta_scale)
10330                    .arg(&freq_scale)
10331                    .arg(t)
10332                    .arg(&eps)
10333                    .arg(&mut *kc)
10334                    .arg(&mut *vc)
10335                    .arg(t_dev)
10336                    .arg(&ktb)
10337                    .arg(&vtb);
10338                unsafe {
10339                    b.launch(cfg)?;
10340                }
10341            }
10342            None => {
10343                let null: u64 = 0;
10344                b.arg(q0)
10345                    .arg(k0)
10346                    .arg(v0)
10347                    .arg(wq)
10348                    .arg(wk)
10349                    .arg(wv)
10350                    .arg(&mut *q)
10351                    .arg(&mut *k)
10352                    .arg(&mut *v)
10353                    .arg(&nc)
10354                    .arg(&rqi)
10355                    .arg(&rki)
10356                    .arg(pos)
10357                    .arg(&nhq)
10358                    .arg(&nhk)
10359                    .arg(&theta_scale)
10360                    .arg(&freq_scale)
10361                    .arg(&null)
10362                    .arg(&eps)
10363                    .arg(&mut *kc)
10364                    .arg(&mut *vc)
10365                    .arg(t_dev)
10366                    .arg(&ktb)
10367                    .arg(&vtb);
10368                unsafe {
10369                    b.launch(cfg)?;
10370                }
10371            }
10372        }
10373        Ok(())
10374    }
10375
10376    /// wave-2 fold: a + b with the sum emitted q8_1 alongside f32.
10377    /// Host-len twin of `rms_norm_qkv_rope_append_dc` (zoo-fusion arc): the eager decode
10378    /// arm tracks kv length on host (`kvl.len`), so the append slot rides the launch arg
10379    /// instead of the device counter. Kernel body is SHARED with the _dc entry (one
10380    /// inlined body — bit-identical to the rms_norm_qkv_rope + append pair it replaces,
10381    /// same law as the dc fold). `n_rot` is the layer's derived rotary width and MUST equal
10382    /// `head_dim` ([`Engine::full_width_rope_only`]).
10383    #[allow(clippy::too_many_arguments)]
10384    pub fn rms_norm_qkv_rope_append(
10385        &self,
10386        q0: &CudaSlice<f32>,
10387        k0: &CudaSlice<f32>,
10388        v0: &CudaSlice<f32>,
10389        wq: &CudaSlice<f32>,
10390        wk: &CudaSlice<f32>,
10391        wv: &CudaSlice<f32>,
10392        q: &mut CudaSlice<f32>,
10393        k: &mut CudaSlice<f32>,
10394        v: &mut CudaSlice<f32>,
10395        head_dim: usize,
10396        n_rot: usize,
10397        rq: usize,
10398        rk: usize,
10399        pos: &CudaSlice<i32>,
10400        nh_q: usize,
10401        nh_k: usize,
10402        base: f32,
10403        freq_scale: f32,
10404        ff: Option<&CudaSlice<f32>>,
10405        eps: f32,
10406        kc: &mut CudaSlice<u8>,
10407        vc: &mut CudaSlice<u8>,
10408        t: usize,
10409        k_tok_bytes: usize,
10410        v_tok_bytes: usize,
10411        g: bool,
10412    ) -> Result<(), Box<dyn std::error::Error>> {
10413        Self::full_width_rope_only("rms_norm_qkv_rope_append", n_rot, head_dim)?;
10414        let rows = rq + rk + rk;
10415        let theta_scale = base.powf(-2.0 / head_dim as f32);
10416        let (nc, rqi, rki, nhq, nhk) = (
10417            head_dim as i32,
10418            rq as i32,
10419            rk as i32,
10420            nh_q as i32,
10421            nh_k as i32,
10422        );
10423        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
10424        let ti = t as i32;
10425        if Self::pdl_on() && Self::pdl_wb_on() {
10426            use cudarc::driver::{DevicePtr, DevicePtrMut};
10427            let s = &self.gpu.stream();
10428            let (p0, _a0) = q0.device_ptr(s);
10429            let (p1, _a1) = k0.device_ptr(s);
10430            let (p2, _a2) = v0.device_ptr(s);
10431            let (pwq, _a3) = wq.device_ptr(s);
10432            let (pwk, _a4) = wk.device_ptr(s);
10433            let (pwv, _a5) = wv.device_ptr(s);
10434            let (pq, _a6) = q.device_ptr_mut(s);
10435            let (pk, _a7) = k.device_ptr_mut(s);
10436            let (pv, _a8) = v.device_ptr_mut(s);
10437            let (pp, _a9) = pos.device_ptr(s);
10438            let pff: u64 = match ff {
10439                Some(t) => {
10440                    let (p, _gg) = t.device_ptr(s);
10441                    p as u64
10442                }
10443                None => 0,
10444            };
10445            let (pkc, _a10) = kc.device_ptr_mut(s);
10446            let (pvc, _a11) = vc.device_ptr_mut(s);
10447            let mut ps = [
10448                &p0 as *const _ as *mut std::ffi::c_void,
10449                &p1 as *const _ as *mut _,
10450                &p2 as *const _ as *mut _,
10451                &pwq as *const _ as *mut _,
10452                &pwk as *const _ as *mut _,
10453                &pwv as *const _ as *mut _,
10454                &pq as *const _ as *mut _,
10455                &pk as *const _ as *mut _,
10456                &pv as *const _ as *mut _,
10457                &nc as *const _ as *mut _,
10458                &rqi as *const _ as *mut _,
10459                &rki as *const _ as *mut _,
10460                &pp as *const _ as *mut _,
10461                &nhq as *const _ as *mut _,
10462                &nhk as *const _ as *mut _,
10463                &theta_scale as *const _ as *mut _,
10464                &freq_scale as *const _ as *mut _,
10465                &pff as *const _ as *mut _,
10466                &eps as *const _ as *mut _,
10467                &pkc as *const _ as *mut _,
10468                &pvc as *const _ as *mut _,
10469                &ti as *const _ as *mut _,
10470                &ktb as *const _ as *mut _,
10471                &vtb as *const _ as *mut _,
10472            ];
10473            unsafe {
10474                self.launch_pdl_flash(
10475                    g,
10476                    "rms_norm_qkv_rope_append_f32",
10477                    (rows as u32, 1, 1),
10478                    (rms_block(), 1, 1),
10479                    0,
10480                    &mut ps,
10481                )?;
10482            }
10483            return Ok(());
10484        }
10485        let f = if g {
10486            self.func_g("rms_norm_qkv_rope_append_f32")
10487        } else {
10488            self.func("rms_norm_qkv_rope_append_f32")
10489        };
10490        let cfg = LaunchConfig {
10491            grid_dim: (rows as u32, 1, 1),
10492            block_dim: (rms_block(), 1, 1),
10493            shared_mem_bytes: 0,
10494        };
10495        let __s_b = self.gpu.stream();
10496        let mut b = __s_b.launch_builder(&f);
10497        let null: u64 = 0;
10498        b.arg(q0)
10499            .arg(k0)
10500            .arg(v0)
10501            .arg(wq)
10502            .arg(wk)
10503            .arg(wv)
10504            .arg(&mut *q)
10505            .arg(&mut *k)
10506            .arg(&mut *v)
10507            .arg(&nc)
10508            .arg(&rqi)
10509            .arg(&rki)
10510            .arg(pos)
10511            .arg(&nhq)
10512            .arg(&nhk)
10513            .arg(&theta_scale)
10514            .arg(&freq_scale);
10515        match ff {
10516            Some(t) => {
10517                b.arg(t);
10518            }
10519            None => {
10520                b.arg(&null);
10521            }
10522        }
10523        b.arg(&eps)
10524            .arg(&mut *kc)
10525            .arg(&mut *vc)
10526            .arg(&ti)
10527            .arg(&ktb)
10528            .arg(&vtb);
10529        unsafe {
10530            b.launch(cfg)?;
10531        }
10532        Ok(())
10533    }
10534
10535    pub fn add_q8_1(
10536        &self,
10537        a: &CudaSlice<f32>,
10538        b: &CudaSlice<f32>,
10539        res: &mut CudaSlice<f32>,
10540        ncols: usize,
10541        nrows: usize,
10542    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10543        debug_assert!(ncols % 128 == 0);
10544        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
10545        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
10546        let f = self.func("add_q8_1_f32");
10547        let cfg = LaunchConfig {
10548            grid_dim: (nrows as u32, 1, 1),
10549            block_dim: (rms_block(), 1, 1),
10550            shared_mem_bytes: 0,
10551        };
10552        let nc = ncols as i32;
10553        let __s_b2 = self.gpu.stream();
10554        let mut b2 = __s_b2.launch_builder(&f);
10555        b2.arg(a)
10556            .arg(b)
10557            .arg(&mut *res)
10558            .arg(&mut out_q)
10559            .arg(&mut out_d)
10560            .arg(&nc);
10561        unsafe {
10562            b2.launch(cfg)?;
10563        }
10564        Ok((out_q, out_d))
10565    }
10566
10567    /// E4B FFN-tail exit fusion (glue wave 5): resid = b + rms(a, wa) emitted f32 + q8_1 pair
10568    /// in ONE launch — replaces rms_norm(a,wa->sn) + add_q8_1(sn,b). Same rms_block() config
10569    /// as both parents (bit-identity: identical reduction + quad-walk quantize).
10570    pub fn rms_pre_add_q8_1(
10571        &self,
10572        a: &CudaSlice<f32>,
10573        wa: &CudaSlice<f32>,
10574        b: &CudaSlice<f32>,
10575        res: &mut CudaSlice<f32>,
10576        ncols: usize,
10577        nrows: usize,
10578        eps: f32,
10579    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10580        debug_assert!(ncols % 128 == 0);
10581        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
10582        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
10583        let f = self.func("rms_pre_add_q8_1_f32");
10584        let cfg = LaunchConfig {
10585            grid_dim: (nrows as u32, 1, 1),
10586            block_dim: (rms_block(), 1, 1),
10587            shared_mem_bytes: 0,
10588        };
10589        let (nc, ep) = (ncols as i32, eps);
10590        let __s_b2 = self.gpu.stream();
10591        let mut b2 = __s_b2.launch_builder(&f);
10592        b2.arg(a)
10593            .arg(wa)
10594            .arg(b)
10595            .arg(&mut *res)
10596            .arg(&mut out_q)
10597            .arg(&mut out_d)
10598            .arg(&nc)
10599            .arg(&ep);
10600        unsafe {
10601            b2.launch(cfg)?;
10602        }
10603        Ok((out_q, out_d))
10604    }
10605
10606    /// L2 norm per row (head_dim), no weight.
10607    /// PREFILL l2 dispatch (round 27): the warp-per-row float4 v2 when the numeric-config
10608    /// seam allows (MEMRA_L2_V2, default ON, d_state==128 only); else the strided kernel.
10609    pub fn l2_v2_on(ncols: usize) -> bool {
10610        ncols == 128 && std::env::var("MEMRA_L2_V2").as_deref() != Ok("0")
10611    }
10612
10613    pub fn l2_norm_pp(
10614        &self,
10615        x: &CudaSlice<f32>,
10616        dst: &mut CudaSlice<f32>,
10617        dst16: Option<&mut CudaSlice<u8>>,
10618        ncols: usize,
10619        nrows: usize,
10620        eps: f32,
10621    ) -> Result<(), Box<dyn std::error::Error>> {
10622        if Self::l2_v2_on(ncols) {
10623            let f = self.func("l2_norm_pp_v2_f32");
10624            let rows_per_block = 8u32; // 256 threads = 8 warps = 8 rows
10625            let cfg = LaunchConfig {
10626                grid_dim: ((nrows as u32).div_ceil(rows_per_block), 1, 1),
10627                block_dim: (256, 1, 1),
10628                shared_mem_bytes: 0,
10629            };
10630            let (nc, nr, e) = (ncols as i32, nrows as i32, eps);
10631            // mirror-fold: bf16 twin address by value (0 = skip; matches the nullable param)
10632            let d16: u64 = match dst16 {
10633                Some(d) => self.addr_u8(d),
10634                None => 0,
10635            };
10636            let __s_b = self.gpu.stream();
10637            let mut b = __s_b.launch_builder(&f);
10638            b.arg(x).arg(dst).arg(&d16).arg(&nc).arg(&nr).arg(&e);
10639            unsafe {
10640                b.launch(cfg)?;
10641            }
10642            return Ok(());
10643        }
10644        self.l2_norm(x, dst, ncols, nrows, eps)
10645    }
10646
10647    pub fn l2_norm(
10648        &self,
10649        x: &CudaSlice<f32>,
10650        dst: &mut CudaSlice<f32>,
10651        ncols: usize,
10652        nrows: usize,
10653        eps: f32,
10654    ) -> Result<(), Box<dyn std::error::Error>> {
10655        let f = self.func("l2_norm_f32");
10656        let cfg = LaunchConfig {
10657            grid_dim: (nrows as u32, 1, 1),
10658            block_dim: (256, 1, 1),
10659            shared_mem_bytes: 0,
10660        };
10661        let (nc, e) = (ncols as i32, eps);
10662        let __s_b = self.gpu.stream();
10663        let mut b = __s_b.launch_builder(&f);
10664        b.arg(x).arg(dst).arg(&nc).arg(&e);
10665        unsafe {
10666            b.launch(cfg)?;
10667        }
10668        Ok(())
10669    }
10670
10671    /// L2-norm with blockDim=32 (warp-tree reduction) — BIT-IDENTICAL to gdn_prep_decode_f32's
10672    /// per-warp L2 norm. The verify path MUST use this to match decode's FP accumulation order:
10673    /// l2_norm at blockDim=256 produces a different shfl-tree reduction of the 128-element
10674    /// squared-sum (pairwise tree vs serial-4-then-warp-tree), causing ULP differences that
10675    /// propagate through gdn_scan and flip argmax on marginal logits.
10676    pub fn l2_norm_decode(
10677        &self,
10678        x: &CudaSlice<f32>,
10679        dst: &mut CudaSlice<f32>,
10680        ncols: usize,
10681        nrows: usize,
10682        eps: f32,
10683    ) -> Result<(), Box<dyn std::error::Error>> {
10684        let f = self.func("l2_norm_f32");
10685        let cfg = LaunchConfig {
10686            grid_dim: (nrows as u32, 1, 1),
10687            block_dim: (32, 1, 1),
10688            shared_mem_bytes: 0,
10689        };
10690        let (nc, e) = (ncols as i32, eps);
10691        let __s_b = self.gpu.stream();
10692        let mut b = __s_b.launch_builder(&f);
10693        b.arg(x).arg(dst).arg(&nc).arg(&e);
10694        unsafe {
10695            b.launch(cfg)?;
10696        }
10697        Ok(())
10698    }
10699
10700    /// RoPE NEOX in-place. x:[head_dim, n_heads, n_tokens], pos:[n_tokens].
10701    pub fn rope_neox(
10702        &self,
10703        x: &mut CudaSlice<f32>,
10704        pos: &CudaSlice<i32>,
10705        head_dim: usize,
10706        n_dims: usize,
10707        n_heads: usize,
10708        n_tokens: usize,
10709        freq_base: f32,
10710        freq_scale: f32,
10711    ) -> Result<(), Box<dyn std::error::Error>> {
10712        let f = self.func("rope_neox_f32");
10713        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
10714        let grid = (n_heads * n_tokens) as u32;
10715        let cfg = LaunchConfig {
10716            grid_dim: (grid, 1, 1),
10717            block_dim: ((head_dim / 2) as u32, 1, 1),
10718            shared_mem_bytes: 0,
10719        };
10720        let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
10721        let __s_b = self.gpu.stream();
10722        let mut b = __s_b.launch_builder(&f);
10723        b.arg(x)
10724            .arg(pos)
10725            .arg(&hd)
10726            .arg(&nd)
10727            .arg(&nh)
10728            .arg(&theta_scale)
10729            .arg(&freq_scale);
10730        unsafe {
10731            b.launch(cfg)?;
10732        }
10733        Ok(())
10734    }
10735
10736    /// RoPE NEOX with per-dim freq factors (gemma4 global layers, rope_freqs.weight [n_dims/2]).
10737    pub fn rope_neox_ff(
10738        &self,
10739        x: &mut CudaSlice<f32>,
10740        pos: &CudaSlice<i32>,
10741        head_dim: usize,
10742        n_dims: usize,
10743        n_heads: usize,
10744        n_tokens: usize,
10745        freq_base: f32,
10746        freq_scale: f32,
10747        ff: &CudaSlice<f32>,
10748    ) -> Result<(), Box<dyn std::error::Error>> {
10749        let f = self.func("rope_neox_ff_f32");
10750        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
10751        let grid = (n_heads * n_tokens) as u32;
10752        let cfg = LaunchConfig {
10753            grid_dim: (grid, 1, 1),
10754            block_dim: ((head_dim / 2) as u32, 1, 1),
10755            shared_mem_bytes: 0,
10756        };
10757        let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
10758        let __s_b = self.gpu.stream();
10759        let mut b = __s_b.launch_builder(&f);
10760        b.arg(x)
10761            .arg(pos)
10762            .arg(&hd)
10763            .arg(&nd)
10764            .arg(&nh)
10765            .arg(&theta_scale)
10766            .arg(&freq_scale)
10767            .arg(ff);
10768        unsafe {
10769            b.launch(cfg)?;
10770        }
10771        Ok(())
10772    }
10773
10774    /// gemma4: rope q and k in one launch (per-row chain = rope_neox / rope_neox_ff verbatim).
10775    #[allow(clippy::too_many_arguments)]
10776    pub fn rope_neox2(
10777        &self,
10778        q: &mut CudaSlice<f32>,
10779        k: &mut CudaSlice<f32>,
10780        pos: &CudaSlice<i32>,
10781        head_dim: usize,
10782        n_dims: usize,
10783        nh_q: usize,
10784        nh_k: usize,
10785        n_tokens: usize,
10786        freq_base: f32,
10787        freq_scale: f32,
10788        ff: Option<&CudaSlice<f32>>,
10789    ) -> Result<(), Box<dyn std::error::Error>> {
10790        let f = self.func("rope_neox2_f32");
10791        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
10792        let grid = ((nh_q + nh_k) * n_tokens) as u32;
10793        let cfg = LaunchConfig {
10794            grid_dim: (grid, 1, 1),
10795            block_dim: ((head_dim / 2) as u32, 1, 1),
10796            shared_mem_bytes: 0,
10797        };
10798        let (hd, nd, nq, nk, nt) = (
10799            head_dim as i32,
10800            n_dims as i32,
10801            nh_q as i32,
10802            nh_k as i32,
10803            n_tokens as i32,
10804        );
10805        let __s_b = self.gpu.stream();
10806        let mut b = __s_b.launch_builder(&f);
10807        b.arg(q)
10808            .arg(k)
10809            .arg(pos)
10810            .arg(&hd)
10811            .arg(&nd)
10812            .arg(&nq)
10813            .arg(&nk)
10814            .arg(&nt)
10815            .arg(&theta_scale)
10816            .arg(&freq_scale);
10817        match ff {
10818            Some(ffv) => {
10819                b.arg(ffv);
10820                unsafe {
10821                    b.launch(cfg)?;
10822                }
10823            }
10824            None => {
10825                let null: u64 = 0;
10826                b.arg(&null);
10827                unsafe {
10828                    b.launch(cfg)?;
10829                }
10830            }
10831        }
10832        Ok(())
10833    }
10834
10835    /// gemma4 R1: dst = GELU_tanh(gate) * up.
10836    pub fn gelu_tanh_mul(
10837        &self,
10838        gate: &CudaSlice<f32>,
10839        up: &CudaSlice<f32>,
10840        dst: &mut CudaSlice<f32>,
10841        n: usize,
10842    ) -> Result<(), Box<dyn std::error::Error>> {
10843        let f = self.func("gelu_tanh_mul_f32");
10844        let cfg = LaunchConfig::for_num_elems(n as u32);
10845        let ni = n as i32;
10846        let __s_b = self.gpu.stream();
10847        let mut b = __s_b.launch_builder(&f);
10848        b.arg(gate).arg(up).arg(dst).arg(&ni);
10849        unsafe {
10850            b.launch(cfg)?;
10851        }
10852        Ok(())
10853    }
10854
10855    pub fn silu_mul(
10856        &self,
10857        gate: &CudaSlice<f32>,
10858        up: &CudaSlice<f32>,
10859        dst: &mut CudaSlice<f32>,
10860        n: usize,
10861    ) -> Result<(), Box<dyn std::error::Error>> {
10862        let f = self.func("silu_mul_f32");
10863        // float4 kernel: one thread per 4 elements (tail handled in-kernel)
10864        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
10865        let ni = n as i32;
10866        let __s_b = self.gpu.stream();
10867        let mut b = __s_b.launch_builder(&f);
10868        b.arg(gate).arg(up).arg(dst).arg(&ni);
10869        unsafe {
10870            b.launch(cfg)?;
10871        }
10872        Ok(())
10873    }
10874
10875    /// SwiGLU twin using Memra's host-matching expf transcription.
10876    pub fn silu_mul_host_expf(
10877        &self,
10878        gate: &CudaSlice<f32>,
10879        up: &CudaSlice<f32>,
10880        dst: &mut CudaSlice<f32>,
10881        n: usize,
10882    ) -> Result<(), Box<dyn std::error::Error>> {
10883        let f = self.func("silu_mul_host_expf_f32");
10884        let cfg = LaunchConfig::for_num_elems(n as u32);
10885        let ni = n as i32;
10886        let __s_b = self.gpu.stream();
10887        let mut b = __s_b.launch_builder(&f);
10888        b.arg(gate).arg(up).arg(dst).arg(&ni);
10889        unsafe {
10890            b.launch(cfg)?;
10891        }
10892        Ok(())
10893    }
10894
10895    /// Step routed-expert clamp twin using Memra's host-matching expf transcription.
10896    pub fn silu_clamped_mul_host_expf(
10897        &self,
10898        gate: &CudaSlice<f32>,
10899        up: &CudaSlice<f32>,
10900        limit: f32,
10901        dst: &mut CudaSlice<f32>,
10902        n: usize,
10903    ) -> Result<(), Box<dyn std::error::Error>> {
10904        if !limit.is_finite() || limit <= 0.0 {
10905            return Err(
10906                format!("Step routed-expert clamp limit must be positive, got {limit}").into(),
10907            );
10908        }
10909        let f = self.func("silu_clamped_mul_host_expf_f32");
10910        let cfg = LaunchConfig::for_num_elems(n as u32);
10911        let ni = n as i32;
10912        let __s_b = self.gpu.stream();
10913        let mut b = __s_b.launch_builder(&f);
10914        b.arg(gate).arg(up).arg(&limit).arg(dst).arg(&ni);
10915        unsafe {
10916            b.launch(cfg)?;
10917        }
10918        Ok(())
10919    }
10920
10921    /// f16out twin of `silu_mul` (task #17): the epilogue also emits the fp16 GEMM operand
10922    /// for the down projection — kills the standalone convert pass. Bit-identical class.
10923    pub fn silu_mul_f16out(
10924        &self,
10925        gate: &CudaSlice<f32>,
10926        up: &CudaSlice<f32>,
10927        dst: &mut CudaSlice<f32>,
10928        dst16: &mut CudaSlice<u8>,
10929        n: usize,
10930    ) -> Result<(), Box<dyn std::error::Error>> {
10931        let f = self.func("silu_mul_f16out_f32");
10932        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
10933        let ni = n as i32;
10934        let __s_b = self.gpu.stream();
10935        let mut b = __s_b.launch_builder(&f);
10936        b.arg(gate).arg(up).arg(dst).arg(dst16).arg(&ni);
10937        unsafe {
10938            b.launch(cfg)?;
10939        }
10940        Ok(())
10941    }
10942
10943    /// FFN SwiGLU epilogue fusion (RANK3 LEVER 2): `dst = silu(gate*gs) * (up*us)` in ONE launch,
10944    /// folding the per-tensor NVFP4 macro-scale (`gs`,`us`) that would otherwise be two separate
10945    /// `scale_inplace` launches on the gate/up matmul outputs. BIT-IDENTICAL to
10946    /// scale_inplace(gate,gs); scale_inplace(up,us); silu_mul(gate,up,dst) — identical float ops in
10947    /// identical order. For non-NVFP4 weights gs==us==1.0 -> identical to `silu_mul`. Net: -2
10948    /// launches per dense FFN layer (the gate+up post-matmul scales).
10949    pub fn silu_mul_scaled(
10950        &self,
10951        gate: &CudaSlice<f32>,
10952        up: &CudaSlice<f32>,
10953        gs: f32,
10954        us: f32,
10955        dst: &mut CudaSlice<f32>,
10956        n: usize,
10957    ) -> Result<(), Box<dyn std::error::Error>> {
10958        let f = self.func("silu_mul_scaled_f32");
10959        let cfg = LaunchConfig::for_num_elems(n as u32);
10960        let ni = n as i32;
10961        let (gsf, usf) = (gs, us);
10962        let __s_b = self.gpu.stream();
10963        let mut b = __s_b.launch_builder(&f);
10964        b.arg(gate).arg(up).arg(&gsf).arg(&usf).arg(dst).arg(&ni);
10965        unsafe {
10966            b.launch(cfg)?;
10967        }
10968        Ok(())
10969    }
10970
10971    /// swigluoai (MiniMax-M3 / GPT-OSS): clamped SwiGLU epilogue, math 1:1 vs llama.cpp
10972    /// ggml_cuda_op_swiglu_oai_single. `dst = swish_alpha(clamp(gate*gs)) * (1 + clamp(up*us))`.
10973    /// gs/us fold the NVFP4 macro-scales exactly like `silu_mul_scaled`.
10974    #[allow(clippy::too_many_arguments)]
10975    pub fn swigluoai_mul_scaled(
10976        &self,
10977        gate: &CudaSlice<f32>,
10978        up: &CudaSlice<f32>,
10979        gs: f32,
10980        us: f32,
10981        alpha: f32,
10982        limit: f32,
10983        dst: &mut CudaSlice<f32>,
10984        n: usize,
10985    ) -> Result<(), Box<dyn std::error::Error>> {
10986        let f = self.func("swigluoai_mul_scaled_f32");
10987        let cfg = LaunchConfig::for_num_elems(n as u32);
10988        let ni = n as i32;
10989        let __s_b = self.gpu.stream();
10990        let mut b = __s_b.launch_builder(&f);
10991        b.arg(gate)
10992            .arg(up)
10993            .arg(&gs)
10994            .arg(&us)
10995            .arg(&alpha)
10996            .arg(&limit)
10997            .arg(dst)
10998            .arg(&ni);
10999        unsafe {
11000            b.launch(cfg)?;
11001        }
11002        Ok(())
11003    }
11004
11005    /// RANK2 LEVER (q8_1 quant-fold): SwiGLU epilogue that EMITS the q8_1 quantization of `act`
11006    /// directly (aq int8 [n] + ad f32 [n/32]), so ffn_down's standalone `quantize_q8_1` launch is
11007    /// removed — the down-proj activation has one consumer, so the quant folds into the producer for
11008    /// free (no extra HBM read; no f32 `act` write). gs/us fold the gate/up NVFP4 macro-scales like
11009    /// `silu_mul_scaled`. BIT-IDENTICAL q8_1 to silu_mul_scaled(...) then quantize_q8_1(...). Only
11010    /// valid when ffn_down uses the q8_1 dp4a/mmvq path; the caller checks `uses_q8_1_fast(ffn_down)`.
11011    /// n must be a multiple of 32 (n_ff always is).
11012    pub fn silu_mul_scaled_q8_1(
11013        &self,
11014        gate: &CudaSlice<f32>,
11015        up: &CudaSlice<f32>,
11016        gs: f32,
11017        us: f32,
11018        n: usize,
11019    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11020        let f = self.func("silu_mul_scaled_q8_1");
11021        let nblk = n / 32;
11022        let mut aq = self.alloc_uninit::<i8>(n)?; // full-overwrite output
11023        let mut ad = self.alloc_uninit::<f32>(nblk)?; // full-overwrite output
11024        // WARP-PER-BLOCK kernel: one warp (32 lanes) per 32-block -> n threads total.
11025        let cfg = LaunchConfig::for_num_elems(n as u32);
11026        let (gsf, usf, ni) = (gs, us, n as i32);
11027        let __s_b = self.gpu.stream();
11028        let mut b = __s_b.launch_builder(&f);
11029        b.arg(gate)
11030            .arg(up)
11031            .arg(&gsf)
11032            .arg(&usf)
11033            .arg(&mut aq)
11034            .arg(&mut ad)
11035            .arg(&ni);
11036        unsafe {
11037            b.launch(cfg)?;
11038        }
11039        Ok((aq, ad))
11040    }
11041
11042    pub fn add(
11043        &self,
11044        a: &CudaSlice<f32>,
11045        b_in: &CudaSlice<f32>,
11046        dst: &mut CudaSlice<f32>,
11047        n: usize,
11048    ) -> Result<(), Box<dyn std::error::Error>> {
11049        let f = self.func("add_f32");
11050        // float4 kernel: one thread per 4 elements (tail handled in-kernel)
11051        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
11052        let ni = n as i32;
11053        let __s_bld = self.gpu.stream();
11054        let mut bld = __s_bld.launch_builder(&f);
11055        bld.arg(a).arg(b_in).arg(dst).arg(&ni);
11056        unsafe {
11057            bld.launch(cfg)?;
11058        }
11059        Ok(())
11060    }
11061
11062    pub fn mul(
11063        &self,
11064        a: &CudaSlice<f32>,
11065        b_in: &CudaSlice<f32>,
11066        dst: &mut CudaSlice<f32>,
11067        n: usize,
11068    ) -> Result<(), Box<dyn std::error::Error>> {
11069        let f = self.func("mul_f32");
11070        let cfg = LaunchConfig::for_num_elems(n as u32);
11071        let ni = n as i32;
11072        let __s_bld = self.gpu.stream();
11073        let mut bld = __s_bld.launch_builder(&f);
11074        bld.arg(a).arg(b_in).arg(dst).arg(&ni);
11075        unsafe {
11076            bld.launch(cfg)?;
11077        }
11078        Ok(())
11079    }
11080
11081    /// Unified weight-tensor matmul: dispatches quant tensors to qmatvec (weights packed) and
11082    /// float tensors to cuBLASLt. y[m,out] = x[m,in] @ W[out,in]^T.
11083    pub fn matmul(
11084        &self,
11085        w: &crate::model::GpuTensor,
11086        x: &CudaSlice<f32>,
11087        m: usize,
11088    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11089        use crate::model::GpuTensor;
11090        let in_f = w.in_features();
11091        let out_f = w.out_features();
11092        // PREFILL (T>1) ROOT FIX: batched tensor-core int8 GEMM. Decodes each weight tile to int8
11093        // in smem ONCE and reuses across all tokens via mma — vs the dp4a matvec's per-token weight
11094        // re-read. Only the 4 daily-hot dtypes; m=1 decode keeps dp4a (it's bandwidth-bound, mma
11095        // gives nothing). Quantize the activation once here then call the GEMM.
11096        // m cutoff FIXED at 16: the m=4 MMA-verify A/B (2026-07-06, was MEMRA_GEMM_M) measured
11097        // NEGATIVE — the MMA tile grid starves at m=4 (BN=256 -> grid.y=1) and its FP order
11098        // shifted verify argmax at tight margins. Do not lower without re-running that battery.
11099        #[allow(non_snake_case)]
11100        // VERIFY-EXACT scope pushes the GEMM crossover out of reach (usize::MAX) — the
11101        // t>=16 dflash verify must ride the decode-exact batched class (parity law).
11102        let GEMM_M_THRESHOLD = if self.verify_exact_on() {
11103            usize::MAX
11104        } else {
11105            16usize
11106        };
11107
11108        // PREFILL GEMM (m>=16). ACCURACY-FIRST dispatch (2026-06-28, prefill-gemm-beat-research wf
11109        // wllbyo6vc step 1): the int8 W4A8 GEMM (qmatvec_gemm, q8_1 activation, s32 accumulate) is
11110        // ACCURATE (prefill logit maxdiff 0.159, < dp4a 0.55) and the default. The FP4 W4A4 mxf4 path
11111        // (try_fp4_gemm) quantizes the ACTIVATION to e2m1 4-bit (8 magnitude levels) -> maxdiff 1.0
11112        // when combined — a real accuracy loss, NOT a math bug. So FP4-W4A4 is taken ONLY under the
11113        // explicit MEMRA_FP4 opt-in AND it must come SECOND (int8 W4A8 is the correct default for NVFP4).
11114        // The workflow plan rebuilds the FP4 path (kill per-K repack, widen K, deepen pipeline, TMA) to
11115        // be both fast AND accurate; until then NVFP4 prefill defaults to the accurate int8 GEMM.
11116        // TINY-OUT_F GUARD (2026-06-28, ncu trace): the tiling GEMM's grid is (ceil(out_f/BM=64),
11117        // ceil(m/BN=256)). For tiny out_f (ssm_beta/ssm_alpha out_f=num_v_heads~32), grid.x=1 -> only
11118        // ceil(m/256) CTAs (e.g. 2 for m=512) on 82 SMs = 0.39% SM throughput, 852us EACH (measured
11119        // worst offender). The dp4a path grids (out_f, m) = far more CTAs, filling the GPU. So route
11120        // out_f < 2*BM to dp4a (skip the tiling GEMM which structurally can't fill the SMs here).
11121        const GEMM_MIN_OUT_F: usize = 128; // 2*BM; below this the GEMM grid.x starves the 82 SMs
11122        // VENDORED llama MMQ prefill GEMMs. NVFP4 W4A8 is DEFAULT-ON (2026-07-05 flip: same int8
11123        // accuracy class as the int8 GEMM below at ~1.9x pp512, rp-loader coexists with the A6
11124        // repack; MEMRA_MMQ_W4A8=0 = escape hatch). W4A4 mxf4nvf4 + Q4_K/Q5_K stay behind MEMRA_MMQ=1.
11125        // The env policy lives in mmq_supports/qmatvec_mmq. Feeds raw f32 activation `x` (the
11126        // launcher quantizes internally). out_f>=MMQ_Y/2 keeps the tile grid from starving the SMs.
11127        // FP8-ACT PREFILL (MEMRA_PP_FP8=1, probe verdict 2026-07-08): F8-E4M3-origin projections
11128        // carry their raw e4m3 device bytes (the `fp8` operand stashed at load next to the Q8_0
11129        // re-encode) — cuBLASLt FP8 TN at 620-795 TF vs 47-72 TF for this class's int8 GEMM.
11130        // Weight side EXACT (checkpoint bytes); activation rides ONE per-batch e4m3 scale
11131        // (amax/448) folded with weight_scale in-GEMM. Prefill only; decode keeps Q8_0 untouched.
11132        if m >= GEMM_M_THRESHOLD {
11133            if let Some(y) = self.try_fp8_gemm(w, x, m)? {
11134                return Ok(y);
11135            }
11136            // PER-BLOCK FP8 MMQ (lane/fp8-mmq): the block-128 class try_fp8_gemm skips (cuBLASLt
11137            // takes no block grid on sm_120). Exact per block — the checkpoint's e4m3 bytes and its
11138            // f32 grid go into the tile unchanged. TWO SOURCES, TWO DEFAULTS: the load-time stash is
11139            // opt-in (MEMRA_FP8_MMQ=1), the native-resident QT_F8_E4M3_BLK grid is DEFAULT ON
11140            // (MEMRA_FP8_MMQ=0 reverts it to dequant-per-call) — see fp8_ffi.rs for why the same
11141            // tile defaults differently by operand source.
11142            if let Some(y) = self.try_fp8_blk_mmq(w, x, m)? {
11143                return Ok(y);
11144            }
11145            // FP16-mirror prefill (MEMRA_PP_F16=1, probe 2026-07-26: 3.2-3.7x the MMQ class).
11146            // Mirror presence IS the gate (only built under the env). Decode never reaches here.
11147            if let Some(y) = self.try_f16_gemm(w, x, m)? {
11148                return Ok(y);
11149            }
11150        }
11151        // F8-E4M3 BLOCK-128 (QT_F8_E4M3_BLK, lane/fp8-blk128-decode). TWO arms, split at the SAME
11152        // m threshold the rest of this method uses:
11153        //   * m >= threshold (prefill): dequant-per-call to the ARM B' Q8_0 slab and recurse, so
11154        //     prefill keeps the floor's kernels AND the floor's bits (try_e4m3_blk_prefill).
11155        //   * m <  threshold: the native per-block GEMV — m=1 decode and the m=2..15 verify tiers.
11156        //     grid.y=m runs the exact m=1 program per (token,row), so the decode-parity law holds
11157        //     across every tier by construction with no batched twin needed.
11158        //
11159        // NOT gated on `fast`: this dtype has no dp4a twin and no Stage-A f32-dequant oracle (the
11160        // generic `deq()` switch has no block-scale input), exactly as QT_F8_E4M3 has none, so
11161        // MEMRA_FAST=0 cannot route it anywhere else. Placed before every GEMM/MMQ arm below
11162        // because gemm_supports/mmq_supports/mmvq_supports all deliberately REFUSE this qtype —
11163        // reaching the generic tail would panic rather than produce wrong numbers, and this pair of
11164        // arms is what makes sure it never gets there.
11165        if let GpuTensor::Quant { qtype, .. } = w {
11166            if *qtype == QT_F8_E4M3_BLK {
11167                if m >= GEMM_M_THRESHOLD {
11168                    if let Some(y) = self.try_e4m3_blk_prefill(w, x, m)? {
11169                        return Ok(y);
11170                    }
11171                }
11172                let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
11173                if let Some(y) = self.try_e4m3_blk_pre(w, &aq, &ad, m)? {
11174                    return Ok(y);
11175                }
11176            }
11177        }
11178        if m >= GEMM_M_THRESHOLD && out_f >= GEMM_MIN_OUT_F && self.mmq_supports(w) {
11179            return self.qmatvec_mmq(w, x, m);
11180        }
11181        if m >= GEMM_M_THRESHOLD && out_f >= GEMM_MIN_OUT_F && self.gemm_supports(w) {
11182            let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
11183            return self.qmatvec_gemm(w, &aq, &ad, m);
11184        }
11185        // FP4 W4A4 only as an explicit speed/accuracy tradeoff opt-in, and only if the int8 GEMM
11186        // above didn't already handle this weight (e.g. NVFP4 with in_f%64!=0, or MEMRA_NO_GEMM set).
11187        if m >= GEMM_M_THRESHOLD {
11188            if let Some(y) = self.try_fp4_gemm(w, x, m, in_f, out_f)? {
11189                return Ok(y);
11190            }
11191        }
11192        // Stage-B fast int8 dp4a is the DEFAULT since 2026-07-08 (it has been the daily path
11193        // for weeks; the old opt-in flag was a silent-slow-path landmine). MEMRA_FAST=0 reverts
11194        // to Stage-A f32-dequant (the correctness oracle path).
11195        let fast = std::env::var("MEMRA_FAST").as_deref() != Ok("0");
11196        // PERF-3 decode-GEMV: m=1 warp-per-row MMVQ (MEMRA_MMVQ). The big decode matvecs reach
11197        // `matmul` directly (ffn_down, lm_head output, wo), so route them here too — not only the
11198        // matmul_pre siblings. qmatvec_mmvq_raw quantizes the activation internally (q8_1) like the
11199        // _fast paths; the NVFP4 macro-scale is applied by the `scale != 1.0` block below.
11200        if m == 1 && fast {
11201            if let GpuTensor::Quant {
11202                bytes,
11203                qtype,
11204                row_bytes,
11205                rp,
11206                rp4,
11207                scale,
11208                ..
11209            } = w
11210            {
11211                if self.mmvq_supports(*qtype) {
11212                    // NVFP4 macro-scale rides the kernel's fused epilogue arg (one launch total);
11213                    // non-NVFP4 has scale==1.0 so qmatvec_mmvq skips scale_inplace either way.
11214                    // Q4_0 split-plane mirror (rp4): the decode arm reads it via the _rp twins.
11215                    let (bytes, rp) = match rp4 {
11216                        Some(m4) => (m4, true),
11217                        None => (bytes, *rp),
11218                    };
11219                    let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
11220                    return self.qmatvec_mmvq(
11221                        bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, *scale, rp,
11222                    );
11223                }
11224            }
11225        }
11226        // BATCHED weight-resident matvec for the m=2-4 band (the MTP/verify forward's ffn_down, wo, and
11227        // lm_head `output` reach `matmul` directly at m=T=2..4). Walks the weight ONCE, dp4a vs all m
11228        // activation columns -> 1 weight read for m tokens (vs grid.y=m re-reading m times below). Quant
11229        // the activation once here (q8_1) like the _fast paths; macro-scale applied via the scale!=1.0
11230        // block below. MEMRA_NO_BATCHED -> per-m path.
11231        //
11232        // DECODE-PARITY GATE (2026-07-07, the 9B synth K=3/4/6 spec FAIL root cause): the batched
11233        // kernels are bit-identical per (token,row) to MMVQ's 32-thread warp reduce, NOT to the
11234        // dp4a kernels' 128-thread two-level reduce. Without MEMRA_MMVQ the m=1 decode chain rides
11235        // dp4a, so a verify riding batched here has a DIFFERENT FP order than the decode it must
11236        // match bit-for-bit — greedy spec flips at tight-margin tokens (the old HANDOVER "ENV LAW:
11237        // FAST+MMVQ both required" footgun, closed here). Parity law: the m>1 kernel CLASS must be
11238        // a pure function of (dtype, env) equal to the m=1 class — batched iff MMVQ. Without MMVQ
11239        // the verify falls to the per-m grid.y=m dp4a path below (each column = the exact m=1
11240        // dp4a program). MEMRA_MMVQ=1 (the daily config) is dispatch-unchanged.
11241        if (2..=16).contains(&m)
11242            && fast
11243            && std::env::var("MEMRA_NO_BATCHED").is_err()
11244            && (m <= 4 || Self::b8_enabled())
11245        {
11246            // b16 tier (2026-07-11, spec K>7): Q4_0/Q6_K have base+_rp b16 kernels; Q8_0's
11247            // b16 exists only as the split-plane _rp twin, so it joins iff the q8rp mirror
11248            // is present (rp4) — the mirror pick below then routes to the _rp family.
11249            // QT_F8_E4M3 joins unconditionally (lane/rp-on-st): its b16 IS the base kernel,
11250            // because the native e4m3 row layout is already aligned and needs no mirror.
11251            // NVFP4/Q4_K/Q8_0 all join unconditionally now (lane/rp-on-st): each has base + _rp
11252            // b16 twins, so either residency layout has its aligned form at this width. Q8_0's
11253            // old `rp4.is_some()` precondition is GONE — the mirror is a bandwidth lever, not the
11254            // exact tier's admission ticket (it was refusing FP8-ST over 23.9 MiB of ssm_beta).
11255            let m_ok = m <= 8
11256                || matches!(w, GpuTensor::Quant { qtype, .. }
11257                if *qtype == QT_Q4_0 || *qtype == QT_Q6_K || *qtype == QT_F8_E4M3
11258                    || *qtype == QT_NVFP4 || *qtype == QT_Q4_K || *qtype == QT_Q5_K || *qtype == QT_Q8_0);
11259            if m_ok {
11260                if let GpuTensor::Quant {
11261                    bytes,
11262                    qtype,
11263                    row_bytes,
11264                    rp,
11265                    rp4,
11266                    ..
11267                } = w
11268                {
11269                    if self.batched_supports(*qtype) && self.mmvq_supports(*qtype) {
11270                        let (bytes, rp) = match rp4 {
11271                            Some(m4) => (m4, true),
11272                            None => (bytes, *rp),
11273                        };
11274                        let mcols = Self::batched_mcols(m);
11275                        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
11276                        let mut y = self.qmatvec_mmvq_batched(
11277                            bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, mcols, 1.0, rp,
11278                        )?;
11279                        if let GpuTensor::Quant { scale, .. } = w {
11280                            if *scale != 1.0 {
11281                                self.scale_inplace(&mut y, *scale, m * out_f)?;
11282                            }
11283                        }
11284                        return Ok(y);
11285                    }
11286                }
11287            }
11288        }
11289        // F8-E4M3 (MEMRA_ST_E4M3) catch-all for the m<16 band the arms above didn't take (m=9..15,
11290        // the K=8 verify tier; or m=2..8 under MEMRA_NO_BATCHED/MEMRA_B8=0): grid.y=m e4m3 mmvq —
11291        // the SAME per-(token,row) program as the m=1 decode launch (bit-identical by construction),
11292        // weight re-read m times (rare tier; exactness over bandwidth here). There is no _dp4a twin
11293        // for this dtype, so the generic match below must never see it under `fast`.
11294        if fast {
11295            if let GpuTensor::Quant {
11296                bytes,
11297                qtype,
11298                row_bytes,
11299                scale,
11300                ..
11301            } = w
11302            {
11303                if *qtype == QT_F8_E4M3 {
11304                    let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
11305                    return self.qmatvec_mmvq(
11306                        bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, *scale, false,
11307                    );
11308                }
11309            }
11310        }
11311        let mut y = match w {
11312            GpuTensor::Quant {
11313                bytes,
11314                qtype,
11315                row_bytes,
11316                ..
11317            } if fast && *qtype == QT_Q8_0 => {
11318                self.qmatvec_q8_0_fast(bytes, x, m, in_f, out_f, *row_bytes)?
11319            }
11320            GpuTensor::Quant {
11321                bytes,
11322                qtype,
11323                row_bytes,
11324                ..
11325            } if fast && *qtype == QT_Q4_K => {
11326                self.qmatvec_q4_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
11327            }
11328            GpuTensor::Quant {
11329                bytes,
11330                qtype,
11331                row_bytes,
11332                ..
11333            } if fast && *qtype == QT_Q6_K => {
11334                self.qmatvec_q6_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
11335            }
11336            GpuTensor::Quant {
11337                bytes,
11338                qtype,
11339                row_bytes,
11340                ..
11341            } if fast && *qtype == QT_Q5_K => {
11342                self.qmatvec_q5_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
11343            }
11344            GpuTensor::Quant {
11345                bytes,
11346                qtype,
11347                row_bytes,
11348                ..
11349            } if fast && *qtype == QT_Q3_K => {
11350                self.qmatvec_q3_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
11351            }
11352            GpuTensor::Quant {
11353                bytes,
11354                qtype,
11355                row_bytes,
11356                rp,
11357                ..
11358            } if fast && *qtype == QT_NVFP4 => self.qmatvec_dp4a_named(
11359                if *rp {
11360                    "qmatvec_nvfp4_dp4a_rp"
11361                } else {
11362                    "qmatvec_nvfp4_dp4a"
11363                },
11364                &bytes.slice(0..bytes.len()),
11365                x,
11366                m,
11367                in_f,
11368                out_f,
11369                *row_bytes,
11370            )?,
11371            // IQ4_XS trunk fast path — DEFAULT ON since 2026-08-02 (MEMRA_IQ_FAST=0 reverts to
11372            // Stage-A; see iq_fast_enabled). The old opt-in default was the KAT-Coder decode
11373            // anomaly (research/kat-anomaly-20260802/).
11374            GpuTensor::Quant {
11375                bytes,
11376                qtype,
11377                row_bytes,
11378                ..
11379            } if fast && *qtype == QT_IQ4_XS && Self::iq_fast_enabled() => {
11380                self.qmatvec_iq4_XS_fast(bytes, x, m, in_f, out_f, *row_bytes)?
11381            }
11382            // B3: IQ3_S uses the Stage-A f32 dequant-in-kernel path. There is NO
11383            // qmatvec_iq3_s_dp4a kernel — do NOT add a `*qtype == QT_IQ3_S` fast guard here
11384            // without first writing the matching kernel, or func() will panic
11385            // "kernel ... not in any fatbin".
11386            GpuTensor::Quant {
11387                bytes,
11388                qtype,
11389                row_bytes,
11390                rp,
11391                ..
11392            } =>
11393            // Stage-A generic: repacked NVFP4 uses the device-side split-plane tag (the
11394            // deq(row,j) form cannot address the planes; same value/product order).
11395            {
11396                self.qmatvec(
11397                    bytes,
11398                    x,
11399                    m,
11400                    in_f,
11401                    out_f,
11402                    if *rp && *qtype == QT_NVFP4 {
11403                        QT_NVFP4_RP
11404                    } else {
11405                        *qtype
11406                    },
11407                    *row_bytes,
11408                )?
11409            }
11410            GpuTensor::Float { data, .. } => self.linear(x, data, m, in_f, out_f)?,
11411            // MEMRA_FULL_PREC bf16-resident weight: dequant-on-use to f32 scratch, then the same
11412            // cuBLASLt f32 GEMV as the Float arm.
11413            GpuTensor::FloatBf16 { data, .. } => {
11414                self.linear_bf16_chunked(x, data, m, in_f, out_f, false, None)?
11415            }
11416        };
11417        // NVFP4 per-tensor macro-scale (post-matmul). scale==1.0 for all other quants/float -> no-op.
11418        if let GpuTensor::Quant { scale, .. } = w {
11419            if *scale != 1.0 {
11420                self.scale_inplace(&mut y, *scale, m * out_f)?;
11421            }
11422        }
11423        Ok(y)
11424    }
11425
11426    /// True when `MEMRA_FAST=0`, i.e. the Stage-A f32 oracle is the requested arithmetic and every
11427    /// `matmul_pre` call will take the raw-f32 escape rather than the q8_1 pair.
11428    ///
11429    /// WHY THIS EXISTS AS ITS OWN PREDICATE: `uses_q8_1_fast` needs a weight, but the callers that
11430    /// have to DECIDE WHETHER TO MATERIALIZE an f32 activation sit one level above any weight — the
11431    /// gemma-4 decode trunk emits a q8_1 pair per layer and hands it down, so the f32 has to be
11432    /// produced (or not) by the loop that owns the residual. Cached in a OnceLock like every other
11433    /// `*_on()` flag here: this is read once per layer per token on the decode path, and the daily
11434    /// path must not pay an env lookup for a flag that is off.
11435    pub fn stage_a_raw_needed() -> bool {
11436        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11437        *ON.get_or_init(|| std::env::var("MEMRA_FAST").as_deref() == Ok("0"))
11438    }
11439
11440    /// True if `w` would take the int8-dp4a fast path under MEMRA_FAST (so its activation can be
11441    /// pre-quantized once and shared across sibling matmuls via `matmul_pre`).
11442    pub fn uses_q8_1_fast(&self, w: &crate::model::GpuTensor) -> bool {
11443        use crate::model::GpuTensor;
11444        if std::env::var("MEMRA_FAST").as_deref() == Ok("0") {
11445            return false;
11446        }
11447        match w {
11448            // QT_F8_E4M3_BLK is admitted for the same reason QT_F8_E4M3 is: its ONLY kernel class
11449            // takes the shared q8_1 activation, so callers may pre-quantize once and share it
11450            // across siblings. It is NOT admitted to any of the fused/dual epilogue doors those
11451            // siblings can then open (`q8_fused_params`, `e4m3_fused_params` and
11452            // `matmul_pre_dual_noscale` all match on their own qtype and refuse this one) — the
11453            // block class has no fused twin yet, so each of its projections takes its own launch.
11454            GpuTensor::Quant { qtype, .. } => {
11455                matches!(
11456                    *qtype,
11457                    QT_Q8_0
11458                        | QT_Q4_K
11459                        | QT_Q6_K
11460                        | QT_Q5_K
11461                        | QT_Q3_K
11462                        | QT_NVFP4
11463                        | QT_F8_E4M3
11464                        | QT_F8_E4M3_BLK
11465                        | QT_Q4_0
11466                ) || (*qtype == QT_IQ4_XS && Self::iq_fast_enabled())
11467            }
11468            GpuTensor::Float { .. } | GpuTensor::FloatBf16 { .. } => false,
11469        }
11470    }
11471
11472    /// matmul with a PRE-QUANTIZED q8_1 activation (aq,ad from `quantize_q8_1`). Skips the
11473    /// per-matmul re-quantize so sibling matmuls that share an input (gate+up share `z`;
11474    /// q/k/v + wqkv/gate/beta/alpha share `h`) quantize ONCE. Caller MUST have checked
11475    /// `uses_q8_1_fast(w)`; falls back to plain `matmul` otherwise (Stage-A / Float / non-fast).
11476    pub fn matmul_pre(
11477        &self,
11478        w: &crate::model::GpuTensor,
11479        aq: &CudaSlice<i8>,
11480        ad: &CudaSlice<f32>,
11481        x_fallback: &CudaSlice<f32>,
11482        m: usize,
11483    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11484        use crate::model::GpuTensor;
11485        // Every raw-f32 arm below (fp8/f16/MMQ/fp4) reads m*in_f from x_fallback. Callers that
11486        // pre-quantized and dropped the f32 input pass an EMPTY x_fallback (E4B's fusion port:
11487        // h = zeros(0)) — the length guard keeps those on the aq/ad GEMM instead of feeding a
11488        // 0-byte buffer to a convert kernel (illegal address -> cublasLt status 13; the E4B
11489        // rc=30013 dig, 2026-07-31).
11490        let x_raw_ok = x_fallback.len() >= m * w.in_features();
11491        // FP8-ACT PREFILL (MEMRA_PP_FP8=1): same arm as `matmul` — the fp8 operand needs the RAW
11492        // f32 activation (per-batch e4m3 quant differs from q8_1), so x_fallback not aq/ad.
11493        if m >= 16 && x_raw_ok && !self.verify_exact_on() {
11494            if let Some(y) = self.try_fp8_gemm(w, x_fallback, m)? {
11495                return Ok(y);
11496            }
11497            // PER-BLOCK FP8 MMQ — same arm as `matmul` (stash opt-in, native-resident default ON);
11498            // its own quantizer wants the RAW f32 activation, so x_fallback not aq/ad.
11499            if let Some(y) = self.try_fp8_blk_mmq(w, x_fallback, m)? {
11500                return Ok(y);
11501            }
11502            // FP16-mirror prefill (same arm as `matmul` — fp16 wants the RAW f32 activation).
11503            if let Some(y) = self.try_f16_gemm(w, x_fallback, m)? {
11504                return Ok(y);
11505            }
11506        }
11507        // BLOCK-128 e4m3 (QT_F8_E4M3_BLK) — the same two arms as `matmul`, split at the same m, and
11508        // placed at the same point in the order (after the prefill GEMM hooks, before every arm
11509        // that refuses this qtype). The prefill arm needs the RAW f32 activation for the Q8_0
11510        // dispatch it recurses into, so it takes x_fallback and is skipped when that is empty
11511        // (a pre-quantized caller that dropped its f32 input never runs at prefill m anyway).
11512        if m >= 16 && x_raw_ok && !self.verify_exact_on() {
11513            if let Some(y) = self.try_e4m3_blk_prefill(w, x_fallback, m)? {
11514                return Ok(y);
11515            }
11516        }
11517        if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
11518            return Ok(y);
11519        }
11520        // VENDORED llama MMQ prefill GEMMs (NVFP4 W4A8 default-on; W4A4/k-quant behind MEMRA_MMQ=1
11521        // — policy in mmq_supports) — use the RAW f32 activation (their own internal quant:
11522        // q8_1 D4 for NVFP4 W4A8, FP8/UE4M3 for W4A4, q8_1 DS4 for Q4_K/Q5_K), so x_fallback not
11523        // aq/ad.
11524        if m >= 16
11525            && w.out_features() >= 128
11526            && self.mmq_supports(w)
11527            && !self.verify_exact_on()
11528            && x_raw_ok
11529        {
11530            return self.qmatvec_mmq(w, x_fallback, m);
11531        }
11532        // Stage-C FP4 prefill (MEMRA_FP4): native mxf4 GEMM needs the f32 activation (FP4-quant differs
11533        // from q8_1), so re-quantize from x_fallback rather than reuse aq/ad. NVFP4 only, m>=16.
11534        if m >= 16 && x_raw_ok && !self.verify_exact_on() {
11535            if let Some(y) =
11536                self.try_fp4_gemm(w, x_fallback, m, w.in_features(), w.out_features())?
11537            {
11538                return Ok(y);
11539            }
11540        }
11541        // Prefill GEMM root fix: if T>1 and the dtype has a GEMM kernel, batch via tensor cores
11542        // (reuses the already-quantized aq/ad — no extra quantize). m=1 falls through to dp4a.
11543        if m >= 16 && self.gemm_supports(w) && !self.verify_exact_on() {
11544            return self.qmatvec_gemm(w, aq, ad, m);
11545        }
11546        // THE SEVENTH ARM. Every raw-f32 arm above is guarded by `x_raw_ok`; this one — the
11547        // Stage-A / Float escape — was not, and it is the ONLY one `MEMRA_FAST=0` opens. So the
11548        // 2026-07-31 E4B rc=30013 fix (add the length guard, keep empty-fallback callers off the
11549        // raw-f32 arms) protected the six arms the FAST path can reach and missed the one the
11550        // ORACLE path reaches. Consequence before this guard: the gemma-4 decode arms, which pass
11551        // `e.zeros(0)` because their trunk emits only a q8_1 pair and never materializes an f32
11552        // attn-normed activation, fell through to Stage-A `qmatvec_f32` (cu/qmatvec.cu:5442),
11553        // which reads `m * in_f` floats out of a 0-byte allocation ->
11554        // CUDA_ERROR_ILLEGAL_ADDRESS at layer 0 of the first decode token. The fault is STICKY:
11555        // it poisons the context, so every LATER request in that process fails with an unrelated
11556        // message ("cache alloc failed: ...") and the true cause appears exactly once, in the
11557        // first failure. That is what made `MEMRA_FAST=0` — the reference named in
11558        // `tools/argmax-margin-gate.sh`'s own header — return an opaque HTTP 500 on the gemma
11559        // dense artifact and left the arm with no working truth instrument.
11560        //
11561        // Refuse loudly instead of reading out of bounds. A named error at the true call site is
11562        // strictly better than an illegal address surfacing later at an unrelated sync point, and
11563        // an oracle that cannot run must say so rather than corrupt the context it runs in.
11564        if !self.uses_q8_1_fast(w) {
11565            if !x_raw_ok {
11566                return Err(format!(
11567                    "matmul_pre: q8_1-fast is off for this weight but x_fallback holds {} f32 \
11568                     (need m*in_f = {}*{} = {}). This call site pre-quantized its activation and \
11569                     dropped the f32, so there is nothing to fall back to — pass the real f32 \
11570                     activation (see Engine::rms_norm_decode, which is bit-identical to \
11571                     rms_norm_q8_1's reduction) or keep the weight on the q8_1 path.",
11572                    x_fallback.len(),
11573                    m,
11574                    w.in_features(),
11575                    m * w.in_features()
11576                )
11577                .into());
11578            }
11579            return self.matmul(w, x_fallback, m);
11580        }
11581        let in_f = w.in_features();
11582        let out_f = w.out_features();
11583        let (bytes, qtype, row_bytes, scale, rp) = match w {
11584            GpuTensor::Quant {
11585                bytes,
11586                qtype,
11587                row_bytes,
11588                scale,
11589                rp,
11590                ..
11591            } => (bytes, *qtype, *row_bytes, *scale, *rp),
11592            _ => unreachable!("uses_q8_1_fast guaranteed Quant"),
11593        };
11594        // Q4_0 split-plane mirror: only the mmvq/batched decode arms read it (the _rp twins);
11595        // the dp4a/oracle tails below keep the raw GGUF bytes.
11596        let (mbytes, mrp) = match w {
11597            GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
11598            _ => (bytes, rp),
11599        };
11600        // PERF-3 decode-GEMV: warp-per-row MMVQ for the m=1 decode arm, gated behind MEMRA_MMVQ.
11601        // Only the 4 daily-hot dtypes have an _mmvq kernel (Q8_0/Q4_K/Q6_K/NVFP4); Q5_K/Q3_K/IQ4_XS
11602        // keep _dp4a (the oracle/fallback). Bit-equivalent to _dp4a up to f32 reduction order.
11603        if m == 1 && self.mmvq_supports(qtype) {
11604            return self.qmatvec_mmvq(mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, mrp);
11605        }
11606        // BATCHED weight-resident matvec for the m=2-4 band (the MTP/verify forward: full_attn_verify
11607        // and decode_step_t run their projections at m=T=k=2..4). The plain _dp4a path below launches
11608        // grid.y=m INDEPENDENT blocks per output row -> the weight row is re-read m times from HBM/L2.
11609        // The _b2/_b4 kernels walk the weight ONCE and dp4a vs all m activation columns, so m tokens
11610        // cost ~1 weight read instead of m (decode is weight-BW-bound). BIT-IDENTICAL per (token,row)
11611        // to the _mmvq path (32-thread warp reduce — NOT the dp4a 128-thread reduce below).
11612        // m=2 -> mcols=2; m∈{3,4} -> mcols=4; m∈{5..8} -> mcols=8 (kernel guards c>=m).
11613        // MEMRA_NO_BATCHED forces the per-m grid.y=m path (the A/B reference); MEMRA_B8=0 keeps
11614        // m=5..8 on the old per-m path (b8-tier-only seam).
11615        // DECODE-PARITY GATE (2026-07-07): batched iff mmvq_supports — see matmul's parity note.
11616        // Without MEMRA_MMVQ, m=1 decode rides dp4a (the arm below at m=1); the verify must ride
11617        // the SAME class per column (grid.y=m dp4a = the exact m=1 dp4a program per column).
11618        if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
11619            && std::env::var("MEMRA_NO_BATCHED").is_err()
11620            && (m <= 4 || Self::b8_enabled())
11621            // b16 tier: every class routed here now has base + _rp b16 kernels (Q4_0/Q6_K
11622            // pre-existing; NVFP4/Q4_K/Q8_0-base/F8_E4M3 added lane/rp-on-st 2026-08-06), so
11623            // there is no mirror precondition left — `mrp` still selects the LAYOUT below.
11624            && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || qtype == QT_NVFP4
11625                || qtype == QT_Q4_K || qtype == QT_Q5_K || qtype == QT_F8_E4M3 || qtype == QT_Q8_0)
11626        {
11627            let mcols = Self::batched_mcols(m);
11628            return self.qmatvec_mmvq_batched(
11629                mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, mrp,
11630            );
11631        }
11632        // F8-E4M3 catch-all (m=9..15 / batched-disabled seams): grid.y=m e4m3 mmvq — this dtype
11633        // has NO _dp4a twin, and per (token,row) the mmvq body is the exact m=1 decode program.
11634        // Q4_0 joins the catch-all (2026-07-11): adaptive-K cap 8 makes verify t=9 reachable
11635        // for the first time (past the b8 tier) and Q4_0 has no dp4a twin either. The mirror
11636        // (mbytes/mrp) keeps the rp layout consistent with the m=1 decode program.
11637        if qtype == QT_F8_E4M3 || qtype == QT_Q4_0 {
11638            let (b2, r2) = if qtype == QT_Q4_0 {
11639                (mbytes, mrp)
11640            } else {
11641                (bytes, rp)
11642            };
11643            return self.qmatvec_mmvq(b2, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, r2);
11644        }
11645        let name = match qtype {
11646            QT_Q8_0 => "qmatvec_q8_0_dp4a",
11647            QT_Q4_K => "qmatvec_q4_K_dp4a",
11648            QT_Q6_K => "qmatvec_q6_K_dp4a",
11649            QT_Q5_K => "qmatvec_q5_K_dp4a",
11650            QT_Q3_K => "qmatvec_q3_K_dp4a",
11651            QT_NVFP4 => {
11652                if rp {
11653                    "qmatvec_nvfp4_dp4a_rp"
11654                } else {
11655                    "qmatvec_nvfp4_dp4a"
11656                }
11657            }
11658            QT_IQ4_XS => "qmatvec_iq4_XS_dp4a",
11659            _ => unreachable!(),
11660        };
11661        let f = self.func(name);
11662        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output: skip memset
11663        let cfg = LaunchConfig {
11664            grid_dim: (out_f as u32, m as u32, 1),
11665            block_dim: (128, 1, 1),
11666            shared_mem_bytes: 0,
11667        };
11668        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
11669        let __s_b = self.gpu.stream();
11670        let mut b = __s_b.launch_builder(&f);
11671        b.arg(bytes)
11672            .arg(aq)
11673            .arg(ad)
11674            .arg(&mut y)
11675            .arg(&inf)
11676            .arg(&outf)
11677            .arg(&mi)
11678            .arg(&rb);
11679        unsafe {
11680            b.launch(cfg)?;
11681        }
11682        if scale != 1.0 {
11683            self.scale_inplace(&mut y, scale, m * out_f)?;
11684        }
11685        Ok(y)
11686    }
11687
11688    /// DECODE-EXACT matmul at any m: guarantees the SAME warp-per-row (MMVQ, 32-thread) FP
11689    /// accumulation order as the T=1 decode path for EVERY token row. The spec-decode verify MUST
11690    /// use this for linear-attn projections to be bit-identical to greedy decode. The dp4a kernel
11691    /// (128 threads, two-level reduction) used by `matmul`/`matmul_pre` at m>=5 has a different
11692    /// shfl-tree shape that produces ULP differences propagating through gdn_scan into argmax flips.
11693    /// The MMVQ kernel with grid.y=m already processes each row independently (same 32-thread warp
11694    /// reduce as m=1); this method just forces that path unconditionally.
11695    pub fn matmul_decode_exact(
11696        &self,
11697        w: &crate::model::GpuTensor,
11698        x: &CudaSlice<f32>,
11699        m: usize,
11700    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11701        use crate::model::GpuTensor;
11702        // FLOAT tensors (35B ssm_beta/ssm_alpha on every linear layer, F32 ne=[2048,32]): the
11703        // generic path is cuBLASLt, whose reduction splits are n-DEPENDENT — m=1 vs m=2 col-0
11704        // outputs differ in every bit (probe 2026-07-06: 32/32 bit-diff, maxdiff 3.5e-3), which
11705        // shifted 35B verify logits 0.26-0.56 vs eager and flipped greedy at tight margins (the
11706        // p3 spec FAIL). Decode-exact contract: per-COLUMN m=1 cuBLASLt calls — each column's
11707        // reduction is the exact kernel the T=1 decode path runs, so verify==decode bit-for-bit.
11708        // m<=10 here (K+2 verify tier), so the extra launches are a handful of 4us gemvs.
11709        if let GpuTensor::Float { data, .. } = w {
11710            return self.linear_decode_exact(x, data, m, w.in_features(), w.out_features());
11711        }
11712        // MEMRA_FULL_PREC bf16-resident weight: dequant-on-use, then the per-column decode-exact
11713        // float linear (same n-independent reduction contract as the Float arm above).
11714        if let GpuTensor::FloatBf16 { data, .. } = w {
11715            let (in_f, out_f) = (w.in_features(), w.out_features());
11716            return self.linear_bf16_chunked(x, data, m, in_f, out_f, true, None);
11717        }
11718        if !self.uses_q8_1_fast(w) {
11719            return self.matmul(w, x, m);
11720        }
11721        let in_f = w.in_features();
11722        let out_f = w.out_features();
11723        let (bytes, qtype, row_bytes, scale, rp) = match w {
11724            GpuTensor::Quant {
11725                bytes,
11726                qtype,
11727                row_bytes,
11728                scale,
11729                rp,
11730                ..
11731            } => (bytes, *qtype, *row_bytes, *scale, *rp),
11732            _ => return self.matmul(w, x, m),
11733        };
11734        // Q4_0 split-plane mirror for the mmvq/batched arms below (dp4a tail = matmul_pre,
11735        // which does its own mirror pick).
11736        let (bytes, rp) = match w {
11737            GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
11738            _ => (bytes, rp),
11739        };
11740        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
11741        // BLOCK-128 e4m3 (QT_F8_E4M3_BLK): the same single kernel every other entry dispatches, so
11742        // the decode-exact contract needs nothing special — grid.y=m runs the m=1 program per
11743        // (token,row) by construction, which is exactly what this method exists to guarantee.
11744        if let Some(y) = self.try_e4m3_blk_pre(w, &aq, &ad, m)? {
11745            return Ok(y);
11746        }
11747        // Batched weight-resident matvec for m=2-8: BIT-IDENTICAL per (token,row) to MMVQ (exact
11748        // integer dp4a, same warp reduce — kernel-check gate rel=0.00e0), one weight read for m
11749        // tokens. The dispatch the divergence fix must avoid is dp4a's 128-thread two-level
11750        // reduce, NOT this. m=5..8 is the K=4..7 spec-verify tier (b8): pre-b8 T=5 fell to the
11751        // grid.y=m per-row MMVQ below = 5 full weight reads/launch — the measured 27B K=4 cliff.
11752        // DECODE-PARITY GATE (2026-07-07): batched (MMVQ-class order) only when the m=1 decode
11753        // chain rides MMVQ too — without MEMRA_MMVQ decode is dp4a, so the exact-contract here
11754        // must be per-column dp4a (matmul_pre fallthrough), not the MMVQ order.
11755        if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
11756            && std::env::var("MEMRA_NO_BATCHED").is_err()
11757            && (m <= 4 || Self::b8_enabled())
11758            // Every b16 class has base + _rp twins after lane/rp-on-st (see matmul_pre's note):
11759            // no mirror precondition, `rp` selects the layout only.
11760            && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || qtype == QT_F8_E4M3
11761                || qtype == QT_NVFP4 || qtype == QT_Q4_K || qtype == QT_Q5_K || qtype == QT_Q8_0)
11762        {
11763            let mcols = Self::batched_mcols(m);
11764            return self.qmatvec_mmvq_batched(
11765                bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, rp,
11766            );
11767        }
11768        if self.mmvq_supports(qtype) {
11769            // MMVQ at grid.y=m: each row is processed by its own warp independently — same 32-thread
11770            // accumulation + warp_reduce_sum as m=1 decode. Bit-identical per row.
11771            return self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, scale, rp);
11772        }
11773        // Fallback for non-MMVQ quant types (Q5_K, Q3_K): use dp4a (the only available kernel).
11774        // These types are not used in the 27B's linear-attn NVFP4+Q4_K layers.
11775        self.matmul_pre(w, &aq, &ad, x, m)
11776    }
11777
11778    /// DECODE-EXACT matmul from a PRE-QUANTIZED q8_1 activation (batched-verify epilogue
11779    /// re-fuse, lane/vt-fixes fix 2, 2026-08-03): the EXACT `matmul_decode_exact` dispatch for
11780    /// q8_1-fast Quant tensors, with the caller's (aq, ad) replacing the internal
11781    /// `quantize_q8_1`. quantize_q8_1 is deterministic (same input bytes -> same q8 bytes), so
11782    /// sharing one quantize across sibling matmuls of the same activation — or consuming the
11783    /// q8 emitted by a fused epilogue (rms_norm_q8_1 / add_rms_norm_q8_1 /
11784    /// silu_mul_scaled_q8_1 / gated_rmsnorm_q8_1, all kernel-check-pinned bit-identical to
11785    /// their unfused chains) — cannot change any dispatched kernel's input bytes.
11786    /// Caller MUST guarantee `uses_q8_1_fast(w)` (the fused epilogues only exist on that path).
11787    pub fn matmul_decode_exact_pre(
11788        &self,
11789        w: &crate::model::GpuTensor,
11790        aq: &CudaSlice<i8>,
11791        ad: &CudaSlice<f32>,
11792        m: usize,
11793    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11794        use crate::model::GpuTensor;
11795        debug_assert!(
11796            self.uses_q8_1_fast(w),
11797            "matmul_decode_exact_pre: caller must guarantee q8_1-fast"
11798        );
11799        // BLOCK-128 e4m3: same single kernel, all m — see matmul_decode_exact's note.
11800        if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
11801            return Ok(y);
11802        }
11803        let in_f = w.in_features();
11804        let out_f = w.out_features();
11805        let (bytes, qtype, row_bytes, scale, rp) = match w {
11806            GpuTensor::Quant {
11807                bytes,
11808                qtype,
11809                row_bytes,
11810                scale,
11811                rp,
11812                ..
11813            } => (bytes, *qtype, *row_bytes, *scale, *rp),
11814            _ => {
11815                return Err(
11816                    "matmul_decode_exact_pre: Quant tensor required (q8_1-fast contract)".into(),
11817                );
11818            }
11819        };
11820        // Q4_0 split-plane mirror — same pick as matmul_decode_exact.
11821        let (bytes, rp) = match w {
11822            GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
11823            _ => (bytes, rp),
11824        };
11825        // Dispatch mirror of matmul_decode_exact's q8_1-fast tail, condition for condition.
11826        if (2..=16).contains(&m)
11827            && self.batched_supports(qtype)
11828            && self.mmvq_supports(qtype)
11829            && std::env::var("MEMRA_NO_BATCHED").is_err()
11830            && (m <= 4 || Self::b8_enabled())
11831            && (m <= 8
11832                || qtype == QT_Q4_0
11833                || qtype == QT_Q6_K
11834                || qtype == QT_F8_E4M3
11835                || qtype == QT_NVFP4
11836                || qtype == QT_Q4_K
11837                || qtype == QT_Q5_K
11838                || qtype == QT_Q8_0)
11839        {
11840            let mcols = Self::batched_mcols(m);
11841            return self.qmatvec_mmvq_batched(
11842                bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, rp,
11843            );
11844        }
11845        if self.mmvq_supports(qtype) {
11846            return self.qmatvec_mmvq(bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, rp);
11847        }
11848        // Non-MMVQ quant types (Q5_K/Q3_K under MEMRA_MMVQ=0): dp4a via matmul_pre — the same
11849        // fallback matmul_decode_exact takes. m <= 16 on the verify tier never reads x_fallback.
11850        let x0 = self.zeros(0)?;
11851        self.matmul_pre(w, aq, ad, &x0, m)
11852    }
11853
11854    /// DUAL gate+up batched matvec from a PRE-QUANTIZED activation, macro-scales DEFERRED
11855    /// (lane/vt-fixes fix 2): same eligibility as `matmul_decode_exact_dual`, but the caller's
11856    /// (aq, ad) replaces the internal quantize and the NVFP4 per-tensor scales are RETURNED
11857    /// instead of applied via two `scale_inplace` launches — the fused SwiGLU epilogue
11858    /// (`silu_mul_scaled_q8_1`) folds them, exactly like the m=1 decode chain does. Deferring
11859    /// is value-exact: `y[i]*s` inline in the epilogue is the same IEEE multiply scale_inplace
11860    /// would store (f32 store/load round-trips are exact). None -> caller falls back to the
11861    /// per-tensor path.
11862    pub fn matmul_decode_exact_dual_pre(
11863        &self,
11864        w0: &crate::model::GpuTensor,
11865        w1: &crate::model::GpuTensor,
11866        aq: &CudaSlice<i8>,
11867        ad: &CudaSlice<f32>,
11868        m: usize,
11869    ) -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn std::error::Error>>
11870    {
11871        use crate::model::GpuTensor;
11872        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11873        let on = *ON.get_or_init(|| {
11874            std::env::var("MEMRA_SPEC_DUAL_T")
11875                .map(|v| v != "0")
11876                .unwrap_or(true)
11877        });
11878        if !on
11879            || !(2..=7).contains(&m)
11880            || std::env::var("MEMRA_NO_BATCHED").is_ok()
11881            || !self.uses_q8_1_fast(w0)
11882            || !self.uses_q8_1_fast(w1)
11883        {
11884            return Ok(None);
11885        }
11886        // DECODE-PARITY GATE (lane/nvfp4-strict, 2026-08-05): batched iff MMVQ — the dual
11887        // kernels are the MMVQ warp-reduce family, and without MEMRA_MMVQ the m=1 decode
11888        // chain this verify must match bit-for-bit rides dp4a (see matmul_decode_exact's
11889        // note). The singles enforce this via `mmvq_supports`; the dual door skipped it.
11890        if !self.mmvq_supports(QT_NVFP4) {
11891            return Ok(None);
11892        }
11893        let (in_f, out_f) = (w0.in_features(), w0.out_features());
11894        if w1.in_features() != in_f || w1.out_features() != out_f {
11895            return Ok(None);
11896        }
11897        let (b0, b1, row_bytes, s0, s1, rp) = match (w0, w1) {
11898            (
11899                GpuTensor::Quant {
11900                    bytes: b0,
11901                    qtype: q0,
11902                    row_bytes: rb0,
11903                    scale: s0,
11904                    rp: rp0,
11905                    rp4: None,
11906                    ..
11907                },
11908                GpuTensor::Quant {
11909                    bytes: b1,
11910                    qtype: q1,
11911                    row_bytes: rb1,
11912                    scale: s1,
11913                    rp: rp1,
11914                    rp4: None,
11915                    ..
11916                },
11917            ) if *q0 == QT_NVFP4 && *q1 == QT_NVFP4 && rb0 == rb1 && rp0 == rp1 => {
11918                (b0, b1, *rb0, *s0, *s1, *rp0)
11919            }
11920            _ => return Ok(None),
11921        };
11922        // m=5..7: only the exact-width rp duals exist (vt-fixes fix 1b); GGUF layout keeps
11923        // the singles. The b8 dual (MCOLS=8 at m=5..8) measured FLAT and stays dead.
11924        if m > 4 && !(rp && Self::b8_enabled() && std::env::var("MEMRA_B567").as_deref() != Ok("0"))
11925        {
11926            return Ok(None);
11927        }
11928        let (y0, y1) =
11929            self.qmatvec_batched_dual_raw(b0, b1, aq, ad, m, in_f, out_f, row_bytes, rp)?;
11930        Ok(Some(((y0, s0), (y1, s1))))
11931    }
11932
11933    /// GROUP-4 GDN-tuple BATCHED matvec (trunk-kernels slice C): the qwen35 linear-layer
11934    /// in-projection 4-tuple (wqkv / wqkv_gate / ssm_beta / ssm_alpha) from ONE pre-quantized
11935    /// activation in ONE launch. Blocks map to the concatenated row space; every out_f must be
11936    /// a multiple of 8 (rows_per_block) so each warp's row pair resolves to one tensor; per
11937    /// (tensor, token, row) the kernel body is `nvfp4_mmvq_batched_rp` VERBATIM with the
11938    /// tensor's macro-scale fused at the write (== the conditional scale_inplace pass,
11939    /// bit-identical) -> BIT-IDENTICAL to the four single launches. Split-plane rp NVFP4 only,
11940    /// m=2..16 (exact-width MCOLS at m=5..7 mirroring the B567 law; m>4 requires b8_enabled
11941    /// like the singles; m=9..=16 rides the b16 form — the E4 width lift that lets
11942    /// `matmul_nvfp4_fused3/4` delegate that class here, lane/orndecode2).
11943    /// None -> caller runs the four singles. MEMRA_TK_GDN_GROUP=0 rollback.
11944    pub fn matmul_decode_exact_group4_pre(
11945        &self,
11946        ws: [&crate::model::GpuTensor; 4],
11947        aq: &CudaSlice<i8>,
11948        ad: &CudaSlice<f32>,
11949        m: usize,
11950    ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
11951        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11952        let on = *ON.get_or_init(|| {
11953            std::env::var("MEMRA_TK_GDN_GROUP")
11954                .map(|v| v != "0")
11955                .unwrap_or(true)
11956        });
11957        self.matmul_decode_exact_group_pre(&ws, aq, ad, m, on, "GDN group4")
11958    }
11959
11960    /// GROUP-3 twin for the qwen35 full-attention q/k/v triple (trunk-kernels slice D):
11961    /// the SAME group4 kernels with n3=0 (blocks never reach the fourth range; W3/y3 are
11962    /// never dereferenced) — per (tensor, token, row) bit-identical to the three singles
11963    /// exactly as the group4 door is to its four. MEMRA_TK_FA_GROUP=0 rollback.
11964    pub fn matmul_decode_exact_group3_pre(
11965        &self,
11966        ws: [&crate::model::GpuTensor; 3],
11967        aq: &CudaSlice<i8>,
11968        ad: &CudaSlice<f32>,
11969        m: usize,
11970    ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
11971        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11972        let on = *ON.get_or_init(|| {
11973            std::env::var("MEMRA_TK_FA_GROUP")
11974                .map(|v| v != "0")
11975                .unwrap_or(true)
11976        });
11977        self.matmul_decode_exact_group_pre(&ws, aq, ad, m, on, "FA group3")
11978    }
11979
11980    /// Shared core of the group3/group4 doors: eligibility mirror of the singles' batched
11981    /// dispatch, then ONE `qmatvec_nvfp4_mmvq_group4_b*_rp` launch over the concatenated
11982    /// row space (3-tensor callers ride n3=0). Returns one output per input tensor.
11983    fn matmul_decode_exact_group_pre(
11984        &self,
11985        ws: &[&crate::model::GpuTensor],
11986        aq: &CudaSlice<i8>,
11987        ad: &CudaSlice<f32>,
11988        m: usize,
11989        on: bool,
11990        tag: &'static str,
11991    ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
11992        use crate::model::GpuTensor;
11993        if !on
11994            || !(2..=16).contains(&m)
11995            || std::env::var("MEMRA_NO_BATCHED").is_ok()
11996            || (m > 4 && !Self::b8_enabled())
11997            || !self.mmvq_supports(QT_NVFP4)
11998            || !self.batched_supports(QT_NVFP4)
11999        {
12000            return Ok(None);
12001        }
12002        let in_f = ws[0].in_features();
12003        let mut parts: Vec<(&CudaSlice<u8>, usize, f32)> = Vec::with_capacity(4);
12004        for w in ws {
12005            if !self.uses_q8_1_fast(w) || w.in_features() != in_f {
12006                return Ok(None);
12007            }
12008            match w {
12009                GpuTensor::Quant {
12010                    bytes,
12011                    qtype,
12012                    scale,
12013                    rp: true,
12014                    rp4: None,
12015                    ..
12016                } if *qtype == QT_NVFP4 && w.out_features() % 8 == 0 => {
12017                    parts.push((bytes, w.out_features(), *scale));
12018                }
12019                _ => return Ok(None),
12020            }
12021        }
12022        // MCOLS tier mirrors the singles: batched_mcols + the B567 exact-width law at m=5..7.
12023        static B567: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12024        let b567 = *B567.get_or_init(|| std::env::var("MEMRA_B567").as_deref() != Ok("0"));
12025        let mcols = if (5..=7).contains(&m) && b567 {
12026            m
12027        } else {
12028            Self::batched_mcols(m)
12029        };
12030        let kname: &'static str = match mcols {
12031            2 => "qmatvec_nvfp4_mmvq_group4_b2_rp",
12032            4 => "qmatvec_nvfp4_mmvq_group4_b4_rp",
12033            5 => "qmatvec_nvfp4_mmvq_group4_b5_rp",
12034            6 => "qmatvec_nvfp4_mmvq_group4_b6_rp",
12035            7 => "qmatvec_nvfp4_mmvq_group4_b7_rp",
12036            8 => "qmatvec_nvfp4_mmvq_group4_b8_rp",
12037            16 => "qmatvec_nvfp4_mmvq_group4_b16_rp",
12038            _ => return Ok(None),
12039        };
12040        // Engagement receipt PER DOOR (dead-arm lesson): one shared Once here suppressed
12041        // the second door's print on the slice-D battery — key the once-set by tag.
12042        if std::env::var("MEMRA_DEBUG").is_ok() {
12043            use std::sync::Mutex;
12044            static SEEN: Mutex<Vec<&'static str>> = Mutex::new(Vec::new());
12045            let mut seen = SEEN.lock().unwrap();
12046            if !seen.contains(&tag) {
12047                seen.push(tag);
12048                eprintln!("[memra] {tag} batched ENGAGED (m={m})");
12049            }
12050        }
12051        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
12052        let rows_per_block = ROWS_PER_BLOCK * 2; // WROWS=2 in the group kernel
12053        let total: usize = parts.iter().map(|p| p.1).sum();
12054        let three = parts.len() == 3;
12055        let mut y0 = self.alloc_uninit::<f32>(m * parts[0].1)?;
12056        let mut y1 = self.alloc_uninit::<f32>(m * parts[1].1)?;
12057        let mut y2 = self.alloc_uninit::<f32>(m * parts[2].1)?;
12058        // 3-tensor callers: n3=0 means no block ever resolves to the fourth range — W3/y3
12059        // are never dereferenced; a 1-element dummy keeps the launch ABI without aliasing y0.
12060        let mut y3 = self.alloc_uninit::<f32>(if three { 1 } else { m * parts[3].1 })?;
12061        let cfg = LaunchConfig {
12062            grid_dim: ((total as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
12063            block_dim: (32, ROWS_PER_BLOCK, 1),
12064            shared_mem_bytes: 0,
12065        };
12066        let (inf, mi) = (in_f as i32, m as i32);
12067        let (n0, n1, n2) = (parts[0].1 as i32, parts[1].1 as i32, parts[2].1 as i32);
12068        let n3 = if three { 0i32 } else { parts[3].1 as i32 };
12069        let (s0, s1, s2) = (parts[0].2, parts[1].2, parts[2].2);
12070        let s3 = if three { 1.0f32 } else { parts[3].2 };
12071        let w3 = if three { parts[0].0 } else { parts[3].0 };
12072        let f = self.func(kname);
12073        let __s_b = self.gpu.stream();
12074        let mut b = __s_b.launch_builder(&f);
12075        b.arg(parts[0].0)
12076            .arg(parts[1].0)
12077            .arg(parts[2].0)
12078            .arg(w3)
12079            .arg(aq)
12080            .arg(ad)
12081            .arg(&mut y0)
12082            .arg(&mut y1)
12083            .arg(&mut y2)
12084            .arg(&mut y3)
12085            .arg(&inf)
12086            .arg(&n0)
12087            .arg(&n1)
12088            .arg(&n2)
12089            .arg(&n3)
12090            .arg(&mi)
12091            .arg(&s0)
12092            .arg(&s1)
12093            .arg(&s2)
12094            .arg(&s3);
12095        unsafe {
12096            b.launch(cfg)?;
12097        }
12098        Ok(Some(if three {
12099            vec![y0, y1, y2]
12100        } else {
12101            vec![y0, y1, y2, y3]
12102        }))
12103    }
12104
12105    /// DUAL gate+up BATCHED matvec at verify t=2..8 (lane/verify-economics, 2026-08-02): ONE
12106    /// launch computes both FFN projections of a verify batch — same activation, same shape,
12107    /// blockIdx.y selects the tensor. Per (tensor, token, row) the kernel body is the single
12108    /// batched program on the SAME layout (split-plane rp: b2 rp / b4 rpr2 / b8 rpr2; GGUF:
12109    /// b2 base / b4 r2 / b8 r2) -> BIT-IDENTICAL to the two single `matmul_decode_exact`
12110    /// launches (kernel-check gates bitwise on both layouts; run-spec K=1..8 arbitrates e2e).
12111    /// The one activation quantize replaces two IDENTICAL quantizes of the same `x` (same
12112    /// kernel, same input -> same q8_1 bytes), and the two independent weight streams in one
12113    /// grid restore the memory-level parallelism the two-launch form loses to tail drain +
12114    /// launch gap (m=1 dual_mr2 precedent: DRAM 40% -> 47-50% on the 27B pair).
12115    /// `Some((y0, y1))` only when both tensors are NVFP4, the SAME layout (both rp or both
12116    /// GGUF, no rp4 mirror), identical (in_f, out_f, row_bytes), q8_1-fast, and m in 2..=4
12117    /// (the b2/b4 tiers = verify T for K=1..3, the profitable-K window — the b8 dual measured
12118    /// FLAT vs the rpsc singles x3 interleaved, research/verify-economics-20260802, and was
12119    /// killed per doctrine). None -> caller runs the two singles. MEMRA_SPEC_DUAL_T=0 rollback.
12120    pub fn matmul_decode_exact_dual(
12121        &self,
12122        w0: &crate::model::GpuTensor,
12123        w1: &crate::model::GpuTensor,
12124        x: &CudaSlice<f32>,
12125        m: usize,
12126    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
12127        use crate::model::GpuTensor;
12128        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12129        let on = *ON.get_or_init(|| {
12130            std::env::var("MEMRA_SPEC_DUAL_T")
12131                .map(|v| v != "0")
12132                .unwrap_or(true)
12133        });
12134        if !on
12135            || !(2..=4).contains(&m)
12136            || std::env::var("MEMRA_NO_BATCHED").is_ok()
12137            || !self.uses_q8_1_fast(w0)
12138            || !self.uses_q8_1_fast(w1)
12139        {
12140            return Ok(None);
12141        }
12142        // DECODE-PARITY GATE (lane/nvfp4-strict, 2026-08-05): batched iff MMVQ — same law as
12143        // the singles' `batched_supports && mmvq_supports` check in matmul_decode_exact,
12144        // which this dual door bypassed. Without MEMRA_MMVQ the m=1 decode is dp4a; the
12145        // verify must ride the per-column dp4a class, not the MMVQ-family dual.
12146        if !self.mmvq_supports(QT_NVFP4) {
12147            return Ok(None);
12148        }
12149        let (in_f, out_f) = (w0.in_features(), w0.out_features());
12150        if w1.in_features() != in_f || w1.out_features() != out_f {
12151            return Ok(None);
12152        }
12153        let (b0, b1, row_bytes, s0, s1, rp) = match (w0, w1) {
12154            (
12155                GpuTensor::Quant {
12156                    bytes: b0,
12157                    qtype: q0,
12158                    row_bytes: rb0,
12159                    scale: s0,
12160                    rp: rp0,
12161                    rp4: None,
12162                    ..
12163                },
12164                GpuTensor::Quant {
12165                    bytes: b1,
12166                    qtype: q1,
12167                    row_bytes: rb1,
12168                    scale: s1,
12169                    rp: rp1,
12170                    rp4: None,
12171                    ..
12172                },
12173            ) if *q0 == QT_NVFP4 && *q1 == QT_NVFP4 && rb0 == rb1 && rp0 == rp1 => {
12174                (b0, b1, *rb0, *s0, *s1, *rp0)
12175            }
12176            _ => return Ok(None),
12177        };
12178        // Engagement receipt (MEMRA_DEBUG=1): the first dead-arm A/B lesson — a `rp: false`
12179        // gate silently no-op'd the whole experiment; prove the arm is live in the log.
12180        if std::env::var("MEMRA_DEBUG").is_ok() {
12181            static ONCE: std::sync::Once = std::sync::Once::new();
12182            ONCE.call_once(|| eprintln!("[memra] dual gate+up batched ENGAGED (m={m} rp={rp})"));
12183        }
12184        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
12185        let (y0, y1) =
12186            self.qmatvec_batched_dual_raw(b0, b1, &aq, &ad, m, in_f, out_f, row_bytes, rp)?;
12187        let mut y0 = y0;
12188        let mut y1 = y1;
12189        if s0 != 1.0 {
12190            self.scale_inplace(&mut y0, s0, m * out_f)?;
12191        }
12192        if s1 != 1.0 {
12193            self.scale_inplace(&mut y1, s1, m * out_f)?;
12194        }
12195        Ok(Some((y0, y1)))
12196    }
12197
12198    /// Launch body of the dual batched twins from raw NVFP4 weight bytes + a pre-quantized q8_1
12199    /// activation (kernel-check's bit-equivalence entry; matmul_decode_exact_dual's core).
12200    /// mcols tier = batched_mcols(m); macro-scale NOT applied. `rp` selects the split-plane
12201    /// twins (both buffers must be the repacked layout).
12202    #[allow(clippy::too_many_arguments)]
12203    pub fn qmatvec_batched_dual_raw(
12204        &self,
12205        b0: &CudaSlice<u8>,
12206        b1: &CudaSlice<u8>,
12207        aq: &CudaSlice<i8>,
12208        ad: &CudaSlice<f32>,
12209        m: usize,
12210        in_f: usize,
12211        out_f: usize,
12212        row_bytes: usize,
12213        rp: bool,
12214    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12215        const ROWS_PER_BLOCK: u32 = 4;
12216        let mcols = Self::batched_mcols(m);
12217        // EXACT-WIDTH duals at m=5..7 (vt-fixes fix 1b): rp-only; bit-identical to the two
12218        // b5/b6/b7 singles (blockIdx.y selects the tensor, same template body).
12219        let tiny_rp1 = rp
12220            && mcols == 4
12221            && out_f <= 128
12222            && std::env::var("MEMRA_NVFP4_AUX_DUAL").as_deref() != Ok("0");
12223        let (name, rows_per_block) = if tiny_rp1 {
12224            ("qmatvec_nvfp4_mmvq_dual_b4_rp", ROWS_PER_BLOCK)
12225        } else {
12226            match (mcols, rp, m) {
12227                (2, false, _) => ("qmatvec_nvfp4_mmvq_dual_b2", ROWS_PER_BLOCK),
12228                (4, false, _) => ("qmatvec_nvfp4_mmvq_dual_b4_r2", ROWS_PER_BLOCK * 2),
12229                (2, true, _) => ("qmatvec_nvfp4_mmvq_dual_b2_rp", ROWS_PER_BLOCK),
12230                (4, true, _) => ("qmatvec_nvfp4_mmvq_dual_b4_rpr2", ROWS_PER_BLOCK * 2),
12231                (8, true, 5) => ("qmatvec_nvfp4_mmvq_dual_b5_rpr2", ROWS_PER_BLOCK * 2),
12232                (8, true, 6) => ("qmatvec_nvfp4_mmvq_dual_b6_rpr2", ROWS_PER_BLOCK * 2),
12233                (8, true, 7) => ("qmatvec_nvfp4_mmvq_dual_b7_rpr2", ROWS_PER_BLOCK * 2),
12234                _ => {
12235                    return Err(
12236                        format!("qmatvec_batched_dual_raw: no dual kernel for m {m}").into(),
12237                    );
12238                }
12239            }
12240        };
12241        let f = self.func(name);
12242        let mut y0 = self.alloc_uninit::<f32>(m * out_f)?;
12243        let mut y1 = self.alloc_uninit::<f32>(m * out_f)?;
12244        let cfg = LaunchConfig {
12245            grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 2, 1),
12246            block_dim: (32, ROWS_PER_BLOCK, 1),
12247            shared_mem_bytes: 0,
12248        };
12249        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
12250        let __s_b = self.gpu.stream();
12251        let mut b = __s_b.launch_builder(&f);
12252        b.arg(b0)
12253            .arg(b1)
12254            .arg(aq)
12255            .arg(ad)
12256            .arg(&mut y0)
12257            .arg(&mut y1)
12258            .arg(&inf)
12259            .arg(&outf)
12260            .arg(&mi)
12261            .arg(&rb);
12262        unsafe {
12263            b.launch(cfg)?;
12264        }
12265        Ok((y0, y1))
12266    }
12267
12268    /// Like `matmul_pre` but RETURNS THE RAW (un-macro-scaled) matmul output together with the
12269    /// per-tensor NVFP4 scale, instead of applying `scale_inplace` internally. Used by the fused
12270    /// SwiGLU epilogue (RANK3 LEVER 2) so the gate/up scales fold into one `silu_mul_scaled` launch.
12271    /// `Some((y_raw, scale))` only on the m==1 decode fast path (mmvq / dp4a) where the scale is a
12272    /// separate post-launch op we can defer; returns `None` for every other path (prefill GEMM, FP4
12273    /// GEMM, Stage-A, Float) so the caller falls back to the scaled `matmul_pre` + `silu_mul`.
12274    /// DUAL gate+up NVFP4 matvec (mm-fusion): ONE launch computes both projections (same
12275    /// activation, same shape) — grid.y selects the tensor. Bit-identical per element to two
12276    /// mr2 launches at m=1. Returns (gate_raw, up_raw) un-scaled (caller folds the two macro
12277    /// scales into the SwiGLU epilogue, same as the matmul_pre_noscale contract). None unless
12278    /// both tensors are NVFP4 q8_1-fast with identical (in_f, out_f, row_bytes) and m==1.
12279    pub fn matmul_pre_dual_noscale(
12280        &self,
12281        w0: &crate::model::GpuTensor,
12282        w1: &crate::model::GpuTensor,
12283        aq: &CudaSlice<i8>,
12284        ad: &CudaSlice<f32>,
12285        m: usize,
12286    ) -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn std::error::Error>>
12287    {
12288        use crate::model::GpuTensor;
12289        if m != 1 || !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) {
12290            return Ok(None);
12291        }
12292        // FP-ORDER LAW (lane/nvfp4-strict, 2026-08-05): every kernel this door can dispatch
12293        // (q8_0 fused2, nvfp4 dual_mr2) is the MMVQ family — 32-thread warp reduce. Without
12294        // MEMRA_MMVQ the m=1 singles ride dp4a (128-thread two-level reduce), so fusing here
12295        // would mix dispatch families across the pair — the exact class `q8_fused_params`
12296        // already refuses for Q8_0. The NVFP4 arm lacked this check, which is why
12297        // decode-batch-gate `--mode strict`'s equalizing env (MEMRA_MMVQ=0) never pinned
12298        // NVFP4 models: decode_step_h kept riding dual_mr2 while the batched body fell to
12299        // dp4a (gate1 maxdiff 1.639e-1 / gate2 step-8 divergence at the 2026-08-05 train
12300        // HEAD, research/nvfp4-strict-20260805/). Default env (MMVQ on) is dispatch-unchanged.
12301        if !self.mmvq_supports(QT_NVFP4) {
12302            return Ok(None);
12303        }
12304        let (in_f, out_f) = (w0.in_features(), w0.out_features());
12305        if w1.in_features() != in_f || w1.out_features() != out_f {
12306            return Ok(None);
12307        }
12308        // Q8_0 ARM (lane/q27-deepdive, 2026-08-05): the dense-FFN gate+up pair on a Q8_0 trunk fell
12309        // through this NVFP4-only gate to two `matmul_pre_noscale` launches — measured 128 of the
12310        // 1015 launches/token on q27-Q8_0 decode, the single largest un-fused class in the tick
12311        // (nsys `research/q27-deepdive-20260805/nsys/`). `q8_fused2_core` already serves the same
12312        // pair shape for the shared-expert gate/up, and its kernel body is `qmatvec_q8_0_mmvq`
12313        // VERBATIM per (tensor,row) -> BIT-IDENTICAL to the two separate launches. Q8_0 carries no
12314        // macro-scale (q8_fused_params requires scale==1.0), so the noscale contract is satisfied
12315        // by returning 1.0 for both: the SwiGLU epilogue's fold becomes the identity it already is
12316        // on this dtype today. Seam: MEMRA_Q8_FFN_FUSE2=0 rolls back to the two-launch pair.
12317        // rp4 guard: with MEMRA_Q8RP the singles route to the `_rp` split-plane twin over the
12318        // mirror buffer; the fused2 kernel has no `_rp` form, so fusing there would swap
12319        // dispatch families mid-model. Bail and let the two singles run (mirror lane unchanged).
12320        let no_mirror =
12321            |w: &crate::model::GpuTensor| !matches!(w, GpuTensor::Quant { rp4: Some(_), .. });
12322        if self.q8_ffn_fuse2_on()
12323            && no_mirror(w0)
12324            && no_mirror(w1)
12325            && let Some([p0, p1]) = self.q8_fused_params(&[w0, w1])
12326        {
12327            let (y0, y1) = self.q8_fused2_core(p0.0, p1.0, aq, ad, in_f, p0.1, p1.1, p0.2)?;
12328            return Ok(Some(((y0, 1.0), (y1, 1.0))));
12329        }
12330        // F8-E4M3 ARM (lane/fp8-decode-v1, 2026-08-05): with native e4m3 residency the FFN gate+up
12331        // pair (and the ssm beta+alpha dual, which routes through this same entry) fell through
12332        // both the NVFP4 gate below and the Q8_0 arm above to two `matmul_pre_noscale` launches —
12333        // native residency was UN-FUSING the trunk relative to the Q8_0 slab it replaces. The
12334        // fused2 kernel body is `qmatvec_e4m3_mmvq` VERBATIM per (tensor,row). Contract match:
12335        // `matmul_pre_noscale` on e4m3 launches with scale 1.0 and RETURNS the per-tensor
12336        // weight_scale for the caller to fold, so we pass ws=1.0 here and return (s0,s1) — same
12337        // bits, and the two macro-scale multiplies still fold into the SwiGLU epilogue.
12338        // MEMRA_E4M3_DUAL=0 rolls back to the two-launch pair.
12339        if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
12340            let (y0, y1) =
12341                self.e4m3_fused2_core(p0.0, p1.0, aq, ad, in_f, p0.1, p1.1, p0.2, 1.0, 1.0)?;
12342            return Ok(Some(((y0, p0.3), (y1, p1.3))));
12343        }
12344        let (b0, q0, rb0, s0, rp0) = match w0 {
12345            GpuTensor::Quant {
12346                bytes,
12347                qtype,
12348                row_bytes,
12349                scale,
12350                rp,
12351                ..
12352            } => (bytes, *qtype, *row_bytes, *scale, *rp),
12353            _ => return Ok(None),
12354        };
12355        let (b1, q1, rb1, s1, rp1) = match w1 {
12356            GpuTensor::Quant {
12357                bytes,
12358                qtype,
12359                row_bytes,
12360                scale,
12361                rp,
12362                ..
12363            } => (bytes, *qtype, *row_bytes, *scale, *rp),
12364            _ => return Ok(None),
12365        };
12366        if q0 != QT_NVFP4 || q1 != QT_NVFP4 || rb0 != rb1 || rp0 != rp1 {
12367            return Ok(None);
12368        }
12369        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
12370        const RPW: u32 = 2;
12371        let rows_per_block = ROWS_PER_BLOCK * RPW;
12372        let f = self.func(if rp0 {
12373            "qmatvec_nvfp4_mmvq_dual_mr2_rp"
12374        } else {
12375            "qmatvec_nvfp4_mmvq_dual_mr2"
12376        });
12377        let mut y0 = self.alloc_uninit::<f32>(out_f)?;
12378        let mut y1 = self.alloc_uninit::<f32>(out_f)?;
12379        let cfg = LaunchConfig {
12380            grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 2, 1),
12381            block_dim: (32, ROWS_PER_BLOCK, 1),
12382            shared_mem_bytes: 0,
12383        };
12384        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, 1i32, rb0 as i64);
12385        // noscale contract: the caller folds s0/s1 into the SwiGLU epilogue — the kernel's fused
12386        // yscale args stay 1.0 here (they exist for the single-tensor callers).
12387        let one = 1.0f32;
12388        let __s_b = self.gpu.stream();
12389        let mut b = __s_b.launch_builder(&f);
12390        b.arg(b0)
12391            .arg(b1)
12392            .arg(aq)
12393            .arg(ad)
12394            .arg(&mut y0)
12395            .arg(&mut y1)
12396            .arg(&inf)
12397            .arg(&outf)
12398            .arg(&mi)
12399            .arg(&rb)
12400            .arg(&one)
12401            .arg(&one);
12402        unsafe {
12403            b.launch(cfg)?;
12404        }
12405        Ok(Some(((y0, s0), (y1, s1))))
12406    }
12407
12408    /// FUSED NVFP4 matvec TRIPLE with unequal out_f (rig-native decode increment 1,
12409    /// lane/rig-native-nvfp4): wq+wk+wv in ONE launch via the q8_0 fused2 block-offset
12410    /// recipe. Per (tensor,row,t) the kernel body is nvfp4_mmvq_multirow_rp VERBATIM ->
12411    /// bit-identical to three separate `matmul_pre` launches; yscales fold in-kernel exactly
12412    /// as the singles do. grid.y = m (the t-parallel verify rows ride the same launch).
12413    /// None when ineligible (not all rp NVFP4 / in_f mismatch / mmvq off) — callers fall
12414    /// back to the three singles.
12415    #[allow(clippy::too_many_arguments)]
12416    pub fn matmul_nvfp4_fused3(
12417        &self,
12418        w0: &crate::model::GpuTensor,
12419        w1: &crate::model::GpuTensor,
12420        w2: &crate::model::GpuTensor,
12421        aq: &CudaSlice<i8>,
12422        ad: &CudaSlice<f32>,
12423        m: usize,
12424    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
12425    {
12426        use crate::model::GpuTensor;
12427        // m == 1 rides the original fused kernel; m = 2..=8 rides the BATCHED fused twin
12428        // (lane/samplat — the fused4_b8 pattern): the seg body is nvfp4_mmvq_batched_rp_sc
12429        // verbatim, weight rows read once for all m columns, bit-identical per
12430        // (tensor,row,column) to the three bN_rpsc singles. The old "at m>1 the fused
12431        // segments would re-read the weight per row" note described the grid.y=m lift,
12432        // which this twin deliberately is NOT.
12433        if !self.mmvq_supports(QT_NVFP4)
12434            || !self.uses_q8_1_fast(w0)
12435            || !self.uses_q8_1_fast(w1)
12436            || !self.uses_q8_1_fast(w2)
12437        {
12438            return Ok(None);
12439        }
12440        // m = 9..=16 (lane/orndecode2): the exact-16 tier's trio width rides the GROUP3
12441        // door — same family and bit-identity law as the fused4 delegate above.
12442        if (9..=16).contains(&m) {
12443            return Ok(
12444                match self.matmul_decode_exact_group3_pre([w0, w1, w2], aq, ad, m)? {
12445                    Some(mut ys) => {
12446                        let y2 = ys.pop().unwrap();
12447                        let y1 = ys.pop().unwrap();
12448                        let y0 = ys.pop().unwrap();
12449                        Some((y0, y1, y2))
12450                    }
12451                    None => None,
12452                },
12453            );
12454        }
12455        if !(1..=8).contains(&m) {
12456            return Ok(None);
12457        }
12458        if m > 1 {
12459            let in_f = w0.in_features();
12460            if std::env::var("MEMRA_NVFP4_FUSED3B").as_deref() == Ok("0")
12461                || !self.batched_supports(QT_NVFP4)
12462                || std::env::var("MEMRA_NO_BATCHED").is_ok()
12463                || (m > 4 && !Self::b8_enabled())
12464                || in_f % 512 != 0
12465                || in_f / 64 > 272
12466            {
12467                return Ok(None);
12468            }
12469        }
12470        let unpack = |w: &crate::model::GpuTensor| match w {
12471            GpuTensor::Quant {
12472                bytes,
12473                qtype,
12474                scale,
12475                rp,
12476                ..
12477            } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
12478            _ => None,
12479        };
12480        let (Some(p0), Some(p1), Some(p2)) = (unpack(w0), unpack(w1), unpack(w2)) else {
12481            return Ok(None);
12482        };
12483        let in_f = w0.in_features();
12484        if w1.in_features() != in_f || w2.in_features() != in_f {
12485            return Ok(None);
12486        }
12487        let (o0, o1, o2) = (w0.out_features(), w1.out_features(), w2.out_features());
12488        const ROWS_PER_BLOCK: u32 = 4; // MEMRA_MMVQ_ROWS
12489        const RPW: u32 = 2;
12490        let rows_pb = ROWS_PER_BLOCK * RPW;
12491        let nb = |o: usize| (o as u32).div_ceil(rows_pb);
12492        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
12493        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
12494        let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
12495        let (inf, oi0, oi1, oi2, mi) = (in_f as i32, o0 as i32, o1 as i32, o2 as i32, m as i32);
12496        // SAFETY: the raw pointers come straight from the &GpuTensor borrows above and are
12497        // only dereferenced for the launch-arg build inside this call.
12498        let (b0, b1, b2) = unsafe { (&*p0.0, &*p1.0, &*p2.0) };
12499        if m > 1 {
12500            // batched twin has no in-kernel scale — refuse scale carriers (GGUF trunk = 1.0).
12501            if p0.1 != 1.0 || p1.1 != 1.0 || p2.1 != 1.0 {
12502                return Ok(None);
12503            }
12504            let f = self.func("qmatvec_nvfp4_mmvq_fused3_b8_rpsc");
12505            let cfg = LaunchConfig {
12506                grid_dim: (nb(o0) + nb(o1) + nb(o2), 1, 1),
12507                block_dim: (32, ROWS_PER_BLOCK, 1),
12508                shared_mem_bytes: 0,
12509            };
12510            let __s_b = self.gpu.stream();
12511            let mut b = __s_b.launch_builder(&f);
12512            b.arg(b0)
12513                .arg(b1)
12514                .arg(b2)
12515                .arg(aq)
12516                .arg(ad)
12517                .arg(&mut y0)
12518                .arg(&mut y1)
12519                .arg(&mut y2)
12520                .arg(&inf)
12521                .arg(&oi0)
12522                .arg(&oi1)
12523                .arg(&oi2)
12524                .arg(&mi);
12525            unsafe {
12526                b.launch(cfg)?;
12527            }
12528            return Ok(Some((y0, y1, y2)));
12529        }
12530        let f = self.func("qmatvec_nvfp4_mmvq_fused3_rp");
12531        let cfg = LaunchConfig {
12532            grid_dim: (nb(o0) + nb(o1) + nb(o2), m as u32, 1),
12533            block_dim: (32, ROWS_PER_BLOCK, 1),
12534            shared_mem_bytes: 0,
12535        };
12536        let __s_b = self.gpu.stream();
12537        let mut b = __s_b.launch_builder(&f);
12538        b.arg(b0)
12539            .arg(b1)
12540            .arg(b2)
12541            .arg(aq)
12542            .arg(ad)
12543            .arg(&mut y0)
12544            .arg(&mut y1)
12545            .arg(&mut y2)
12546            .arg(&inf)
12547            .arg(&oi0)
12548            .arg(&oi1)
12549            .arg(&oi2)
12550            .arg(&mi)
12551            .arg(&p0.1)
12552            .arg(&p1.1)
12553            .arg(&p2.1);
12554        unsafe {
12555            b.launch(cfg)?;
12556        }
12557        Ok(Some((y0, y1, y2)))
12558    }
12559
12560    /// fused2 twin of `matmul_nvfp4_fused3`, for MIXED-type weight groups: the gemma4
12561    /// dense NVFP4mix recipe keeps `attn_v` and `ffn_down` at Q8_0 (full-NVFP4 was
12562    /// measured garbage on this dense class), so its q/k/v trio and gate/up/down never
12563    /// satisfy an all-NVFP4 fused3 — the pairs that ARE uniformly NVFP4 (q,k and
12564    /// gate,up) fuse here instead. m==1 only, same law as fused3/fused4: per
12565    /// (tensor,row) the kernel seg body is VERBATIM, so the fusion is bit-identical to
12566    /// two separate launches. `MEMRA_NVFP4_FUSED2=0` is the rollback seam and the
12567    /// same-binary interleaved A/B arm.
12568    pub fn matmul_nvfp4_fused2(
12569        &self,
12570        w0: &crate::model::GpuTensor,
12571        w1: &crate::model::GpuTensor,
12572        aq: &CudaSlice<i8>,
12573        ad: &CudaSlice<f32>,
12574        m: usize,
12575    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
12576        use crate::model::GpuTensor;
12577        static FUSED2_OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12578        let off =
12579            *FUSED2_OFF.get_or_init(|| std::env::var("MEMRA_NVFP4_FUSED2").as_deref() == Ok("0"));
12580        // m==1 ONLY: at m>1 the singles ride the _b16 weight-once column program (one weight
12581        // read serves all m rows); the fused segments would re-read the weight per row.
12582        if off
12583            || m != 1
12584            || !self.mmvq_supports(QT_NVFP4)
12585            || !self.uses_q8_1_fast(w0)
12586            || !self.uses_q8_1_fast(w1)
12587        {
12588            return Ok(None);
12589        }
12590        let unpack = |w: &crate::model::GpuTensor| match w {
12591            GpuTensor::Quant {
12592                bytes,
12593                qtype,
12594                scale,
12595                rp,
12596                ..
12597            } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
12598            _ => None,
12599        };
12600        let (Some(p0), Some(p1)) = (unpack(w0), unpack(w1)) else {
12601            return Ok(None);
12602        };
12603        let in_f = w0.in_features();
12604        if w1.in_features() != in_f {
12605            return Ok(None);
12606        }
12607        let (o0, o1) = (w0.out_features(), w1.out_features());
12608        const ROWS_PER_BLOCK: u32 = 4; // MEMRA_MMVQ_ROWS
12609        const RPW: u32 = 2;
12610        let rows_pb = ROWS_PER_BLOCK * RPW;
12611        let nb = |o: usize| (o as u32).div_ceil(rows_pb);
12612        let f = self.func("qmatvec_nvfp4_mmvq_fused2_rp");
12613        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
12614        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
12615        let cfg = LaunchConfig {
12616            grid_dim: (nb(o0) + nb(o1), m as u32, 1),
12617            block_dim: (32, ROWS_PER_BLOCK, 1),
12618            shared_mem_bytes: 0,
12619        };
12620        let (inf, oi0, oi1, mi) = (in_f as i32, o0 as i32, o1 as i32, m as i32);
12621        // SAFETY: the raw pointers come straight from the &GpuTensor borrows above and are
12622        // only dereferenced for the launch-arg build inside this call.
12623        let (b0, b1) = unsafe { (&*p0.0, &*p1.0) };
12624        // PDL wave-B: the kernel carries MEMRA_PDL_ENTRY — launch overlaps the producer's
12625        // drain (quantize/norm emitting aq/ad). Same math, same order, bit-identical.
12626        if Self::pdl_on() && Self::pdl_mmvq_on() && Self::pdl_nvfp4q8_on() {
12627            {
12628                use cudarc::driver::{DevicePtr, DevicePtrMut};
12629                let s = &self.gpu.stream();
12630                let (pw0, _g0) = b0.device_ptr(s);
12631                let (pw1, _g1) = b1.device_ptr(s);
12632                let (paq, _g2) = aq.device_ptr(s);
12633                let (pad, _g3) = ad.device_ptr(s);
12634                let (py0, _g4) = y0.device_ptr_mut(s);
12635                let (py1, _g5) = y1.device_ptr_mut(s);
12636                let (s0, s1) = (p0.1, p1.1);
12637                let mut ps = [
12638                    &pw0 as *const _ as *mut std::ffi::c_void,
12639                    &pw1 as *const _ as *mut _,
12640                    &paq as *const _ as *mut _,
12641                    &pad as *const _ as *mut _,
12642                    &py0 as *const _ as *mut _,
12643                    &py1 as *const _ as *mut _,
12644                    &inf as *const _ as *mut _,
12645                    &oi0 as *const _ as *mut _,
12646                    &oi1 as *const _ as *mut _,
12647                    &mi as *const _ as *mut _,
12648                    &s0 as *const _ as *mut _,
12649                    &s1 as *const _ as *mut _,
12650                ];
12651                unsafe {
12652                    self.launch_pdl(
12653                        "qmatvec_nvfp4_mmvq_fused2_rp",
12654                        cfg.grid_dim,
12655                        cfg.block_dim,
12656                        &mut ps,
12657                    )?;
12658                }
12659            }
12660            return Ok(Some((y0, y1)));
12661        }
12662        let __s_b = self.gpu.stream();
12663        let mut b = __s_b.launch_builder(&f);
12664        b.arg(b0)
12665            .arg(b1)
12666            .arg(aq)
12667            .arg(ad)
12668            .arg(&mut y0)
12669            .arg(&mut y1)
12670            .arg(&inf)
12671            .arg(&oi0)
12672            .arg(&oi1)
12673            .arg(&mi)
12674            .arg(&p0.1)
12675            .arg(&p1.1);
12676        unsafe {
12677            b.launch(cfg)?;
12678        }
12679        Ok(Some((y0, y1)))
12680    }
12681
12682    /// Slot-fed fused2 twin (alloc-free capture lane): identical launch to
12683    /// `matmul_nvfp4_fused2`, caller-owned outputs — the gemma4 dc_slotted graph body
12684    /// needs zero mem nodes, so the allocating wrapper can't serve it. Returns false
12685    /// (decline) on any non-NVFP4/rp pair; the caller chains or refuses.
12686    pub fn matmul_nvfp4_fused2_into(
12687        &self,
12688        w0: &crate::model::GpuTensor,
12689        w1: &crate::model::GpuTensor,
12690        aq: &CudaSlice<i8>,
12691        ad: &CudaSlice<f32>,
12692        y0: &mut CudaSlice<f32>,
12693        y1: &mut CudaSlice<f32>,
12694    ) -> Result<bool, Box<dyn std::error::Error>> {
12695        use crate::model::GpuTensor;
12696        static FUSED2_OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12697        let off =
12698            *FUSED2_OFF.get_or_init(|| std::env::var("MEMRA_NVFP4_FUSED2").as_deref() == Ok("0"));
12699        if off
12700            || !self.mmvq_supports(QT_NVFP4)
12701            || !self.uses_q8_1_fast(w0)
12702            || !self.uses_q8_1_fast(w1)
12703        {
12704            return Ok(false);
12705        }
12706        let unpack = |w: &crate::model::GpuTensor| match w {
12707            GpuTensor::Quant {
12708                bytes,
12709                qtype,
12710                scale,
12711                rp,
12712                ..
12713            } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
12714            _ => None,
12715        };
12716        let (Some(p0), Some(p1)) = (unpack(w0), unpack(w1)) else {
12717            return Ok(false);
12718        };
12719        let in_f = w0.in_features();
12720        if w1.in_features() != in_f {
12721            return Ok(false);
12722        }
12723        let (o0, o1) = (w0.out_features(), w1.out_features());
12724        if y0.len() < o0 || y1.len() < o1 {
12725            return Ok(false);
12726        }
12727        const ROWS_PER_BLOCK: u32 = 4; // MEMRA_MMVQ_ROWS
12728        const RPW: u32 = 2;
12729        let rows_pb = ROWS_PER_BLOCK * RPW;
12730        let nb = |o: usize| (o as u32).div_ceil(rows_pb);
12731        let f = self.func("qmatvec_nvfp4_mmvq_fused2_rp");
12732        let cfg = LaunchConfig {
12733            grid_dim: (nb(o0) + nb(o1), 1, 1),
12734            block_dim: (32, ROWS_PER_BLOCK, 1),
12735            shared_mem_bytes: 0,
12736        };
12737        let (inf, oi0, oi1, mi) = (in_f as i32, o0 as i32, o1 as i32, 1i32);
12738        // SAFETY: the raw pointers come straight from the &GpuTensor borrows above and are
12739        // only dereferenced for the launch-arg build inside this call.
12740        let (b0, b1) = unsafe { (&*p0.0, &*p1.0) };
12741        let __s_b = self.gpu.stream();
12742        let mut b = __s_b.launch_builder(&f);
12743        b.arg(b0)
12744            .arg(b1)
12745            .arg(aq)
12746            .arg(ad)
12747            .arg(&mut *y0)
12748            .arg(&mut *y1)
12749            .arg(&inf)
12750            .arg(&oi0)
12751            .arg(&oi1)
12752            .arg(&mi)
12753            .arg(&p0.1)
12754            .arg(&p1.1);
12755        unsafe {
12756            b.launch(cfg)?;
12757        }
12758        Ok(true)
12759    }
12760
12761    /// fused4 twin of `matmul_nvfp4_fused3`: the Linear-mixer projection quartet
12762    /// (wqkv + wqkv_gate + ssm_beta + ssm_alpha) in one launch, m==1 only. Per
12763    /// (tensor,row) the kernel body is nvfp4_mmvq_multirow_rp VERBATIM — bit-identical
12764    /// to four separate launches (rig-native decode increment 2, RIG-NATIVE-DECODE.md).
12765    #[allow(clippy::type_complexity)]
12766    pub fn matmul_nvfp4_fused4(
12767        &self,
12768        w0: &crate::model::GpuTensor,
12769        w1: &crate::model::GpuTensor,
12770        w2: &crate::model::GpuTensor,
12771        w3: &crate::model::GpuTensor,
12772        aq: &CudaSlice<i8>,
12773        ad: &CudaSlice<f32>,
12774        m: usize,
12775    ) -> Result<
12776        Option<(
12777            CudaSlice<f32>,
12778            CudaSlice<f32>,
12779            CudaSlice<f32>,
12780            CudaSlice<f32>,
12781        )>,
12782        Box<dyn std::error::Error>,
12783    > {
12784        use crate::model::GpuTensor;
12785        // MEMRA_NVFP4_FUSED4=0: rollback seam + the same-binary interleaved A/B arm.
12786        // m == 1 rides the original fused kernel; m = 2..=8 rides the BATCHED fused twin
12787        // (lane/samplat, 2026-08-21): same quartet-in-one-launch shape, seg body =
12788        // nvfp4_mmvq_batched_rp_sc verbatim (weight rows read once for all m columns) —
12789        // bit-identical per (tensor,row,column) to the four bN_rpsc singles it replaces.
12790        // Admission mirrors the singles' batched gates below.
12791        if std::env::var("MEMRA_NVFP4_FUSED4").as_deref() == Ok("0")
12792            || !self.mmvq_supports(QT_NVFP4)
12793            || !self.uses_q8_1_fast(w0)
12794            || !self.uses_q8_1_fast(w1)
12795            || !self.uses_q8_1_fast(w2)
12796            || !self.uses_q8_1_fast(w3)
12797        {
12798            return Ok(None);
12799        }
12800        // m = 9..=16 (lane/orndecode2, the exact-16 tier's trunk): the rp-sc seg twins stop
12801        // at 8; this width class rides the GROUP4 door instead — nvfp4_mmvq_batched_rp<16,_>
12802        // body, the SAME family as the b16_rp singles the tier would otherwise launch four
12803        // times, bit-identical per (tensor,token,row) incl. the fused write-side scale.
12804        if (9..=16).contains(&m) {
12805            return Ok(
12806                match self.matmul_decode_exact_group4_pre([w0, w1, w2, w3], aq, ad, m)? {
12807                    Some(mut ys) => {
12808                        let y3 = ys.pop().unwrap();
12809                        let y2 = ys.pop().unwrap();
12810                        let y1 = ys.pop().unwrap();
12811                        let y0 = ys.pop().unwrap();
12812                        Some((y0, y1, y2, y3))
12813                    }
12814                    None => None,
12815                },
12816            );
12817        }
12818        if !(1..=8).contains(&m) {
12819            return Ok(None);
12820        }
12821        if m > 1 {
12822            // the batched-twin gates: the bN_rpsc program this must stay byte-identical to
12823            // (matmul_pre's batched arm), plus the rp-sc dispatch shape requirements.
12824            let in_f = w0.in_features();
12825            if !self.batched_supports(QT_NVFP4)
12826                || std::env::var("MEMRA_NO_BATCHED").is_ok()
12827                || (m > 4 && !Self::b8_enabled())
12828                || in_f % 512 != 0
12829                || in_f / 64 > 272
12830            {
12831                return Ok(None);
12832            }
12833        }
12834        let unpack = |w: &crate::model::GpuTensor| match w {
12835            GpuTensor::Quant {
12836                bytes,
12837                qtype,
12838                scale,
12839                rp,
12840                ..
12841            } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
12842            _ => None,
12843        };
12844        let (Some(p0), Some(p1), Some(p2), Some(p3)) =
12845            (unpack(w0), unpack(w1), unpack(w2), unpack(w3))
12846        else {
12847            return Ok(None);
12848        };
12849        let in_f = w0.in_features();
12850        if w1.in_features() != in_f || w2.in_features() != in_f || w3.in_features() != in_f {
12851            return Ok(None);
12852        }
12853        let (o0, o1, o2, o3) = (
12854            w0.out_features(),
12855            w1.out_features(),
12856            w2.out_features(),
12857            w3.out_features(),
12858        );
12859        const ROWS_PER_BLOCK: u32 = 4; // MEMRA_MMVQ_ROWS
12860        const RPW: u32 = 2;
12861        let rows_pb = ROWS_PER_BLOCK * RPW;
12862        let nb = |o: usize| (o as u32).div_ceil(rows_pb);
12863        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
12864        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
12865        let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
12866        let mut y3 = self.alloc_uninit::<f32>(m * o3)?;
12867        let (inf, oi0, oi1, oi2, oi3, mi) = (
12868            in_f as i32,
12869            o0 as i32,
12870            o1 as i32,
12871            o2 as i32,
12872            o3 as i32,
12873            m as i32,
12874        );
12875        // SAFETY: the raw pointers come straight from the &GpuTensor borrows above and are
12876        // only dereferenced for the launch-arg build inside this call.
12877        let (b0, b1, b2, b3) = unsafe { (&*p0.0, &*p1.0, &*p2.0, &*p3.0) };
12878        if m > 1 {
12879            // Batched fused twin: no in-kernel scale (the bN_rpsc program has none) — refuse
12880            // scale-carrying tensors so the singles path keeps them (GGUF trunk scales are 1.0).
12881            if p0.1 != 1.0 || p1.1 != 1.0 || p2.1 != 1.0 || p3.1 != 1.0 {
12882                return Ok(None);
12883            }
12884            let f = self.func("qmatvec_nvfp4_mmvq_fused4_b8_rpsc");
12885            let cfg = LaunchConfig {
12886                grid_dim: (nb(o0) + nb(o1) + nb(o2) + nb(o3), 1, 1),
12887                block_dim: (32, ROWS_PER_BLOCK, 1),
12888                shared_mem_bytes: 0,
12889            };
12890            let __s_b = self.gpu.stream();
12891            let mut b = __s_b.launch_builder(&f);
12892            b.arg(b0)
12893                .arg(b1)
12894                .arg(b2)
12895                .arg(b3)
12896                .arg(aq)
12897                .arg(ad)
12898                .arg(&mut y0)
12899                .arg(&mut y1)
12900                .arg(&mut y2)
12901                .arg(&mut y3)
12902                .arg(&inf)
12903                .arg(&oi0)
12904                .arg(&oi1)
12905                .arg(&oi2)
12906                .arg(&oi3)
12907                .arg(&mi);
12908            unsafe {
12909                b.launch(cfg)?;
12910            }
12911            return Ok(Some((y0, y1, y2, y3)));
12912        }
12913        let f = self.func("qmatvec_nvfp4_mmvq_fused4_rp");
12914        let cfg = LaunchConfig {
12915            grid_dim: (nb(o0) + nb(o1) + nb(o2) + nb(o3), m as u32, 1),
12916            block_dim: (32, ROWS_PER_BLOCK, 1),
12917            shared_mem_bytes: 0,
12918        };
12919        let __s_b = self.gpu.stream();
12920        let mut b = __s_b.launch_builder(&f);
12921        b.arg(b0)
12922            .arg(b1)
12923            .arg(b2)
12924            .arg(b3)
12925            .arg(aq)
12926            .arg(ad)
12927            .arg(&mut y0)
12928            .arg(&mut y1)
12929            .arg(&mut y2)
12930            .arg(&mut y3)
12931            .arg(&inf)
12932            .arg(&oi0)
12933            .arg(&oi1)
12934            .arg(&oi2)
12935            .arg(&oi3)
12936            .arg(&mi)
12937            .arg(&p0.1)
12938            .arg(&p1.1)
12939            .arg(&p2.1)
12940            .arg(&p3.1);
12941        unsafe {
12942            b.launch(cfg)?;
12943        }
12944        Ok(Some((y0, y1, y2, y3)))
12945    }
12946
12947    /// FUSED Q8_0 m=1 matvec PAIR with UNEQUAL out_f (trunk launch-fusion, 2026-07-05). Folds two
12948    /// same-input q8_0 projections (35B trunk: wqkv+wqkv_gate 8192/4096, gate_shexp+up_shexp
12949    /// 512/512) into ONE launch via a block-offset split (blocks [0,nb0) -> w0, rest -> w1) — the
12950    /// dual-mr2 recipe with the same-out_f restriction lifted. Per (tensor,row) the kernel body is
12951    /// qmatvec_q8_0_mmvq VERBATIM -> BIT-IDENTICAL to two separate m=1 launches. Returns None when
12952    /// ineligible (not both Q8_0 / in_f mismatch / MEMRA_MMVQ off / MEMRA_Q8_DUAL=0) — caller falls
12953    /// back to the per-tensor path.
12954    pub fn matmul_q8_fused2(
12955        &self,
12956        w0: &crate::model::GpuTensor,
12957        w1: &crate::model::GpuTensor,
12958        aq: &CudaSlice<i8>,
12959        ad: &CudaSlice<f32>,
12960    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
12961        // e4m3 twin (lane/fp8-decode-v1): this entry is the trunk's generic m=1 pair door
12962        // (wqkv+wqkv_gate, ssm_beta+alpha, gate_shexp+up_shexp), so admitting QT_F8_E4M3 here
12963        // fuses the NATIVE-RESIDENCY FP8 trunk at every existing call site with no call-site
12964        // change. Scale is folded in-kernel per range -> the returned buffers are already scaled,
12965        // exactly like the per-tensor `matmul_pre` e4m3 dispatch this replaces.
12966        if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
12967            return Ok(Some(self.e4m3_fused2_core(
12968                p0.0,
12969                p1.0,
12970                aq,
12971                ad,
12972                w0.in_features(),
12973                p0.1,
12974                p1.1,
12975                p0.2,
12976                p0.3,
12977                p1.3,
12978            )?));
12979        }
12980        let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
12981            return Ok(None);
12982        };
12983        Ok(Some(self.q8_fused2_core(
12984            p0.0,
12985            p1.0,
12986            aq,
12987            ad,
12988            w0.in_features(),
12989            p0.1,
12990            p1.1,
12991            p0.2,
12992        )?))
12993    }
12994
12995    #[allow(clippy::too_many_arguments)]
12996    fn q8_fused2_core(
12997        &self,
12998        b0: &CudaSlice<u8>,
12999        b1: &CudaSlice<u8>,
13000        aq: &CudaSlice<i8>,
13001        ad: &CudaSlice<f32>,
13002        in_f: usize,
13003        out0: usize,
13004        out1: usize,
13005        row_bytes: usize,
13006    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13007        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
13008        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
13009        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
13010        let f = self.func("qmatvec_q8_0_mmvq_fused2");
13011        let mut y0 = self.alloc_uninit::<f32>(out0)?;
13012        let mut y1 = self.alloc_uninit::<f32>(out1)?;
13013        let cfg = LaunchConfig {
13014            grid_dim: (nb0 + nb1, 1, 1),
13015            block_dim: (32, ROWS_PER_BLOCK, 1),
13016            shared_mem_bytes: 0,
13017        };
13018        let (inf, o0, o1, rbl) = (in_f as i32, out0 as i32, out1 as i32, row_bytes as i64);
13019        let __s_b = self.gpu.stream();
13020        let mut b = __s_b.launch_builder(&f);
13021        b.arg(b0)
13022            .arg(b1)
13023            .arg(aq)
13024            .arg(ad)
13025            .arg(&mut y0)
13026            .arg(&mut y1)
13027            .arg(&inf)
13028            .arg(&o0)
13029            .arg(&o1)
13030            .arg(&rbl);
13031        unsafe {
13032            b.launch(cfg)?;
13033        }
13034        Ok((y0, y1))
13035    }
13036
13037    /// f32-activation entry for the fused2 pair: quantizes x to q8_1 ONCE then runs the fused
13038    /// launch — replaces two `matmul(w, x, 1)` calls that would each re-quantize the same x
13039    /// (35B shared-expert gate+up per MoE layer per token). Same bits: quantize_q8_1 is
13040    /// deterministic, the fused body is the MMVQ kernel verbatim. None when ineligible (the
13041    /// callers' m==1-under-MEMRA_FAST dispatch would take MMVQ; anything else falls back).
13042    pub fn matmul_q8_fused2_x(
13043        &self,
13044        w0: &crate::model::GpuTensor,
13045        w1: &crate::model::GpuTensor,
13046        x: &CudaSlice<f32>,
13047    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
13048        if !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) {
13049            return Ok(None);
13050        }
13051        if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
13052            let (aq, ad) = self.quantize_q8_1(x, 1, w0.in_features())?;
13053            return Ok(Some(self.e4m3_fused2_core(
13054                p0.0,
13055                p1.0,
13056                &aq,
13057                &ad,
13058                w0.in_features(),
13059                p0.1,
13060                p1.1,
13061                p0.2,
13062                p0.3,
13063                p1.3,
13064            )?));
13065        }
13066        let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
13067            return Ok(None);
13068        };
13069        let (aq, ad) = self.quantize_q8_1(x, 1, w0.in_features())?;
13070        Ok(Some(self.q8_fused2_core(
13071            p0.0,
13072            p1.0,
13073            &aq,
13074            &ad,
13075            w0.in_features(),
13076            p0.1,
13077            p1.1,
13078            p0.2,
13079        )?))
13080    }
13081
13082    /// Test entry for the kernel_check gate: launch the fused2 kernel from raw weight bytes,
13083    /// quantizing the f32 activation internally (mirrors qmatvec_mmvq_raw; no env gating).
13084    #[allow(clippy::too_many_arguments)]
13085    pub fn qmatvec_q8_fused2_raw(
13086        &self,
13087        b0: &CudaSlice<u8>,
13088        b1: &CudaSlice<u8>,
13089        x: &CudaSlice<f32>,
13090        in_f: usize,
13091        out0: usize,
13092        out1: usize,
13093        row_bytes: usize,
13094    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13095        let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
13096        self.q8_fused2_core(b0, b1, &aq, &ad, in_f, out0, out1, row_bytes)
13097    }
13098
13099    /// FUSED Q8_0 m=1 matvec TRIPLE (wq+wk+wv on the 35B full-attn layers: out_f 8192/512/512).
13100    /// Same block-offset recipe as `matmul_q8_fused2` with three ranges. BIT-IDENTICAL per
13101    /// (tensor,row) to three separate m=1 MMVQ launches.
13102    /// FUSED Q4_0 m=1 TRIPLE (gemma q/k/v — same quantized input; per (tensor,row) chain
13103    /// identical to the mr2 kernel). Returns None unless all three are Q4_0 with equal in_f.
13104    pub fn matmul_q4_fused3(
13105        &self,
13106        w0: &crate::model::GpuTensor,
13107        w1: &crate::model::GpuTensor,
13108        w2: &crate::model::GpuTensor,
13109        aq: &CudaSlice<i8>,
13110        ad: &CudaSlice<f32>,
13111    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
13112    {
13113        use crate::model::GpuTensor;
13114        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
13115            match w {
13116                GpuTensor::Quant {
13117                    qtype, row_bytes, ..
13118                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
13119                _ => None,
13120            }
13121        };
13122        let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (q4(w0), q4(w1), q4(w2)) else {
13123            return Ok(None);
13124        };
13125        if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
13126            return Ok(None);
13127        }
13128        // Effective (bytes, rp) per tensor: mirror (rp4) OR the in-place swap (rp flag,
13129        // bytes already split). Mixed layouts cannot share one fused launch -> fall back to
13130        // the separate matvecs (each routes its own rp).
13131        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
13132            match w {
13133                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
13134                    Some(m) => (m, true),
13135                    None => (bytes, *rp),
13136                },
13137                _ => unreachable!(),
13138            }
13139        }
13140        let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
13141        if rp0 != rp1 || rp1 != rp2 {
13142            return Ok(None);
13143        }
13144        let rp = rp0;
13145        let rpb: u32 = 4;
13146        // mr1 (one row/warp, 2026-07-14): follows the singles' MEMRA_Q40_MR default — the
13147        // fused t=1 kernels were left on mr2 when the singles flipped (DRAM-duty map:
13148        // fused3 57% / fused2 86%; small qkv segments starve under mr2's half grid).
13149        let mr1 = rp && Self::q40_mr1_on();
13150        let nb = |o: usize| {
13151            if mr1 {
13152                (o as u32).div_ceil(rpb)
13153            } else {
13154                (o as u32).div_ceil(2).div_ceil(rpb)
13155            }
13156        };
13157        let grid = nb(o0) + nb(o1) + nb(o2);
13158        let mut y0 = self.alloc_uninit::<f32>(o0)?;
13159        let mut y1 = self.alloc_uninit::<f32>(o1)?;
13160        let mut y2 = self.alloc_uninit::<f32>(o2)?;
13161        let f = self.func(if mr1 {
13162            "qmatvec_q4_0_mmvq_fused3_mr1_rp"
13163        } else if rp {
13164            "qmatvec_q4_0_mmvq_fused3_rp"
13165        } else {
13166            "qmatvec_q4_0_mmvq_fused3"
13167        });
13168        let cfg = LaunchConfig {
13169            grid_dim: (grid, 1, 1),
13170            block_dim: (32, rpb, 1),
13171            shared_mem_bytes: 0,
13172        };
13173        let inf = w0.in_features() as i32;
13174        let (oo0, oo1, oo2) = (o0 as i32, o1 as i32, o2 as i32);
13175        let (r0, r1, r2) = (rb0 as i64, rb1 as i64, rb2 as i64);
13176        // PDL wave-A (2026-07-23): the mr1 kernel carries MEMRA_PDL_ENTRY; only that
13177        // variant may take the programmatic-serialization launch.
13178        if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
13179            {
13180                use cudarc::driver::{DevicePtr, DevicePtrMut};
13181                let s = &self.gpu.stream();
13182                let (p0, _g0) = b0.device_ptr(s);
13183                let (p1, _g1) = b1.device_ptr(s);
13184                let (p2, _g2) = b2.device_ptr(s);
13185                let (paq, _g3) = aq.device_ptr(s);
13186                let (pad, _g4) = ad.device_ptr(s);
13187                let (py0, _g5) = y0.device_ptr_mut(s);
13188                let (py1, _g6) = y1.device_ptr_mut(s);
13189                let (py2, _g7) = y2.device_ptr_mut(s);
13190                let mut ps = [
13191                    &p0 as *const _ as *mut std::ffi::c_void,
13192                    &p1 as *const _ as *mut _,
13193                    &p2 as *const _ as *mut _,
13194                    &paq as *const _ as *mut _,
13195                    &pad as *const _ as *mut _,
13196                    &py0 as *const _ as *mut _,
13197                    &py1 as *const _ as *mut _,
13198                    &py2 as *const _ as *mut _,
13199                    &inf as *const _ as *mut _,
13200                    &oo0 as *const _ as *mut _,
13201                    &oo1 as *const _ as *mut _,
13202                    &oo2 as *const _ as *mut _,
13203                    &r0 as *const _ as *mut _,
13204                    &r1 as *const _ as *mut _,
13205                    &r2 as *const _ as *mut _,
13206                ];
13207                unsafe {
13208                    self.launch_pdl(
13209                        "qmatvec_q4_0_mmvq_fused3_mr1_rp",
13210                        (grid, 1, 1),
13211                        (32, rpb, 1),
13212                        &mut ps,
13213                    )?;
13214                }
13215            }
13216            return Ok(Some((y0, y1, y2)));
13217        }
13218        let __s_b = self.gpu.stream();
13219        let mut b = __s_b.launch_builder(&f);
13220        b.arg(b0)
13221            .arg(b1)
13222            .arg(b2)
13223            .arg(aq)
13224            .arg(ad)
13225            .arg(&mut y0)
13226            .arg(&mut y1)
13227            .arg(&mut y2)
13228            .arg(&inf)
13229            .arg(&oo0)
13230            .arg(&oo1)
13231            .arg(&oo2)
13232            .arg(&r0)
13233            .arg(&r1)
13234            .arg(&r2);
13235        unsafe {
13236            b.launch(cfg)?;
13237        }
13238        Ok(Some((y0, y1, y2)))
13239    }
13240
13241    /// Slot-fed fused3 twin (alloc-free capture lane): identical launch, caller-owned outputs.
13242    /// Returns Ok(false) when the fused path is unavailable (caller falls back).
13243    #[allow(clippy::too_many_arguments)]
13244    pub fn matmul_q4_fused3_into(
13245        &self,
13246        w0: &crate::model::GpuTensor,
13247        w1: &crate::model::GpuTensor,
13248        w2: &crate::model::GpuTensor,
13249        aq: &CudaSlice<i8>,
13250        ad: &CudaSlice<f32>,
13251        y0: &mut CudaSlice<f32>,
13252        y1: &mut CudaSlice<f32>,
13253        y2: &mut CudaSlice<f32>,
13254    ) -> Result<bool, Box<dyn std::error::Error>> {
13255        use crate::model::GpuTensor;
13256        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
13257            match w {
13258                GpuTensor::Quant {
13259                    qtype, row_bytes, ..
13260                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
13261                _ => None,
13262            }
13263        };
13264        let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (q4(w0), q4(w1), q4(w2)) else {
13265            return Ok(false);
13266        };
13267        if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
13268            return Ok(false);
13269        }
13270        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
13271            match w {
13272                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
13273                    Some(m) => (m, true),
13274                    None => (bytes, *rp),
13275                },
13276                _ => unreachable!(),
13277            }
13278        }
13279        let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
13280        if rp0 != rp1 || rp1 != rp2 {
13281            return Ok(false);
13282        }
13283        let rp = rp0;
13284        let rpb: u32 = 4;
13285        let mr1 = rp && Self::q40_mr1_on();
13286        let nb = |o: usize| {
13287            if mr1 {
13288                (o as u32).div_ceil(rpb)
13289            } else {
13290                (o as u32).div_ceil(2).div_ceil(rpb)
13291            }
13292        };
13293        let grid = nb(o0) + nb(o1) + nb(o2);
13294        debug_assert!(y0.len() >= o0 && y1.len() >= o1 && y2.len() >= o2);
13295        let f = self.func(if mr1 {
13296            "qmatvec_q4_0_mmvq_fused3_mr1_rp"
13297        } else if rp {
13298            "qmatvec_q4_0_mmvq_fused3_rp"
13299        } else {
13300            "qmatvec_q4_0_mmvq_fused3"
13301        });
13302        let cfg = LaunchConfig {
13303            grid_dim: (grid, 1, 1),
13304            block_dim: (32, rpb, 1),
13305            shared_mem_bytes: 0,
13306        };
13307        let inf = w0.in_features() as i32;
13308        let (oo0, oo1, oo2) = (o0 as i32, o1 as i32, o2 as i32);
13309        let (r0, r1, r2) = (rb0 as i64, rb1 as i64, rb2 as i64);
13310        // PDL wave-A: identical to the owned twin (capture-lane parity).
13311        if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
13312            use cudarc::driver::{DevicePtr, DevicePtrMut};
13313            let s = &self.gpu.stream();
13314            let (p0, _g0) = b0.device_ptr(s);
13315            let (p1, _g1) = b1.device_ptr(s);
13316            let (p2, _g2) = b2.device_ptr(s);
13317            let (paq, _g3) = aq.device_ptr(s);
13318            let (pad, _g4) = ad.device_ptr(s);
13319            let (py0, _g5) = y0.device_ptr_mut(s);
13320            let (py1, _g6) = y1.device_ptr_mut(s);
13321            let (py2, _g7) = y2.device_ptr_mut(s);
13322            let mut ps = [
13323                &p0 as *const _ as *mut std::ffi::c_void,
13324                &p1 as *const _ as *mut _,
13325                &p2 as *const _ as *mut _,
13326                &paq as *const _ as *mut _,
13327                &pad as *const _ as *mut _,
13328                &py0 as *const _ as *mut _,
13329                &py1 as *const _ as *mut _,
13330                &py2 as *const _ as *mut _,
13331                &inf as *const _ as *mut _,
13332                &oo0 as *const _ as *mut _,
13333                &oo1 as *const _ as *mut _,
13334                &oo2 as *const _ as *mut _,
13335                &r0 as *const _ as *mut _,
13336                &r1 as *const _ as *mut _,
13337                &r2 as *const _ as *mut _,
13338            ];
13339            unsafe {
13340                self.launch_pdl(
13341                    "qmatvec_q4_0_mmvq_fused3_mr1_rp",
13342                    (grid, 1, 1),
13343                    (32, rpb, 1),
13344                    &mut ps,
13345                )?;
13346            }
13347            return Ok(true);
13348        }
13349        let __s_b = self.gpu.stream();
13350        let mut b = __s_b.launch_builder(&f);
13351        b.arg(b0)
13352            .arg(b1)
13353            .arg(b2)
13354            .arg(aq)
13355            .arg(ad)
13356            .arg(&mut *y0)
13357            .arg(&mut *y1)
13358            .arg(&mut *y2)
13359            .arg(&inf)
13360            .arg(&oo0)
13361            .arg(&oo1)
13362            .arg(&oo2)
13363            .arg(&r0)
13364            .arg(&r1)
13365            .arg(&r2);
13366        unsafe {
13367            b.launch(cfg)?;
13368        }
13369        Ok(true)
13370    }
13371
13372    /// FUSED Q4_0 m=1 PAIR (gemma shared gate+up).
13373    pub fn matmul_q4_fused2(
13374        &self,
13375        w0: &crate::model::GpuTensor,
13376        w1: &crate::model::GpuTensor,
13377        aq: &CudaSlice<i8>,
13378        ad: &CudaSlice<f32>,
13379    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
13380        use crate::model::GpuTensor;
13381        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
13382            match w {
13383                GpuTensor::Quant {
13384                    qtype, row_bytes, ..
13385                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
13386                _ => None,
13387            }
13388        };
13389        let (Some((rb0, o0)), Some((rb1, o1))) = (q4(w0), q4(w1)) else {
13390            return Ok(None);
13391        };
13392        if w0.in_features() != w1.in_features() {
13393            return Ok(None);
13394        }
13395        // Effective (bytes, rp) per tensor (mirror or in-place swap); mixed -> separate matvecs.
13396        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
13397            match w {
13398                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
13399                    Some(m) => (m, true),
13400                    None => (bytes, *rp),
13401                },
13402                _ => unreachable!(),
13403            }
13404        }
13405        let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
13406        if rp0 != rp1 {
13407            return Ok(None);
13408        }
13409        let rp = rp0;
13410        let rpb: u32 = 4;
13411        // mr1 twin — see matmul_q4_fused3.
13412        let mr1 = rp && Self::q40_mr1_on();
13413        let nb = |o: usize| {
13414            if mr1 {
13415                (o as u32).div_ceil(rpb)
13416            } else {
13417                (o as u32).div_ceil(2).div_ceil(rpb)
13418            }
13419        };
13420        let grid = nb(o0) + nb(o1);
13421        let mut y0 = self.alloc_uninit::<f32>(o0)?;
13422        let mut y1 = self.alloc_uninit::<f32>(o1)?;
13423        let f = self.func(if mr1 {
13424            "qmatvec_q4_0_mmvq_fused2_mr1_rp"
13425        } else if rp {
13426            "qmatvec_q4_0_mmvq_fused2_rp"
13427        } else {
13428            "qmatvec_q4_0_mmvq_fused2"
13429        });
13430        let cfg = LaunchConfig {
13431            grid_dim: (grid, 1, 1),
13432            block_dim: (32, rpb, 1),
13433            shared_mem_bytes: 0,
13434        };
13435        let inf = w0.in_features() as i32;
13436        let (oo0, oo1) = (o0 as i32, o1 as i32);
13437        let (r0, r1) = (rb0 as i64, rb1 as i64);
13438        // PDL wave-A: mr1 kernel carries MEMRA_PDL_ENTRY.
13439        if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
13440            {
13441                use cudarc::driver::{DevicePtr, DevicePtrMut};
13442                let s = &self.gpu.stream();
13443                let (p0, _g0) = b0.device_ptr(s);
13444                let (p1, _g1) = b1.device_ptr(s);
13445                let (paq, _g2) = aq.device_ptr(s);
13446                let (pad, _g3) = ad.device_ptr(s);
13447                let (py0, _g4) = y0.device_ptr_mut(s);
13448                let (py1, _g5) = y1.device_ptr_mut(s);
13449                let mut ps = [
13450                    &p0 as *const _ as *mut std::ffi::c_void,
13451                    &p1 as *const _ as *mut _,
13452                    &paq as *const _ as *mut _,
13453                    &pad as *const _ as *mut _,
13454                    &py0 as *const _ as *mut _,
13455                    &py1 as *const _ as *mut _,
13456                    &inf as *const _ as *mut _,
13457                    &oo0 as *const _ as *mut _,
13458                    &oo1 as *const _ as *mut _,
13459                    &r0 as *const _ as *mut _,
13460                    &r1 as *const _ as *mut _,
13461                ];
13462                unsafe {
13463                    self.launch_pdl(
13464                        "qmatvec_q4_0_mmvq_fused2_mr1_rp",
13465                        (grid, 1, 1),
13466                        (32, rpb, 1),
13467                        &mut ps,
13468                    )?;
13469                }
13470            }
13471            return Ok(Some((y0, y1)));
13472        }
13473        let __s_b = self.gpu.stream();
13474        let mut b = __s_b.launch_builder(&f);
13475        b.arg(b0)
13476            .arg(b1)
13477            .arg(aq)
13478            .arg(ad)
13479            .arg(&mut y0)
13480            .arg(&mut y1)
13481            .arg(&inf)
13482            .arg(&oo0)
13483            .arg(&oo1)
13484            .arg(&r0)
13485            .arg(&r1);
13486        unsafe {
13487            b.launch(cfg)?;
13488        }
13489        Ok(Some((y0, y1)))
13490    }
13491
13492    /// Slot-fed fused2 twin (alloc-free capture lane): identical launch, caller-owned outputs.
13493    pub fn matmul_q4_fused2_into(
13494        &self,
13495        w0: &crate::model::GpuTensor,
13496        w1: &crate::model::GpuTensor,
13497        aq: &CudaSlice<i8>,
13498        ad: &CudaSlice<f32>,
13499        y0: &mut CudaSlice<f32>,
13500        y1: &mut CudaSlice<f32>,
13501    ) -> Result<bool, Box<dyn std::error::Error>> {
13502        use crate::model::GpuTensor;
13503        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
13504            match w {
13505                GpuTensor::Quant {
13506                    qtype, row_bytes, ..
13507                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
13508                _ => None,
13509            }
13510        };
13511        let (Some((rb0, o0)), Some((rb1, o1))) = (q4(w0), q4(w1)) else {
13512            return Ok(false);
13513        };
13514        if w0.in_features() != w1.in_features() {
13515            return Ok(false);
13516        }
13517        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
13518            match w {
13519                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
13520                    Some(m) => (m, true),
13521                    None => (bytes, *rp),
13522                },
13523                _ => unreachable!(),
13524            }
13525        }
13526        let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
13527        if rp0 != rp1 {
13528            return Ok(false);
13529        }
13530        let rp = rp0;
13531        let rpb: u32 = 4;
13532        let mr1 = rp && Self::q40_mr1_on();
13533        let nb = |o: usize| {
13534            if mr1 {
13535                (o as u32).div_ceil(rpb)
13536            } else {
13537                (o as u32).div_ceil(2).div_ceil(rpb)
13538            }
13539        };
13540        let grid = nb(o0) + nb(o1);
13541        debug_assert!(y0.len() >= o0 && y1.len() >= o1);
13542        let f = self.func(if mr1 {
13543            "qmatvec_q4_0_mmvq_fused2_mr1_rp"
13544        } else if rp {
13545            "qmatvec_q4_0_mmvq_fused2_rp"
13546        } else {
13547            "qmatvec_q4_0_mmvq_fused2"
13548        });
13549        let cfg = LaunchConfig {
13550            grid_dim: (grid, 1, 1),
13551            block_dim: (32, rpb, 1),
13552            shared_mem_bytes: 0,
13553        };
13554        let inf = w0.in_features() as i32;
13555        let (oo0, oo1) = (o0 as i32, o1 as i32);
13556        let (r0, r1) = (rb0 as i64, rb1 as i64);
13557        // PDL wave-A: identical to the owned twin (capture-lane parity).
13558        if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
13559            use cudarc::driver::{DevicePtr, DevicePtrMut};
13560            let s = &self.gpu.stream();
13561            let (p0, _g0) = b0.device_ptr(s);
13562            let (p1, _g1) = b1.device_ptr(s);
13563            let (paq, _g2) = aq.device_ptr(s);
13564            let (pad, _g3) = ad.device_ptr(s);
13565            let (py0, _g4) = y0.device_ptr_mut(s);
13566            let (py1, _g5) = y1.device_ptr_mut(s);
13567            let mut ps = [
13568                &p0 as *const _ as *mut std::ffi::c_void,
13569                &p1 as *const _ as *mut _,
13570                &paq as *const _ as *mut _,
13571                &pad as *const _ as *mut _,
13572                &py0 as *const _ as *mut _,
13573                &py1 as *const _ as *mut _,
13574                &inf as *const _ as *mut _,
13575                &oo0 as *const _ as *mut _,
13576                &oo1 as *const _ as *mut _,
13577                &r0 as *const _ as *mut _,
13578                &r1 as *const _ as *mut _,
13579            ];
13580            unsafe {
13581                self.launch_pdl(
13582                    "qmatvec_q4_0_mmvq_fused2_mr1_rp",
13583                    (grid, 1, 1),
13584                    (32, rpb, 1),
13585                    &mut ps,
13586                )?;
13587            }
13588            return Ok(true);
13589        }
13590        let __s_b = self.gpu.stream();
13591        let mut b = __s_b.launch_builder(&f);
13592        b.arg(b0)
13593            .arg(b1)
13594            .arg(aq)
13595            .arg(ad)
13596            .arg(&mut *y0)
13597            .arg(&mut *y1)
13598            .arg(&inf)
13599            .arg(&oo0)
13600            .arg(&oo1)
13601            .arg(&r0)
13602            .arg(&r1);
13603        unsafe {
13604            b.launch(cfg)?;
13605        }
13606        Ok(true)
13607    }
13608
13609    /// BATCHED fused2 (2026-07-13, megakernel-microcosm probe): gate+up b-tier matvecs in
13610    /// ONE segmented-grid launch — the up segment fills SMs as the gate segment drains
13611    /// (the per-launch tail waves behind the 6x-falsified b-tier plateau). Bit-identical
13612    /// per row to two mr2_rp launches. rp layout required; m in 2..=8 (b16 has no twin).
13613    pub fn matmul_q4_fused2_batched(
13614        &self,
13615        w0: &crate::model::GpuTensor,
13616        w1: &crate::model::GpuTensor,
13617        aq: &CudaSlice<i8>,
13618        ad: &CudaSlice<f32>,
13619        m: usize,
13620    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
13621        use crate::model::GpuTensor;
13622        if m < 2 || m > 8 {
13623            return Ok(None);
13624        }
13625        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
13626            match w {
13627                GpuTensor::Quant {
13628                    qtype, row_bytes, ..
13629                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
13630                _ => None,
13631            }
13632        };
13633        let (Some((rb0, o0)), Some((_rb1, o1))) = (q4(w0), q4(w1)) else {
13634            return Ok(None);
13635        };
13636        if w0.in_features() != w1.in_features() {
13637            return Ok(None);
13638        }
13639        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
13640            match w {
13641                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
13642                    Some(mr) => (mr, true),
13643                    None => (bytes, *rp),
13644                },
13645                _ => unreachable!(),
13646            }
13647        }
13648        let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
13649        if !rp0 || !rp1 {
13650            return Ok(None);
13651        }
13652        let mcols = Self::batched_mcols(m);
13653        let rpb: u32 = 4;
13654        let nb = |o: usize| (o as u32).div_ceil(2 * rpb);
13655        let grid = nb(o0) + nb(o1);
13656        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
13657        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
13658        let f = self.func(match mcols {
13659            2 => "qmatvec_q4_0_mmvq_b2_f2_rp",
13660            4 => "qmatvec_q4_0_mmvq_b4_f2_rp",
13661            _ => "qmatvec_q4_0_mmvq_b8_f2_rp",
13662        });
13663        let cfg = LaunchConfig {
13664            grid_dim: (grid, 1, 1),
13665            block_dim: (32, rpb, 1),
13666            shared_mem_bytes: 0,
13667        };
13668        let inf = w0.in_features() as i32;
13669        let (oo0, oo1, mi) = (o0 as i32, o1 as i32, m as i32);
13670        let rb = rb0 as i64;
13671        let __s_b = self.gpu.stream();
13672        let mut b = __s_b.launch_builder(&f);
13673        b.arg(b0)
13674            .arg(b1)
13675            .arg(aq)
13676            .arg(ad)
13677            .arg(&mut y0)
13678            .arg(&mut y1)
13679            .arg(&inf)
13680            .arg(&oo0)
13681            .arg(&oo1)
13682            .arg(&mi)
13683            .arg(&rb);
13684        unsafe {
13685            b.launch(cfg)?;
13686        }
13687        Ok(Some((y0, y1)))
13688    }
13689
13690    /// BATCHED fused3 (see matmul_q4_fused2_batched): three-segment single launch for the
13691    /// verify qkv triple. Same-in_f q4_0 rp tensors, m in 2..=8. Bit-identical per row.
13692    #[allow(clippy::too_many_arguments)]
13693    pub fn matmul_q4_fused3_batched(
13694        &self,
13695        w0: &crate::model::GpuTensor,
13696        w1: &crate::model::GpuTensor,
13697        w2: &crate::model::GpuTensor,
13698        aq: &CudaSlice<i8>,
13699        ad: &CudaSlice<f32>,
13700        m: usize,
13701    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
13702    {
13703        use crate::model::GpuTensor;
13704        if m < 2 || m > 8 {
13705            return Ok(None);
13706        }
13707        let q4 = |w: &GpuTensor| -> Option<usize> {
13708            match w {
13709                GpuTensor::Quant { qtype, .. } if *qtype == QT_Q4_0 => Some(w.out_features()),
13710                _ => None,
13711            }
13712        };
13713        let (Some(o0), Some(o1), Some(o2)) = (q4(w0), q4(w1), q4(w2)) else {
13714            return Ok(None);
13715        };
13716        if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
13717            return Ok(None);
13718        }
13719        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
13720            match w {
13721                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
13722                    Some(mr) => (mr, true),
13723                    None => (bytes, *rp),
13724                },
13725                _ => unreachable!(),
13726            }
13727        }
13728        let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
13729        if !rp0 || !rp1 || !rp2 {
13730            return Ok(None);
13731        }
13732        let mcols = Self::batched_mcols(m);
13733        let rpb: u32 = 4;
13734        let nb = |o: usize| (o as u32).div_ceil(2 * rpb);
13735        let grid = nb(o0) + nb(o1) + nb(o2);
13736        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
13737        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
13738        let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
13739        let f = self.func(match mcols {
13740            2 => "qmatvec_q4_0_mmvq_b2_f3_rp",
13741            4 => "qmatvec_q4_0_mmvq_b4_f3_rp",
13742            _ => "qmatvec_q4_0_mmvq_b8_f3_rp",
13743        });
13744        let cfg = LaunchConfig {
13745            grid_dim: (grid, 1, 1),
13746            block_dim: (32, rpb, 1),
13747            shared_mem_bytes: 0,
13748        };
13749        let inf = w0.in_features() as i32;
13750        let (oo0, oo1, oo2, mi) = (o0 as i32, o1 as i32, o2 as i32, m as i32);
13751        let rb = 0i64;
13752        let __s_b = self.gpu.stream();
13753        let mut b = __s_b.launch_builder(&f);
13754        b.arg(b0)
13755            .arg(b1)
13756            .arg(b2)
13757            .arg(aq)
13758            .arg(ad)
13759            .arg(&mut y0)
13760            .arg(&mut y1)
13761            .arg(&mut y2)
13762            .arg(&inf)
13763            .arg(&oo0)
13764            .arg(&oo1)
13765            .arg(&oo2)
13766            .arg(&mi)
13767            .arg(&rb);
13768        unsafe {
13769            b.launch(cfg)?;
13770        }
13771        Ok(Some((y0, y1, y2)))
13772    }
13773
13774    pub fn matmul_q8_fused3(
13775        &self,
13776        w0: &crate::model::GpuTensor,
13777        w1: &crate::model::GpuTensor,
13778        w2: &crate::model::GpuTensor,
13779        aq: &CudaSlice<i8>,
13780        ad: &CudaSlice<f32>,
13781    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
13782    {
13783        // e4m3 twin (lane/fp8-decode-v1): the full-attn wq/wk/wv triple — on the NV-27B those three
13784        // are per-tensor FP8, so native residency without this arm meant three separate launches.
13785        if let Some([p0, p1, p2]) = self.e4m3_fused_params(&[w0, w1, w2]) {
13786            return Ok(Some(self.e4m3_fused3_core(
13787                p0.0,
13788                p1.0,
13789                p2.0,
13790                aq,
13791                ad,
13792                w0.in_features(),
13793                p0.1,
13794                p1.1,
13795                p2.1,
13796                p0.2,
13797                p0.3,
13798                p1.3,
13799                p2.3,
13800            )?));
13801        }
13802        let Some([p0, p1, p2]) = self.q8_fused_params(&[w0, w1, w2]) else {
13803            return Ok(None);
13804        };
13805        Ok(Some(self.q8_fused3_core(
13806            p0.0,
13807            p1.0,
13808            p2.0,
13809            aq,
13810            ad,
13811            w0.in_features(),
13812            p0.1,
13813            p1.1,
13814            p2.1,
13815            p0.2,
13816        )?))
13817    }
13818
13819    #[allow(clippy::too_many_arguments)]
13820    fn q8_fused3_core(
13821        &self,
13822        b0: &CudaSlice<u8>,
13823        b1: &CudaSlice<u8>,
13824        b2: &CudaSlice<u8>,
13825        aq: &CudaSlice<i8>,
13826        ad: &CudaSlice<f32>,
13827        in_f: usize,
13828        out0: usize,
13829        out1: usize,
13830        out2: usize,
13831        row_bytes: usize,
13832    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13833        const ROWS_PER_BLOCK: u32 = 4;
13834        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
13835        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
13836        let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
13837        let f = self.func("qmatvec_q8_0_mmvq_fused3");
13838        let mut y0 = self.alloc_uninit::<f32>(out0)?;
13839        let mut y1 = self.alloc_uninit::<f32>(out1)?;
13840        let mut y2 = self.alloc_uninit::<f32>(out2)?;
13841        let cfg = LaunchConfig {
13842            grid_dim: (nb0 + nb1 + nb2, 1, 1),
13843            block_dim: (32, ROWS_PER_BLOCK, 1),
13844            shared_mem_bytes: 0,
13845        };
13846        let (inf, o0, o1, o2, rbl) = (
13847            in_f as i32,
13848            out0 as i32,
13849            out1 as i32,
13850            out2 as i32,
13851            row_bytes as i64,
13852        );
13853        let __s_b = self.gpu.stream();
13854        let mut b = __s_b.launch_builder(&f);
13855        b.arg(b0)
13856            .arg(b1)
13857            .arg(b2)
13858            .arg(aq)
13859            .arg(ad)
13860            .arg(&mut y0)
13861            .arg(&mut y1)
13862            .arg(&mut y2)
13863            .arg(&inf)
13864            .arg(&o0)
13865            .arg(&o1)
13866            .arg(&o2)
13867            .arg(&rbl);
13868        unsafe {
13869            b.launch(cfg)?;
13870        }
13871        Ok((y0, y1, y2))
13872    }
13873
13874    /// Test entry for the kernel_check gate: fused3 from raw weight bytes (internal q8_1 quant).
13875    #[allow(clippy::too_many_arguments)]
13876    pub fn qmatvec_q8_fused3_raw(
13877        &self,
13878        b0: &CudaSlice<u8>,
13879        b1: &CudaSlice<u8>,
13880        b2: &CudaSlice<u8>,
13881        x: &CudaSlice<f32>,
13882        in_f: usize,
13883        out0: usize,
13884        out1: usize,
13885        out2: usize,
13886        row_bytes: usize,
13887    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13888        let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
13889        self.q8_fused3_core(b0, b1, b2, &aq, &ad, in_f, out0, out1, out2, row_bytes)
13890    }
13891
13892    /// BATCHED twin of `matmul_q8_fused2` for the verify t=2-4 tier (MEMRA_SPEC_FUSED_T call
13893    /// sites, lane/close35b): ONE launch computes both same-input Q8_0 projections for m tokens.
13894    /// Per (tensor,token,row) the kernel body is q8_0_mmvq_batched VERBATIM with the identical
13895    /// row mapping (Q8_0's batched_variant is always "base") -> BIT-IDENTICAL to the two
13896    /// per-tensor _b2/_b4 launches `matmul_decode_exact` dispatches at m=2-4, with the caller's
13897    /// single shared q8_1 activation replacing two per-call re-quantizes (quantize_q8_1 is
13898    /// deterministic -> same bytes). None when ineligible (m outside 2..=4 / not both Q8_0 /
13899    /// in_f mismatch / MEMRA_MMVQ=0 / MEMRA_Q8_DUAL=0 / MEMRA_NO_BATCHED set — the last keeps
13900    /// dispatch parity: without batched kernels decode-exact runs grid.y=m MMVQ, and the fused
13901    /// twin must not introduce a batched program the reference path would not run).
13902    pub fn matmul_q8_fused2_t(
13903        &self,
13904        w0: &crate::model::GpuTensor,
13905        w1: &crate::model::GpuTensor,
13906        aq: &CudaSlice<i8>,
13907        ad: &CudaSlice<f32>,
13908        m: usize,
13909    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
13910        // m<=8 (lane/q27-deepdive, 2026-08-05): was 2..=4 (the verify tier's mcols 2/4). The
13911        // serving tick's mcols-8 tier now has its fused2_b8 wrapper, so c=5..8 batched decode
13912        // fuses too — same template body, still bit-identical to the two _b8 launches.
13913        if !(2..=8).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok() {
13914            return Ok(None);
13915        }
13916        // e4m3 twin: MEMRA_B8 parity — without it m=5..8 e4m3 decode runs the per-m grid.y=m path,
13917        // so the fused b8 launch would introduce a batched program the reference path would not run.
13918        if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
13919            if m > 4 && !Self::b8_enabled() {
13920                return Ok(None);
13921            }
13922            return Ok(Some(self.e4m3_fused2_t_core(
13923                p0.0,
13924                p1.0,
13925                aq,
13926                ad,
13927                m,
13928                w0.in_features(),
13929                p0.1,
13930                p1.1,
13931                p0.2,
13932                p0.3,
13933                p1.3,
13934            )?));
13935        }
13936        let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
13937            return Ok(None);
13938        };
13939        Ok(Some(self.q8_fused2_t_core(
13940            p0.0,
13941            p1.0,
13942            aq,
13943            ad,
13944            m,
13945            w0.in_features(),
13946            p0.1,
13947            p1.1,
13948            p0.2,
13949        )?))
13950    }
13951
13952    #[allow(clippy::too_many_arguments)]
13953    fn q8_fused2_t_core(
13954        &self,
13955        b0: &CudaSlice<u8>,
13956        b1: &CudaSlice<u8>,
13957        aq: &CudaSlice<i8>,
13958        ad: &CudaSlice<f32>,
13959        m: usize,
13960        in_f: usize,
13961        out0: usize,
13962        out1: usize,
13963        row_bytes: usize,
13964    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13965        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
13966        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
13967        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
13968        let f = self.func(match Self::batched_mcols(m) {
13969            2 => "qmatvec_q8_0_mmvq_fused2_b2",
13970            4 => "qmatvec_q8_0_mmvq_fused2_b4",
13971            // b8 = the SERVING tier (lane/q27-deepdive): c=5..8 batched decode.
13972            _ => "qmatvec_q8_0_mmvq_fused2_b8",
13973        });
13974        let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
13975        let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
13976        let cfg = LaunchConfig {
13977            grid_dim: (nb0 + nb1, 1, 1),
13978            block_dim: (32, ROWS_PER_BLOCK, 1),
13979            shared_mem_bytes: 0,
13980        };
13981        let (inf, o0, o1, mi, rbl) = (
13982            in_f as i32,
13983            out0 as i32,
13984            out1 as i32,
13985            m as i32,
13986            row_bytes as i64,
13987        );
13988        let __s_b = self.gpu.stream();
13989        let mut b = __s_b.launch_builder(&f);
13990        b.arg(b0)
13991            .arg(b1)
13992            .arg(aq)
13993            .arg(ad)
13994            .arg(&mut y0)
13995            .arg(&mut y1)
13996            .arg(&inf)
13997            .arg(&o0)
13998            .arg(&o1)
13999            .arg(&mi)
14000            .arg(&rbl);
14001        unsafe {
14002            b.launch(cfg)?;
14003        }
14004        Ok((y0, y1))
14005    }
14006
14007    /// Test entry for the kernel_check gate: fused2 batched from raw weight bytes (internal
14008    /// q8_1 quant of the [m, in_f] activation), no env gating.
14009    #[allow(clippy::too_many_arguments)]
14010    pub fn qmatvec_q8_fused2_t_raw(
14011        &self,
14012        b0: &CudaSlice<u8>,
14013        b1: &CudaSlice<u8>,
14014        x: &CudaSlice<f32>,
14015        m: usize,
14016        in_f: usize,
14017        out0: usize,
14018        out1: usize,
14019        row_bytes: usize,
14020    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14021        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
14022        self.q8_fused2_t_core(b0, b1, &aq, &ad, m, in_f, out0, out1, row_bytes)
14023    }
14024
14025    /// BATCHED twin of `matmul_q8_fused3` (wq+wk+wv at verify t=2-4). Same contract as
14026    /// `matmul_q8_fused2_t` with three ranges.
14027    #[allow(clippy::too_many_arguments)]
14028    pub fn matmul_q8_fused3_t(
14029        &self,
14030        w0: &crate::model::GpuTensor,
14031        w1: &crate::model::GpuTensor,
14032        w2: &crate::model::GpuTensor,
14033        aq: &CudaSlice<i8>,
14034        ad: &CudaSlice<f32>,
14035        m: usize,
14036    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
14037    {
14038        if !(2..=4).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok() {
14039            return Ok(None);
14040        }
14041        if let Some([p0, p1, p2]) = self.e4m3_fused_params(&[w0, w1, w2]) {
14042            return Ok(Some(self.e4m3_fused3_t_core(
14043                p0.0,
14044                p1.0,
14045                p2.0,
14046                aq,
14047                ad,
14048                m,
14049                w0.in_features(),
14050                p0.1,
14051                p1.1,
14052                p2.1,
14053                p0.2,
14054                p0.3,
14055                p1.3,
14056                p2.3,
14057            )?));
14058        }
14059        let Some([p0, p1, p2]) = self.q8_fused_params(&[w0, w1, w2]) else {
14060            return Ok(None);
14061        };
14062        Ok(Some(self.q8_fused3_t_core(
14063            p0.0,
14064            p1.0,
14065            p2.0,
14066            aq,
14067            ad,
14068            m,
14069            w0.in_features(),
14070            p0.1,
14071            p1.1,
14072            p2.1,
14073            p0.2,
14074        )?))
14075    }
14076
14077    #[allow(clippy::too_many_arguments)]
14078    fn q8_fused3_t_core(
14079        &self,
14080        b0: &CudaSlice<u8>,
14081        b1: &CudaSlice<u8>,
14082        b2: &CudaSlice<u8>,
14083        aq: &CudaSlice<i8>,
14084        ad: &CudaSlice<f32>,
14085        m: usize,
14086        in_f: usize,
14087        out0: usize,
14088        out1: usize,
14089        out2: usize,
14090        row_bytes: usize,
14091    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14092        const ROWS_PER_BLOCK: u32 = 4;
14093        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14094        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14095        let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
14096        let f = self.func(if Self::batched_mcols(m) == 2 {
14097            "qmatvec_q8_0_mmvq_fused3_b2"
14098        } else {
14099            "qmatvec_q8_0_mmvq_fused3_b4"
14100        });
14101        let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
14102        let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
14103        let mut y2 = self.alloc_uninit::<f32>(m * out2)?;
14104        let cfg = LaunchConfig {
14105            grid_dim: (nb0 + nb1 + nb2, 1, 1),
14106            block_dim: (32, ROWS_PER_BLOCK, 1),
14107            shared_mem_bytes: 0,
14108        };
14109        let (inf, o0, o1, o2, mi, rbl) = (
14110            in_f as i32,
14111            out0 as i32,
14112            out1 as i32,
14113            out2 as i32,
14114            m as i32,
14115            row_bytes as i64,
14116        );
14117        let __s_b = self.gpu.stream();
14118        let mut b = __s_b.launch_builder(&f);
14119        b.arg(b0)
14120            .arg(b1)
14121            .arg(b2)
14122            .arg(aq)
14123            .arg(ad)
14124            .arg(&mut y0)
14125            .arg(&mut y1)
14126            .arg(&mut y2)
14127            .arg(&inf)
14128            .arg(&o0)
14129            .arg(&o1)
14130            .arg(&o2)
14131            .arg(&mi)
14132            .arg(&rbl);
14133        unsafe {
14134            b.launch(cfg)?;
14135        }
14136        Ok((y0, y1, y2))
14137    }
14138
14139    /// Test entry for the kernel_check gate: fused3 batched from raw weight bytes.
14140    #[allow(clippy::too_many_arguments)]
14141    pub fn qmatvec_q8_fused3_t_raw(
14142        &self,
14143        b0: &CudaSlice<u8>,
14144        b1: &CudaSlice<u8>,
14145        b2: &CudaSlice<u8>,
14146        x: &CudaSlice<f32>,
14147        m: usize,
14148        in_f: usize,
14149        out0: usize,
14150        out1: usize,
14151        out2: usize,
14152        row_bytes: usize,
14153    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14154        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
14155        self.q8_fused3_t_core(b0, b1, b2, &aq, &ad, m, in_f, out0, out1, out2, row_bytes)
14156    }
14157
14158    /// Rollback seam for the Q8_0 dense-FFN gate+up fusion arm in `matmul_pre_dual_noscale`
14159    /// (lane/q27-deepdive, 2026-08-05). Default ON; `MEMRA_Q8_FFN_FUSE2=0` restores the
14160    /// two-`matmul_pre_noscale` pair. Read once — the dispatch must not vary within a run.
14161    pub fn q8_ffn_fuse2_on(&self) -> bool {
14162        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
14163        *ON.get_or_init(|| std::env::var("MEMRA_Q8_FFN_FUSE2").as_deref() != Ok("0"))
14164    }
14165
14166    /// Eligibility + param extraction for the fused q8_0 launches: every tensor must be Quant Q8_0
14167    /// with macro-scale 1.0 (always true for GGUF q8_0; only NVFP4 carries scale) and share w[0]'s
14168    /// in_f (q8_0 row_bytes is a pure function of in_f, so equal in_f => equal row_bytes). MEMRA_MMVQ
14169    /// must be on: the fused body is the MMVQ kernel; without it decode m=1 runs dp4a and fusing
14170    /// would mix dispatch families (FP-order law). MEMRA_Q8_DUAL=0 = rollback seam.
14171    #[allow(clippy::type_complexity)]
14172    fn q8_fused_params<'w, const N: usize>(
14173        &self,
14174        ws: &[&'w crate::model::GpuTensor; N],
14175    ) -> Option<[(&'w CudaSlice<u8>, usize, usize); N]> {
14176        use crate::model::GpuTensor;
14177        if std::env::var("MEMRA_MMVQ").as_deref() == Ok("0") {
14178            return None;
14179        }
14180        if std::env::var("MEMRA_Q8_DUAL").is_ok_and(|v| v == "0") {
14181            return None;
14182        }
14183        let in_f = ws[0].in_features();
14184        let mut out: [Option<(&CudaSlice<u8>, usize, usize)>; N] = [None; N];
14185        for (i, w) in ws.iter().enumerate() {
14186            match w {
14187                GpuTensor::Quant {
14188                    bytes,
14189                    qtype,
14190                    row_bytes,
14191                    scale,
14192                    ..
14193                } if *qtype == QT_Q8_0 && *scale == 1.0 && w.in_features() == in_f => {
14194                    out[i] = Some((bytes, w.out_features(), *row_bytes))
14195                }
14196                _ => return None,
14197            }
14198        }
14199        Some(out.map(|o| o.unwrap()))
14200    }
14201
14202    /// Rollback seam for the F8-E4M3 launch-fusion arm (lane/fp8-decode-v1, 2026-08-05).
14203    /// Default ON; `MEMRA_E4M3_DUAL=0` restores the per-tensor m=1/batched launches.
14204    pub fn e4m3_dual_on(&self) -> bool {
14205        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
14206        *ON.get_or_init(|| std::env::var("MEMRA_E4M3_DUAL").as_deref() != Ok("0"))
14207    }
14208
14209    /// Eligibility + param extraction for the FUSED e4m3 launches — the QT_F8_E4M3 twin of
14210    /// `q8_fused_params`. Differences that are inherent to the dtype, not policy:
14211    ///   * each tensor carries its OWN per-tensor `weight_scale` (returned as the 4th field);
14212    ///     Q8_0 hard-requires scale==1.0 because it has no macro-scale at all.
14213    ///   * no MEMRA_MMVQ gate: `mmvq_supports` exempts QT_F8_E4M3 (the e4m3 mmvq family is that
14214    ///     dtype's ONLY int8-act kernel class), so the per-tensor fallback these fused kernels
14215    ///     replace is ALWAYS the same mmvq body under every env — the FP-order law holds.
14216    ///   * `row_bytes == in_f` is asserted rather than derived: the native-residency load arm keeps
14217    ///     the checkpoint's raw [out_f, in_f] rows, and a re-encoded slab must never reach here.
14218    /// Rejects any split-plane mirror (`rp`/`rp4`): there is no `_rp` e4m3 fused form, so fusing
14219    /// there would swap dispatch families mid-model. MEMRA_E4M3_DUAL=0 = rollback seam.
14220    #[allow(clippy::type_complexity)]
14221    fn e4m3_fused_params<'w, const N: usize>(
14222        &self,
14223        ws: &[&'w crate::model::GpuTensor; N],
14224    ) -> Option<[(&'w CudaSlice<u8>, usize, usize, f32); N]> {
14225        use crate::model::GpuTensor;
14226        if !self.e4m3_dual_on() {
14227            return None;
14228        }
14229        let in_f = ws[0].in_features();
14230        let mut out: [Option<(&CudaSlice<u8>, usize, usize, f32)>; N] = [None; N];
14231        for (i, w) in ws.iter().enumerate() {
14232            match w {
14233                GpuTensor::Quant {
14234                    bytes,
14235                    qtype,
14236                    row_bytes,
14237                    scale,
14238                    rp,
14239                    rp4,
14240                    ..
14241                } if *qtype == QT_F8_E4M3
14242                    && w.in_features() == in_f
14243                    && *row_bytes == in_f
14244                    && !*rp
14245                    && rp4.is_none() =>
14246                {
14247                    out[i] = Some((bytes, w.out_features(), *row_bytes, *scale))
14248                }
14249                _ => return None,
14250            }
14251        }
14252        Some(out.map(|o| o.unwrap()))
14253    }
14254
14255    /// FUSED e4m3 m=1 PAIR. Block-offset split (`qmatvec_e4m3_mmvq_fused2`), per-tensor
14256    /// weight_scale folded at the write like the single-tensor `qmatvec_e4m3_mmvq` — so per
14257    /// (tensor,row) this is BIT-IDENTICAL to two separate m=1 launches, scale included.
14258    #[allow(clippy::too_many_arguments)]
14259    fn e4m3_fused2_core(
14260        &self,
14261        b0: &CudaSlice<u8>,
14262        b1: &CudaSlice<u8>,
14263        aq: &CudaSlice<i8>,
14264        ad: &CudaSlice<f32>,
14265        in_f: usize,
14266        out0: usize,
14267        out1: usize,
14268        row_bytes: usize,
14269        ws0: f32,
14270        ws1: f32,
14271    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14272        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
14273        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14274        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14275        let f = self.func("qmatvec_e4m3_mmvq_fused2");
14276        let mut y0 = self.alloc_uninit::<f32>(out0)?;
14277        let mut y1 = self.alloc_uninit::<f32>(out1)?;
14278        let cfg = LaunchConfig {
14279            grid_dim: (nb0 + nb1, 1, 1),
14280            block_dim: (32, ROWS_PER_BLOCK, 1),
14281            shared_mem_bytes: 0,
14282        };
14283        let (inf, o0, o1, rbl) = (in_f as i32, out0 as i32, out1 as i32, row_bytes as i64);
14284        let __s_b = self.gpu.stream();
14285        let mut b = __s_b.launch_builder(&f);
14286        b.arg(b0)
14287            .arg(b1)
14288            .arg(aq)
14289            .arg(ad)
14290            .arg(&mut y0)
14291            .arg(&mut y1)
14292            .arg(&inf)
14293            .arg(&o0)
14294            .arg(&o1)
14295            .arg(&rbl)
14296            .arg(&ws0)
14297            .arg(&ws1);
14298        unsafe {
14299            b.launch(cfg)?;
14300        }
14301        Ok((y0, y1))
14302    }
14303
14304    /// FUSED e4m3 m=1 TRIPLE (`qmatvec_e4m3_mmvq_fused3`). Same contract as the pair.
14305    #[allow(clippy::too_many_arguments)]
14306    fn e4m3_fused3_core(
14307        &self,
14308        b0: &CudaSlice<u8>,
14309        b1: &CudaSlice<u8>,
14310        b2: &CudaSlice<u8>,
14311        aq: &CudaSlice<i8>,
14312        ad: &CudaSlice<f32>,
14313        in_f: usize,
14314        out0: usize,
14315        out1: usize,
14316        out2: usize,
14317        row_bytes: usize,
14318        ws0: f32,
14319        ws1: f32,
14320        ws2: f32,
14321    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14322        const ROWS_PER_BLOCK: u32 = 4;
14323        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14324        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14325        let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
14326        let f = self.func("qmatvec_e4m3_mmvq_fused3");
14327        let mut y0 = self.alloc_uninit::<f32>(out0)?;
14328        let mut y1 = self.alloc_uninit::<f32>(out1)?;
14329        let mut y2 = self.alloc_uninit::<f32>(out2)?;
14330        let cfg = LaunchConfig {
14331            grid_dim: (nb0 + nb1 + nb2, 1, 1),
14332            block_dim: (32, ROWS_PER_BLOCK, 1),
14333            shared_mem_bytes: 0,
14334        };
14335        let (inf, o0, o1, o2, rbl) = (
14336            in_f as i32,
14337            out0 as i32,
14338            out1 as i32,
14339            out2 as i32,
14340            row_bytes as i64,
14341        );
14342        let __s_b = self.gpu.stream();
14343        let mut b = __s_b.launch_builder(&f);
14344        b.arg(b0)
14345            .arg(b1)
14346            .arg(b2)
14347            .arg(aq)
14348            .arg(ad)
14349            .arg(&mut y0)
14350            .arg(&mut y1)
14351            .arg(&mut y2)
14352            .arg(&inf)
14353            .arg(&o0)
14354            .arg(&o1)
14355            .arg(&o2)
14356            .arg(&rbl)
14357            .arg(&ws0)
14358            .arg(&ws1)
14359            .arg(&ws2);
14360        unsafe {
14361            b.launch(cfg)?;
14362        }
14363        Ok((y0, y1, y2))
14364    }
14365
14366    /// BATCHED FUSED e4m3 pair (m=2..8). The batched kernels carry no `ws` arg (every batched
14367    /// kernel in the tree is scale-free), so each output takes its own `scale_inplace` — the
14368    /// SAME post-op the per-tensor batched dispatch applies, hence still bit-identical.
14369    #[allow(clippy::too_many_arguments)]
14370    fn e4m3_fused2_t_core(
14371        &self,
14372        b0: &CudaSlice<u8>,
14373        b1: &CudaSlice<u8>,
14374        aq: &CudaSlice<i8>,
14375        ad: &CudaSlice<f32>,
14376        m: usize,
14377        in_f: usize,
14378        out0: usize,
14379        out1: usize,
14380        row_bytes: usize,
14381        ws0: f32,
14382        ws1: f32,
14383    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14384        const ROWS_PER_BLOCK: u32 = 4;
14385        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14386        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14387        let f = self.func(match Self::batched_mcols(m) {
14388            2 => "qmatvec_e4m3_mmvq_fused2_b2",
14389            4 => "qmatvec_e4m3_mmvq_fused2_b4",
14390            _ => "qmatvec_e4m3_mmvq_fused2_b8",
14391        });
14392        let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
14393        let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
14394        let cfg = LaunchConfig {
14395            grid_dim: (nb0 + nb1, 1, 1),
14396            block_dim: (32, ROWS_PER_BLOCK, 1),
14397            shared_mem_bytes: 0,
14398        };
14399        let (inf, o0, o1, mi, rbl) = (
14400            in_f as i32,
14401            out0 as i32,
14402            out1 as i32,
14403            m as i32,
14404            row_bytes as i64,
14405        );
14406        let __s_b = self.gpu.stream();
14407        let mut b = __s_b.launch_builder(&f);
14408        b.arg(b0)
14409            .arg(b1)
14410            .arg(aq)
14411            .arg(ad)
14412            .arg(&mut y0)
14413            .arg(&mut y1)
14414            .arg(&inf)
14415            .arg(&o0)
14416            .arg(&o1)
14417            .arg(&mi)
14418            .arg(&rbl);
14419        unsafe {
14420            b.launch(cfg)?;
14421        }
14422        if ws0 != 1.0 {
14423            self.scale_inplace(&mut y0, ws0, m * out0)?;
14424        }
14425        if ws1 != 1.0 {
14426            self.scale_inplace(&mut y1, ws1, m * out1)?;
14427        }
14428        Ok((y0, y1))
14429    }
14430
14431    /// BATCHED FUSED e4m3 triple (m=2..4). Same contract as the batched pair.
14432    #[allow(clippy::too_many_arguments)]
14433    fn e4m3_fused3_t_core(
14434        &self,
14435        b0: &CudaSlice<u8>,
14436        b1: &CudaSlice<u8>,
14437        b2: &CudaSlice<u8>,
14438        aq: &CudaSlice<i8>,
14439        ad: &CudaSlice<f32>,
14440        m: usize,
14441        in_f: usize,
14442        out0: usize,
14443        out1: usize,
14444        out2: usize,
14445        row_bytes: usize,
14446        ws0: f32,
14447        ws1: f32,
14448        ws2: f32,
14449    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14450        const ROWS_PER_BLOCK: u32 = 4;
14451        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14452        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14453        let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
14454        let f = self.func(if Self::batched_mcols(m) == 2 {
14455            "qmatvec_e4m3_mmvq_fused3_b2"
14456        } else {
14457            "qmatvec_e4m3_mmvq_fused3_b4"
14458        });
14459        let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
14460        let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
14461        let mut y2 = self.alloc_uninit::<f32>(m * out2)?;
14462        let cfg = LaunchConfig {
14463            grid_dim: (nb0 + nb1 + nb2, 1, 1),
14464            block_dim: (32, ROWS_PER_BLOCK, 1),
14465            shared_mem_bytes: 0,
14466        };
14467        let (inf, o0, o1, o2, mi, rbl) = (
14468            in_f as i32,
14469            out0 as i32,
14470            out1 as i32,
14471            out2 as i32,
14472            m as i32,
14473            row_bytes as i64,
14474        );
14475        let __s_b = self.gpu.stream();
14476        let mut b = __s_b.launch_builder(&f);
14477        b.arg(b0)
14478            .arg(b1)
14479            .arg(b2)
14480            .arg(aq)
14481            .arg(ad)
14482            .arg(&mut y0)
14483            .arg(&mut y1)
14484            .arg(&mut y2)
14485            .arg(&inf)
14486            .arg(&o0)
14487            .arg(&o1)
14488            .arg(&o2)
14489            .arg(&mi)
14490            .arg(&rbl);
14491        unsafe {
14492            b.launch(cfg)?;
14493        }
14494        if ws0 != 1.0 {
14495            self.scale_inplace(&mut y0, ws0, m * out0)?;
14496        }
14497        if ws1 != 1.0 {
14498            self.scale_inplace(&mut y1, ws1, m * out1)?;
14499        }
14500        if ws2 != 1.0 {
14501            self.scale_inplace(&mut y2, ws2, m * out2)?;
14502        }
14503        Ok((y0, y1, y2))
14504    }
14505
14506    /// BLOCK-128 e4m3 MMVQ launcher (`qmatvec_e4m3_blk_mmvq`, lane/fp8-blk128-decode 2026-08-05).
14507    /// The per-block-dequant twin of `qmatvec_mmvq`'s QT_F8_E4M3 arm: same grid/block decomposition
14508    /// (warp per output row, ROWS_PER_BLOCK warps per block, grid.y = m), same q8_1 activation, but
14509    /// the weight scale is a resident [rows, cols] f32 grid read per k128 block inside the kernel
14510    /// instead of one scalar folded at the write. It cannot share `qmatvec_mmvq`'s body because
14511    /// that launcher's arg list is fixed at (bytes, aq, ad, y, in_f, out_f, m, row_bytes [, scale]).
14512    ///
14513    /// `mr` and `rp` have no analogue here (no split-plane e4m3 layout exists), so there is exactly
14514    /// one kernel and no name table — a shape this cannot serve must be refused at LOAD, not here.
14515    pub fn qmatvec_e4m3_blk_mmvq(
14516        &self,
14517        bytes: &CudaSlice<u8>,
14518        aq: &CudaSlice<i8>,
14519        ad: &CudaSlice<f32>,
14520        scales: &CudaSlice<f32>,
14521        m: usize,
14522        in_f: usize,
14523        out_f: usize,
14524        row_bytes: usize,
14525        scale_cols: usize,
14526    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14527        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
14528        self.qmatvec_e4m3_blk_mmvq_into(
14529            bytes, aq, ad, scales, m, in_f, out_f, row_bytes, scale_cols, &mut y,
14530        )?;
14531        Ok(y)
14532    }
14533
14534    /// Slot-fed twin of `qmatvec_e4m3_blk_mmvq` (caller-owned output; the alloc-free capture lane).
14535    #[allow(clippy::too_many_arguments)]
14536    pub fn qmatvec_e4m3_blk_mmvq_into(
14537        &self,
14538        bytes: &CudaSlice<u8>,
14539        aq: &CudaSlice<i8>,
14540        ad: &CudaSlice<f32>,
14541        scales: &CudaSlice<f32>,
14542        m: usize,
14543        in_f: usize,
14544        out_f: usize,
14545        row_bytes: usize,
14546        scale_cols: usize,
14547        y: &mut CudaSlice<f32>,
14548    ) -> Result<(), Box<dyn std::error::Error>> {
14549        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
14550        let f = self.func("qmatvec_e4m3_blk_mmvq");
14551        let cfg = LaunchConfig {
14552            grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), m as u32, 1),
14553            block_dim: (32, ROWS_PER_BLOCK, 1), // warp-per-row
14554            shared_mem_bytes: 0,                // warp-only reduce
14555        };
14556        let (inf, outf, mi, rb, sc) = (
14557            in_f as i32,
14558            out_f as i32,
14559            m as i32,
14560            row_bytes as i64,
14561            scale_cols as i32,
14562        );
14563        let __s_b = self.gpu.stream();
14564        let mut b = __s_b.launch_builder(&f);
14565        b.arg(bytes)
14566            .arg(aq)
14567            .arg(ad)
14568            .arg(scales)
14569            .arg(&mut *y)
14570            .arg(&inf)
14571            .arg(&outf)
14572            .arg(&mi)
14573            .arg(&rb)
14574            .arg(&sc);
14575        unsafe {
14576            b.launch(cfg)?;
14577        }
14578        Ok(())
14579    }
14580
14581    /// BLOCK-128 e4m3 BATCHED matvec (lane/rp-on-st, 2026-08-06): the weight-read-once twin of
14582    /// `qmatvec_e4m3_blk_mmvq` for m=2..16. Per (token,row) BIT-IDENTICAL to the grid.y=m launch
14583    /// (same fmaf chain, same per-k32 `s * ad` fold, same warp reduce), so it inherits the
14584    /// decode-exactness contract while reading the weight ONCE for up to `mcols` columns instead
14585    /// of `m` times. `mcols` must be one of {2,4,8,16} and satisfy `mcols >= m`.
14586    #[allow(clippy::too_many_arguments)]
14587    pub fn qmatvec_e4m3_blk_mmvq_batched(
14588        &self,
14589        bytes: &CudaSlice<u8>,
14590        aq: &CudaSlice<i8>,
14591        ad: &CudaSlice<f32>,
14592        scales: &CudaSlice<f32>,
14593        m: usize,
14594        in_f: usize,
14595        out_f: usize,
14596        row_bytes: usize,
14597        scale_cols: usize,
14598        mcols: usize,
14599    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14600        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
14601        debug_assert!(mcols >= m, "blk batched: mcols {mcols} < m {m}");
14602        let name = match mcols {
14603            2 => "qmatvec_e4m3_blk_mmvq_b2",
14604            4 => "qmatvec_e4m3_blk_mmvq_b4",
14605            8 => "qmatvec_e4m3_blk_mmvq_b8",
14606            16 => "qmatvec_e4m3_blk_mmvq_b16",
14607            _ => {
14608                return Err(
14609                    format!("qmatvec_e4m3_blk_mmvq_batched: no kernel for mcols {mcols}").into(),
14610                );
14611            }
14612        };
14613        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
14614        let f = self.func(name);
14615        let cfg = LaunchConfig {
14616            grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
14617            block_dim: (32, ROWS_PER_BLOCK, 1),
14618            shared_mem_bytes: 0,
14619        };
14620        let (inf, outf, mi, rb, sc) = (
14621            in_f as i32,
14622            out_f as i32,
14623            m as i32,
14624            row_bytes as i64,
14625            scale_cols as i32,
14626        );
14627        let __s_b = self.gpu.stream();
14628        let mut b = __s_b.launch_builder(&f);
14629        b.arg(bytes)
14630            .arg(aq)
14631            .arg(ad)
14632            .arg(scales)
14633            .arg(&mut y)
14634            .arg(&inf)
14635            .arg(&outf)
14636            .arg(&mi)
14637            .arg(&rb)
14638            .arg(&sc);
14639        unsafe {
14640            b.launch(cfg)?;
14641        }
14642        Ok(y)
14643    }
14644
14645    /// Test entry for the kernel_check exactness gate: the block-128 e4m3 batched MMVQ from raw
14646    /// bytes with an internal q8_1 quantize (mirrors `qmatvec_batched_raw`).
14647    #[allow(clippy::too_many_arguments)]
14648    pub fn qmatvec_e4m3_blk_batched_raw(
14649        &self,
14650        bytes: &CudaSlice<u8>,
14651        x: &CudaSlice<f32>,
14652        scales: &CudaSlice<f32>,
14653        m: usize,
14654        in_f: usize,
14655        out_f: usize,
14656        row_bytes: usize,
14657        scale_cols: usize,
14658        mcols: usize,
14659    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14660        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
14661        self.qmatvec_e4m3_blk_mmvq_batched(
14662            bytes, &aq, &ad, scales, m, in_f, out_f, row_bytes, scale_cols, mcols,
14663        )
14664    }
14665
14666    /// Test entry for the kernel_check exactness gate: the block-128 e4m3 MMVQ from raw bytes with
14667    /// an internal q8_1 quantize (mirrors `qmatvec_mmvq_raw`).
14668    #[allow(clippy::too_many_arguments)]
14669    pub fn qmatvec_e4m3_blk_mmvq_raw(
14670        &self,
14671        bytes: &CudaSlice<u8>,
14672        x: &CudaSlice<f32>,
14673        scales: &CudaSlice<f32>,
14674        m: usize,
14675        in_f: usize,
14676        out_f: usize,
14677        row_bytes: usize,
14678        scale_cols: usize,
14679    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14680        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
14681        self.qmatvec_e4m3_blk_mmvq(
14682            bytes, &aq, &ad, scales, m, in_f, out_f, row_bytes, scale_cols,
14683        )
14684    }
14685
14686    /// Test entries for the kernel_check bit-parity gate: fused e4m3 launches from raw weight
14687    /// bytes with internal q8_1 quantize, no env gating (mirrors `qmatvec_q8_fused*_raw`).
14688    #[allow(clippy::too_many_arguments)]
14689    pub fn qmatvec_e4m3_fused2_raw(
14690        &self,
14691        b0: &CudaSlice<u8>,
14692        b1: &CudaSlice<u8>,
14693        x: &CudaSlice<f32>,
14694        in_f: usize,
14695        out0: usize,
14696        out1: usize,
14697        row_bytes: usize,
14698        ws0: f32,
14699        ws1: f32,
14700    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14701        let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
14702        self.e4m3_fused2_core(b0, b1, &aq, &ad, in_f, out0, out1, row_bytes, ws0, ws1)
14703    }
14704
14705    #[allow(clippy::too_many_arguments)]
14706    pub fn qmatvec_e4m3_fused3_raw(
14707        &self,
14708        b0: &CudaSlice<u8>,
14709        b1: &CudaSlice<u8>,
14710        b2: &CudaSlice<u8>,
14711        x: &CudaSlice<f32>,
14712        in_f: usize,
14713        out0: usize,
14714        out1: usize,
14715        out2: usize,
14716        row_bytes: usize,
14717        ws0: f32,
14718        ws1: f32,
14719        ws2: f32,
14720    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14721        let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
14722        self.e4m3_fused3_core(
14723            b0, b1, b2, &aq, &ad, in_f, out0, out1, out2, row_bytes, ws0, ws1, ws2,
14724        )
14725    }
14726
14727    #[allow(clippy::too_many_arguments)]
14728    pub fn qmatvec_e4m3_fused2_t_raw(
14729        &self,
14730        b0: &CudaSlice<u8>,
14731        b1: &CudaSlice<u8>,
14732        x: &CudaSlice<f32>,
14733        m: usize,
14734        in_f: usize,
14735        out0: usize,
14736        out1: usize,
14737        row_bytes: usize,
14738        ws0: f32,
14739        ws1: f32,
14740    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14741        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
14742        self.e4m3_fused2_t_core(b0, b1, &aq, &ad, m, in_f, out0, out1, row_bytes, ws0, ws1)
14743    }
14744
14745    #[allow(clippy::too_many_arguments)]
14746    pub fn qmatvec_e4m3_fused3_t_raw(
14747        &self,
14748        b0: &CudaSlice<u8>,
14749        b1: &CudaSlice<u8>,
14750        b2: &CudaSlice<u8>,
14751        x: &CudaSlice<f32>,
14752        m: usize,
14753        in_f: usize,
14754        out0: usize,
14755        out1: usize,
14756        out2: usize,
14757        row_bytes: usize,
14758        ws0: f32,
14759        ws1: f32,
14760        ws2: f32,
14761    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14762        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
14763        self.e4m3_fused3_t_core(
14764            b0, b1, b2, &aq, &ad, m, in_f, out0, out1, out2, row_bytes, ws0, ws1, ws2,
14765        )
14766    }
14767
14768    /// THE single dispatch point for `QT_F8_E4M3_BLK` from a PRE-QUANTIZED q8_1 activation
14769    /// (lane/fp8-blk128-decode). Every `matmul_pre`-family entry calls this first, so the block-128
14770    /// class has exactly ONE code path across `matmul`, `matmul_pre`, `matmul_pre_noscale`,
14771    /// `matmul_decode_exact` and `matmul_decode_exact_pre` — the same kernel at the same grid for
14772    /// every m, which is what makes verify == decode bit-for-bit at every tier for free.
14773    ///
14774    /// Returns None for any other qtype (the caller continues its normal dispatch). The `blk: Some`
14775    /// pattern is part of the match, not an unwrap: qtype and grid presence are set together in the
14776    /// one residency arm that builds this tensor, and a qtype-without-grid would be a construction
14777    /// bug — better to fall through and hit a loud refusal than to unwrap a None here.
14778    fn try_e4m3_blk_pre(
14779        &self,
14780        w: &crate::model::GpuTensor,
14781        aq: &CudaSlice<i8>,
14782        ad: &CudaSlice<f32>,
14783        m: usize,
14784    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
14785        use crate::model::GpuTensor;
14786        if let GpuTensor::Quant {
14787            bytes,
14788            qtype,
14789            row_bytes,
14790            blk: Some(g),
14791            ..
14792        } = w
14793        {
14794            if *qtype == QT_F8_E4M3_BLK {
14795                // BATCHED tier m=2..16 (lane/rp-on-st): weight read ONCE for up to mcols columns
14796                // instead of m grid.y re-reads. Bit-identical per (token,row) to the grid.y=m form
14797                // below, so the decode-exactness contract is preserved at every width. Gated by
14798                // the same seams the other batched families honor (MEMRA_NO_BATCHED, MEMRA_B8) so
14799                // one rollback door covers every dtype's batched tier.
14800                if (2..=16).contains(&m)
14801                    && std::env::var("MEMRA_NO_BATCHED").is_err()
14802                    && (m <= 4 || Self::b8_enabled())
14803                {
14804                    let mcols = Self::batched_mcols(m);
14805                    return Ok(Some(self.qmatvec_e4m3_blk_mmvq_batched(
14806                        bytes,
14807                        aq,
14808                        ad,
14809                        &g.scales,
14810                        m,
14811                        w.in_features(),
14812                        w.out_features(),
14813                        *row_bytes,
14814                        g.cols,
14815                        mcols,
14816                    )?));
14817                }
14818                return Ok(Some(self.qmatvec_e4m3_blk_mmvq(
14819                    bytes,
14820                    aq,
14821                    ad,
14822                    &g.scales,
14823                    m,
14824                    w.in_features(),
14825                    w.out_features(),
14826                    *row_bytes,
14827                    g.cols,
14828                )?));
14829            }
14830        }
14831        Ok(None)
14832    }
14833
14834    /// PREFILL (m >= GEMM_M_THRESHOLD) for `QT_F8_E4M3_BLK` — DEQUANT-PER-CALL to the Q8_0 slab
14835    /// this class's residency replaced, then the ordinary Q8_0 prefill dispatch on the transient.
14836    ///
14837    /// WHY THIS EXISTS AT ALL, i.e. the regression it prevents: the decode kernel is a warp-per-row
14838    /// GEMV. At grid.y=m it re-reads the whole weight once PER TOKEN, so letting a 512-token prefill
14839    /// chunk reach it would be a ~500x weight-traffic blowup on the single most bandwidth-bound part
14840    /// of the forward. Native residency is a DECODE win and must not be paid for in prefill, so
14841    /// prefill keeps the floor's arithmetic and the floor's kernels.
14842    ///
14843    /// WHY DEQUANT-PER-CALL rather than a second resident slab: a resident slab is dual residency —
14844    /// it gives back the entire 1.0-vs-1.0625 B/weight win this lane exists to capture (and then
14845    /// some, since the e4m3 copy stays too). The transient costs one linear device pass per
14846    /// (projection, prefill call) and frees immediately.
14847    ///
14848    /// NUMERICALLY IT IS THE FLOOR, EXACTLY: `fp8_blk_dequant_q8_0` is the merged ARM B' kernel,
14849    /// gate-proven BYTE-IDENTICAL to the host dequant+re-encode (kernel-check `fp8-blk-gpu`). So the
14850    /// slab these bytes form is bit-for-bit the slab the `MEMRA_ST_E4M3_BLK=0` arm makes resident,
14851    /// and every prefill kernel downstream sees identical input — prefill logits under this lane are
14852    /// bit-identical to prefill logits under the floor, which is what makes the decode A/B a clean
14853    /// single-variable comparison instead of a two-variable one.
14854    ///
14855    /// WHAT IT COSTS, MEASURED, AND WHY THAT COST IS MOSTLY STRUCTURAL (27B block-128 ckpt, pp512,
14856    /// this rig = RTX 5090 Laptop, ~896 GB/s GDDR7). This arm makes prefill move the weight THREE
14857    /// times instead of once: read 6.88 GB of e4m3, write 7.31 GB of Q8_0, then the MMQ reads that
14858    /// 7.31 GB back. The two extra passes are 14.19 GB = 15.8 ms at this card's roofline against a
14859    /// ~332 ms pp512, i.e. **~-4.5% pp is a floor no kernel tuning can remove** — only deleting the
14860    /// dequant can. Measured: the dequant kernel costs 27.9 ms/pass (nsys, 208 projections) after
14861    /// the 2026-08-05 vector rewrite (was 66.5 ms at one byte per thread), and e2e pp512 is
14862    /// 1451.4 vs the slab arm's 1541.6 tok/s = -5.8% (N=3 interleaved pairs). So ~1.3pp of the
14863    /// -5.8% is residual kernel inefficiency and ~4.5pp is the extra traffic itself.
14864    ///
14865    /// SO THE DEQUANT IS NO LONGER THE DEFAULT ROUTE — it is the FALLBACK. The per-block FP8 MMQ
14866    /// tile (`try_fp8_blk_mmq`) consumes the resident e4m3 bytes + grid DIRECTLY, deleting both extra
14867    /// passes, and since 2026-08-05 it runs FIRST and by default for the native-resident source
14868    /// (`fp8_blk_mmq_native_enabled`; `MEMRA_FP8_MMQ=0` is the seam back to this dequant). On paper
14869    /// the trade was unassumable — lane/fp8-mmq-v2 measured that tile at 0.85-1.09x the Q8_0 MMQ
14870    /// floor GEMM-only, so it swapped a -4.5% traffic cost for a 0-to-15% GEMM cost of unknown sign.
14871    /// Measured on the 27B (3 arms interleaved, N=3, research/fp8blk-20260805/VERDICT.md): slab
14872    /// 1540.5 / this dequant 1449.1 / the tile 1553.3 tok/s, min(tile) > max(slab). The tile wins
14873    /// because v2's denominator had its slab already resident while this class's floor must build it
14874    /// every call; same tile, opposite sign, because the question changed.
14875    ///
14876    /// THIS ARM STILL RUNS, and is not dead code: every `try_fp8_blk_mmq` precondition (in_f % 16,
14877    /// grid dims vs shape, per-tensor scale == 1.0, the e4m3-NaN scan) refuses by falling through to
14878    /// here, so a checkpoint the tile cannot take keeps exact prefill on the floor's own bits rather
14879    /// than losing the class. It is also what `MEMRA_FP8_MMQ=0` reverts to.
14880    fn try_e4m3_blk_prefill(
14881        &self,
14882        w: &crate::model::GpuTensor,
14883        x: &CudaSlice<f32>,
14884        m: usize,
14885    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
14886        use crate::model::GpuTensor;
14887        let GpuTensor::Quant {
14888            bytes,
14889            qtype,
14890            blk: Some(g),
14891            ..
14892        } = w
14893        else {
14894            return Ok(None);
14895        };
14896        if *qtype != QT_F8_E4M3_BLK {
14897            return Ok(None);
14898        }
14899        // NO-DEQUANT ROUTE, THE DEFAULT (MEMRA_FP8_MMQ=0 reverts): the per-block MMQ tile eats the
14900        // resident e4m3 bytes and grid as-is, so neither extra weight pass happens. Its own
14901        // preconditions (in_f % 16, grid dims, scale == 1.0, no e4m3 NaN code) can refuse — fall
14902        // through to the dequant below when they do, never silently produce nothing.
14903        if let Some(y) = self.try_fp8_blk_mmq(w, x, m)? {
14904            return Ok(Some(y));
14905        }
14906        let (in_f, out_f) = (w.in_features(), w.out_features());
14907        let slab = self.fp8_blk_dequant_q8_0_dev(bytes, &g.scales, out_f, in_f)?;
14908        let tmp = GpuTensor::Quant {
14909            bytes: slab,
14910            qtype: QT_Q8_0,
14911            row_bytes: in_f / 32 * 34,
14912            ne: vec![in_f as u64, out_f as u64],
14913            scale: 1.0,
14914            rp: false,
14915            #[cfg(memra_cutlass)]
14916            cutlass: None,
14917            fp8: None,
14918            blk: None,
14919            f16: None,
14920            rp4: None,
14921        };
14922        // Recursion terminates: `tmp` is QT_Q8_0 with `blk: None`, so it cannot re-enter this arm.
14923        Ok(Some(self.matmul(&tmp, x, m)?))
14924    }
14925
14926    pub fn matmul_pre_noscale(
14927        &self,
14928        w: &crate::model::GpuTensor,
14929        aq: &CudaSlice<i8>,
14930        ad: &CudaSlice<f32>,
14931        m: usize,
14932    ) -> Result<Option<(CudaSlice<f32>, f32)>, Box<dyn std::error::Error>> {
14933        use crate::model::GpuTensor;
14934        // BLOCK-128 e4m3: every scale factor is folded inside the kernel per k128, so the
14935        // "separable post-op scale" this entry exists to defer is 1.0 — return it explicitly
14936        // rather than let the tail below refuse and cost the caller a re-dispatch.
14937        if m == 1 {
14938            if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
14939                return Ok(Some((y, 1.0)));
14940            }
14941        }
14942        // Only the m==1 fast path applies the scale as a separable post-op; bail everywhere else.
14943        if m != 1 || !self.uses_q8_1_fast(w) {
14944            return Ok(None);
14945        }
14946        let in_f = w.in_features();
14947        let out_f = w.out_features();
14948        let (bytes, qtype, row_bytes, scale, rp) = match w {
14949            GpuTensor::Quant {
14950                bytes,
14951                qtype,
14952                row_bytes,
14953                scale,
14954                rp,
14955                ..
14956            } => (bytes, *qtype, *row_bytes, *scale, *rp),
14957            _ => return Ok(None),
14958        };
14959        // MMVQ warp-per-row (scale==1.0 passed -> kernel skips its internal scale; we return scale).
14960        if self.mmvq_supports(qtype) {
14961            // Q4_0 split-plane mirror (dp4a fallback below keeps the raw GGUF bytes).
14962            let (mbytes, mrp) = match w {
14963                GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
14964                _ => (bytes, rp),
14965            };
14966            let y = self.qmatvec_mmvq(
14967                mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, /*scale*/ 1.0, mrp,
14968            )?;
14969            return Ok(Some((y, scale)));
14970        }
14971        // dp4a fallback: same launch as matmul_pre but WITHOUT the post scale_inplace.
14972        let name = match qtype {
14973            QT_Q8_0 => "qmatvec_q8_0_dp4a",
14974            QT_Q4_K => "qmatvec_q4_K_dp4a",
14975            QT_Q6_K => "qmatvec_q6_K_dp4a",
14976            QT_Q5_K => "qmatvec_q5_K_dp4a",
14977            QT_Q3_K => "qmatvec_q3_K_dp4a",
14978            QT_NVFP4 => {
14979                if rp {
14980                    "qmatvec_nvfp4_dp4a_rp"
14981                } else {
14982                    "qmatvec_nvfp4_dp4a"
14983                }
14984            }
14985            QT_IQ4_XS => "qmatvec_iq4_XS_dp4a",
14986            _ => return Ok(None),
14987        };
14988        let f = self.func(name);
14989        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
14990        let cfg = LaunchConfig {
14991            grid_dim: (out_f as u32, m as u32, 1),
14992            block_dim: (128, 1, 1),
14993            shared_mem_bytes: 0,
14994        };
14995        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
14996        let __s_b = self.gpu.stream();
14997        let mut b = __s_b.launch_builder(&f);
14998        b.arg(bytes)
14999            .arg(aq)
15000            .arg(ad)
15001            .arg(&mut y)
15002            .arg(&inf)
15003            .arg(&outf)
15004            .arg(&mi)
15005            .arg(&rb);
15006        unsafe {
15007            b.launch(cfg)?;
15008        }
15009        Ok(Some((y, scale)))
15010    }
15011
15012    /// True if `qtype` has a warp-per-row MMVQ decode kernel AND MEMRA_MMVQ is set. Only the 4
15013    /// daily-hot dtypes (Q8_0, Q4_K, Q6_K, NVFP4) — others keep the _dp4a matvec (oracle/fallback).
15014    pub fn mmvq_supports(&self, qtype: i32) -> bool {
15015        // DEFAULT ON since 2026-07-08 (MEMRA_MMVQ=0 reverts to the _dp4a matvec class).
15016        // QT_F8_E4M3 is exempt from the MEMRA_MMVQ=0 escape: the e4m3 mmvq family is that dtype's
15017        // ONLY int8-act kernel class (there is no _dp4a twin), so its m=1/verify/batched dispatch
15018        // is a pure function of the dtype — the decode-parity law holds under every env.
15019        if qtype == QT_F8_E4M3 {
15020            return true;
15021        }
15022        if std::env::var("MEMRA_MMVQ").as_deref() == Ok("0") {
15023            return false;
15024        }
15025        matches!(
15026            qtype,
15027            QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q6_K | QT_NVFP4 | QT_Q4_0
15028        )
15029    }
15030
15031    /// PERF-3 warp-per-row MMVQ launcher (decode m=1 hot path). block=(32,ROWS_PER_BLOCK,1):
15032    /// one warp owns one output row, warp-only __shfl reduction (no smem barrier). Bit-equivalent
15033    /// to qmatvec_*_dp4a up to f32 reduction order. Pre-quantized q8_1 activation (aq,ad). NVFP4
15034    /// per-tensor macro-scale applied post (scale==1.0 for other dtypes -> no-op).
15035    pub fn qmatvec_mmvq(
15036        &self,
15037        bytes: &CudaSlice<u8>,
15038        aq: &CudaSlice<i8>,
15039        ad: &CudaSlice<f32>,
15040        m: usize,
15041        in_f: usize,
15042        out_f: usize,
15043        qtype: i32,
15044        row_bytes: usize,
15045        scale: f32,
15046        rp: bool,
15047    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15048        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output: skip memset
15049        self.qmatvec_mmvq_into(
15050            bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, rp, &mut y,
15051        )?;
15052        Ok(y)
15053    }
15054
15055    /// Slot-fed MMVQ twin (alloc-free capture lane): full policy body, caller-owned output.
15056    #[allow(clippy::too_many_arguments)]
15057    pub fn qmatvec_mmvq_into(
15058        &self,
15059        bytes: &CudaSlice<u8>,
15060        aq: &CudaSlice<i8>,
15061        ad: &CudaSlice<f32>,
15062        m: usize,
15063        in_f: usize,
15064        out_f: usize,
15065        qtype: i32,
15066        row_bytes: usize,
15067        scale: f32,
15068        rp: bool,
15069        y: &mut CudaSlice<f32>,
15070    ) -> Result<(), Box<dyn std::error::Error>> {
15071        debug_assert!(y.len() >= m * out_f);
15072        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
15073        // SMALL-SHAPE GRID FILL (H100 lane, 2026-07-26 microbench: attn qkv out_f=2048 =
15074        // 0.97 waves at the 4-warp block -> 66% of peak). The g2 twin (2 warps/block)
15075        // doubles the grid when the 4-warp launch would be sub-wave; per-row program
15076        // identical -> bit-identical. MEMRA_Q80_G2=0 reverts.
15077        if qtype == QT_Q8_0
15078            && rp
15079            && m == 1
15080            && out_f >= 64
15081            && (out_f as u32).div_ceil(ROWS_PER_BLOCK) < 4 * self.sm_count() as u32
15082            && {
15083                static G2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15084                *G2.get_or_init(|| std::env::var("MEMRA_Q80_G2").as_deref() != Ok("0"))
15085            }
15086        {
15087            let f = self.func("qmatvec_q8_0_mmvq_rp_g2");
15088            let cfg = LaunchConfig {
15089                grid_dim: ((out_f as u32).div_ceil(2), 1, 1),
15090                block_dim: (32, 2, 1),
15091                shared_mem_bytes: 0,
15092            };
15093            let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, 1i32, row_bytes as i64);
15094            let __s_b = self.gpu.stream();
15095            let mut b = __s_b.launch_builder(&f);
15096            b.arg(bytes)
15097                .arg(aq)
15098                .arg(ad)
15099                .arg(&mut *y)
15100                .arg(&inf)
15101                .arg(&outf)
15102                .arg(&mi)
15103                .arg(&rb);
15104            unsafe {
15105                b.launch(cfg)?;
15106            }
15107            if scale != 1.0 {
15108                self.scale_inplace(y, scale, out_f)?;
15109            }
15110            return Ok(());
15111        }
15112        // Multi-row-per-warp (mr2) policy, fixed since the 2026-07 sweeps (the MEMRA_MMVQ_MR
15113        // override + mr4 kernel were retired 2026-07-08 — mr4 regressed on register pressure and
15114        // crashed under rp; q4_K/q6_K mr2 measured flat, "no gain = no change"):
15115        //   NVFP4 m=1 -> mr2 (clean +1-2% on 9B: RPW acc chains hide the weight-load latency
15116        //     that pins the single-row kernel at 30-46% DRAM). Bit-identical per row.
15117        //   Q5_K m=1 -> mr2 (2026-07-05: the FR-Spec trimmed draft head is Q5_K 32768 rows = 8%
15118        //     of the 27B p3 spec wall; latency-bound like the other k-quants pre-fix).
15119        //   Q4_K/Q6_K m=1 -> single-row (mr2 measured +0.7% / flat — weight-bandwidth-bound).
15120        let mut mr: u32 = if m == 1 && (qtype == QT_NVFP4 || qtype == QT_Q5_K) {
15121            2
15122        } else {
15123            1
15124        };
15125        // Q4_0 mr (gemma trunk): DEFAULT 1 since 2026-07-13 (MEMRA_Q40_MR=2 reverts) — the
15126        // mr1 rp twin doubles the block count and wins the tail-quantization/latency battle
15127        // on every gemma model (E4B +3.75%: 198.9 vs 191.7; 26B +0.7%; 31B +0.9%; N=2-3
15128        // valid-window interleaved, bit-identical per row — same dot program).
15129        if m == 1 && qtype == QT_Q4_0 {
15130            static Q40MR: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
15131            // shape policy PROBED NEGATIVE (2026-07-13): tall-only mr1 197.2 vs
15132            // mr1-everywhere 198.7 — mr1 wins wide-output shapes too; arm removed.
15133            mr = *Q40MR.get_or_init(|| {
15134                std::env::var("MEMRA_Q40_MR")
15135                    .ok()
15136                    .and_then(|v| v.parse().ok())
15137                    .unwrap_or(1)
15138            });
15139        }
15140        // q5issue lane (2026-07-08): MEMRA_Q5K_ISSUE swaps the q5_K m=1 mmvq kernels for the
15141        // issue-reduced `_il` bodies (uint4 header/qh/qs loads + branchless scale decode —
15142        // cuts ~34 LDG.U16 + ~5 LDG.U8 + a warp-divergent scale branch per 32-elem group-row
15143        // to 5 LDG.128). Bit-identical per (token,row) to the reference kernels.
15144        // `1` = shape-aware policy (N=3 clock-locked micro-bench, mem P0, synthetic real shapes):
15145        //   out_f <= 65536 (trunk/frspec regime): il at the default mr — mr2_il -9.5%/-10.5%
15146        //     on 4096x4096/4096x8192, -3.1% on the 32768 frspec head vs the mr2-ref default;
15147        //   out_f > 65536 (the 248320-row 27B lm_head, already ~97% of the mem wall): mr2_il
15148        //     REGRESSES +22% there but mr1_il wins -2.1% vs the mr2-ref default -> force mr=1.
15149        // `2` = force il at the current mr for EVERY shape (A/B probe seam). Default OFF.
15150        let q5_mode = std::env::var("MEMRA_Q5K_ISSUE").ok();
15151        let q5_force = q5_mode.as_deref() == Some("2");
15152        // DEFAULT ON since 2026-07-08 (MEMRA_Q5K_ISSUE=0 reverts): +1.8% 9B plain e2e N=3
15153        // (128.2 -> 130.4), 27B flat (its big head is already at the mem wall), all gates green.
15154        let q5_il = qtype == QT_Q5_K
15155            && m == 1
15156            && (q5_force || q5_mode.as_deref().map(|v| v != "0").unwrap_or(true));
15157        if q5_il && !q5_force && out_f > 65536 {
15158            mr = 1;
15159        }
15160        // Q4_0 split-plane rp: mr2 default; MEMRA_Q40_MR=1 reaches the mr1 rp twin
15161        // (2026-07-13 — the tall-input/short-output tail-quantization probe).
15162        if qtype == QT_Q4_0 && rp && mr != 1 {
15163            mr = 2;
15164        }
15165        // Q8_0 rp (H100 lane): mr1 default — the q4_0 mr2 recipe MEASURED NEGATIVE on H100
15166        // (2026-07-26 N=3: mr1 186.2 vs mr2 171.5 tok/s; halving the grid on 132 SMs costs
15167        // more than 2-row ILP buys). mr2 kernel stays behind MEMRA_Q80_MR=2 for the corpus.
15168        if qtype == QT_Q8_0 && rp {
15169            static Q80MR: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
15170            mr = *Q80MR.get_or_init(|| {
15171                std::env::var("MEMRA_Q80_MR")
15172                    .ok()
15173                    .and_then(|v| v.parse().ok())
15174                    .unwrap_or(1)
15175            });
15176        }
15177        let name = match (qtype, mr, rp) {
15178            (QT_NVFP4, 2, false) => "qmatvec_nvfp4_mmvq_mr2",
15179            (QT_NVFP4, 2, true) => "qmatvec_nvfp4_mmvq_mr2_rp",
15180            (QT_NVFP4, _, true) => "qmatvec_nvfp4_mmvq_rp",
15181            (QT_Q4_0, 1, true) => "qmatvec_q4_0_mmvq_rp",
15182            (QT_Q4_0, _, true) => "qmatvec_q4_0_mmvq_mr2_rp",
15183            (QT_Q5_K, 2, _) => {
15184                if q5_il {
15185                    "qmatvec_q5_K_mmvq_mr2_il"
15186                } else {
15187                    "qmatvec_q5_K_mmvq_mr2"
15188                }
15189            }
15190            (QT_Q8_0, 2, true) => "qmatvec_q8_0_mmvq_mr2_rp",
15191            // rpca (cp.async-staged weight ring): MEASURED NEGATIVE on H100 for Q8_0
15192            // (2026-07-26 N=3: 181.8 vs plain rp 185.5 — the smem round-trip exceeds the
15193            // latency it hides for 8-bit direct-dp4a; the NVFP4 win case overlaps table
15194            // decode with half the bytes). OPT-IN via MEMRA_Q80_CA=1 for the corpus.
15195            (QT_Q8_0, _, true)
15196                if in_f % 1024 == 0 && {
15197                    static CA: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15198                    *CA.get_or_init(|| std::env::var("MEMRA_Q80_CA").as_deref() == Ok("1"))
15199                } =>
15200            {
15201                "qmatvec_q8_0_mmvq_rpca"
15202            }
15203            (QT_Q8_0, _, true) => "qmatvec_q8_0_mmvq_rp",
15204            (QT_Q8_0, _, _) => "qmatvec_q8_0_mmvq",
15205            // K-quant split-plane twins (H100 K-quant coalescing fix, 2026-08-01): the rp4
15206            // mirror routes here; GGUF layout keeps the plain kernels. rp bytes MUST never
15207            // reach a GGUF-layout kernel or vice versa.
15208            (QT_Q4_K, _, true) => "qmatvec_q4_K_mmvq_rp",
15209            (QT_Q6_K, _, true) => "qmatvec_q6_K_mmvq_rp",
15210            (QT_Q4_K, _, _) => "qmatvec_q4_K_mmvq",
15211            (QT_Q4_0, 2, false) => "qmatvec_q4_0_mmvq_mr2",
15212            (QT_Q4_0, _, false) => "qmatvec_q4_0_mmvq",
15213            (QT_Q5_K, _, _) => {
15214                if q5_il {
15215                    "qmatvec_q5_K_mmvq_il"
15216                } else {
15217                    "qmatvec_q5_K_mmvq"
15218                }
15219            }
15220            (QT_Q6_K, _, _) => "qmatvec_q6_K_mmvq",
15221            (QT_NVFP4, _, false) => "qmatvec_nvfp4_mmvq",
15222            (QT_F8_E4M3, _, _) => "qmatvec_e4m3_mmvq",
15223            _ => panic!("qmatvec_mmvq: qtype {qtype} has no MMVQ kernel"),
15224        };
15225        let f = self.func(name);
15226        // each block still has ROWS_PER_BLOCK warps; with mr rows/warp it covers ROWS_PER_BLOCK*mr rows.
15227        let rows_per_block = ROWS_PER_BLOCK * mr;
15228        let cfg = LaunchConfig {
15229            grid_dim: (
15230                (out_f as u32 + rows_per_block - 1) / rows_per_block,
15231                m as u32,
15232                1,
15233            ),
15234            block_dim: (32, ROWS_PER_BLOCK, 1), // warp-per-row (x mr rows each)
15235            shared_mem_bytes: 0,                // warp-only reduce at m=1
15236        };
15237        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
15238        let __s_b = self.gpu.stream();
15239        let mut b = __s_b.launch_builder(&f);
15240        // NVFP4 + e4m3 mmvq kernels take the macro-scale as a fused epilogue arg (applied at the
15241        // write — bit-identical to the old separate scale_inplace pass, minus one launch per matvec:
15242        // 53 scale launches/token on the 9B; for e4m3 the scale is the checkpoint's per-tensor f32
15243        // weight_scale). Other mmvq kernels keep the 8-arg signature.
15244        if qtype == QT_NVFP4 || qtype == QT_F8_E4M3 {
15245            // PDL wave-B: the nvfp4 mr2_rp single (gemma wo / generic rp singles) joins
15246            // the wave-A launch class — 9-arg flavor (fused macro-scale epilogue).
15247            if Self::pdl_on()
15248                && Self::pdl_mmvq_on()
15249                && Self::pdl_nvfp4q8_on()
15250                && name == "qmatvec_nvfp4_mmvq_mr2_rp"
15251            {
15252                use cudarc::driver::{DevicePtr, DevicePtrMut};
15253                let s = &self.gpu.stream();
15254                let (pw, _g0) = bytes.device_ptr(s);
15255                let (paq, _g1) = aq.device_ptr(s);
15256                let (pad, _g2) = ad.device_ptr(s);
15257                let (py, _g3) = y.device_ptr_mut(s);
15258                let mut ps = [
15259                    &pw as *const _ as *mut std::ffi::c_void,
15260                    &paq as *const _ as *mut _,
15261                    &pad as *const _ as *mut _,
15262                    &py as *const _ as *mut _,
15263                    &inf as *const _ as *mut _,
15264                    &outf as *const _ as *mut _,
15265                    &mi as *const _ as *mut _,
15266                    &rb as *const _ as *mut _,
15267                    &scale as *const _ as *mut _,
15268                ];
15269                unsafe {
15270                    self.launch_pdl(name, cfg.grid_dim, cfg.block_dim, &mut ps)?;
15271                }
15272                return Ok(());
15273            }
15274            b.arg(bytes)
15275                .arg(aq)
15276                .arg(ad)
15277                .arg(&mut *y)
15278                .arg(&inf)
15279                .arg(&outf)
15280                .arg(&mi)
15281                .arg(&rb)
15282                .arg(&scale);
15283            unsafe {
15284                b.launch(cfg)?;
15285            }
15286        } else if Self::pdl_on()
15287            && Self::pdl_mmvq_on()
15288            && (matches!(
15289                name,
15290                "qmatvec_q4_0_mmvq_rp" | "qmatvec_q6_K_mmvq" | "qmatvec_q6_K_mmvq_rp"
15291            ) || (Self::pdl_nvfp4q8_on()
15292                && matches!(name, "qmatvec_q8_0_mmvq_rp" | "qmatvec_q8_0_mmvq_mr2_rp")))
15293        {
15294            // PDL wave-A (2026-07-23): the two decode-hot single-matvec kernels carry
15295            // MEMRA_PDL_ENTRY — grid launches while the producer drains. ONLY the marked
15296            // names may take this launch (unmarked kernels would read unordered).
15297            {
15298                use cudarc::driver::{DevicePtr, DevicePtrMut};
15299                let s = &self.gpu.stream();
15300                let (pw, _g0) = bytes.device_ptr(s);
15301                let (paq, _g1) = aq.device_ptr(s);
15302                let (pad, _g2) = ad.device_ptr(s);
15303                let (py, _g3) = y.device_ptr_mut(s);
15304                let mut ps = [
15305                    &pw as *const _ as *mut std::ffi::c_void,
15306                    &paq as *const _ as *mut _,
15307                    &pad as *const _ as *mut _,
15308                    &py as *const _ as *mut _,
15309                    &inf as *const _ as *mut _,
15310                    &outf as *const _ as *mut _,
15311                    &mi as *const _ as *mut _,
15312                    &rb as *const _ as *mut _,
15313                ];
15314                unsafe {
15315                    self.launch_pdl(name, cfg.grid_dim, cfg.block_dim, &mut ps)?;
15316                }
15317            }
15318            if scale != 1.0 {
15319                self.scale_inplace(y, scale, m * out_f)?;
15320            }
15321        } else {
15322            b.arg(bytes)
15323                .arg(aq)
15324                .arg(ad)
15325                .arg(&mut *y)
15326                .arg(&inf)
15327                .arg(&outf)
15328                .arg(&mi)
15329                .arg(&rb);
15330            unsafe {
15331                b.launch(cfg)?;
15332            }
15333            if scale != 1.0 {
15334                self.scale_inplace(y, scale, m * out_f)?;
15335            }
15336        }
15337        Ok(())
15338    }
15339
15340    /// Test entry for the kernel_check bit-equivalence gate: run the warp-per-row MMVQ directly
15341    /// from raw weight bytes (quantize the f32 activation `x` to q8_1 internally). NVFP4 per-tensor
15342    /// macro-scale is NOT applied (caller compares bare, like qmatvec_*_fast). Mirrors qmatvec_gemm_raw.
15343    pub fn qmatvec_mmvq_raw(
15344        &self,
15345        bytes: &CudaSlice<u8>,
15346        x: &CudaSlice<f32>,
15347        m: usize,
15348        in_f: usize,
15349        out_f: usize,
15350        qtype: i32,
15351        row_bytes: usize,
15352        rp: bool,
15353    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15354        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15355        self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, 1.0, rp)
15356    }
15357
15358    /// True if `qtype` has a batched weight-resident (`_b2`/`_b4`) matvec kernel. These mirror the
15359    /// `_mmvq` kernels but iterate the m token columns INSIDE one warp/row, so the weight bytes leave
15360    /// HBM/L2 once for m tokens (vs grid.y=m re-reading m times). The 5 daily-hot dtypes have them.
15361    pub fn batched_supports(&self, qtype: i32) -> bool {
15362        matches!(
15363            qtype,
15364            QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q6_K | QT_NVFP4 | QT_F8_E4M3 | QT_Q4_0
15365        )
15366    }
15367
15368    /// IQ4_XS trunk fast seam: MEMRA_IQ_FAST=0 reverts non-expert IQ4_XS matmuls to the Stage-A
15369    /// f32 oracle path. Default ON since 2026-08-02 (research/kat-anomaly-20260802/): the old
15370    /// opt-in default left every IQ4_XS-trunk artifact (KAT-Coder IQ4_XS: attn_qkv/attn_gate/
15371    /// ssm_out/shexp, ~0.52GB re-read per decode tick) on the oracle kernel — decode 106.7 ->
15372    /// 193.4 tok/s (x5 interleaved), pp512 228 -> 697, same bytes, via qmatvec_iq4_XS_dp4a. The
15373    /// supported artifacts carry IQ4_XS only in EXPERT banks (their own dispatch, not this seam),
15374    /// so this admission is dispatch-unchanged for every non-IQ4_XS-trunk model.
15375    pub fn iq_fast_enabled() -> bool {
15376        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15377        *ON.get_or_init(|| {
15378            std::env::var("MEMRA_IQ_FAST")
15379                .map(|v| v != "0")
15380                .unwrap_or(true)
15381        })
15382    }
15383
15384    /// b8 tier seam: MEMRA_B8=0 keeps m=5..8 on the per-m grid.y=m path (m=2..4 batched dispatch
15385    /// unaffected). Default ON — the K=4..7 spec-verify weight-read-once fix.
15386    pub fn b8_enabled() -> bool {
15387        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15388        *ON.get_or_init(|| std::env::var("MEMRA_B8").map(|v| v != "0").unwrap_or(true))
15389    }
15390
15391    /// Compile-time column batch for a runtime m: 2 -> b2, 3..4 -> b4, 5..8 -> b8.
15392    pub fn batched_mcols(m: usize) -> usize {
15393        if m == 2 {
15394            2
15395        } else if m <= 4 {
15396            4
15397        } else if m <= 8 {
15398            8
15399        } else {
15400            16
15401        }
15402    }
15403
15404    /// Kernel name for the batched matvec of `(qtype, mcols)`. mcols ∈ {2,4,8}. The b8 tier is the
15405    /// K=4..7 spec-verify fix (T=5..8): pre-b8 those T fell to grid.y=m per-row MMVQ = m full
15406    /// weight reads/launch — the measured 27B K=4 cliff (101 -> 73 tok/s at p3 despite acceptance
15407    /// holding 54%). One b8 launch reads the weight ONCE for up to 8 columns (c >= m masked).
15408    fn batched_kernel_name(qtype: i32, mcols: usize) -> Option<&'static str> {
15409        Some(match (qtype, mcols) {
15410            (QT_Q8_0, 2) => "qmatvec_q8_0_mmvq_b2",
15411            (QT_Q8_0, 4) => "qmatvec_q8_0_mmvq_b4",
15412            (QT_Q8_0, 8) => "qmatvec_q8_0_mmvq_b8",
15413            // b16 now has BOTH forms (lane/rp-on-st, 2026-08-06). It used to be rp-ONLY, which
15414            // made the q8rp mirror the exact-16 tier's admission ticket for any model carrying a
15415            // single Q8_0 matmul — measured as the FP8-ST refusal (`L0.ssm_beta qtype=0
15416            // rp4=false`, 96 t / 23.9 MiB = 0.143% of resident weight). The mirror stays a
15417            // BANDWIDTH lever on Q8_0-dominant GGUFs; it is no longer a correctness prerequisite.
15418            (QT_Q8_0, 16) => "qmatvec_q8_0_mmvq_b16",
15419            (QT_Q4_K, 2) => "qmatvec_q4_K_mmvq_b2",
15420            (QT_Q4_K, 4) => "qmatvec_q4_K_mmvq_b4",
15421            (QT_Q4_K, 8) => "qmatvec_q4_K_mmvq_b8",
15422            // b16 base + _rp (lane/rp-on-st): the 9B NVFP4 GGUF's blocker — real NVFP4 GGUFs keep
15423            // Q4_K attention next to NVFP4 MLP, and the tier's predicate is an ALL.
15424            (QT_Q4_K, 16) => "qmatvec_q4_K_mmvq_b16",
15425            (QT_Q5_K, 2) => "qmatvec_q5_K_mmvq_b2",
15426            (QT_Q5_K, 4) => "qmatvec_q5_K_mmvq_b4",
15427            (QT_Q5_K, 8) => "qmatvec_q5_K_mmvq_b8",
15428            // b16 base only (lane/rp-on-st): Q5_K has no rp twins at any width, so there is
15429            // nothing to mirror. Named by the diagnostic as `L0.wqkv_gate qtype=3` on the 9B.
15430            (QT_Q5_K, 16) => "qmatvec_q5_K_mmvq_b16",
15431            (QT_Q6_K, 2) => "qmatvec_q6_K_mmvq_b2",
15432            (QT_Q6_K, 4) => "qmatvec_q6_K_mmvq_b4",
15433            (QT_Q6_K, 8) => "qmatvec_q6_K_mmvq_b8",
15434            (QT_Q6_K, 16) => "qmatvec_q6_K_mmvq_b16",
15435            (QT_NVFP4, 2) => "qmatvec_nvfp4_mmvq_b2",
15436            (QT_NVFP4, 4) => "qmatvec_nvfp4_mmvq_b4",
15437            (QT_NVFP4, 8) => "qmatvec_nvfp4_mmvq_b8",
15438            // b16 (lane/rp-on-st): no mirror needed — NVFP4's 36 B/k32 block is already the
15439            // aligned form its own kernel walks. Unlocks the exact-16 tier for every NVFP4 model
15440            // AND for the mixed FP8-ST artifact, whose 193 NVFP4 tensors were refusing it.
15441            (QT_NVFP4, 16) => "qmatvec_nvfp4_mmvq_b16",
15442            (QT_F8_E4M3, 2) => "qmatvec_e4m3_mmvq_b2",
15443            (QT_F8_E4M3, 4) => "qmatvec_e4m3_mmvq_b4",
15444            (QT_F8_E4M3, 8) => "qmatvec_e4m3_mmvq_b8",
15445            // b16 tier (lane/rp-on-st): e4m3 needs NO split-plane mirror to reach it — its native
15446            // row-major layout is already 32B-aligned per k32 block, so the base kernel IS the
15447            // aligned form. Contrast Q8_0, whose b16 exists only as the `_rp` twin (hence q8rp).
15448            (QT_F8_E4M3, 16) => "qmatvec_e4m3_mmvq_b16",
15449            (QT_Q4_0, 2) => "qmatvec_q4_0_mmvq_b2",
15450            (QT_Q4_0, 4) => "qmatvec_q4_0_mmvq_b4",
15451            (QT_Q4_0, 8) => "qmatvec_q4_0_mmvq_b8",
15452            (QT_Q4_0, 16) => "qmatvec_q4_0_mmvq_b16",
15453            _ => return None,
15454        })
15455    }
15456
15457    /// BATCHED weight-tile-resident matvec from a PRE-QUANTIZED q8_1 activation (the m=2-8 verify/MTP
15458    /// win). One warp walks the weight row ONCE, dp4a vs all m activation columns -> weight HBM/L2
15459    /// traffic 1x for m tokens (vs grid.y=m re-reading it m times). `mcols` ∈ {2,4,8} is the
15460    /// compile-time batch; m must be <= mcols (the c >= m columns are masked in-kernel). y is
15461    /// [m, out_f] token-major. NVFP4 per-tensor macro-scale applied post
15462    /// (scale==1.0 for other dtypes -> no-op). BIT-IDENTICAL per (token,row) to qmatvec_*_mmvq.
15463    ///
15464    /// NVFP4 VARIANT DISPATCH: the batched NVFP4 kernel measured memory-LATENCY bound on the real
15465    /// 27B verify (ncu --set full, 12 steady launches: long_scoreboard 18-30 stalls/issue vs <=1.7
15466    /// for every other reason, DRAM only 41-51% active, lg_throttle 0.7, L1 hit 94% — ONE 6-LDG
15467    /// weight wavefront in flight per warp is the binding constraint, NOT bandwidth and NOT the
15468    /// column-unroll break). Two exactness-free fixes, chosen PER SHAPE from the DRAM-cold 8-copy
15469    /// msweep on all six 27B shapes (2026-07-03):
15470    ///   `pf` = next-g weight-prefetch double-buffer (48 regs, occupancy intact) — wins everywhere
15471    ///          it applies for b4 (-3..-14%), never loses;
15472    ///   `r2` = two rows/warp (67 regs -> 7 resident blocks/SM) — the bigger win (-8.5..-30%) but
15473    ///          wave-quantization-sensitive: with the grid halved to ceil(out_f/8) blocks, a
15474    ///          fractional straggler wave (waves in ~1.05-1.5) costs a full extra latency round on
15475    ///          a latency-bound kernel (27B ffn_down 640 blocks / 574 resident = 1.11 waves: +17%),
15476    ///          while <=1 wave (9B ffn_down 0.89: -30%) or >=2 waves (tail amortized; qkv 2.2:
15477    ///          -8.5%, ffn_gate 3.8: -12.5%) win. For b2, r2 wins on DEEP k-loops (in_f>=6144:
15478    ///          -8..-19%) where the 2-col body starves weight MLP hardest; pf measured negative.
15479    /// b4: r2 when waves(out_f) <= 1 (and grid fills >=half the SMs) or >= 2, else pf.
15480    /// b2: in_f>=6144 -> r2, else base.
15481    /// MEMRA_MMVQ_BV=base|pf|r2|pfr2 forces one variant everywhere (A/B + rollback seam).
15482    /// All variants BIT-IDENTICAL per (token,row): same dp4a order, scales, adg factor, reduce —
15483    /// only load issue time and the row->warp mapping change (kernel-check gates all of them).
15484    /// `rp` = the weight buffer is the A6 SPLIT-PLANE repacked layout (NVFP4 only): the same
15485    /// wave-aware auto rule applies, mapped onto the `_rp` twins (rp/rpr2/rpr2w8 mirror
15486    /// pf/r2/r2w8 — regs 44/67/64 land in the same residency classes).
15487    /// The variant the batched dispatch will pick for this (shape, m, mcols, layout) — exposed so
15488    /// gates can distinguish bit-identical variants (bit-bad==0 required) from the k-split family
15489    /// (deterministic but k-reduce-order-shifted: rel<1e-3 + run-to-run bit-identity required).
15490    /// Device SM count (cached) — grid-fill policy input.
15491    pub fn sm_count(&self) -> i32 {
15492        static SMS: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
15493        *SMS.get_or_init(|| {
15494            use cudarc::driver::sys::CUdevice_attribute_enum as A;
15495            self.gpu
15496                .ctx
15497                .attribute(A::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)
15498                .unwrap_or(82)
15499        })
15500    }
15501
15502    pub fn batched_variant(
15503        &self,
15504        _m: usize,
15505        in_f: usize,
15506        out_f: usize,
15507        qtype: i32,
15508        row_bytes: usize,
15509        mcols: usize,
15510        rp: bool,
15511    ) -> &'static str {
15512        // Q8_0 never joined the auto variant machinery (on sm_120 its only batched shapes
15513        // were tiny aux tensors). On Q8_0-trunk models the layout is the whole game: the
15514        // split-plane mirror (rp) routes to the _rp twins (H100 coalescing fix, 2026-07-26);
15515        // GGUF layout stays "base". rp bytes MUST never reach the base kernel or vice versa.
15516        if qtype == QT_Q8_0 {
15517            return if rp { "rp" } else { "base" };
15518        }
15519        static BV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
15520        let bv = *BV.get_or_init(|| match std::env::var("MEMRA_MMVQ_BV").as_deref() {
15521            Ok("base") => "base",
15522            Ok("pf") => "pf",
15523            Ok("r2") => "r2",
15524            Ok("r2w8") => "r2w8",
15525            Ok("pfr2") => "pfr2",
15526            Ok("ca") => "ca",
15527            Ok("car2") => "car2",
15528            // rp* = SPLIT-PLANE REPACKED layout kernels (A6 prototype): W must already be the
15529            // repacked buffer (msweep MSWEEP_RP harness) — never valid on GGUF-layout weights.
15530            Ok("rp") => "rp",
15531            Ok("rpr2") => "rpr2",
15532            Ok("rpr2w8") => "rpr2w8",
15533            // rpca* = cp.async software-pipelined split-plane (2026-07-05): hides the _rp
15534            // long_scoreboard load stall. rp-layout only; b4/b2 (no b8 twin).
15535            Ok("rpca") => "rpca",
15536            Ok("rpcar2") => "rpcar2",
15537            // 2026-07-06 m-small latency arc: rpsc = rpr2 + per-warp smem scale prestage (kills
15538            // the scale-plane global dependency, zero reg growth); rpms/rpmsc = m-split x2
15539            // across warp pairs (2x blocks of rpr2, column halves per warp, BIT-identical to
15540            // _rp); rpks/rpksc = k-split x2 (fastest microbench cells but k-reduce-order-shifted:
15541            // run-spec self-consistency FAILED on the 27B daily driver — verify logits must be
15542            // bit-identical to the decode path — measurement corpus ONLY, never auto).
15543            Ok("rpsc") => "rpsc",
15544            Ok("rpms") => "rpms",
15545            Ok("rpmsc") => "rpmsc",
15546            Ok("rpks") => "rpks",
15547            Ok("rpksc") => "rpksc",
15548            _ => "auto",
15549        });
15550        // cp.async ring variants need 16B-aligned rows (in_f%256==0 -> (in_f/64)*36 % 16 == 0)
15551        // and whole 32-group warp iterations (nsb%32==0 <=> in_f%1024==0). All 27B/9B trunk
15552        // shapes qualify; anything else falls back to the register variants.
15553        let ca_ok = qtype == QT_NVFP4 && (row_bytes % 16 == 0) && (in_f % 1024 == 0);
15554        // rpsc: smem scale plane fits (nsb64 <= 272) + int4-aligned staging (nsb64 % 4 == 0).
15555        // rpks/rpksc: half-plane staging alignment needs nsb64 % 8 == 0 (in_f % 512 == 0).
15556        // MEMRA_KS=0 removes the 2026-07-06 rpsc/rpks/rpksc entries from AUTO (rollback seam;
15557        // forced MEMRA_MMVQ_BV values still work).
15558        static KS_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15559        let ks_on = *KS_ON.get_or_init(|| std::env::var("MEMRA_KS").as_deref() != Ok("0"));
15560        let sc_ok = ks_on && qtype == QT_NVFP4 && (in_f % 256 == 0) && (in_f / 64 <= 272);
15561        let ks_ok = ks_on && qtype == QT_NVFP4 && (in_f % 512 == 0) && (in_f / 64 <= 272);
15562        static SMS: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
15563        let sms = *SMS.get_or_init(|| {
15564            use cudarc::driver::sys::CUdevice_attribute_enum as A;
15565            self.gpu
15566                .ctx
15567                .attribute(A::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)
15568                .unwrap_or(82)
15569        });
15570        // k-quant r2 port (2026-07-04): q4_K/q5_K/q6_K have _r2/_r2w8 twins. ncu on the DRAM-cold
15571        // 9B msweep showed q4_K/q5_K b4 memory-latency bound like NVFP4 pre-fix (long_scoreboard
15572        // 19.6/16.4 per issue, DRAM 47.7/38.2%, L2 weight hit ~13%); q6_K lm_head is the exception
15573        // at DRAM 90-91% = wall-bound (yet r2 still wins -8%: deeper MLP raises achieved DRAM).
15574        // No _pf port (a k-quant group stages 10+ words vs NVFP4's 5 — register cost outweighs;
15575        // r2 covers the same MLP) and no rp (GGUF layout only). Q8_0 stays base: its only real
15576        // batched shapes are the tiny out_f=32 ssm_alpha/beta (8-block grids never fill one SM).
15577        // AUTO RULE = the measured winners table (differs from NVFP4's!):
15578        //   r2w8 NEVER in auto — the reg squeeze (72 -> 64 regs = stack spill) loses to unbounded
15579        //     r2 on every measured k-quant cell, incl. the wave-crossing lm_heads (q6_K 1316 vs
15580        //     r2 1258us) — kernels kept behind the force seam for the corpus;
15581        //   q4_K: r2 whenever the halved grid fills the SMs (blocks >= 4*SMs), INCLUDING the
15582        //     1.05-2.0 straggler window where NVFP4's r2 lost (qkv 1.78 waves: r2 -15% here; the
15583        //     k-quant base kernel leaves more latency on the table than a straggler wave costs);
15584        //   q5_K/q6_K: r2 only at waves >= 2 (the 248320-row lm_heads, 48+ waves: q6_K -8%, q5_K
15585        //     -2%); mid shapes measured base-or-flat (q5_K qkv 49.1 base vs 49.7 r2, attn_gate
15586        //     flat, attn_k base) — the 5/6-bit two-stream unpack makes r2's staging pricier.
15587        //   b2 same table with 8-row blocks: q4_K r2 when filled (-3..-22% all measured shapes),
15588        //     q5_K/q6_K r2 at waves >= 2 (27B lm_head -2.9%; 9B q6_K flat, harmless).
15589        let kq_r2 = matches!(qtype, QT_Q4_K | QT_Q5_K | QT_Q6_K);
15590        // MEMRA_KQ_BV=base|r2|r2w8 forces the k-quant variant WITHOUT touching the NVFP4 dispatch
15591        // (MEMRA_MMVQ_BV is global — an interleaved k-quant-only e2e A/B needs this narrower seam).
15592        static KQBV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
15593        let kq_bv = *KQBV.get_or_init(|| match std::env::var("MEMRA_KQ_BV").as_deref() {
15594            Ok("base") => "base",
15595            Ok("r2") => "r2",
15596            Ok("r2w8") => "r2w8",
15597            _ => "auto",
15598        });
15599        let variant: &'static str = if qtype == QT_Q4_0 {
15600            // Q4_0 r2 (gemma verify trunk, 2026-07-10): shared activation loads + the
15601            // row-independent ones-sum computed once per (col,group) for 2 rows. Same
15602            // fill rule as q4_K: r2 when the halved grid still fills the SMs.
15603            static Q40BV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
15604            let q40 = *Q40BV.get_or_init(|| match std::env::var("MEMRA_Q40_BV").as_deref() {
15605                // ms/sm/la = force-only measurement seams (ALL FLAT/NEGATIVE 2026-07-13,
15606                // never auto): m-split flat (nvcc keeps 72 regs); smem-slab −11% (staging
15607                // + syncs cost more than the stalls, bank-pad made no difference);
15608                // register load-ahead flat (nvcc already reorders). The b-tier limiter
15609                // is still unidentified — see the jsonl row.
15610                Ok("base") => "base",
15611                Ok("r2") => "r2",
15612                Ok("ms") => "ms",
15613                Ok("sm") => "sm",
15614                Ok("la") => "la",
15615                _ => "auto",
15616            });
15617            let v = if q40 != "auto" {
15618                q40
15619            } else if (out_f as u32).div_ceil(8) >= 4 * sms as u32 {
15620                "r2"
15621            } else {
15622                "base"
15623            };
15624            // split-plane mirror twins (2026-07-10): same fill rule, _rp names.
15625            // (m-split r2 pair twin PROBED FLAT 2026-07-13 — nvcc kept 72 regs either way
15626            // and the limiter is the per-column activation load chain (long_scoreboard
15627            // 42.5%), not occupancy; arm killed per doctrine, jsonl row is the record.)
15628            if rp {
15629                match v {
15630                    "ms" => "r2ms_rp",
15631                    "sm" => "r2sm_rp",
15632                    "la" => "r2la_rp",
15633                    "r2" => "r2_rp",
15634                    _ => "rp",
15635                }
15636            } else if matches!(v, "ms" | "sm" | "la") {
15637                "r2"
15638            } else {
15639                v
15640            }
15641        } else if qtype != QT_NVFP4 && !kq_r2 {
15642            "base"
15643        } else if kq_r2 && rp {
15644            // K-quant split-plane mirror (2026-08-01): only the plain _rp batched twins are
15645            // compiled for q4_K/q6_K — rp is a LAYOUT, it must survive every heuristic
15646            // (split-plane bytes through a GGUF-layout kernel = NaN). q5_K never mirrors.
15647            "rp"
15648        } else if kq_r2 {
15649            // k-quant r2w8 only exists at b4 (b2_r2 already 8-resident; b8 has no w8 twin) ->
15650            // mcols != 4 forced r2w8 falls to unbounded r2.
15651            if kq_bv != "auto" {
15652                if kq_bv == "r2w8" && mcols != 4 {
15653                    "r2"
15654                } else {
15655                    kq_bv
15656                }
15657            } else if bv != "auto" {
15658                match bv {
15659                    "r2" | "pfr2" | "rpr2" | "car2" => "r2",
15660                    "r2w8" | "rpr2w8" => {
15661                        if mcols != 4 {
15662                            "r2"
15663                        } else {
15664                            "r2w8"
15665                        }
15666                    }
15667                    _ => "base", // base/pf/ca/rp forced -> base (no such k-quant kernels)
15668                }
15669            } else {
15670                let blocks = (out_f + 7) / 8;
15671                let waves = blocks as f64 / (7 * sms as usize) as f64;
15672                let filled = blocks >= 4 * sms as usize;
15673                let use_r2 = if qtype == QT_Q4_K {
15674                    filled
15675                } else {
15676                    waves >= 2.0
15677                };
15678                if use_r2 { "r2" } else { "base" }
15679            }
15680        } else if bv != "auto" {
15681            // r2w8 only exists for b4/b8 (the b2_r2 kernel is already 8-blocks-resident at 60 regs).
15682            // ca/car2 need the alignment gate AND have no b8 twins; pfr2 has no b8 twin either —
15683            // unsupported (shape, mcols) combos fall back to pf/r2.
15684            // On rp buffers, forced legacy names map to their rp twins (layout law).
15685            let v = if bv == "r2w8" && mcols == 2 {
15686                "r2"
15687            } else if bv == "ca" && (!ca_ok || mcols == 8) {
15688                "pf"
15689            } else if bv == "car2" && (!ca_ok || mcols == 8) {
15690                "r2"
15691            } else if bv == "pfr2" && mcols == 8 {
15692                "r2"
15693            } else if (bv == "rpr2w8" || bv == "rpr2") && mcols == 2 {
15694                "rpr2"
15695            }
15696            // rpca* has no b8 twin (falls to rpr2w8/rpr2); needs the ca alignment gate.
15697            else if (bv == "rpca" || bv == "rpcar2") && (!ca_ok || mcols == 8) {
15698                if mcols == 8 { "rpr2w8" } else { "rpr2" }
15699            } else if bv == "rpcar2" && mcols == 2 {
15700                "rpca"
15701            }
15702            // rpsc/rpmsc/rpks* gate on smem-fit + alignment; fall to rpr2 outside it
15703            // (rpms has no smem and no alignment need — always valid on rp buffers).
15704            else if (bv == "rpsc" || bv == "rpmsc") && !sc_ok {
15705                "rpr2"
15706            } else if (bv == "rpks" || bv == "rpksc") && !ks_ok {
15707                "rpr2"
15708            } else {
15709                bv
15710            };
15711            if rp {
15712                match v {
15713                    "base" | "pf" | "ca" | "rp" => "rp",
15714                    "r2" | "pfr2" | "car2" | "rpr2" => "rpr2",
15715                    "r2w8" | "rpr2w8" => {
15716                        if mcols == 2 {
15717                            "rpr2"
15718                        } else {
15719                            "rpr2w8"
15720                        }
15721                    }
15722                    other => other, // rpca/rpcar2/rpsc/rpks/rpksc pass through (already rp-layout)
15723                }
15724            } else {
15725                v
15726            }
15727        } else if mcols == 8 {
15728            // b8 AUTO (2026-07-06 m-small latency arc, rtx6000 DRAM-cold rp msweep m=5/6/8 all five
15729            // 27B shapes): rpsc — the rpr2w8 schedule with the warp's scale rows prestaged to
15730            // smem, leaving ONE global dependency (the quant stream) in the k-loop at zero reg
15731            // growth. BIT-identical to rpr2w8 and wins or ties EVERY b8 cell: ffn_gate m5
15732            // 50.7->46.9 m8 64.1->57.1 (-11%), qkv m8 34.6->33.0, ssm_out m8 29.7->28.8,
15733            // attn_gate m8 26.9->26.1, ffn_down m5 58.2->56.9. The faster split-grid twins are
15734            // OUT: rpksc (k-split, ffn_down m5 -21%) broke run-spec self-consistency (k-reduce
15735            // order shifts verify argmax at tie margins — verify must stay bit-identical to the
15736            // m=1 decode chain); rpmsc (m-split, bit-identical) measured NEGATIVE everywhere
15737            // (twin warp's duplicated weight stream: ffn_down m5 85.7 vs 56.9).
15738            if rp {
15739                if sc_ok { "rpsc" } else { "rpr2w8" }
15740            } else {
15741                "r2w8"
15742            }
15743        } else if mcols >= 4 {
15744            // r2 runs 7 resident blocks/SM (67 regs); its __launch_bounds__(128,8) twin `r2w8`
15745            // (64 regs) runs 8. grid = ceil(out_f/8) for both. rp twins land in the same
15746            // residency classes (rp 44 regs ~ pf-class occupancy, rpr2 67, rpr2w8 64).
15747            let blocks = (out_f + 7) / 8;
15748            let r7 = 7 * sms as usize;
15749            let r8 = 8 * sms as usize;
15750            let waves = blocks as f64 / r7 as f64;
15751            let filled = blocks >= 4 * sms as usize;
15752            // 2026-07-06 m-small latency arc: b4 keeps the wave rule (rpms/rpmsc measured
15753            // flat-to-negative at m=3/4 on every shape — the m-split twin duplicates the weight
15754            // stream; rpsc b4 also negative on r2-class picks, ffn_down m4 51.1 vs 46.5).
15755            if filled && blocks.div_ceil(r8) < blocks.div_ceil(r7) {
15756                // the extra residency drops the INTEGER wave count -> the straggler wave a
15757                // latency-bound kernel pays in full disappears (ffn_down 1.11 -> 0.98 waves:
15758                // 112.5 -> 81.6us, beats pf 90.1; qkv 2.23 -> 1.95: 58.1 -> 51.1).
15759                if rp { "rpr2w8" } else { "r2w8" }
15760            } else if waves >= 2.0 || (waves <= 1.0 && filled) {
15761                // tail amortized (>=2 waves) or single wave: unbounded r2 (no reg-squeeze tax —
15762                // gate/up 81.1 vs 83.9 bounded, attn_q 61.0 vs 63.4).
15763                if rp { "rpr2" } else { "r2" }
15764            } else {
15765                // fractional straggler-wave window with no crossing, or grid too small to fill
15766                // the SMs (tiny out_f<=1024 shapes want max row-parallelism): prefetch variant
15767                // (rp = the r1 split-plane twin — measured the attn_gate winner, 35.4 vs pf 36.4).
15768                if rp { "rp" } else { "pf" }
15769            }
15770        } else if in_f >= 6144 {
15771            // b2 deep-k (2026-07-06): every new twin measured flat-to-negative here (rpms 44.1
15772            // vs rpr2 40.8 ffn_down; rpsc 43.6; the winning rpks is banned on k-order) — rpr2
15773            // stays.
15774            if rp { "rpr2" } else { "r2" }
15775        } else if rp {
15776            // b2 shallow-k: qkv (out_f=10240, 0.97 waves at 7-resident) is the one measured cell
15777            // where the r2-schedule scale-prestage twin beats the r1 rp pick (24.7 vs 28.9us
15778            // -15%); the wider (ffn_gate 1.65 waves) and smaller (attn_gate 0.58) shapes LOSE
15779            // (41.8 vs 38.2 / 16.6 vs 14.6) — gate on the single-wave window.
15780            let waves = ((out_f + 7) / 8) as f64 / (7 * sms as usize) as f64;
15781            if sc_ok && waves >= 0.9 && waves <= 1.1 {
15782                "rpsc"
15783            } else {
15784                "rp"
15785            }
15786        } else {
15787            "base"
15788        };
15789        variant
15790    }
15791
15792    pub fn qmatvec_mmvq_batched(
15793        &self,
15794        bytes: &CudaSlice<u8>,
15795        aq: &CudaSlice<i8>,
15796        ad: &CudaSlice<f32>,
15797        m: usize,
15798        in_f: usize,
15799        out_f: usize,
15800        qtype: i32,
15801        row_bytes: usize,
15802        mcols: usize,
15803        scale: f32,
15804        rp: bool,
15805    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15806        const ROWS_PER_BLOCK: u32 = 4;
15807        // TUNE SEAM (H100 lane): MEMRA_BVAR forces the batched-variant pick for the whole
15808        // process — the auto heuristics were tuned on sm_120 (82 SMs / 858 GB/s) and the
15809        // sm_90a re-tune sweeps this seam empirically. Layout variants stay safe: an rp
15810        // weight keeps its rp-layout kernel family regardless of the override.
15811        let forced: Option<&'static str> = {
15812            static V: std::sync::OnceLock<Option<String>> = std::sync::OnceLock::new();
15813            V.get_or_init(|| std::env::var("MEMRA_BVAR").ok())
15814                .as_deref()
15815                .map(|s| Box::leak(s.to_string().into_boxed_str()) as &'static str)
15816        };
15817        let variant = match forced {
15818            Some(v) if !rp || v.contains("rp") => v,
15819            _ => self.batched_variant(m, in_f, out_f, qtype, row_bytes, mcols, rp),
15820        };
15821        let base_name = Self::batched_kernel_name(qtype, mcols).ok_or_else(|| {
15822            format!("qmatvec_mmvq_batched: no kernel for qtype {qtype} mcols {mcols}")
15823        })?;
15824        // b16 tier (t=9..16 verify): only base/_rp b16 kernels are compiled — the b2..b8
15825        // per-shape perf variants (r2/pf/...) do not apply at this width. rp is a LAYOUT,
15826        // not a perf variant: it must survive (base kernel on split-plane bytes = NaN).
15827        let variant = if mcols == 16 {
15828            if rp { "rp" } else { "base" }
15829        } else {
15830            variant
15831        };
15832        // EXACT-WIDTH b5/b6/b7 twins (lane/vt-fixes fix 1, 2026-08-03): the b8 kernels
15833        // allocate acc[WROWS][8] at ANY m, so T=5..7 verify paid the full 8-wide register
15834        // tax — the measured T=4->5 cliff. The same template at MCOLS=m runs the identical
15835        // per-(token,row) chain (columns c >= m never execute in either form) ->
15836        // BIT-IDENTICAL to the b8 launch. NVFP4 split-plane only (the sm_120 default trunk);
15837        // covers both b8 auto schedules (rpsc, rpr2w8). MEMRA_B567=0 rollback.
15838        static B567: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15839        let b567 = *B567.get_or_init(|| std::env::var("MEMRA_B567").as_deref() != Ok("0"));
15840        if b567
15841            && qtype == QT_NVFP4
15842            && rp
15843            && mcols == 8
15844            && (5..=7).contains(&m)
15845            && matches!(variant, "rpsc" | "rpr2w8")
15846        {
15847            let f = self.func(&format!("qmatvec_nvfp4_mmvq_b{m}_{variant}"));
15848            let rows_per_block = ROWS_PER_BLOCK * 2; // r2-class schedules: 2 rows/warp
15849            let mut y = self.alloc_uninit::<f32>(m * out_f)?;
15850            let cfg = LaunchConfig {
15851                grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
15852                block_dim: (32, ROWS_PER_BLOCK, 1),
15853                shared_mem_bytes: 0,
15854            };
15855            let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
15856            let __s_b = self.gpu.stream();
15857            let mut b = __s_b.launch_builder(&f);
15858            b.arg(bytes)
15859                .arg(aq)
15860                .arg(ad)
15861                .arg(&mut y)
15862                .arg(&inf)
15863                .arg(&outf)
15864                .arg(&mi)
15865                .arg(&rb);
15866            unsafe {
15867                b.launch(cfg)?;
15868            }
15869            if scale != 1.0 {
15870                self.scale_inplace(&mut y, scale, m * out_f)?;
15871            }
15872            return Ok(y);
15873        }
15874        let (name, rows_per_block): (std::borrow::Cow<'static, str>, u32) = match variant {
15875            "base" => (base_name.into(), ROWS_PER_BLOCK),
15876            "pf" => (format!("{base_name}_pf").into(), ROWS_PER_BLOCK),
15877            "ca" => (format!("{base_name}_ca").into(), ROWS_PER_BLOCK),
15878            "rp" => (format!("{base_name}_rp").into(), ROWS_PER_BLOCK),
15879            "rpca" => (format!("{base_name}_rpca").into(), ROWS_PER_BLOCK), // 1 row/warp cp.async
15880            // split families: 2 warp-pairs x 2 rows = 4 rows/block (the k-range or column set
15881            // splits across the pair's two warps; grid.x doubles vs rpr2 at the same regs).
15882            "rpks" => (format!("{base_name}_rpks").into(), ROWS_PER_BLOCK),
15883            "rpksc" => (format!("{base_name}_rpksc").into(), ROWS_PER_BLOCK),
15884            "rpms" => (format!("{base_name}_rpms").into(), ROWS_PER_BLOCK),
15885            "rpmsc" => (format!("{base_name}_rpmsc").into(), ROWS_PER_BLOCK),
15886            "r2ms_rp" => (format!("{base_name}_r2ms_rp").into(), ROWS_PER_BLOCK),
15887            "r2sm_rp" => (format!("{base_name}_r2sm_rp").into(), ROWS_PER_BLOCK * 2),
15888            "r2la_rp" => (format!("{base_name}_r2la_rp").into(), ROWS_PER_BLOCK * 2),
15889            v => (format!("{base_name}_{v}").into(), ROWS_PER_BLOCK * 2), // r2-class: 2 rows/warp
15890        };
15891        debug_assert!(
15892            !rp || name.contains("_rp"),
15893            "rp weight dispatched to a GGUF-layout kernel"
15894        );
15895        let f = self.func(&name);
15896        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
15897        // r2sm_rp: [MCOLS][32 blk][8 int] activation slab + [MCOLS][32] f32 scales.
15898        let smem = if name.contains("_r2sm_rp") {
15899            (mcols * 32 * 9 * 4 + mcols * 32 * 4) as u32
15900        } else {
15901            0
15902        };
15903        let cfg = LaunchConfig {
15904            grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
15905            block_dim: (32, ROWS_PER_BLOCK, 1),
15906            shared_mem_bytes: smem,
15907        };
15908        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
15909        let __s_b = self.gpu.stream();
15910        let mut b = __s_b.launch_builder(&f);
15911        b.arg(bytes)
15912            .arg(aq)
15913            .arg(ad)
15914            .arg(&mut y)
15915            .arg(&inf)
15916            .arg(&outf)
15917            .arg(&mi)
15918            .arg(&rb);
15919        unsafe {
15920            b.launch(cfg)?;
15921        }
15922        if scale != 1.0 {
15923            self.scale_inplace(&mut y, scale, m * out_f)?;
15924        }
15925        Ok(y)
15926    }
15927
15928    /// BATCHED weight-tile-resident matvec from raw weight bytes (quantizes the f32 activation `x` to
15929    /// q8_1 internally; macro-scale NOT applied — caller compares bare, like qmatvec_*_fast). For the
15930    /// kernel_check bit-equivalence gate. `mcols` ∈ {2,4,8}. Works for Q8_0/Q4_K/Q5_K/Q6_K/NVFP4.
15931    pub fn qmatvec_batched_raw(
15932        &self,
15933        bytes: &CudaSlice<u8>,
15934        x: &CudaSlice<f32>,
15935        m: usize,
15936        in_f: usize,
15937        out_f: usize,
15938        qtype: i32,
15939        row_bytes: usize,
15940        mcols: usize,
15941        rp: bool,
15942    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15943        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15944        self.qmatvec_mmvq_batched(
15945            bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, mcols, 1.0, rp,
15946        )
15947    }
15948
15949    /// Back-compat NVFP4-only batched raw launcher (used by older gates). Delegates to the generic one.
15950    pub fn qmatvec_nvfp4_batched_raw(
15951        &self,
15952        bytes: &CudaSlice<u8>,
15953        x: &CudaSlice<f32>,
15954        m: usize,
15955        in_f: usize,
15956        out_f: usize,
15957        row_bytes: usize,
15958        mcols: usize,
15959        rp: bool,
15960    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15961        self.qmatvec_batched_raw(bytes, x, m, in_f, out_f, QT_NVFP4, row_bytes, mcols, rp)
15962    }
15963
15964    /// Stage-C FP4 gate (MEMRA_FP4): if `w` is an NVFP4 weight with in_f%64==0, run the native mxf4
15965    /// block-scale GEMM and apply the per-tensor macro-scale, returning Some(y). Else None (caller
15966    /// falls through to the int8 GEMM / dp4a). Strict opt-in over the proven int8 path; m>=16 only.
15967    fn try_fp4_gemm(
15968        &self,
15969        w: &crate::model::GpuTensor,
15970        x: &CudaSlice<f32>,
15971        m: usize,
15972        in_f: usize,
15973        out_f: usize,
15974    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
15975        use crate::model::GpuTensor;
15976        if cfg!(memra_portable_cuda) {
15977            return Ok(None);
15978        }
15979        if std::env::var("MEMRA_FP4").is_err() {
15980            return Ok(None);
15981        }
15982        // CUTLASS prefill branch (m>=128 + MEMRA_FP4_CUTLASS + a repacked CutlassWeight present): route
15983        // to the CUTLASS sm120 NVFP4 GEMM, folding the per-tensor macro-scale into the epilogue alpha
15984        // (1/scale) — no post-matmul scale_inplace. Decode (m<128) and the m∈[16,128) middle band keep
15985        // the hand-roll below: CUTLASS's 128-row M-tile wastes work under 128.
15986        // The hand-roll applies the per-tensor macro-scale as a POST-matmul MULTIPLY (scale_inplace(y,
15987        // scale)); CUTLASS's epilogue does D = alpha * (A@B^T), so alpha == scale reproduces it exactly
15988        // (NOT 1/scale — the plan sketch had this inverted; the kernel_check arm gates it). scale==1.0
15989        // for the common no-macro-scale case.
15990        #[cfg(memra_cutlass)]
15991        if m >= 128 && std::env::var("MEMRA_FP4_CUTLASS").is_ok() {
15992            if let GpuTensor::Quant {
15993                bytes,
15994                qtype,
15995                scale,
15996                row_bytes,
15997                cutlass,
15998                ..
15999            } = w
16000            {
16001                if *qtype == QT_NVFP4 && in_f % 64 == 0 {
16002                    if let Some(cw) = cutlass {
16003                        // Resident fast path: load-time-repacked B + swizzled SFB (no per-call repack).
16004                        let y = self.cutlass_fp4_gemm(
16005                            &cw.b_packed,
16006                            &cw.sfb_swizzled,
16007                            x,
16008                            *scale,
16009                            m,
16010                            out_f,
16011                            in_f,
16012                        )?;
16013                        return Ok(Some(y));
16014                    } else if std::env::var("MEMRA_FP4_CUTLASS_OTF").is_ok() {
16015                        // On-the-fly repack (MEMRA_FP4_CUTLASS_OTF): de-interleave + swizzle the B operand
16016                        // from raw bytes per prefill call. No resident doubling of the NVFP4 weight VRAM
16017                        // (the load-time repack ~doubles it) — needed for models that don't fit the
16018                        // resident path (e.g. the 27B on 24GB). Slower (per-call repack) but argmax-exact.
16019                        let (b_packed, sfb_sw) =
16020                            self.build_cutlass_weight(bytes, out_f, in_f, *row_bytes)?;
16021                        let y =
16022                            self.cutlass_fp4_gemm(&b_packed, &sfb_sw, x, *scale, m, out_f, in_f)?;
16023                        return Ok(Some(y));
16024                    }
16025                }
16026            }
16027        }
16028        if let GpuTensor::Quant {
16029            bytes,
16030            qtype,
16031            row_bytes,
16032            scale,
16033            rp,
16034            ..
16035        } = w
16036        {
16037            // A6: the hand-rolled W4A4 mxf4 GEMM reads 36B GGUF blocks — no rp port (MEMRA_FP4 is
16038            // an opt-in accuracy tradeoff); repacked tensors fall through to the int8 GEMM.
16039            if *qtype == QT_NVFP4 && in_f % 64 == 0 && !*rp {
16040                let y =
16041                    self.qmatvec_gemm_nvfp4_fp4(bytes, x, m, in_f, out_f, *row_bytes, *scale)?;
16042                return Ok(Some(y));
16043            }
16044        }
16045        Ok(None)
16046    }
16047
16048    /// rms_norm + fused fp16 twin (task #14): f32 output verbatim `rms_norm` + the fp16
16049    /// copy the f16-mirror GEMM group would otherwise produce with a standalone convert
16050    /// launch. BIT-IDENTICAL end-to-end (same reduction, same __float2half values).
16051    pub fn rms_norm_f16out(
16052        &self,
16053        x: &CudaSlice<f32>,
16054        w: &CudaSlice<f32>,
16055        dst: &mut CudaSlice<f32>,
16056        dst16: &mut CudaSlice<u8>,
16057        ncols: usize,
16058        nrows: usize,
16059        eps: f32,
16060    ) -> Result<(), Box<dyn std::error::Error>> {
16061        let f = self.func("rms_norm_f16out_f32");
16062        let cfg = LaunchConfig {
16063            grid_dim: (nrows as u32, 1, 1),
16064            block_dim: (rms_block(), 1, 1),
16065            shared_mem_bytes: 0,
16066        };
16067        let (nc, e) = (ncols as i32, eps);
16068        let __s_b = self.gpu.stream();
16069        let mut b = __s_b.launch_builder(&f);
16070        b.arg(x).arg(w).arg(dst).arg(dst16).arg(&nc).arg(&e);
16071        unsafe {
16072            b.launch(cfg)?;
16073        }
16074        Ok(())
16075    }
16076
16077    /// add+norm(+f16out) fusion for the prefill trunk (round 28; add_rms_norm precedent —
16078    /// bit-identical to add_f32 -> rms_norm_f16out). block_dim matches rms_norm_f16out's.
16079    #[allow(clippy::too_many_arguments)]
16080    pub fn add_rms_norm_f16out(
16081        &self,
16082        a: &CudaSlice<f32>,
16083        b: &CudaSlice<f32>,
16084        w: &CudaSlice<f32>,
16085        res: &mut CudaSlice<f32>,
16086        dst: &mut CudaSlice<f32>,
16087        dst16: &mut CudaSlice<u8>,
16088        ncols: usize,
16089        nrows: usize,
16090        eps: f32,
16091    ) -> Result<(), Box<dyn std::error::Error>> {
16092        let f = self.func("add_rms_norm_f16out_f32");
16093        let cfg = LaunchConfig {
16094            grid_dim: (nrows as u32, 1, 1),
16095            block_dim: (rms_block(), 1, 1),
16096            shared_mem_bytes: 0,
16097        };
16098        let (nc, e) = (ncols as i32, eps);
16099        let __s_lb = self.gpu.stream();
16100        let mut lb = __s_lb.launch_builder(&f);
16101        lb.arg(a)
16102            .arg(b)
16103            .arg(w)
16104            .arg(res)
16105            .arg(dst)
16106            .arg(dst16)
16107            .arg(&nc)
16108            .arg(&e);
16109        unsafe {
16110            lb.launch(cfg)?;
16111        }
16112        Ok(())
16113    }
16114
16115    /// matmul_group with a PRE-EMITTED fp16 activation (task #14: the producer norm fused
16116    /// the convert). Mirror-less members fall back to `matmul` on the f32 activation.
16117    pub fn matmul_group_xh(
16118        &self,
16119        ws: &[&crate::model::GpuTensor],
16120        x: &CudaSlice<f32>,
16121        xh: &CudaSlice<u8>,
16122        m: usize,
16123    ) -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
16124        let mut out = Vec::with_capacity(ws.len());
16125        let in_f = ws[0].in_features();
16126        for w in ws {
16127            if w.in_features() == in_f && m >= 16 && !self.verify_exact_on() {
16128                if let Some(y) = self.try_f16_gemm_pre(w, xh, m)? {
16129                    out.push(y);
16130                    continue;
16131                }
16132            }
16133            out.push(self.matmul(w, x, m)?);
16134        }
16135        Ok(out)
16136    }
16137
16138    /// task #14 pad-proofing: zero beta/g_log at rows >= len_d[0] (pads become identity
16139    /// GDN steps). Layouts [T, H].
16140    pub fn gdn_pad_mask(
16141        &self,
16142        beta: &mut CudaSlice<f32>,
16143        g_log: &mut CudaSlice<f32>,
16144        len_d: &CudaSlice<i32>,
16145        h: usize,
16146        t: usize,
16147    ) -> Result<(), Box<dyn std::error::Error>> {
16148        let f = self.func("gdn_pad_mask_f32");
16149        let cfg = LaunchConfig::for_num_elems((t * h) as u32);
16150        let (hi, ti) = (h as i32, t as i32);
16151        let __s_b = self.gpu.stream();
16152        let mut b = __s_b.launch_builder(&f);
16153        b.arg(beta).arg(g_log).arg(len_d).arg(&hi).arg(&ti);
16154        unsafe {
16155            b.launch(cfg)?;
16156        }
16157        Ok(())
16158    }
16159
16160    /// task #14 pad-proofing: dst[ncols] = src row (len_d[0]-1) — device-indexed last-row
16161    /// gather for the padded prime graph's h_seed/hlast.
16162    pub fn row_gather_dev(
16163        &self,
16164        src: &CudaSlice<f32>,
16165        dst: &mut CudaSlice<f32>,
16166        len_d: &CudaSlice<i32>,
16167        ncols: usize,
16168    ) -> Result<(), Box<dyn std::error::Error>> {
16169        let f = self.func("row_gather_dev_f32");
16170        let cfg = LaunchConfig::for_num_elems(ncols as u32);
16171        let nc = ncols as i32;
16172        let __s_b = self.gpu.stream();
16173        let mut b = __s_b.launch_builder(&f);
16174        b.arg(src).arg(dst).arg(len_d).arg(&nc);
16175        unsafe {
16176            b.launch(cfg)?;
16177        }
16178        Ok(())
16179    }
16180
16181    /// Grouped matmul: several weights consuming ONE activation (hybrid layers: the GDN
16182    /// 4-tuple wqkv/gate/beta/alpha, attention q/k/v, ffn gate/up). Semantics identical to
16183    /// calling `matmul` per weight; the f16-mirror arm converts the activation ONCE for the
16184    /// whole group instead of once per GEMM (the standalone converts were ~250 launches/prime
16185    /// of small-kernel gap fuel — nsys 2026-07-26). Any member without a mirror (or with a
16186    /// different in_f) falls back to its own `matmul` — behavior unchanged.
16187    pub fn matmul_group(
16188        &self,
16189        ws: &[&crate::model::GpuTensor],
16190        x: &CudaSlice<f32>,
16191        m: usize,
16192    ) -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
16193        use crate::model::GpuTensor;
16194        let mut out = Vec::with_capacity(ws.len());
16195        let any_mirror = ws
16196            .iter()
16197            .any(|w| matches!(w, GpuTensor::Quant { f16: Some(_), .. }));
16198        if m >= 16 && any_mirror && !self.verify_exact_on() {
16199            let in_f = ws[0].in_features();
16200            let xh = self.f16_act(x, m * in_f, in_f)?;
16201            for w in ws {
16202                if w.in_features() == in_f {
16203                    if let Some(y) = self.try_f16_gemm_pre(w, &xh, m)? {
16204                        out.push(y);
16205                        continue;
16206                    }
16207                }
16208                out.push(self.matmul(w, x, m)?);
16209            }
16210            return Ok(out);
16211        }
16212        for w in ws {
16213            out.push(self.matmul(w, x, m)?);
16214        }
16215        Ok(out)
16216    }
16217
16218    /// Cross-request grouped matmul (task #13): run ONE projection group over the
16219    /// CONCATENATION of several sequences' activations (m = sum of per-seq rows — the
16220    /// GEMM-batch win vLLM gets from continuous batching), then split each output back
16221    /// into per-seq buffers. Zero view plumbing: gather/scatter are stream-ordered D2D
16222    /// copies (~us at prime sizes). NUMERIC CONFIG NOTE: a GEMM at m=sum tiles K
16223    /// differently than per-seq GEMMs — argmax-gated like every prefill GEMM change.
16224    pub fn matmul_group_multi(
16225        &self,
16226        ws: &[&crate::model::GpuTensor],
16227        xs: &[&CudaSlice<f32>],
16228        ms: &[usize],
16229    ) -> Result<Vec<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
16230        assert_eq!(xs.len(), ms.len());
16231        let in_f = ws[0].in_features();
16232        let total: usize = ms.iter().sum();
16233        let mut xcat = self.uninit(total * in_f)?;
16234        let mut off = 0usize;
16235        for (x, &m) in xs.iter().zip(ms) {
16236            self.copy_into(&mut xcat, off * in_f, x, m * in_f)?;
16237            off += m;
16238        }
16239        let ys = self.matmul_group(ws, &xcat, total)?;
16240        let mut out: Vec<Vec<CudaSlice<f32>>> = (0..xs.len()).map(|_| Vec::new()).collect();
16241        for (w, y) in ws.iter().zip(ys) {
16242            let out_f = w.out_features();
16243            let mut off = 0usize;
16244            for (s, &m) in ms.iter().enumerate() {
16245                let mut ys_s = self.uninit(m * out_f)?;
16246                let src = y.slice(off * out_f..(off + m) * out_f);
16247                self.gpu.stream().memcpy_dtod(&src, &mut ys_s)?;
16248                out[s].push(ys_s);
16249                off += m;
16250            }
16251        }
16252        Ok(out)
16253    }
16254
16255    /// True if `w`'s qtype has a batched tensor-core GEMM kernel (the prefill T>1 root fix).
16256    /// Only the 4 daily-hot dtypes: Q8_0, Q4_K, Q6_K, NVFP4. NVFP4 needs in_f % 64 == 0.
16257    /// DEFAULT-ON (2026-06-28): measured pp512 9B-NVFP4 = 1413 tok/s WITH this GEMM vs 298 with the
16258    /// dp4a fallback (4.7x) AND MORE accurate (prefill logit maxdiff 0.159 vs dp4a 0.55, both argmax
16259    /// MATCH). The int8 tensor-core GEMM is unconditional (its historical MEMRA_GEMM opt-in gate
16260    /// shipped with Phase 0 — mma + smem swizzle + cp.async — and was removed). Prefill-only
16261    /// (m>=GEMM_M_THRESHOLD); m=1 decode keeps dp4a/MMVQ (this returns true but matmul only calls it
16262    /// at m>=threshold). Portable CUDA targets always use the correctness fallback; on sm_120a,
16263    /// MEMRA_NO_GEMM forces that same dp4a fallback (the bit-reference).
16264    pub fn gemm_supports(&self, w: &crate::model::GpuTensor) -> bool {
16265        use crate::model::GpuTensor;
16266        if !legacy_quant_gemm_allowed(
16267            cfg!(memra_portable_cuda),
16268            cfg!(memra_hopper_mma),
16269            std::env::var_os("MEMRA_NO_GEMM").is_some(),
16270        ) {
16271            return false;
16272        }
16273        match w {
16274            GpuTensor::Quant { qtype, .. } => {
16275                matches!(*qtype, QT_Q8_0 | QT_Q4_K | QT_Q6_K | QT_Q5_K | QT_Q4_0)
16276                    || (*qtype == QT_NVFP4 && w.in_features() % 64 == 0)
16277            }
16278            GpuTensor::Float { .. } | GpuTensor::FloatBf16 { .. } => false,
16279        }
16280    }
16281
16282    /// Batched tensor-core int8 GEMM with a PRE-QUANTIZED q8_1 activation (aq,ad). The prefill
16283    /// (T>1) root fix: decode each weight 32-block to int8 in shared memory ONCE per (row-tile,
16284    /// K-step) and reuse it across all BN tokens via mma.sync.m16n8k32.s8 — amortizing the weight
16285    /// read/decode N-fold (vs the dp4a matvec's per-token re-read). s32 accumulate is exact vs
16286    /// dp4a; only the final f32 block-scale rounding differs. Caller MUST have checked
16287    /// `gemm_supports(w)`. y[m,out] token-major. NVFP4 per-tensor macro-scale applied post.
16288    pub fn qmatvec_gemm(
16289        &self,
16290        w: &crate::model::GpuTensor,
16291        aq: &CudaSlice<i8>,
16292        ad: &CudaSlice<f32>,
16293        m: usize,
16294    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16295        use crate::model::GpuTensor;
16296        let in_f = w.in_features();
16297        let out_f = w.out_features();
16298        let (bytes, qtype, row_bytes, scale, rp) = match w {
16299            GpuTensor::Quant {
16300                bytes,
16301                qtype,
16302                row_bytes,
16303                scale,
16304                rp,
16305                ..
16306            } => (bytes, *qtype, *row_bytes, *scale, *rp),
16307            _ => unreachable!("gemm_supports guaranteed Quant"),
16308        };
16309        // wgmma arm (sm_90a, task 8): the m64n64k32 warpgroup kernel reads the rp4 split-plane
16310        // mirror AS-IS (qplane rows = its A operand, the half dplane its scales) and the same
16311        // (aq, ad) activation planes. Same numeric class as the mma kernel below (exact s32 per
16312        // 32-block, one f32 scale fold per block, ascending K) — argmax/tolerance gated like
16313        // every prefill GEMM, not bit-gated. MEMRA_WGMMA=0 restores the portable kernel.
16314        if cfg!(memra_hopper_mma) && qtype == QT_Q8_0 && out_f % 64 == 0 && wgmma_gemm_enabled() {
16315            if let GpuTensor::Quant { rp4: Some(m4), .. } = w {
16316                let mut y = self.qmatvec_gemm_q8_0_wgmma_raw(m4, aq, ad, m, in_f, out_f)?;
16317                if scale != 1.0 {
16318                    self.scale_inplace(&mut y, scale, m * out_f)?;
16319                }
16320                return Ok(y);
16321            }
16322        }
16323        let name = match qtype {
16324            QT_Q8_0 => "qmatvec_gemm_q8_0",
16325            QT_Q4_K => "qmatvec_gemm_q4_K",
16326            QT_Q4_0 => {
16327                if rp {
16328                    "qmatvec_gemm_q4_0_rp"
16329                } else {
16330                    "qmatvec_gemm_q4_0"
16331                }
16332            }
16333            QT_Q5_K => "qmatvec_gemm_q5_K",
16334            QT_Q6_K => "qmatvec_gemm_q6_K",
16335            QT_NVFP4 => {
16336                if rp {
16337                    "qmatvec_gemm_nvfp4_rp"
16338                } else {
16339                    "qmatvec_gemm_nvfp4"
16340                }
16341            }
16342            _ => unreachable!(),
16343        };
16344        let f = self.func(name);
16345        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output: skip memset
16346        // CTA tile MUST match the .cu per-kernel tile. MMQ-PORT: kernel1 (Q8_0/Q4_K/Q5_K) runs llama's
16347        // 128x128 SQUARE tile (K1_BM=128 x K1_BN=128, 8 warps); kernel2 (Q6_K/NVFP4) keeps 64x256, 4 warps
16348        // (the macro BM/BN in the .cu). Grid dims are selected by qtype so each launches its own tile.
16349        let is_k1 = matches!(qtype, QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q4_0);
16350        // TUNE SEAM: MEMRA_GEMM_K1_LAUNCH overrides kernel1's launch tile to match a -D-swept fatbin.
16351        let k1_tile = if is_k1 {
16352            k1_launch_override().unwrap_or((128, 128, 8))
16353        } else {
16354            (128, 128, 8)
16355        };
16356        let (bm, bn): (u32, u32) = if is_k1 {
16357            (k1_tile.0, k1_tile.1)
16358        } else {
16359            (64, 256)
16360        };
16361        let warps: u32 = if is_k1 {
16362            k1_tile.2
16363        } else {
16364            match qtype {
16365                QT_NVFP4 => 8,
16366                _ => 4,
16367            }
16368        };
16369        let cfg = LaunchConfig {
16370            grid_dim: ((out_f as u32 + bm - 1) / bm, (m as u32 + bn - 1) / bn, 1),
16371            block_dim: (32, warps, 1),
16372            shared_mem_bytes: 0,
16373        };
16374        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
16375        let __s_b = self.gpu.stream();
16376        let mut b = __s_b.launch_builder(&f);
16377        b.arg(bytes)
16378            .arg(aq)
16379            .arg(ad)
16380            .arg(&mut y)
16381            .arg(&inf)
16382            .arg(&outf)
16383            .arg(&mi)
16384            .arg(&rb);
16385        unsafe {
16386            b.launch(cfg)?;
16387        }
16388        if scale != 1.0 {
16389            self.scale_inplace(&mut y, scale, m * out_f)?;
16390        }
16391        Ok(y)
16392    }
16393
16394    /// Test entry: run the GEMM directly from raw weight bytes + qtype (no GpuTensor). Quantizes
16395    /// the f32 activation `x` to q8_1 internally then launches the tensor-core GEMM. NVFP4 per-tensor
16396    /// macro-scale is NOT applied here (caller passes it separately, like the dp4a path). Used by
16397    /// kernel_check for the bit-equivalence gate vs qmatvec_*_dp4a.
16398    pub fn qmatvec_gemm_raw(
16399        &self,
16400        bytes: &CudaSlice<u8>,
16401        x: &CudaSlice<f32>,
16402        m: usize,
16403        in_f: usize,
16404        out_f: usize,
16405        qtype: i32,
16406        row_bytes: usize,
16407    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16408        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
16409        let name = match qtype {
16410            QT_Q8_0 => "qmatvec_gemm_q8_0",
16411            QT_Q4_K => "qmatvec_gemm_q4_K",
16412            QT_Q4_0 => "qmatvec_gemm_q4_0",
16413            QT_Q5_K => "qmatvec_gemm_q5_K",
16414            QT_Q6_K => "qmatvec_gemm_q6_K",
16415            QT_NVFP4 => "qmatvec_gemm_nvfp4",
16416            QT_NVFP4_RP => "qmatvec_gemm_nvfp4_rp",
16417            _ => panic!("qmatvec_gemm_raw: qtype {qtype} has no GEMM kernel"),
16418        };
16419        let f = self.func(name);
16420        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output: skip memset
16421        // MMQ-PORT: kernel1 (Q8_0/Q4_K/Q5_K) = llama 128x128 tile, 8 warps; kernel2 (Q6_K/NVFP4) = 64x256,
16422        // 4/8 warps. Grid tile per qtype (must match the .cu K1_BM/K1_BN vs BM/BN). KEEP IN SYNC w/ qmatvec_gemm.
16423        let is_k1 = matches!(qtype, QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q4_0);
16424        // TUNE SEAM: MEMRA_GEMM_K1_LAUNCH overrides kernel1's launch tile to match a -D-swept fatbin.
16425        let k1_tile = if is_k1 {
16426            k1_launch_override().unwrap_or((128, 128, 8))
16427        } else {
16428            (128, 128, 8)
16429        };
16430        let (bm, bn): (u32, u32) = if is_k1 {
16431            (k1_tile.0, k1_tile.1)
16432        } else {
16433            (64, 256)
16434        };
16435        let warps: u32 = if is_k1 {
16436            k1_tile.2
16437        } else {
16438            match qtype {
16439                QT_NVFP4 | QT_NVFP4_RP => 8,
16440                _ => 4,
16441            }
16442        };
16443        let cfg = LaunchConfig {
16444            grid_dim: ((out_f as u32 + bm - 1) / bm, (m as u32 + bn - 1) / bn, 1),
16445            block_dim: (32, warps, 1),
16446            shared_mem_bytes: 0,
16447        };
16448        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
16449        let __s_b = self.gpu.stream();
16450        let mut b = __s_b.launch_builder(&f);
16451        b.arg(bytes)
16452            .arg(&aq)
16453            .arg(&ad)
16454            .arg(&mut y)
16455            .arg(&inf)
16456            .arg(&outf)
16457            .arg(&mi)
16458            .arg(&rb);
16459        unsafe {
16460            b.launch(cfg)?;
16461        }
16462        Ok(y)
16463    }
16464
16465    /// H100 warpgroup GEMM raw entry (task 8): launch `qmatvec_gemm_q8_0_wgmma` on an rp4
16466    /// split-plane mirror + pre-quantized (aq, ad) activation planes. One warpgroup (128 thr)
16467    /// owns a 64x64 C tile; grid (out_f/64, ceil(m/64)). out_f % 64 == 0 REQUIRED (row loads
16468    /// and dplane scale reads are unguarded); the token edge is guarded in-kernel.
16469    /// Standalone harness verdict (tools/bench_q8_gemm_wgmma.cu, 4096x4096x512): rel 1.6e-05
16470    /// vs CPU ref, 179us vs the portable mma kernel's 688us (3.84x, unpipelined).
16471    pub fn qmatvec_gemm_q8_0_wgmma_raw(
16472        &self,
16473        rp4: &CudaSlice<u8>,
16474        aq: &CudaSlice<i8>,
16475        ad: &CudaSlice<f32>,
16476        m: usize,
16477        in_f: usize,
16478        out_f: usize,
16479    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16480        assert!(
16481            out_f % 64 == 0 && in_f % 32 == 0,
16482            "wgmma GEMM needs out_f%64==0, in_f%32==0"
16483        );
16484        let f = self.func("qmatvec_gemm_q8_0_wgmma");
16485        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output
16486        let cfg = LaunchConfig {
16487            grid_dim: ((out_f / 64) as u32, (m as u32).div_ceil(64), 1),
16488            block_dim: (128, 1, 1),
16489            shared_mem_bytes: 0,
16490        };
16491        let (inf, outf, mi) = (in_f as i32, out_f as i32, m as i32);
16492        let __s_b = self.gpu.stream();
16493        let mut b = __s_b.launch_builder(&f);
16494        b.arg(rp4)
16495            .arg(aq)
16496            .arg(ad)
16497            .arg(&mut y)
16498            .arg(&inf)
16499            .arg(&outf)
16500            .arg(&mi);
16501        unsafe {
16502            b.launch(cfg)?;
16503        }
16504        Ok(y)
16505    }
16506
16507    /// y[i] *= s. NVFP4 per-tensor macro-scale broadcast over the whole output.
16508    pub fn scale_inplace(
16509        &self,
16510        y: &mut CudaSlice<f32>,
16511        s: f32,
16512        n: usize,
16513    ) -> Result<(), Box<dyn std::error::Error>> {
16514        let f = self.func("scale_f32");
16515        let cfg = LaunchConfig::for_num_elems(n as u32);
16516        let (sf, ni) = (s, n as i32);
16517        let __s_b = self.gpu.stream();
16518        let mut b = __s_b.launch_builder(&f);
16519        b.arg(y).arg(&sf).arg(&ni);
16520        unsafe {
16521            b.launch(cfg)?;
16522        }
16523        Ok(())
16524    }
16525
16526    /// MEMRA_FULL_PREC dequant-on-use: expand a bf16-resident weight (`GpuTensor::FloatBf16`, raw
16527    /// bf16 bytes) to a transient f32 scratch of `n` elements, which then feeds the existing f32
16528    /// cuBLASLt GEMV. The scratch is freed when the caller drops it, so peak VRAM = resident bf16
16529    /// weights + ONE (largest) weight's f32 expansion + activations. SLOW IS FINE (research mode).
16530    pub fn bf16_to_f32(
16531        &self,
16532        data: &cudarc::driver::CudaView<'_, u8>,
16533        n: usize,
16534    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16535        let mut out = self.alloc_uninit::<f32>(n)?;
16536        let f = self.func("bf16_to_f32");
16537        let cfg = LaunchConfig::for_num_elems(n as u32);
16538        let ni = n as i32;
16539        let __s_b = self.gpu.stream();
16540        let mut b = __s_b.launch_builder(&f);
16541        b.arg(data).arg(&mut out).arg(&ni);
16542        unsafe {
16543            b.launch(cfg)?;
16544        }
16545        Ok(out)
16546    }
16547
16548    /// Chunked bf16 linear (MEMRA_FULL_PREC): y[m,out] = x @ W_bf16^T with the f32 dequant scratch
16549    /// bounded to CHUNK_ROWS rows (256MB at in_f=4096) instead of the whole weight — the 4GB
16550    /// lm_head expansion OOM'd the 24GB budget. Row-chunking partitions OUTPUT rows; each row's
16551    /// dot is computed by the identical kernel on identical bytes, so per-(token,row) results are
16552    /// bit-identical to the unchunked form. `exact` selects linear_decode_exact (per-column m=1
16553    /// calls, the spec-verify contract) vs plain linear.
16554    fn linear_bf16_chunked(
16555        &self,
16556        x: &CudaSlice<f32>,
16557        data: &CudaSlice<u8>,
16558        m: usize,
16559        in_f: usize,
16560        out_f: usize,
16561        exact: bool,
16562        canonical_chunk_rows: Option<usize>,
16563    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16564        // MEMRA_STEP_TP_TIMING=1: cumulative cost of the per-call BF16->F32 expansion class
16565        // (alloc + convert kernel + f32 cuBLASLt = ~5x weight traffic). Prints every 1024 calls.
16566        // The stream sync per call perturbs wall; diagnostic only, never in a receipts run.
16567        static EXP_NS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
16568        static EXP_CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
16569        static EXP_WBYTES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
16570        let timing = std::env::var("MEMRA_STEP_TP_TIMING").as_deref() == Ok("1");
16571        let started = timing.then(std::time::Instant::now);
16572        let result =
16573            self.linear_bf16_chunked_inner(x, data, m, in_f, out_f, exact, canonical_chunk_rows);
16574        if let Some(started) = started {
16575            use std::sync::atomic::Ordering;
16576            self.stream().synchronize()?;
16577            let ns = EXP_NS.fetch_add(started.elapsed().as_nanos() as u64, Ordering::Relaxed)
16578                + started.elapsed().as_nanos() as u64;
16579            let wb = EXP_WBYTES.fetch_add((in_f * out_f * 2) as u64, Ordering::Relaxed)
16580                + (in_f * out_f * 2) as u64;
16581            let calls = EXP_CALLS.fetch_add(1, Ordering::Relaxed) + 1;
16582            if calls % 1024 == 0 {
16583                eprintln!(
16584                    "[bf16-expand-timing] calls={calls} total_ms={:.1} avg_us={:.1} \
16585                     weight_gb={:.2}",
16586                    ns as f64 / 1.0e6,
16587                    ns as f64 / calls as f64 / 1.0e3,
16588                    wb as f64 / 1.0e9,
16589                );
16590            }
16591        }
16592        result
16593    }
16594
16595    /// MEMRA_BF16_MMV=1: decode-time (m=1) BF16-resident matvec door. Numeric class change vs
16596    /// the expansion path (single-kernel deterministic tree reduce vs f32 cuBLASLt), so it is
16597    /// default OFF and gated by the run-gen argmax gate + boot battery like the other
16598    /// numeric-class doors (DEV_ROUTES precedent).
16599    pub(crate) fn bf16_mmv_on() -> bool {
16600        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16601        *ON.get_or_init(|| std::env::var("MEMRA_BF16_MMV").as_deref() == Ok("1"))
16602    }
16603
16604    /// One-block-per-row BF16 matvec: y[out_f] = W_bf16[out_f, in_f] @ x[in_f], f32 accumulate.
16605    /// Weights read once as raw bf16 (same bits<<16 expansion contract as `deq`'s QT_BF16 arm).
16606    fn matvec_bf16(
16607        &self,
16608        data: &CudaSlice<u8>,
16609        x: &CudaSlice<f32>,
16610        in_f: usize,
16611        out_f: usize,
16612    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16613        if data.len() != in_f * out_f * 2 || x.len() < in_f || in_f % 8 != 0 {
16614            return Err(format!(
16615                "matvec_bf16 geometry bytes={} x={} in={in_f} out={out_f}",
16616                data.len(),
16617                x.len()
16618            )
16619            .into());
16620        }
16621        let mut y = self.alloc_uninit::<f32>(out_f)?;
16622        let f = self.func("matvec_bf16_f32acc");
16623        let cfg = LaunchConfig {
16624            grid_dim: (out_f as u32, 1, 1),
16625            block_dim: (mmv_block(), 1, 1),
16626            shared_mem_bytes: 0,
16627        };
16628        let ini = in_f as i32;
16629        let __s_bld = self.gpu.stream();
16630        let mut bld = __s_bld.launch_builder(&f);
16631        bld.arg(data).arg(x).arg(&mut y).arg(&ini);
16632        unsafe {
16633            bld.launch(cfg)?;
16634        }
16635        Ok(y)
16636    }
16637
16638    /// Fused QK rms-norm + neox rope (t=1): one launch per rank replaces two rms_norm
16639    /// launches, a position upload, and the rope launch; the position is read directly from
16640    /// the caller's device counter (UVA). Numeric-class door (see the kernel doc).
16641    #[allow(clippy::too_many_arguments)]
16642    /// FUSION #1: qk norms + rope + dcw KV append + last-block len inc, one launch
16643    /// (replaces qk_norm_rope_into + append_kv_quantized_dcw + inc_i32 on the dcw path).
16644    /// Bit-identical to the split kernels; requires head_dim == 128 and
16645    /// kv_dim_v == kv_dim_k == nh_k * head_dim (caller-guarded fallback otherwise).
16646    #[allow(clippy::too_many_arguments)]
16647    pub fn qk_norm_rope_append_inc_dcw(
16648        &self,
16649        q_raw: &CudaSlice<f32>,
16650        k_raw: &CudaSlice<f32>,
16651        v_raw: &CudaSlice<f32>,
16652        qw: &CudaSlice<f32>,
16653        kw: &CudaSlice<f32>,
16654        q_out: &mut CudaSlice<f32>,
16655        k_out: &mut CudaSlice<f32>,
16656        pos: &CudaSlice<i32>,
16657        k_plane: &mut CudaSlice<u8>,
16658        v_plane: &mut CudaSlice<u8>,
16659        // Shared ref by the planes_and_counters_mut split-borrow contract; the kernel is the
16660        // (single) writer, exactly like the split append+inc pair it replaces.
16661        len_dev: &CudaSlice<i32>,
16662        base_dev: Option<&CudaSlice<i32>>,
16663        done_ctr: &mut CudaSlice<u32>,
16664        kv_dim_k: usize,
16665        kv_dim_v: usize,
16666        k_tok_bytes: usize,
16667        v_tok_bytes: usize,
16668        head_dim: usize,
16669        n_dims: usize,
16670        nh_q: usize,
16671        nh_k: usize,
16672        eps: f32,
16673        freq_base: f32,
16674        freq_scale: f32,
16675        ff: Option<&CudaSlice<f32>>,
16676    ) -> Result<(), Box<dyn std::error::Error>> {
16677        if head_dim != 128
16678            || kv_dim_v != kv_dim_k
16679            || kv_dim_k != nh_k * head_dim
16680            || q_raw.len() < nh_q * head_dim
16681            || k_raw.len() < nh_k * head_dim
16682            || v_raw.len() < kv_dim_v
16683            || q_out.len() < nh_q * head_dim
16684            || k_out.len() < nh_k * head_dim
16685            || pos.is_empty()
16686            || done_ctr.is_empty()
16687        {
16688            return Err(format!(
16689                "qk_norm_rope_append_inc geometry head_dim={head_dim} nh_q={nh_q} nh_k={nh_k}                  kv_k={kv_dim_k} kv_v={kv_dim_v}"
16690            )
16691            .into());
16692        }
16693        let f = self.func("qk_norm_rope_append_inc_dcw");
16694        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
16695        let cfg = LaunchConfig {
16696            grid_dim: ((nh_q + nh_k) as u32, 1, 1),
16697            block_dim: (128, 1, 1),
16698            shared_mem_bytes: 0,
16699        };
16700        let (kvk, kvv) = (kv_dim_k as i32, kv_dim_v as i32);
16701        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
16702        let (hd, nd, nq) = (head_dim as i32, n_dims as i32, nh_q as i32);
16703        let null: u64 = 0;
16704        let __s_b = self.gpu.stream();
16705        let mut b = __s_b.launch_builder(&f);
16706        b.arg(q_raw)
16707            .arg(k_raw)
16708            .arg(v_raw)
16709            .arg(qw)
16710            .arg(kw)
16711            .arg(q_out)
16712            .arg(k_out)
16713            .arg(pos)
16714            .arg(&mut *k_plane)
16715            .arg(&mut *v_plane)
16716            .arg(len_dev);
16717        match base_dev {
16718            Some(base) => {
16719                b.arg(base);
16720            }
16721            None => {
16722                b.arg(&null);
16723            }
16724        }
16725        b.arg(&mut *done_ctr)
16726            .arg(&kvk)
16727            .arg(&kvv)
16728            .arg(&ktb)
16729            .arg(&vtb)
16730            .arg(&hd)
16731            .arg(&nd)
16732            .arg(&nq)
16733            .arg(&eps)
16734            .arg(&theta_scale)
16735            .arg(&freq_scale);
16736        match ff {
16737            Some(freqs) => {
16738                b.arg(freqs);
16739            }
16740            None => {
16741                b.arg(&null);
16742            }
16743        }
16744        unsafe {
16745            b.launch(cfg)?;
16746        }
16747        Ok(())
16748    }
16749
16750    pub fn qk_norm_rope_into(
16751        &self,
16752        q_raw: &CudaSlice<f32>,
16753        k_raw: &CudaSlice<f32>,
16754        qw: &CudaSlice<f32>,
16755        kw: &CudaSlice<f32>,
16756        q_out: &mut CudaSlice<f32>,
16757        k_out: &mut CudaSlice<f32>,
16758        pos: &CudaSlice<i32>,
16759        head_dim: usize,
16760        n_dims: usize,
16761        nh_q: usize,
16762        nh_k: usize,
16763        eps: f32,
16764        freq_base: f32,
16765        freq_scale: f32,
16766        ff: Option<&CudaSlice<f32>>,
16767    ) -> Result<(), Box<dyn std::error::Error>> {
16768        if head_dim > 512
16769            || q_raw.len() < nh_q * head_dim
16770            || k_raw.len() < nh_k * head_dim
16771            || q_out.len() < nh_q * head_dim
16772            || k_out.len() < nh_k * head_dim
16773            || qw.len() < head_dim
16774            || kw.len() < head_dim
16775            || pos.is_empty()
16776        {
16777            return Err(format!(
16778                "qk_norm_rope geometry head_dim={head_dim} nh_q={nh_q} nh_k={nh_k}"
16779            )
16780            .into());
16781        }
16782        let f = self.func("qk_norm_rope_f32");
16783        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
16784        let cfg = LaunchConfig {
16785            grid_dim: ((nh_q + nh_k) as u32, 1, 1),
16786            block_dim: (128, 1, 1),
16787            shared_mem_bytes: 0,
16788        };
16789        let (hd, nd, nq) = (head_dim as i32, n_dims as i32, nh_q as i32);
16790        let __s_b = self.gpu.stream();
16791        let mut b = __s_b.launch_builder(&f);
16792        b.arg(q_raw)
16793            .arg(k_raw)
16794            .arg(qw)
16795            .arg(kw)
16796            .arg(q_out)
16797            .arg(k_out)
16798            .arg(pos)
16799            .arg(&hd)
16800            .arg(&nd)
16801            .arg(&nq)
16802            .arg(&eps)
16803            .arg(&theta_scale)
16804            .arg(&freq_scale);
16805        match ff {
16806            Some(ffv) => {
16807                b.arg(ffv);
16808                unsafe {
16809                    b.launch(cfg)?;
16810                }
16811            }
16812            None => {
16813                let null: u64 = 0;
16814                b.arg(&null);
16815                unsafe {
16816                    b.launch(cfg)?;
16817                }
16818            }
16819        }
16820        Ok(())
16821    }
16822
16823    /// Four-block F32 matvec with in-order block accumulation (see the kernel doc): one
16824    /// launch computes a rank's whole O partial from its four canonical column blocks.
16825    #[allow(clippy::too_many_arguments)]
16826    pub fn matvec_f32_b4_into(
16827        &self,
16828        w: [&CudaSlice<f32>; 4],
16829        x: &CudaSlice<f32>,
16830        y: &mut CudaSlice<f32>,
16831        block_cols: usize,
16832        out_f: usize,
16833    ) -> Result<(), Box<dyn std::error::Error>> {
16834        if block_cols % 4 != 0
16835            || x.len() < 4 * block_cols
16836            || y.len() < out_f
16837            || w.iter().any(|w| w.len() != out_f * block_cols)
16838        {
16839            return Err(format!(
16840                "matvec_f32_b4 geometry block_cols={block_cols} out={out_f} x={}",
16841                x.len()
16842            )
16843            .into());
16844        }
16845        let f = self.func("matvec_f32_b4");
16846        let cfg = LaunchConfig {
16847            grid_dim: (out_f as u32, 1, 1),
16848            block_dim: (128, 1, 1),
16849            shared_mem_bytes: 0,
16850        };
16851        let (bc, of) = (block_cols as i32, out_f as i32);
16852        let __s_b = self.gpu.stream();
16853        let mut b = __s_b.launch_builder(&f);
16854        b.arg(w[0])
16855            .arg(w[1])
16856            .arg(w[2])
16857            .arg(w[3])
16858            .arg(x)
16859            .arg(y)
16860            .arg(&bc)
16861            .arg(&of);
16862        unsafe {
16863            b.launch(cfg)?;
16864        }
16865        Ok(())
16866    }
16867
16868    /// Sequential weighted row-sum: y[i] = sum_p w[p] * x[p*width+i] in row order — the exact
16869    /// per-element FP chain of a zero-reset plus n_rows sequential axpy launches.
16870    pub fn axpy_rows_seq_into(
16871        &self,
16872        x: &CudaSlice<f32>,
16873        w: &CudaSlice<f32>,
16874        y: &mut CudaSlice<f32>,
16875        width: usize,
16876        n_rows: usize,
16877    ) -> Result<(), Box<dyn std::error::Error>> {
16878        if x.len() < n_rows * width || w.len() < n_rows || y.len() < width {
16879            return Err(format!(
16880                "axpy_rows_seq geometry x={} w={} y={} width={width} rows={n_rows}",
16881                x.len(),
16882                w.len(),
16883                y.len()
16884            )
16885            .into());
16886        }
16887        let f = self.func("axpy_rows_seq_f32");
16888        let cfg = LaunchConfig::for_num_elems(width as u32);
16889        let (wi, nr) = (width as i32, n_rows as i32);
16890        let __s_b = self.gpu.stream();
16891        let mut b = __s_b.launch_builder(&f);
16892        b.arg(x).arg(w).arg(y).arg(&wi).arg(&nr);
16893        unsafe {
16894            b.launch(cfg)?;
16895        }
16896        Ok(())
16897    }
16898
16899    /// Row-offset twin of `axpy_rows_seq_md_into` (spec verify t-column combine): the
16900    /// accumulation runs over rows [row0, row0+n_rows) of a taller partial slab — the
16901    /// exact sequential FP chain of the base kernel over that window.
16902    #[allow(clippy::too_many_arguments)]
16903    pub fn axpy_rows_seq_md_off_into(
16904        &self,
16905        x: &CudaSlice<f32>,
16906        w_route: &CudaSlice<f32>,
16907        md: &CudaSlice<f32>,
16908        sel: &CudaSlice<i32>,
16909        y: &mut CudaSlice<f32>,
16910        width: usize,
16911        n_rows: usize,
16912        row0: usize,
16913    ) -> Result<(), Box<dyn std::error::Error>> {
16914        if x.len() < (row0 + n_rows) * width
16915            || w_route.len() < row0 + n_rows
16916            || sel.len() < row0 + n_rows
16917            || y.len() < width
16918        {
16919            return Err(format!(
16920                "axpy_rows_seq_md_off geometry x={} w={} sel={} y={} width={width} \
16921                 rows={n_rows} row0={row0}",
16922                x.len(),
16923                w_route.len(),
16924                sel.len(),
16925                y.len()
16926            )
16927            .into());
16928        }
16929        let f = self.func("axpy_rows_seq_md_off_f32");
16930        let cfg = LaunchConfig::for_num_elems(width as u32);
16931        let (wi, nr, r0) = (width as i32, n_rows as i32, row0 as i32);
16932        let __s_b = self.gpu.stream();
16933        let mut b = __s_b.launch_builder(&f);
16934        b.arg(x)
16935            .arg(w_route)
16936            .arg(md)
16937            .arg(sel)
16938            .arg(y)
16939            .arg(&wi)
16940            .arg(&nr)
16941            .arg(&r0);
16942        unsafe {
16943            b.launch(cfg)?;
16944        }
16945        Ok(())
16946    }
16947
16948    /// T-COLUMN twin of `qmatvec_nvfp4_sel_gu_into` (spec verify, MEMRA_TCOL_FFN):
16949    /// 2*n_sel_col selection pairs over TWO activation rows (pair t reads row
16950    /// t/n_sel_col). Per-(pair,row) FP program == the t=1 gu kernel: each column's
16951    /// outputs are bit-equal to its own t=1 launch.
16952    #[allow(clippy::too_many_arguments)]
16953    pub fn qmatvec_nvfp4_sel_gu_tcol_into(
16954        &self,
16955        gate_bank: &CudaSlice<u8>,
16956        up_bank: &CudaSlice<u8>,
16957        sel: &CudaSlice<i32>,
16958        aq: &CudaSlice<i8>,
16959        ad: &CudaSlice<f32>,
16960        yg: &mut CudaSlice<f32>,
16961        yu: &mut CudaSlice<f32>,
16962        n_sel: usize,
16963        n_sel_col: usize,
16964        in_f: usize,
16965        out_f: usize,
16966        row_bytes: usize,
16967        expert_stride: usize,
16968        act_row_stride: usize,
16969        ad_row_stride: usize,
16970    ) -> Result<(), Box<dyn std::error::Error>> {
16971        assert!(in_f % 64 == 0, "NVFP4 dp4a requires in_f % 64 == 0");
16972        if yg.len() < n_sel * out_f
16973            || yu.len() < n_sel * out_f
16974            || sel.len() < n_sel
16975            || n_sel_col == 0
16976            || n_sel % n_sel_col != 0
16977        {
16978            return Err("NVFP4 gu tcol geometry".into());
16979        }
16980        let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_gu_tcol");
16981        let cfg = LaunchConfig {
16982            grid_dim: ((2 * out_f) as u32, n_sel as u32, 1),
16983            block_dim: (128, 1, 1),
16984            shared_mem_bytes: 0,
16985        };
16986        let (inf, outf, ns, nsc) = (in_f as i32, out_f as i32, n_sel as i32, n_sel_col as i32);
16987        let (rb, es) = (row_bytes as i64, expert_stride as i64);
16988        let (ars, adrs) = (act_row_stride as i64, ad_row_stride as i64);
16989        let __s_b = self.gpu.stream();
16990        let mut b = __s_b.launch_builder(&f);
16991        b.arg(gate_bank)
16992            .arg(up_bank)
16993            .arg(sel)
16994            .arg(aq)
16995            .arg(ad)
16996            .arg(yg)
16997            .arg(yu)
16998            .arg(&inf)
16999            .arg(&outf)
17000            .arg(&ns)
17001            .arg(&rb)
17002            .arg(&es)
17003            .arg(&ars)
17004            .arg(&adrs)
17005            .arg(&nsc);
17006        unsafe {
17007            b.launch(cfg)?;
17008        }
17009        Ok(())
17010    }
17011
17012    /// Device-routed twin of `axpy_rows_seq_into`: the per-row weight folds in-kernel
17013    /// (w_route[p] * md[sel[p]] — the same single f32 multiply the host fold performs).
17014    #[allow(clippy::too_many_arguments)]
17015    pub fn axpy_rows_seq_md_into(
17016        &self,
17017        x: &CudaSlice<f32>,
17018        w_route: &CudaSlice<f32>,
17019        md: &CudaSlice<f32>,
17020        sel: &CudaSlice<i32>,
17021        y: &mut CudaSlice<f32>,
17022        width: usize,
17023        n_rows: usize,
17024    ) -> Result<(), Box<dyn std::error::Error>> {
17025        if x.len() < n_rows * width
17026            || w_route.len() < n_rows
17027            || sel.len() < n_rows
17028            || y.len() < width
17029        {
17030            return Err(format!(
17031                "axpy_rows_seq_md geometry x={} w={} sel={} y={} width={width} rows={n_rows}",
17032                x.len(),
17033                w_route.len(),
17034                sel.len(),
17035                y.len()
17036            )
17037            .into());
17038        }
17039        let f = self.func("axpy_rows_seq_md_f32");
17040        let cfg = LaunchConfig::for_num_elems(width as u32);
17041        let (wi, nr) = (width as i32, n_rows as i32);
17042        let __s_b = self.gpu.stream();
17043        let mut b = __s_b.launch_builder(&f);
17044        b.arg(x)
17045            .arg(w_route)
17046            .arg(md)
17047            .arg(sel)
17048            .arg(y)
17049            .arg(&wi)
17050            .arg(&nr);
17051        unsafe {
17052            b.launch(cfg)?;
17053        }
17054        Ok(())
17055    }
17056
17057    /// BF16 twin of `matvec_f32_qkv_into` (weights as raw checkpoint bf16 bytes).
17058    #[allow(clippy::too_many_arguments)]
17059    /// T-COLUMN twin of `matvec_bf16_qkvg_into` (spec verify): weights read once, T input
17060    /// columns accumulated with per-column FP order identical to the t=1 kernel. Outputs
17061    /// land column-major-of-rows: yq[c*out_q + row] etc.
17062    #[allow(clippy::too_many_arguments)]
17063    pub fn matvec_bf16_qkvg_tcol_into(
17064        &self,
17065        wq: &CudaSlice<u8>,
17066        wk: &CudaSlice<u8>,
17067        wv: &CudaSlice<u8>,
17068        wg: &CudaSlice<u8>,
17069        x_t: &CudaSlice<f32>,
17070        yq: &mut CudaSlice<f32>,
17071        yk: &mut CudaSlice<f32>,
17072        yv: &mut CudaSlice<f32>,
17073        yg: &mut CudaSlice<f32>,
17074        in_f: usize,
17075        out_q: usize,
17076        out_kv: usize,
17077        out_g: usize,
17078        t: usize,
17079    ) -> Result<(), Box<dyn std::error::Error>> {
17080        if t == 0
17081            || t > 8
17082            || in_f % 8 != 0
17083            || x_t.len() < t * in_f
17084            || yq.len() < t * out_q
17085            || yk.len() < t * out_kv
17086            || yv.len() < t * out_kv
17087            || (out_g > 0 && yg.len() < t * out_g)
17088        {
17089            return Err("matvec_bf16_qkvg_tcol geometry".into());
17090        }
17091        let f = self.func("matvec_bf16_qkvg_tcol");
17092        let grid = out_q + 2 * out_kv + out_g;
17093        let cfg = LaunchConfig {
17094            grid_dim: (grid as u32, 1, 1),
17095            block_dim: (mmv_block(), 1, 1),
17096            shared_mem_bytes: 0,
17097        };
17098        let (ini, oq, okv, og, ti) = (
17099            in_f as i32,
17100            out_q as i32,
17101            out_kv as i32,
17102            out_g as i32,
17103            t as i32,
17104        );
17105        let __s_b = self.gpu.stream();
17106        let mut b = __s_b.launch_builder(&f);
17107        b.arg(wq)
17108            .arg(wk)
17109            .arg(wv)
17110            .arg(wg)
17111            .arg(x_t)
17112            .arg(yq)
17113            .arg(yk)
17114            .arg(yv)
17115            .arg(yg)
17116            .arg(&ini)
17117            .arg(&oq)
17118            .arg(&okv)
17119            .arg(&og)
17120            .arg(&ti);
17121        unsafe {
17122            b.launch(cfg)?;
17123        }
17124        Ok(())
17125    }
17126
17127    pub fn matvec_bf16_qkvg_into(
17128        &self,
17129        wq: &CudaSlice<u8>,
17130        wk: &CudaSlice<u8>,
17131        wv: &CudaSlice<u8>,
17132        wg: &CudaSlice<u8>,
17133        x: &CudaSlice<f32>,
17134        yq: &mut CudaSlice<f32>,
17135        yk: &mut CudaSlice<f32>,
17136        yv: &mut CudaSlice<f32>,
17137        yg: &mut CudaSlice<f32>,
17138        in_f: usize,
17139        out_q: usize,
17140        out_kv: usize,
17141        out_g: usize,
17142    ) -> Result<(), Box<dyn std::error::Error>> {
17143        if in_f % 8 != 0
17144            || wq.len() != out_q * in_f * 2
17145            || wk.len() != out_kv * in_f * 2
17146            || wv.len() != out_kv * in_f * 2
17147            || wg.len() < out_g * in_f * 2
17148            || x.len() < in_f
17149            || yq.len() < out_q
17150            || yk.len() < out_kv
17151            || yv.len() < out_kv
17152            || (out_g > 0 && yg.len() < out_g)
17153        {
17154            return Err(format!(
17155                "fused bf16 QKV geometry in={in_f} out_q={out_q} out_kv={out_kv} out_g={out_g}"
17156            )
17157            .into());
17158        }
17159        let f = self.func("matvec_bf16_qkvg");
17160        let cfg = LaunchConfig {
17161            grid_dim: ((out_q + 2 * out_kv + out_g) as u32, 1, 1),
17162            block_dim: (mmv_block(), 1, 1),
17163            shared_mem_bytes: 0,
17164        };
17165        let (inf, oq, okv, og) = (in_f as i32, out_q as i32, out_kv as i32, out_g as i32);
17166        let __s_b = self.gpu.stream();
17167        let mut b = __s_b.launch_builder(&f);
17168        b.arg(wq)
17169            .arg(wk)
17170            .arg(wv)
17171            .arg(wg)
17172            .arg(x)
17173            .arg(yq)
17174            .arg(yk)
17175            .arg(yv)
17176            .arg(yg)
17177            .arg(&inf)
17178            .arg(&oq)
17179            .arg(&okv)
17180            .arg(&og);
17181        unsafe {
17182            b.launch(cfg)?;
17183        }
17184        Ok(())
17185    }
17186
17187    /// BF16 twin of `matvec_f32_b4_into` (weights as raw checkpoint bf16 bytes).
17188    pub fn matvec_bf16_b4_into(
17189        &self,
17190        w: [&CudaSlice<u8>; 4],
17191        x: &CudaSlice<f32>,
17192        y: &mut CudaSlice<f32>,
17193        block_cols: usize,
17194        out_f: usize,
17195    ) -> Result<(), Box<dyn std::error::Error>> {
17196        if block_cols % 8 != 0
17197            || x.len() < 4 * block_cols
17198            || y.len() < out_f
17199            || w.iter().any(|w| w.len() != out_f * block_cols * 2)
17200        {
17201            return Err(format!(
17202                "bf16 b4 geometry block_cols={block_cols} out={out_f} x={}",
17203                x.len()
17204            )
17205            .into());
17206        }
17207        // MEMRA_B4_X2=1: the #2b grid-halving twin — half the blocks, two rows each,
17208        // bit-identical per row (the second row's stream hides the first's reduce tail).
17209        static B4_X2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
17210        let x2 = *B4_X2.get_or_init(|| std::env::var("MEMRA_B4_X2").as_deref() == Ok("1"));
17211        let f = self.func(if x2 {
17212            "matvec_bf16_b4_x2"
17213        } else {
17214            "matvec_bf16_b4"
17215        });
17216        let grid = if x2 { out_f.div_ceil(2) } else { out_f };
17217        let cfg = LaunchConfig {
17218            grid_dim: (grid as u32, 1, 1),
17219            block_dim: (mmv_block(), 1, 1),
17220            shared_mem_bytes: 0,
17221        };
17222        let (bc, of) = (block_cols as i32, out_f as i32);
17223        let __s_b = self.gpu.stream();
17224        let mut b = __s_b.launch_builder(&f);
17225        b.arg(w[0])
17226            .arg(w[1])
17227            .arg(w[2])
17228            .arg(w[3])
17229            .arg(x)
17230            .arg(y)
17231            .arg(&bc)
17232            .arg(&of);
17233        unsafe {
17234            b.launch(cfg)?;
17235        }
17236        Ok(())
17237    }
17238
17239    /// T-COLUMN twin of `matvec_bf16_b4_into` (spec verify o_proj): weights read once, T
17240    /// gated rows (each 4*block_cols wide) accumulated with per-column FP order identical
17241    /// to the t=1 kernel. Outputs land y[c*out_f + row]. Same blockDim as the t=1 launch —
17242    /// the shared-memory reduce order depends on it. Refuses under MEMRA_B4_X2 (different
17243    /// t=1 program).
17244    pub fn matvec_bf16_b4_tcol_into(
17245        &self,
17246        w: [&CudaSlice<u8>; 4],
17247        x_t: &CudaSlice<f32>,
17248        y_t: &mut CudaSlice<f32>,
17249        block_cols: usize,
17250        out_f: usize,
17251        t: usize,
17252    ) -> Result<(), Box<dyn std::error::Error>> {
17253        if block_cols % 8 != 0
17254            || t == 0
17255            || t > 8
17256            || x_t.len() < t * 4 * block_cols
17257            || y_t.len() < t * out_f
17258            || w.iter().any(|w| w.len() != out_f * block_cols * 2)
17259        {
17260            return Err(format!(
17261                "bf16 b4 tcol geometry block_cols={block_cols} out={out_f} t={t} x={}",
17262                x_t.len()
17263            )
17264            .into());
17265        }
17266        if std::env::var("MEMRA_B4_X2").as_deref() == Ok("1") {
17267            return Err(
17268                "b4 tcol verify is qualified against the plain b4 kernel only \
17269                        (MEMRA_B4_X2=1 is a different t=1 program)"
17270                    .into(),
17271            );
17272        }
17273        let f = self.func("matvec_bf16_b4_tcol");
17274        let cfg = LaunchConfig {
17275            grid_dim: (out_f as u32, 1, 1),
17276            block_dim: (mmv_block(), 1, 1),
17277            shared_mem_bytes: 0,
17278        };
17279        let (bc, of, ti) = (block_cols as i32, out_f as i32, t as i32);
17280        let __s_b = self.gpu.stream();
17281        let mut b = __s_b.launch_builder(&f);
17282        b.arg(w[0])
17283            .arg(w[1])
17284            .arg(w[2])
17285            .arg(w[3])
17286            .arg(x_t)
17287            .arg(y_t)
17288            .arg(&bc)
17289            .arg(&of)
17290            .arg(&ti);
17291        unsafe {
17292            b.launch(cfg)?;
17293        }
17294        Ok(())
17295    }
17296
17297    /// `matvec_bf16` writing into a caller-owned output (persistent-workspace form).
17298    pub fn matvec_bf16_into(
17299        &self,
17300        data: &CudaSlice<u8>,
17301        x: &CudaSlice<f32>,
17302        y: &mut CudaSlice<f32>,
17303        in_f: usize,
17304        out_f: usize,
17305    ) -> Result<(), Box<dyn std::error::Error>> {
17306        if data.len() != in_f * out_f * 2 || x.len() < in_f || in_f % 8 != 0 || y.len() < out_f {
17307            return Err(format!(
17308                "matvec_bf16_into geometry bytes={} x={} y={} in={in_f} out={out_f}",
17309                data.len(),
17310                x.len(),
17311                y.len()
17312            )
17313            .into());
17314        }
17315        // MEMRA_DOWN_X4=1 (short-row shapes, in_f<=2048): four sequential rows per
17316        // block, exact f32acc per-row program — cures the 1-iteration latency
17317        // starvation (shexp down measured 420GB/s at in_f=1280).
17318        static X4: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
17319        let x4 = *X4.get_or_init(|| std::env::var("MEMRA_DOWN_X4").as_deref() == Ok("1"))
17320            && in_f <= 2048;
17321        if x4 {
17322            let f = self.func("matvec_bf16_f32acc_x4");
17323            let cfg = LaunchConfig {
17324                grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
17325                block_dim: (mmv_block(), 1, 1),
17326                shared_mem_bytes: 0,
17327            };
17328            let (ini, outi) = (in_f as i32, out_f as i32);
17329            let __s_b = self.gpu.stream();
17330            let mut b = __s_b.launch_builder(&f);
17331            b.arg(data).arg(x).arg(y).arg(&ini).arg(&outi);
17332            unsafe {
17333                b.launch(cfg)?;
17334            }
17335            return Ok(());
17336        }
17337        let f = self.func("matvec_bf16_f32acc");
17338        let cfg = LaunchConfig {
17339            grid_dim: (out_f as u32, 1, 1),
17340            block_dim: (mmv_block(), 1, 1),
17341            shared_mem_bytes: 0,
17342        };
17343        let ini = in_f as i32;
17344        let __s_b = self.gpu.stream();
17345        let mut b = __s_b.launch_builder(&f);
17346        b.arg(data).arg(x).arg(y).arg(&ini);
17347        unsafe {
17348            b.launch(cfg)?;
17349        }
17350        Ok(())
17351    }
17352
17353    /// `matvec_bf16_into` over a WEIGHT VIEW (row-range slice of a bf16 tensor): the head-split
17354    /// door feeds each device its half of the lm-head rows. Same kernel, same per-row program.
17355    pub fn matvec_bf16_view_into(
17356        &self,
17357        data: &cudarc::driver::CudaView<'_, u8>,
17358        x: &CudaSlice<f32>,
17359        y: &mut CudaSlice<f32>,
17360        in_f: usize,
17361        out_f: usize,
17362    ) -> Result<(), Box<dyn std::error::Error>> {
17363        if data.len() != in_f * out_f * 2 || x.len() < in_f || in_f % 8 != 0 || y.len() < out_f {
17364            return Err(format!(
17365                "matvec_bf16_view_into geometry bytes={} x={} y={} in={in_f} out={out_f}",
17366                data.len(),
17367                x.len(),
17368                y.len()
17369            )
17370            .into());
17371        }
17372        let f = self.func("matvec_bf16_f32acc");
17373        let cfg = LaunchConfig {
17374            grid_dim: (out_f as u32, 1, 1),
17375            block_dim: (mmv_block(), 1, 1),
17376            shared_mem_bytes: 0,
17377        };
17378        let ini = in_f as i32;
17379        let __s_b = self.gpu.stream();
17380        let mut b = __s_b.launch_builder(&f);
17381        b.arg(data).arg(x).arg(y).arg(&ini);
17382        unsafe {
17383            b.launch(cfg)?;
17384        }
17385        Ok(())
17386    }
17387
17388    /// `matvec_bf16_into` with a RAW u64 output pointer (UVA — the dev1-shexp down row
17389    /// lands root-resident over P2P). Same kernel, same per-row program: bit-identical.
17390    pub fn matvec_bf16_raw_out(
17391        &self,
17392        w: &CudaSlice<u8>,
17393        x: &CudaSlice<f32>,
17394        y_raw: u64,
17395        in_f: usize,
17396        out_f: usize,
17397    ) -> Result<(), Box<dyn std::error::Error>> {
17398        if w.len() != in_f * out_f * 2 || x.len() < in_f || in_f % 8 != 0 || y_raw == 0 {
17399            return Err("matvec_bf16_raw_out geometry".into());
17400        }
17401        let f = self.func("matvec_bf16_f32acc");
17402        let cfg = LaunchConfig {
17403            grid_dim: (out_f as u32, 1, 1),
17404            block_dim: (mmv_block(), 1, 1),
17405            shared_mem_bytes: 0,
17406        };
17407        let ini = in_f as i32;
17408        let __s_b = self.gpu.stream();
17409        let mut b = __s_b.launch_builder(&f);
17410        b.arg(w).arg(x).arg(&y_raw).arg(&ini);
17411        unsafe {
17412            b.launch(cfg)?;
17413        }
17414        Ok(())
17415    }
17416
17417    /// MOE TAIL FUSION M1: dst = (a + b) + sh*scale[0] in one launch (sh/scale as RAW
17418    /// UVA pointers so the caller passes persistent-static rows without holding locks).
17419    /// Exact per-element sequence of the split add + add_scaled_rows pair.
17420    pub fn add3_raw(
17421        &self,
17422        a: &CudaSlice<f32>,
17423        b: &CudaSlice<f32>,
17424        sh_raw: u64,
17425        scale_raw: u64,
17426        dst: &mut CudaSlice<f32>,
17427        n: usize,
17428    ) -> Result<(), Box<dyn std::error::Error>> {
17429        if a.len() < n || b.len() < n || dst.len() < n || sh_raw == 0 || scale_raw == 0 {
17430            return Err("add3_raw geometry".into());
17431        }
17432        let f = self.func("add3_f32");
17433        let cfg = LaunchConfig {
17434            grid_dim: ((n as u32).div_ceil(256), 1, 1),
17435            block_dim: (256, 1, 1),
17436            shared_mem_bytes: 0,
17437        };
17438        let ni = n as i32;
17439        let __s_b = self.gpu.stream();
17440        let mut bld = __s_b.launch_builder(&f);
17441        bld.arg(a)
17442            .arg(b)
17443            .arg(&sh_raw)
17444            .arg(&scale_raw)
17445            .arg(dst)
17446            .arg(&ni);
17447        unsafe {
17448            bld.launch(cfg)?;
17449        }
17450        Ok(())
17451    }
17452
17453    /// FUSION #2e: shexp down matvec + scaled accumulate (dst[r] += dot_r * scale[0]),
17454    /// one launch replacing matvec_bf16_into + the ownership copy + add_scaled_rows.
17455    pub fn matvec_bf16_down_addscale_into(
17456        &self,
17457        w: &CudaSlice<u8>,
17458        x: &CudaSlice<f32>,
17459        scale: &CudaSlice<f32>,
17460        dst: &mut CudaSlice<f32>,
17461        in_f: usize,
17462        out_f: usize,
17463    ) -> Result<(), Box<dyn std::error::Error>> {
17464        if w.len() != in_f * out_f * 2
17465            || x.len() < in_f
17466            || in_f % 8 != 0
17467            || dst.len() < out_f
17468            || scale.is_empty()
17469        {
17470            return Err("matvec_bf16_down_addscale geometry".into());
17471        }
17472        let f = self.func("matvec_bf16_down_addscale");
17473        let cfg = LaunchConfig {
17474            grid_dim: (out_f as u32, 1, 1),
17475            block_dim: (mmv_block(), 1, 1),
17476            shared_mem_bytes: 0,
17477        };
17478        let ini = in_f as i32;
17479        let __s_b = self.gpu.stream();
17480        let mut b = __s_b.launch_builder(&f);
17481        b.arg(w).arg(x).arg(scale).arg(dst).arg(&ini);
17482        unsafe {
17483            b.launch(cfg)?;
17484        }
17485        Ok(())
17486    }
17487
17488    /// FUSION #2b: shexp dual matvec + SwiGLU act, one launch (bit-identical to
17489    /// matvec_bf16_dual_into + ffn_act_lim at gs=us=1; limit=None takes plain silu).
17490    pub fn matvec_bf16_dual_silu_into(
17491        &self,
17492        wg: &CudaSlice<u8>,
17493        wu: &CudaSlice<u8>,
17494        x: &CudaSlice<f32>,
17495        act: &mut CudaSlice<f32>,
17496        in_f: usize,
17497        out_f: usize,
17498        limit: Option<f32>,
17499    ) -> Result<(), Box<dyn std::error::Error>> {
17500        if wg.len() != in_f * out_f * 2
17501            || wu.len() != in_f * out_f * 2
17502            || x.len() < in_f
17503            || in_f % 8 != 0
17504            || act.len() < out_f
17505        {
17506            return Err("matvec_bf16_dual_silu geometry".into());
17507        }
17508        let f = self.func("matvec_bf16_dual_silu");
17509        let cfg = LaunchConfig {
17510            grid_dim: (out_f as u32, 1, 1),
17511            block_dim: (mmv_block(), 1, 1),
17512            shared_mem_bytes: 0,
17513        };
17514        let (ini, outi) = (in_f as i32, out_f as i32);
17515        let lim = limit.unwrap_or(0.0);
17516        let __s_b = self.gpu.stream();
17517        let mut b = __s_b.launch_builder(&f);
17518        b.arg(wg)
17519            .arg(wu)
17520            .arg(x)
17521            .arg(act)
17522            .arg(&ini)
17523            .arg(&outi)
17524            .arg(&lim);
17525        unsafe {
17526            b.launch(cfg)?;
17527        }
17528        Ok(())
17529    }
17530
17531    /// `matvec_bf16_dual_into` over WEIGHT VIEWS (row-range slices): the shexp row-split
17532    /// door feeds each device its half of the gate/up rows. Same kernel, same per-row program.
17533    #[allow(clippy::too_many_arguments)]
17534    pub fn matvec_bf16_dual_view_into(
17535        &self,
17536        wg: &cudarc::driver::CudaView<'_, u8>,
17537        wu: &cudarc::driver::CudaView<'_, u8>,
17538        x: &CudaSlice<f32>,
17539        yg: &mut CudaSlice<f32>,
17540        yu: &mut CudaSlice<f32>,
17541        in_f: usize,
17542        out_f: usize,
17543    ) -> Result<(), Box<dyn std::error::Error>> {
17544        if wg.len() != in_f * out_f * 2
17545            || wu.len() != in_f * out_f * 2
17546            || x.len() < in_f
17547            || in_f % 8 != 0
17548            || yg.len() < out_f
17549            || yu.len() < out_f
17550        {
17551            return Err(format!(
17552                "matvec_bf16_dual_view_into geometry wg={} wu={} x={} in={in_f} out={out_f}",
17553                wg.len(),
17554                wu.len(),
17555                x.len()
17556            )
17557            .into());
17558        }
17559        let f = self.func("matvec_bf16_dual");
17560        let cfg = LaunchConfig {
17561            grid_dim: ((2 * out_f) as u32, 1, 1),
17562            block_dim: (mmv_block(), 1, 1),
17563            shared_mem_bytes: 0,
17564        };
17565        let (ini, outi) = (in_f as i32, out_f as i32);
17566        let __s_b = self.gpu.stream();
17567        let mut b = __s_b.launch_builder(&f);
17568        b.arg(wg)
17569            .arg(wu)
17570            .arg(x)
17571            .arg(yg)
17572            .arg(yu)
17573            .arg(&ini)
17574            .arg(&outi);
17575        unsafe {
17576            b.launch(cfg)?;
17577        }
17578        Ok(())
17579    }
17580
17581    /// `matvec_bf16_dual` writing into caller-owned outputs (persistent-workspace form).
17582    #[allow(clippy::too_many_arguments)]
17583    pub fn matvec_bf16_dual_into(
17584        &self,
17585        wg: &CudaSlice<u8>,
17586        wu: &CudaSlice<u8>,
17587        x: &CudaSlice<f32>,
17588        yg: &mut CudaSlice<f32>,
17589        yu: &mut CudaSlice<f32>,
17590        in_f: usize,
17591        out_f: usize,
17592    ) -> Result<(), Box<dyn std::error::Error>> {
17593        if wg.len() != in_f * out_f * 2
17594            || wu.len() != in_f * out_f * 2
17595            || x.len() < in_f
17596            || in_f % 8 != 0
17597            || yg.len() < out_f
17598            || yu.len() < out_f
17599        {
17600            return Err(format!(
17601                "matvec_bf16_dual_into geometry wg={} wu={} x={} in={in_f} out={out_f}",
17602                wg.len(),
17603                wu.len(),
17604                x.len()
17605            )
17606            .into());
17607        }
17608        let f = self.func("matvec_bf16_dual");
17609        let cfg = LaunchConfig {
17610            grid_dim: ((2 * out_f) as u32, 1, 1),
17611            block_dim: (mmv_block(), 1, 1),
17612            shared_mem_bytes: 0,
17613        };
17614        let (ini, outi) = (in_f as i32, out_f as i32);
17615        let __s_b = self.gpu.stream();
17616        let mut b = __s_b.launch_builder(&f);
17617        b.arg(wg)
17618            .arg(wu)
17619            .arg(x)
17620            .arg(yg)
17621            .arg(yu)
17622            .arg(&ini)
17623            .arg(&outi);
17624        unsafe {
17625            b.launch(cfg)?;
17626        }
17627        Ok(())
17628    }
17629
17630    /// Dual bf16 matvec: gate/up (same shape) from one shared input in one launch. Per row
17631    /// bit-identical to two `matvec_bf16` launches. Returns (gate, up).
17632    pub(crate) fn matvec_bf16_dual(
17633        &self,
17634        wg: &CudaSlice<u8>,
17635        wu: &CudaSlice<u8>,
17636        x: &CudaSlice<f32>,
17637        in_f: usize,
17638        out_f: usize,
17639    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
17640        if wg.len() != in_f * out_f * 2
17641            || wu.len() != in_f * out_f * 2
17642            || x.len() < in_f
17643            || in_f % 8 != 0
17644        {
17645            return Err(format!(
17646                "matvec_bf16_dual geometry wg={} wu={} x={} in={in_f} out={out_f}",
17647                wg.len(),
17648                wu.len(),
17649                x.len()
17650            )
17651            .into());
17652        }
17653        let mut yg = self.alloc_uninit::<f32>(out_f)?;
17654        let mut yu = self.alloc_uninit::<f32>(out_f)?;
17655        let f = self.func("matvec_bf16_dual");
17656        let cfg = LaunchConfig {
17657            grid_dim: ((2 * out_f) as u32, 1, 1),
17658            block_dim: (mmv_block(), 1, 1),
17659            shared_mem_bytes: 0,
17660        };
17661        let (ini, outi) = (in_f as i32, out_f as i32);
17662        let __s_b = self.gpu.stream();
17663        let mut b = __s_b.launch_builder(&f);
17664        b.arg(wg)
17665            .arg(wu)
17666            .arg(x)
17667            .arg(&mut yg)
17668            .arg(&mut yu)
17669            .arg(&ini)
17670            .arg(&outi);
17671        unsafe {
17672            b.launch(cfg)?;
17673        }
17674        Ok((yg, yu))
17675    }
17676
17677    #[allow(clippy::too_many_arguments)]
17678    fn linear_bf16_chunked_inner(
17679        &self,
17680        x: &CudaSlice<f32>,
17681        data: &CudaSlice<u8>,
17682        m: usize,
17683        in_f: usize,
17684        out_f: usize,
17685        exact: bool,
17686        canonical_chunk_rows: Option<usize>,
17687    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17688        const CHUNK_BYTES: usize = 256 << 20;
17689        // canonical_chunk_rows callers are the step TP resident ranks whose cross-topology
17690        // equality program is defined in cuBLASLt chunk shapes — the door leaves them alone.
17691        if m == 1
17692            && !exact
17693            && canonical_chunk_rows.is_none()
17694            && in_f % 8 == 0
17695            && Self::bf16_mmv_on()
17696        {
17697            return self.matvec_bf16(data, x, in_f, out_f);
17698        }
17699        let row_bytes = in_f
17700            .checked_mul(std::mem::size_of::<f32>())
17701            .ok_or("BF16 chunk row byte count overflow")?;
17702        if row_bytes == 0 || out_f == 0 {
17703            return Err("BF16 chunk dimensions must be nonzero".into());
17704        }
17705        let max_chunk_rows = (CHUNK_BYTES / row_bytes).max(1).min(out_f);
17706        let chunk_rows = match canonical_chunk_rows {
17707            Some(rows) if rows == 0 => {
17708                return Err("canonical BF16 chunk rows must be nonzero".into());
17709            }
17710            Some(rows) if rows > max_chunk_rows => {
17711                return Err(format!(
17712                    "canonical BF16 chunk rows {rows} exceed the {max_chunk_rows}-row scratch limit"
17713                )
17714                .into());
17715            }
17716            Some(rows) if out_f % rows != 0 => {
17717                return Err(format!(
17718                    "BF16 output width {out_f} is not divisible by canonical {rows}-row chunks"
17719                )
17720                .into());
17721            }
17722            Some(rows) => rows,
17723            None => max_chunk_rows,
17724        };
17725        if chunk_rows >= out_f {
17726            let wf32 = self.bf16_to_f32(&data.slice(0..in_f * out_f * 2), in_f * out_f)?;
17727            return if exact {
17728                self.linear_decode_exact(x, &wf32, m, in_f, out_f)
17729            } else {
17730                self.linear(x, &wf32, m, in_f, out_f)
17731            };
17732        }
17733        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
17734        let mut r0 = 0usize;
17735        while r0 < out_f {
17736            let rows = chunk_rows.min(out_f - r0);
17737            let wslice = data.slice(r0 * in_f * 2..(r0 + rows) * in_f * 2);
17738            let wf32 = self.bf16_to_f32(&wslice, in_f * rows)?;
17739            let yc = if exact {
17740                self.linear_decode_exact(x, &wf32, m, in_f, rows)?
17741            } else {
17742                self.linear(x, &wf32, m, in_f, rows)?
17743            };
17744            // scatter [m, rows] into y[m, out_f] at column offset r0 (m is tiny in decode/verify)
17745            for mi in 0..m {
17746                let src = yc.slice(mi * rows..(mi + 1) * rows);
17747                let mut dst = y.slice_mut(mi * out_f + r0..mi * out_f + r0 + rows);
17748                self.gpu.stream().memcpy_dtod(&src, &mut dst)?;
17749            }
17750            r0 += rows;
17751        }
17752        Ok(y)
17753    }
17754
17755    /// Execute an already resident BF16 projection. This is the model-faithful substrate used by
17756    /// Step tensor-parallel correctness ranks; it preserves checkpoint bytes and the existing
17757    /// chunked BF16 numerical program instead of re-encoding the weight.
17758    pub fn linear_bf16_resident(
17759        &self,
17760        x: &CudaSlice<f32>,
17761        data: &CudaSlice<u8>,
17762        m: usize,
17763        in_f: usize,
17764        out_f: usize,
17765    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17766        if data.len() != in_f * out_f * 2 {
17767            return Err(format!("resident BF16 bytes {} != {out_f}x{in_f}x2", data.len()).into());
17768        }
17769        self.linear_bf16_chunked(x, data, m, in_f, out_f, false, None)
17770    }
17771
17772    /// Execute a resident BF16 projection as fixed-width output-row chunks.
17773    ///
17774    /// Tensor-parallel ranks use this to give TP1/TP2/TP4/TP8 the same cuBLASLt problem shape
17775    /// for every checkpoint row. Callers must derive `canonical_chunk_rows` from the registered
17776    /// model topology rather than the active rank count.
17777    pub fn linear_bf16_resident_canonical_rows(
17778        &self,
17779        x: &CudaSlice<f32>,
17780        data: &CudaSlice<u8>,
17781        m: usize,
17782        in_f: usize,
17783        out_f: usize,
17784        canonical_chunk_rows: usize,
17785    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17786        if data.len() != in_f * out_f * 2 {
17787            return Err(format!("resident BF16 bytes {} != {out_f}x{in_f}x2", data.len()).into());
17788        }
17789        self.linear_bf16_chunked(x, data, m, in_f, out_f, false, Some(canonical_chunk_rows))
17790    }
17791
17792    /// Execute a load-time F32 mirror with the same fixed output-row chunks as the BF16 path.
17793    ///
17794    /// Expanding the checkpoint bytes once changes residency, not arithmetic: every cuBLASLt
17795    /// call receives the same F32 values and problem shape as `linear_bf16_chunked`.
17796    pub fn linear_f32_resident_canonical_rows(
17797        &self,
17798        x: &CudaSlice<f32>,
17799        data: &CudaSlice<f32>,
17800        m: usize,
17801        in_f: usize,
17802        out_f: usize,
17803        canonical_chunk_rows: usize,
17804    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17805        self.linear_f32_resident_canonical_rows_inner(
17806            x,
17807            data,
17808            m,
17809            in_f,
17810            out_f,
17811            canonical_chunk_rows,
17812            false,
17813        )
17814    }
17815
17816    /// Execute fixed output-row chunks and assemble them with one strided placement per chunk.
17817    ///
17818    /// The projection shapes and values are identical to
17819    /// [`Self::linear_f32_resident_canonical_rows`]. Only the byte-preserving output layout step
17820    /// changes, replacing one device copy per token with one placement kernel per output chunk.
17821    pub fn linear_f32_resident_canonical_rows_strided(
17822        &self,
17823        x: &CudaSlice<f32>,
17824        data: &CudaSlice<f32>,
17825        m: usize,
17826        in_f: usize,
17827        out_f: usize,
17828        canonical_chunk_rows: usize,
17829    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17830        self.linear_f32_resident_canonical_rows_inner(
17831            x,
17832            data,
17833            m,
17834            in_f,
17835            out_f,
17836            canonical_chunk_rows,
17837            true,
17838        )
17839    }
17840
17841    fn linear_f32_resident_canonical_rows_inner(
17842        &self,
17843        x: &CudaSlice<f32>,
17844        data: &CudaSlice<f32>,
17845        m: usize,
17846        in_f: usize,
17847        out_f: usize,
17848        canonical_chunk_rows: usize,
17849        strided_output: bool,
17850    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17851        if data.len() != in_f * out_f {
17852            return Err(format!("resident F32 values {} != {out_f}x{in_f}", data.len()).into());
17853        }
17854        if canonical_chunk_rows == 0
17855            || canonical_chunk_rows > out_f
17856            || out_f % canonical_chunk_rows != 0
17857        {
17858            return Err(format!(
17859                "invalid canonical F32 chunk rows {canonical_chunk_rows} for output width {out_f}"
17860            )
17861            .into());
17862        }
17863        if canonical_chunk_rows == out_f {
17864            return self.linear(x, data, m, in_f, out_f);
17865        }
17866
17867        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
17868        let input = x.slice(0..x.len());
17869        for r0 in (0..out_f).step_by(canonical_chunk_rows) {
17870            let weights = data.slice(r0 * in_f..(r0 + canonical_chunk_rows) * in_f);
17871            if m == 1 {
17872                let mut destination = y.slice_mut(r0..r0 + canonical_chunk_rows);
17873                self.linear_device_into(
17874                    &input,
17875                    &weights,
17876                    &mut destination,
17877                    1,
17878                    in_f,
17879                    canonical_chunk_rows,
17880                )?;
17881                continue;
17882            }
17883            let chunk = self.linear_device(&input, &weights, m, in_f, canonical_chunk_rows)?;
17884            if strided_output {
17885                self.place_rows_strided(&chunk, &mut y, canonical_chunk_rows, m, out_f, r0)?;
17886            } else {
17887                for token in 0..m {
17888                    let source = chunk
17889                        .slice(token * canonical_chunk_rows..(token + 1) * canonical_chunk_rows);
17890                    let mut destination =
17891                        y.slice_mut(token * out_f + r0..token * out_f + r0 + canonical_chunk_rows);
17892                    self.gpu.stream().memcpy_dtod(&source, &mut destination)?;
17893                }
17894            }
17895        }
17896        Ok(y)
17897    }
17898
17899    /// One-token twin of `linear_f32_resident_canonical_rows` writing into a caller-owned
17900    /// output. Same cuBLASLt calls, values, and chunk order as the allocating variant at
17901    /// `m == 1`; only the output residency changes (persistent workspace instead of a fresh
17902    /// allocation per call). This is the projection substrate of the v2 Step TP decode driver.
17903    pub fn linear_f32_resident_canonical_rows_t1_into(
17904        &self,
17905        x: &CudaSlice<f32>,
17906        data: &CudaSlice<f32>,
17907        y: &mut CudaSlice<f32>,
17908        in_f: usize,
17909        out_f: usize,
17910        canonical_chunk_rows: usize,
17911    ) -> Result<(), Box<dyn std::error::Error>> {
17912        if data.len() != in_f * out_f {
17913            return Err(format!("resident F32 values {} != {out_f}x{in_f}", data.len()).into());
17914        }
17915        if y.len() != out_f || x.len() != in_f {
17916            return Err(format!(
17917                "resident F32 t1 shapes x={} y={} != in {in_f} out {out_f}",
17918                x.len(),
17919                y.len()
17920            )
17921            .into());
17922        }
17923        if canonical_chunk_rows == 0
17924            || canonical_chunk_rows > out_f
17925            || out_f % canonical_chunk_rows != 0
17926        {
17927            return Err(format!(
17928                "invalid canonical F32 chunk rows {canonical_chunk_rows} for output width {out_f}"
17929            )
17930            .into());
17931        }
17932        let input = x.slice(0..x.len());
17933        for r0 in (0..out_f).step_by(canonical_chunk_rows) {
17934            let weights = data.slice(r0 * in_f..(r0 + canonical_chunk_rows) * in_f);
17935            let mut destination = y.slice_mut(r0..r0 + canonical_chunk_rows);
17936            self.linear_device_into(
17937                &input,
17938                &weights,
17939                &mut destination,
17940                1,
17941                in_f,
17942                canonical_chunk_rows,
17943            )?;
17944        }
17945        Ok(())
17946    }
17947
17948    /// One-token view-to-view linear into a caller-owned destination — the `linear` twin
17949    /// without the allocation, for workspace-resident operands.
17950    pub fn linear_t1_into(
17951        &self,
17952        x: &cudarc::driver::CudaView<'_, f32>,
17953        w: &cudarc::driver::CudaView<'_, f32>,
17954        y: &mut cudarc::driver::CudaViewMut<'_, f32>,
17955        in_f: usize,
17956        out_f: usize,
17957    ) -> Result<(), Box<dyn std::error::Error>> {
17958        self.linear_device_into(x, w, y, 1, in_f, out_f)
17959    }
17960
17961    /// On-device linear: y[m,out] = x[m,in] @ W[out,in]^T, weights row-major [out,in] (ggml).
17962    /// cuBLASLt col-major mapping (see memra_runtime::Gpu::linear_f32 for the derivation).
17963    /// DECODE-EXACT float linear: per-column m=1 cuBLASLt calls. cuBLASLt's reduction split is
17964    /// n-dependent (lt_ndep probe: m=1 vs m=2 col0 differs every bit), so spec-verify batches
17965    /// must not batch float matmuls the T=1 decode chain runs at m=1. Used by the small-t MoE
17966    /// router/shexp sites and matmul_decode_exact's Float arm.
17967    pub fn linear_decode_exact(
17968        &self,
17969        x: &CudaSlice<f32>,
17970        w: &CudaSlice<f32>,
17971        m_tokens: usize,
17972        in_f: usize,
17973        out_f: usize,
17974    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17975        if m_tokens == 1 {
17976            return self.linear(x, w, 1, in_f, out_f);
17977        }
17978        let xv = self.view(x, m_tokens * in_f);
17979        let mut y = self.alloc_uninit::<f32>(m_tokens * out_f)?;
17980        for t in 0..m_tokens {
17981            let row = xv.slice(t * in_f..(t + 1) * in_f);
17982            let mut xr = self.alloc_uninit::<f32>(in_f)?;
17983            self.copy_view_into(&mut xr, 0, &row, in_f)?;
17984            let yr = self.linear(&xr, w, 1, in_f, out_f)?;
17985            self.copy_into(&mut y, t * out_f, &yr, out_f)?;
17986        }
17987        Ok(y)
17988    }
17989
17990    pub fn linear(
17991        &self,
17992        x: &CudaSlice<f32>,
17993        w: &CudaSlice<f32>,
17994        m_tokens: usize,
17995        in_f: usize,
17996        out_f: usize,
17997    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17998        self.linear_device(x, w, m_tokens, in_f, out_f)
17999    }
18000
18001    fn linear_device<I>(
18002        &self,
18003        x: &I,
18004        w: &I,
18005        m_tokens: usize,
18006        in_f: usize,
18007        out_f: usize,
18008    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>>
18009    where
18010        I: cudarc::driver::DevicePtr<f32>,
18011    {
18012        let mut c = self.alloc_uninit::<f32>(m_tokens * out_f)?; // cuBLASLt beta=0: C fully written
18013        self.linear_device_into(x, w, &mut c, m_tokens, in_f, out_f)?;
18014        Ok(c)
18015    }
18016
18017    fn linear_device_into<I, O>(
18018        &self,
18019        x: &I,
18020        w: &I,
18021        c: &mut O,
18022        m_tokens: usize,
18023        in_f: usize,
18024        out_f: usize,
18025    ) -> Result<(), Box<dyn std::error::Error>>
18026    where
18027        I: cudarc::driver::DevicePtr<f32>,
18028        O: cudarc::driver::DevicePtrMut<f32>,
18029    {
18030        use cudarc::cublaslt::{Matmul, MatmulConfig};
18031        let cfg = MatmulConfig {
18032            transa: true,
18033            transb: false,
18034            transc: false,
18035            m: out_f as u64,
18036            n: m_tokens as u64,
18037            k: in_f as u64,
18038            alpha: 1.0,
18039            lda: in_f as i64,
18040            ldb: in_f as i64,
18041            beta: 0.0,
18042            ldc: out_f as i64,
18043            stride_a: None,
18044            stride_b: None,
18045            stride_c: None,
18046            stride_bias: None,
18047            batch_size: None,
18048        };
18049        let blas = self.gpu.blas();
18050        unsafe {
18051            blas.matmul(cfg, w, x, c, None, None)?;
18052        }
18053        Ok(())
18054    }
18055
18056    /// Naive SDPA. Q:[head_dim,n_head,T], K/V:[head_dim,n_head_kv,T_kv] -> O:[head_dim,n_head,T].
18057    pub fn sdpa_naive(
18058        &self,
18059        q: &CudaSlice<f32>,
18060        k: &CudaSlice<f32>,
18061        v: &CudaSlice<f32>,
18062        o: &mut CudaSlice<f32>,
18063        head_dim: usize,
18064        n_head: usize,
18065        n_head_kv: usize,
18066        t: usize,
18067        t_kv: usize,
18068        scale: f32,
18069        causal: bool,
18070    ) -> Result<(), Box<dyn std::error::Error>> {
18071        let f = self.func("sdpa_naive_f32");
18072        let cfg = LaunchConfig {
18073            grid_dim: (n_head as u32, t as u32, 1),
18074            block_dim: (128, 1, 1),
18075            shared_mem_bytes: (t_kv * 4) as u32,
18076        };
18077        let (hd, nh, nhkv, ti, tkvi, cz) = (
18078            head_dim as i32,
18079            n_head as i32,
18080            n_head_kv as i32,
18081            t as i32,
18082            t_kv as i32,
18083            causal as i32,
18084        );
18085        let __s_b = self.gpu.stream();
18086        let mut b = __s_b.launch_builder(&f);
18087        b.arg(q)
18088            .arg(k)
18089            .arg(v)
18090            .arg(o)
18091            .arg(&hd)
18092            .arg(&nh)
18093            .arg(&nhkv)
18094            .arg(&ti)
18095            .arg(&tkvi)
18096            .arg(&scale)
18097            .arg(&cz);
18098        unsafe {
18099            b.launch(cfg)?;
18100        }
18101        Ok(())
18102    }
18103
18104    /// Island twin (lane/gemma-vision): causal + sliding-window attention with
18105    /// bidirectional image islands. `span_id` labels each absolute kv position
18106    /// (-1 text, >=0 island id); same-island keys are visible unconditionally,
18107    /// reproducing the reference's non-causal image batch. window 0 = no window.
18108    #[allow(clippy::too_many_arguments)]
18109    pub fn sdpa_naive_island(
18110        &self,
18111        q: &CudaSlice<f32>,
18112        k: &CudaSlice<f32>,
18113        v: &CudaSlice<f32>,
18114        o: &mut CudaSlice<f32>,
18115        span_id: &CudaSlice<i32>,
18116        head_dim: usize,
18117        n_head: usize,
18118        n_head_kv: usize,
18119        t: usize,
18120        t_kv: usize,
18121        scale: f32,
18122        window: usize,
18123    ) -> Result<(), Box<dyn std::error::Error>> {
18124        let f = self.func("sdpa_naive_island_f32");
18125        let cfg = LaunchConfig {
18126            grid_dim: (n_head as u32, t as u32, 1),
18127            block_dim: (128, 1, 1),
18128            shared_mem_bytes: (t_kv * 4) as u32,
18129        };
18130        let (hd, nh, nhkv, ti, tkvi, wi) = (
18131            head_dim as i32,
18132            n_head as i32,
18133            n_head_kv as i32,
18134            t as i32,
18135            t_kv as i32,
18136            window as i32,
18137        );
18138        let __s_b = self.gpu.stream();
18139        let mut b = __s_b.launch_builder(&f);
18140        b.arg(q)
18141            .arg(k)
18142            .arg(v)
18143            .arg(o)
18144            .arg(span_id)
18145            .arg(&hd)
18146            .arg(&nh)
18147            .arg(&nhkv)
18148            .arg(&ti)
18149            .arg(&tkvi)
18150            .arg(&scale)
18151            .arg(&wi);
18152        unsafe {
18153            b.launch(cfg)?;
18154        }
18155        Ok(())
18156    }
18157
18158    /// Windowed sdpa_naive twin (gemma4 R6): masks keys older than q_pos-(window-1).
18159    #[allow(clippy::too_many_arguments)]
18160    pub fn sdpa_naive_w(
18161        &self,
18162        q: &CudaSlice<f32>,
18163        k: &CudaSlice<f32>,
18164        v: &CudaSlice<f32>,
18165        o: &mut CudaSlice<f32>,
18166        head_dim: usize,
18167        n_head: usize,
18168        n_head_kv: usize,
18169        t: usize,
18170        t_kv: usize,
18171        scale: f32,
18172        causal: bool,
18173        window: usize,
18174    ) -> Result<(), Box<dyn std::error::Error>> {
18175        let f = self.func("sdpa_naive_w_f32");
18176        let cfg = LaunchConfig {
18177            grid_dim: (n_head as u32, t as u32, 1),
18178            block_dim: (128, 1, 1),
18179            shared_mem_bytes: (t_kv * 4) as u32,
18180        };
18181        let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
18182            head_dim as i32,
18183            n_head as i32,
18184            n_head_kv as i32,
18185            t as i32,
18186            t_kv as i32,
18187            causal as i32,
18188            window as i32,
18189        );
18190        let __s_b = self.gpu.stream();
18191        let mut b = __s_b.launch_builder(&f);
18192        b.arg(q)
18193            .arg(k)
18194            .arg(v)
18195            .arg(o)
18196            .arg(&hd)
18197            .arg(&nh)
18198            .arg(&nhkv)
18199            .arg(&ti)
18200            .arg(&tkvi)
18201            .arg(&scale)
18202            .arg(&cz)
18203            .arg(&wi);
18204        unsafe {
18205            b.launch(cfg)?;
18206        }
18207        Ok(())
18208    }
18209
18210    /// Lo-clipped windowed sdpa_naive twin (lane/dflash2-longctx, DFLASH2-EVAL §10.6(c)).
18211    /// Same mask law as `sdpa_naive_w`, but keys below every query's window floor are never
18212    /// read: kv_lo = max(0, (t_kv - t) + 1 - window) — the oldest key visible to the OLDEST
18213    /// query row (q_pos = t_kv - t). Dynamic shared memory shrinks from t_kv*4 bytes (which
18214    /// blows the 48KB launch bound at ~12k rows — the B2 ctx crash) to (t_kv - kv_lo)*4 =
18215    /// (window - 1 + t)*4, and the key scan drops from O(t_kv) to O(window + t). Output is
18216    /// byte-identical to `sdpa_naive_w` (masked keys contribute exact zeros to same-order
18217    /// reductions; kernel_check `sdpa_naive_w_lo` pins it). window == 0 (no window) keeps
18218    /// kv_lo = 0 and is then shape-identical to the legacy kernel, including its bound.
18219    #[allow(clippy::too_many_arguments)]
18220    pub fn sdpa_naive_w_lo(
18221        &self,
18222        q: &CudaSlice<f32>,
18223        k: &CudaSlice<f32>,
18224        v: &CudaSlice<f32>,
18225        o: &mut CudaSlice<f32>,
18226        head_dim: usize,
18227        n_head: usize,
18228        n_head_kv: usize,
18229        t: usize,
18230        t_kv: usize,
18231        scale: f32,
18232        causal: bool,
18233        window: usize,
18234    ) -> Result<(), Box<dyn std::error::Error>> {
18235        let kv_lo = if window > 0 {
18236            (t_kv - t + 1).saturating_sub(window)
18237        } else {
18238            0
18239        };
18240        let smem = (t_kv - kv_lo) * 4;
18241        if smem > 48 * 1024 {
18242            return Err(format!(
18243                "sdpa_naive_w_lo: window {window} + T {t} rows need {smem} bytes of dynamic \
18244                 shared memory (> 48KB launch bound) — this kernel clips the OLD side only; \
18245                 a window this wide needs the multi-pass long-ctx kernel"
18246            )
18247            .into());
18248        }
18249        let f = self.func("sdpa_naive_w_lo_f32");
18250        let cfg = LaunchConfig {
18251            grid_dim: (n_head as u32, t as u32, 1),
18252            block_dim: (128, 1, 1),
18253            shared_mem_bytes: smem as u32,
18254        };
18255        let (hd, nh, nhkv, ti, tkvi, cz, wi, lo) = (
18256            head_dim as i32,
18257            n_head as i32,
18258            n_head_kv as i32,
18259            t as i32,
18260            t_kv as i32,
18261            causal as i32,
18262            window as i32,
18263            kv_lo as i32,
18264        );
18265        let __s_b = self.gpu.stream();
18266        let mut b = __s_b.launch_builder(&f);
18267        b.arg(q)
18268            .arg(k)
18269            .arg(v)
18270            .arg(o)
18271            .arg(&hd)
18272            .arg(&nh)
18273            .arg(&nhkv)
18274            .arg(&ti)
18275            .arg(&tkvi)
18276            .arg(&scale)
18277            .arg(&cz)
18278            .arg(&wi)
18279            .arg(&lo);
18280        unsafe {
18281            b.launch(cfg)?;
18282        }
18283        Ok(())
18284    }
18285
18286    /// SDPA where K/V are CudaViews into a resident KV cache (decode hot path, no host round-trip).
18287    pub fn sdpa_naive_view(
18288        &self,
18289        q: &CudaSlice<f32>,
18290        k: &cudarc::driver::CudaView<f32>,
18291        v: &cudarc::driver::CudaView<f32>,
18292        o: &mut CudaSlice<f32>,
18293        head_dim: usize,
18294        n_head: usize,
18295        n_head_kv: usize,
18296        t: usize,
18297        t_kv: usize,
18298        scale: f32,
18299        causal: bool,
18300    ) -> Result<(), Box<dyn std::error::Error>> {
18301        let f = self.func("sdpa_naive_f32");
18302        let cfg = LaunchConfig {
18303            grid_dim: (n_head as u32, t as u32, 1),
18304            block_dim: (128, 1, 1),
18305            shared_mem_bytes: (t_kv * 4) as u32,
18306        };
18307        let (hd, nh, nhkv, ti, tkvi, cz) = (
18308            head_dim as i32,
18309            n_head as i32,
18310            n_head_kv as i32,
18311            t as i32,
18312            t_kv as i32,
18313            causal as i32,
18314        );
18315        let __s_b = self.gpu.stream();
18316        let mut b = __s_b.launch_builder(&f);
18317        b.arg(q)
18318            .arg(k)
18319            .arg(v)
18320            .arg(o)
18321            .arg(&hd)
18322            .arg(&nh)
18323            .arg(&nhkv)
18324            .arg(&ti)
18325            .arg(&tkvi)
18326            .arg(&scale)
18327            .arg(&cz);
18328        unsafe {
18329            b.launch(cfg)?;
18330        }
18331        Ok(())
18332    }
18333
18334    /// Correctness fallback for quantized resident K/V views. Dequantizes K and V once into f32
18335    /// workspaces, then calls `sdpa_naive`. This is an explicit API: the optimized prefill view
18336    /// dispatch remains unchanged, so callers can use it as a reference or compatibility path.
18337    /// Dequant a quantized KV view into caller-owned f32 buffers (one grid-stride launch).
18338    /// `g` picks the kf8vf8-module stamp for e4m3 caches (same flag contract as fa_decode/
18339    /// fa_prefill_view). Used by the E4B shared-KV prefill arms (2026-07-31) to feed the
18340    /// f32 fa_prefill_w / fa_prefill_hd512 twins from the target layer's quantized rows.
18341    #[allow(clippy::too_many_arguments)]
18342    pub fn fa_dequant_kv_view_f32(
18343        &self,
18344        k: &cudarc::driver::CudaView<u8>,
18345        v: &cudarc::driver::CudaView<u8>,
18346        kf: &mut CudaSlice<f32>,
18347        vf: &mut CudaSlice<f32>,
18348        kv_dim_k: usize,
18349        kv_dim_v: usize,
18350        t_kv: usize,
18351        k_tok_bytes: usize,
18352        v_tok_bytes: usize,
18353        g: bool,
18354    ) -> Result<(), Box<dyn std::error::Error>> {
18355        let f = if g {
18356            self.func_g("fa_dequant_kv_ws_f32")
18357        } else {
18358            self.func("fa_dequant_kv_ws_f32")
18359        };
18360        let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
18361        let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
18362        let cfg = LaunchConfig {
18363            grid_dim: (nblk.max(1), 1, 1),
18364            block_dim: (256, 1, 1),
18365            shared_mem_bytes: 0,
18366        };
18367        let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
18368        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
18369        let __s_b = self.gpu.stream();
18370        let mut b = __s_b.launch_builder(&f);
18371        b.arg(k)
18372            .arg(v)
18373            .arg(&mut *kf)
18374            .arg(&mut *vf)
18375            .arg(&kdk)
18376            .arg(&kdv)
18377            .arg(&tkvi)
18378            .arg(&ktb)
18379            .arg(&vtb);
18380        unsafe {
18381            b.launch(cfg)?;
18382        }
18383        Ok(())
18384    }
18385
18386    #[allow(clippy::too_many_arguments)]
18387    pub fn sdpa_naive_quantized_view(
18388        &self,
18389        q: &CudaSlice<f32>,
18390        k: &cudarc::driver::CudaView<u8>,
18391        v: &cudarc::driver::CudaView<u8>,
18392        o: &mut CudaSlice<f32>,
18393        head_dim: usize,
18394        n_head: usize,
18395        n_head_kv: usize,
18396        t: usize,
18397        t_kv: usize,
18398        scale: f32,
18399        causal: bool,
18400        k_tok_bytes: usize,
18401        v_tok_bytes: usize,
18402    ) -> Result<(), Box<dyn std::error::Error>> {
18403        let kv_dim = n_head_kv * head_dim;
18404        let mut kf = self.uninit(t_kv * kv_dim)?;
18405        let mut vf = self.uninit(t_kv * kv_dim)?;
18406        let f = self.func("fa_dequant_kv_ws_f32");
18407        let total = (2 * t_kv * kv_dim) as u64;
18408        let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
18409        let cfg = LaunchConfig {
18410            grid_dim: (nblk.max(1), 1, 1),
18411            block_dim: (256, 1, 1),
18412            shared_mem_bytes: 0,
18413        };
18414        let (kv_dim_i, t_kv_i) = (kv_dim as i32, t_kv as i32);
18415        let (k_tok_bytes_i, v_tok_bytes_i) = (k_tok_bytes as i64, v_tok_bytes as i64);
18416        let __s_b = self.gpu.stream();
18417        let mut b = __s_b.launch_builder(&f);
18418        b.arg(k)
18419            .arg(v)
18420            .arg(&mut kf)
18421            .arg(&mut vf)
18422            .arg(&kv_dim_i)
18423            .arg(&kv_dim_i)
18424            .arg(&t_kv_i)
18425            .arg(&k_tok_bytes_i)
18426            .arg(&v_tok_bytes_i);
18427        unsafe { b.launch(cfg)? };
18428        self.sdpa_naive(
18429            q, &kf, &vf, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
18430        )
18431    }
18432
18433    /// WINDOWED twin of `sdpa_naive_quantized_view` (step35 SWA prefill): dequant the KV byte
18434    /// view into f32 workspaces with the SAME `fa_dequant_kv_ws_f32` launch, then run
18435    /// `sdpa_naive_w` instead of `sdpa_naive`. `window == 0` is the unwindowed form (the kernel
18436    /// treats a non-positive window as "no window mask"), so this is a strict superset of the
18437    /// unwindowed function above and produces bit-identical output at window == 0.
18438    ///
18439    /// Why this exists: EVERY windowed FlashAttention stamp in flash_attn.cu is head_dim-256
18440    /// only (`fa_prefill_w_f32` == `fa_prefill_f32_body<256>`, and the quantized-view windowed
18441    /// twins likewise), while step35 is head_dim 128. Its SWA layers therefore have no windowed
18442    /// FA path and take this f32 floor in v0 — same cache bytes, same numeric class as the
18443    /// unwindowed quantized-view fallback, so the chunk-invariance contract holds on both.
18444    #[allow(clippy::too_many_arguments)]
18445    pub fn sdpa_naive_w_quantized_view(
18446        &self,
18447        q: &CudaSlice<f32>,
18448        k: &cudarc::driver::CudaView<u8>,
18449        v: &cudarc::driver::CudaView<u8>,
18450        o: &mut CudaSlice<f32>,
18451        head_dim: usize,
18452        n_head: usize,
18453        n_head_kv: usize,
18454        t: usize,
18455        t_kv: usize,
18456        scale: f32,
18457        causal: bool,
18458        window: usize,
18459        k_tok_bytes: usize,
18460        v_tok_bytes: usize,
18461    ) -> Result<(), Box<dyn std::error::Error>> {
18462        let kv_dim = n_head_kv * head_dim;
18463        let mut kf = self.uninit(t_kv * kv_dim)?;
18464        let mut vf = self.uninit(t_kv * kv_dim)?;
18465        let f = self.func("fa_dequant_kv_ws_f32");
18466        let total = (2 * t_kv * kv_dim) as u64;
18467        let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
18468        let cfg = LaunchConfig {
18469            grid_dim: (nblk.max(1), 1, 1),
18470            block_dim: (256, 1, 1),
18471            shared_mem_bytes: 0,
18472        };
18473        let (kv_dim_i, t_kv_i) = (kv_dim as i32, t_kv as i32);
18474        let (k_tok_bytes_i, v_tok_bytes_i) = (k_tok_bytes as i64, v_tok_bytes as i64);
18475        let __s_b = self.gpu.stream();
18476        let mut b = __s_b.launch_builder(&f);
18477        b.arg(k)
18478            .arg(v)
18479            .arg(&mut kf)
18480            .arg(&mut vf)
18481            .arg(&kv_dim_i)
18482            .arg(&kv_dim_i)
18483            .arg(&t_kv_i)
18484            .arg(&k_tok_bytes_i)
18485            .arg(&v_tok_bytes_i);
18486        unsafe { b.launch(cfg)? };
18487        self.sdpa_naive_w(
18488            q, &kf, &vf, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal, window,
18489        )
18490    }
18491
18492    /// Hand-written FlashAttention prefill (sm_120, FA-2 online softmax on validated mma.sync,
18493    /// head_dim 256 or 128 (template-stamped twins), GQA, causal). Replaces sdpa_naive for T>1.
18494    /// Q/K/V/O [head_dim, n_head(_kv), T].
18495    pub fn fa_prefill(
18496        &self,
18497        q: &CudaSlice<f32>,
18498        k: &CudaSlice<f32>,
18499        v: &CudaSlice<f32>,
18500        o: &mut CudaSlice<f32>,
18501        head_dim: usize,
18502        n_head: usize,
18503        n_head_kv: usize,
18504        t: usize,
18505        t_kv: usize,
18506        scale: f32,
18507        causal: bool,
18508    ) -> Result<(), Box<dyn std::error::Error>> {
18509        if portable_mma_gated() {
18510            return self.sdpa_naive(
18511                q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
18512            );
18513        }
18514        // FA3 v10 arm (task #20, OPT-IN MEMRA_FA3=1 — harness-proven 883us vs the shipped
18515        // kernel's 993us at T=2048): TMA-swizzled wgmma FA, fresh causal hd256 only.
18516        // NEW NUMERIC CONFIG (GDN-mma precedent): online softmax / bf16-P class — the
18517        // run-gen argmax + greedy-stream batteries arbitrate; not bit-paired.
18518        // PROMOTED default-ON hopper (2026-07-27): 3-seed 2048-prime -> 128-decode
18519        // streams MATCH vs mma, full battery green, lane interleaved 5/5 (+2.4%).
18520        // MEMRA_FA3=0 reverts; kernel-check pins the mma config regardless.
18521        let fa3_on = head_dim == 256
18522            && causal
18523            && t == t_kv
18524            && match std::env::var("MEMRA_FA3").as_deref() {
18525                Ok("0") => false,
18526                Ok("1") => true,
18527                _ => cfg!(memra_hopper_mma),
18528            };
18529        if fa3_on {
18530            let n = t * n_head * head_dim;
18531            let nkv = t * n_head_kv * head_dim;
18532            let mut q16 = self.alloc_u8_uninit(n * 2)?;
18533            let mut k16 = self.alloc_u8_uninit(nkv * 2)?;
18534            let mut v16 = self.alloc_u8_uninit(nkv * 2)?;
18535            self.f32_to_bf16_into(q, &mut q16, n)?;
18536            self.f32_to_bf16_into(k, &mut k16, nkv)?;
18537            self.f32_to_bf16_into(v, &mut v16, nkv)?;
18538            let rc = {
18539                use cudarc::driver::{DevicePtr, DevicePtrMut};
18540                let stream = self.gpu.stream();
18541                let (qp, _g1) = q16.device_ptr(&stream);
18542                let (kp, _g2) = k16.device_ptr(&stream);
18543                let (vp, _g3) = v16.device_ptr(&stream);
18544                let (op, _g4) = o.device_ptr_mut(&stream);
18545                unsafe {
18546                    memra_fa3_prefill(
18547                        qp as *const core::ffi::c_void,
18548                        kp as *const core::ffi::c_void,
18549                        vp as *const core::ffi::c_void,
18550                        op as *mut f32,
18551                        t as i32,
18552                        n_head as i32,
18553                        n_head_kv as i32,
18554                        head_dim as i32,
18555                        scale,
18556                        stream.cu_stream() as *mut core::ffi::c_void,
18557                    )
18558                }
18559            };
18560            if rc != 0 {
18561                return Err(format!("memra_fa3_prefill rc={rc}").into());
18562            }
18563            return Ok(());
18564        }
18565        // FLOOR PORT (P2+P0a+P0b+P1): 4 warps/CTA, BLOCK_Q=64 query rows, BK=32 KV tile,
18566        // Q-in-reg + register-O, grid.y=n_head_kv (4 Q-heads share staged K/V).
18567        // P1 plain arm (MEMRA_FA_P1=1 opt-in until the qwen battery): the engine-study body
18568        // (FA2 schedule + boundary split + swizzle) on the non-windowed lane. bf16 pre-convert.
18569        static FA_P1: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
18570        let fa_p1 = *FA_P1.get_or_init(|| std::env::var("MEMRA_FA_P1").as_deref() == Ok("1"));
18571        if fa_p1 && head_dim == 256 && !std::env::var("MEMRA_FA_FLOOR").is_ok() {
18572            const BLOCK_Q: usize = 64;
18573            const BKX: usize = 32;
18574            let f = self.func("fa_prefill_bf16_p1");
18575            let shmem = (2 * (2 * BKX * head_dim + BLOCK_Q * BKX)
18576                + 4 * (BLOCK_Q * BKX + 2 * BLOCK_Q)) as u32;
18577            use cudarc::driver::sys::CUfunction_attribute_enum as A;
18578            f.set_attribute(
18579                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18580                shmem as i32,
18581            )?;
18582            let cfg = LaunchConfig {
18583                grid_dim: (
18584                    (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
18585                    n_head as u32,
18586                    1,
18587                ),
18588                block_dim: (32, 4, 1),
18589                shared_mem_bytes: shmem,
18590            };
18591            let (hd, nh, nhkv, ti, tkvi, cz) = (
18592                head_dim as i32,
18593                n_head as i32,
18594                n_head_kv as i32,
18595                t as i32,
18596                t_kv as i32,
18597                causal as i32,
18598            );
18599            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
18600            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
18601            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
18602            let __s_b = self.gpu.stream();
18603            let mut b = __s_b.launch_builder(&f);
18604            b.arg(&qb)
18605                .arg(&kb)
18606                .arg(&vb)
18607                .arg(o)
18608                .arg(&hd)
18609                .arg(&nh)
18610                .arg(&nhkv)
18611                .arg(&ti)
18612                .arg(&tkvi)
18613                .arg(&scale)
18614                .arg(&cz);
18615            unsafe {
18616                b.launch(cfg)?;
18617            }
18618            return Ok(());
18619        }
18620        // Edge 5a (DEFAULT): fa_prefill_f32_pp — register-resident softmax (no sSw smem
18621        // round-trip), the FA3 softmax-GEMM overlap variant. ncu (pp512): short_scoreboard
18622        // 4.32->3.47, wait 1.99->1.45, per-call ~577us->~440us (1.31x) at flat 12.1% warps /
18623        // 255 regs / 2 CTAs (occupancy preserved). Bit-safe: 9B+27B argmax MATCH, rel 2.55e-3
18624        // vs floor 3.03e-3. MEMRA_FA_FLOOR reverts to the serialized-softmax floor kernel.
18625        const BK: usize = 32;
18626        // W2 lane (MEMRA_FA_PP_W2=1, ncu 2026-07-26): 2-warp/32-row CTA tile doubles grid.x —
18627        // bit-identical per-row math, pure coverage trade for the 6.25%-occupancy starvation.
18628        let w2 = std::env::var("MEMRA_FA_PP_W2").as_deref() == Ok("1");
18629        let (block_q, warps, w2_sfx): (usize, u32, &str) =
18630            if w2 { (32, 2, "_w2") } else { (64, 4, "") };
18631        // hd128 twins (2026-07-07): the prefill kernels are template-stamped at 256 (original
18632        // names, dispatch unchanged) and 128 (`_hd128`, the MiniMax-M3 class). Callers gate
18633        // other head_dims to sdpa_naive before reaching here.
18634        let hd_sfx = fa_hd_suffix(head_dim)?;
18635        let floor = std::env::var("MEMRA_FA_FLOOR").is_ok();
18636        // BF16-KV staging lane (2026-07-26, default ON): the kernel converts K/V to bf16
18637        // during staging anyway — pre-converting to bf16 mirrors is BIT-IDENTICAL (same
18638        // __float2bfloat16 values into the same mma) and turns the 67%-of-stalls scalar
18639        // staging into int4 vector copies. MEMRA_FA_BF16KV=0 reverts.
18640        let bf16kv = !floor && !w2 && std::env::var("MEMRA_FA_BF16KV").as_deref() != Ok("0");
18641        let (kb16, vb16) = if bf16kv {
18642            let n = t_kv * n_head_kv * head_dim;
18643            let mut kb = self.alloc_u8_uninit(n * 2)?;
18644            let mut vb = self.alloc_u8_uninit(n * 2)?;
18645            let fcv = self.func("f32_to_bf16_bulk");
18646            let ni = n as i64;
18647            let cfgc = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
18648            let __s_b = self.gpu.stream();
18649            let mut b = __s_b.launch_builder(&fcv);
18650            b.arg(k).arg(&mut kb).arg(&ni);
18651            unsafe {
18652                b.launch(cfgc)?;
18653            }
18654            let __s_b = self.gpu.stream();
18655            let mut b = __s_b.launch_builder(&fcv);
18656            b.arg(v).arg(&mut vb).arg(&ni);
18657            unsafe {
18658                b.launch(cfgc)?;
18659            }
18660            (Some(kb), Some(vb))
18661        } else {
18662            (None, None)
18663        };
18664        let f = self.func(&if bf16kv {
18665            format!("fa_prefill_bf16kv_pp{hd_sfx}")
18666        } else {
18667            format!(
18668                "fa_prefill_f32{}{}{hd_sfx}",
18669                if floor { "" } else { "_pp" },
18670                if floor { "" } else { w2_sfx }
18671            )
18672        });
18673        // persistent smem: bf16*(KV_STAGES*(sK + sV) + sP) + f32*(sS + sM + sL);
18674        // the bf16kv ring doubles the K/V stages (KV_STAGES=2).
18675        let kv_stages = if bf16kv { 2 } else { 1 };
18676        let shmem = (2 * (kv_stages * 2 * BK * head_dim + block_q * BK)
18677            + 4 * (block_q * BK + 2 * block_q)) as u32;
18678        use cudarc::driver::sys::CUfunction_attribute_enum as A;
18679        f.set_attribute(
18680            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18681            shmem as i32,
18682        )?;
18683        let cfg = LaunchConfig {
18684            grid_dim: (
18685                (t as u32 + block_q as u32 - 1) / block_q as u32,
18686                n_head as u32,
18687                1,
18688            ),
18689            block_dim: (32, warps, 1),
18690            shared_mem_bytes: shmem,
18691        };
18692        let (hd, nh, nhkv, ti, tkvi, cz) = (
18693            head_dim as i32,
18694            n_head as i32,
18695            n_head_kv as i32,
18696            t as i32,
18697            t_kv as i32,
18698            causal as i32,
18699        );
18700        let __s_b = self.gpu.stream();
18701        let mut b = __s_b.launch_builder(&f);
18702        b.arg(q);
18703        match (&kb16, &vb16) {
18704            (Some(kb), Some(vb)) => {
18705                b.arg(kb).arg(vb);
18706            }
18707            _ => {
18708                b.arg(k).arg(v);
18709            }
18710        }
18711        b.arg(o)
18712            .arg(&hd)
18713            .arg(&nh)
18714            .arg(&nhkv)
18715            .arg(&ti)
18716            .arg(&tkvi)
18717            .arg(&scale)
18718            .arg(&cz);
18719        unsafe {
18720            b.launch(cfg)?;
18721        }
18722        Ok(())
18723    }
18724
18725    /// Windowed FA prefill (gemma4 SWA layers past the sliding window, hd256): fa_prefill's
18726    /// exact dispatch (pp default, MEMRA_FA_FLOOR seam) with the sliding-window mask + tile
18727    /// skip in-kernel. Replaces the O(T*T_kv) scalar sdpa_naive_w on the prime path.
18728    #[allow(clippy::too_many_arguments)]
18729    pub fn fa_prefill_w(
18730        &self,
18731        q: &CudaSlice<f32>,
18732        k: &CudaSlice<f32>,
18733        v: &CudaSlice<f32>,
18734        o: &mut CudaSlice<f32>,
18735        head_dim: usize,
18736        n_head: usize,
18737        n_head_kv: usize,
18738        t: usize,
18739        t_kv: usize,
18740        scale: f32,
18741        causal: bool,
18742        window: usize,
18743    ) -> Result<(), Box<dyn std::error::Error>> {
18744        // sm_90a rides the mma twins (portable_mma_gated, 2026-07-31 — the raw
18745        // portable_cuda gate was stale-conservative on Hopper; fa_prefill already flipped).
18746        if portable_mma_gated() {
18747            return self.sdpa_naive_w(
18748                q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal, window,
18749            );
18750        }
18751        // Default: bf16-prestaged twin (same treatment as hd512 — Q/K/V pre-converted once,
18752        // int4 stage copies; bit-identical, kernel_check-gated). MEMRA_FAW_STAGE=f32 reverts;
18753        // MEMRA_FA_FLOOR keeps the f32 floor stamp untouched.
18754        static FAW_F32: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
18755        let faw_f32 =
18756            *FAW_F32.get_or_init(|| std::env::var("MEMRA_FAW_STAGE").as_deref() == Ok("f32"));
18757        let floor = std::env::var("MEMRA_FA_FLOOR").is_ok();
18758        self.fa_prefill_w_arm(
18759            q,
18760            k,
18761            v,
18762            o,
18763            head_dim,
18764            n_head,
18765            n_head_kv,
18766            t,
18767            t_kv,
18768            scale,
18769            causal,
18770            window,
18771            floor || faw_f32,
18772            floor,
18773        )
18774    }
18775
18776    /// Windowed FA prefill with PRE-CONVERTED bf16 operands (producer-emitted; 31B glue lane).
18777    /// Launches the P1 stamp directly — callers guarantee qb/kb/vb hold the exact bf16 of q/k/v.
18778    #[allow(clippy::too_many_arguments)]
18779    pub fn fa_prefill_w_pre(
18780        &self,
18781        qb: &CudaSlice<u8>,
18782        kb: &CudaSlice<u8>,
18783        vb: &CudaSlice<u8>,
18784        o: &mut CudaSlice<f32>,
18785        head_dim: usize,
18786        n_head: usize,
18787        n_head_kv: usize,
18788        t: usize,
18789        t_kv: usize,
18790        scale: f32,
18791        causal: bool,
18792        window: usize,
18793        v_f16: bool,
18794    ) -> Result<(), Box<dyn std::error::Error>> {
18795        const BLOCK_Q: usize = 64;
18796        const BK: usize = 32;
18797        debug_assert_eq!(head_dim, 256);
18798        let hp = fa_f16pv_on() && faw_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
18799        debug_assert!(!v_f16 || hp, "f16 V emitted but the SWA hp arm is off");
18800        if hp {
18801            const BLOCK_QH: usize = 32;
18802            // V bytes must be f16 for the h2 stamp; producer normally emits f16 (v_f16),
18803            // else re-encode through the pooled scratch (stream-ordered reuse).
18804            let mut vguard = self.fa_vf16_scratch.lock().unwrap();
18805            let vh: &CudaSlice<u8> = if v_f16 {
18806                vb
18807            } else {
18808                let n = t_kv * n_head_kv * head_dim;
18809                if vguard.as_ref().map(|b| b.len() < n * 2).unwrap_or(true) {
18810                    *vguard = Some(self.alloc_uninit::<u8>(n * 2)?);
18811                }
18812                self.bf16_to_f16_into(vb, n, vguard.as_mut().unwrap())?;
18813                vguard.as_ref().unwrap()
18814            };
18815            let f = self.func("fa_prefill_w_bf16_p1h2");
18816            let shmem = (2 * (2 * BK * head_dim + 2 * BLOCK_QH * BK) + 4 * (2 * BLOCK_QH)) as u32;
18817            use cudarc::driver::sys::CUfunction_attribute_enum as A;
18818            f.set_attribute(
18819                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18820                shmem as i32,
18821            )?;
18822            let cfg = LaunchConfig {
18823                grid_dim: ((t as u32).div_ceil(BLOCK_QH as u32), (n_head / 2) as u32, 1),
18824                block_dim: (32, 4, 1),
18825                shared_mem_bytes: shmem,
18826            };
18827            let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
18828                head_dim as i32,
18829                n_head as i32,
18830                n_head_kv as i32,
18831                t as i32,
18832                t_kv as i32,
18833                causal as i32,
18834                window as i32,
18835            );
18836            let __s_b = self.gpu.stream();
18837            let mut b = __s_b.launch_builder(&f);
18838            b.arg(qb)
18839                .arg(kb)
18840                .arg(vh)
18841                .arg(o)
18842                .arg(&hd)
18843                .arg(&nh)
18844                .arg(&nhkv)
18845                .arg(&ti)
18846                .arg(&tkvi)
18847                .arg(&scale)
18848                .arg(&cz)
18849                .arg(&wi);
18850            unsafe {
18851                b.launch(cfg)?;
18852            }
18853            return Ok(());
18854        }
18855        let f = self.func("fa_prefill_w_bf16_p1");
18856        let shmem =
18857            (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
18858        use cudarc::driver::sys::CUfunction_attribute_enum as A;
18859        f.set_attribute(
18860            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18861            shmem as i32,
18862        )?;
18863        let cfg = LaunchConfig {
18864            grid_dim: (
18865                (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
18866                n_head as u32,
18867                1,
18868            ),
18869            block_dim: (32, 4, 1),
18870            shared_mem_bytes: shmem,
18871        };
18872        let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
18873            head_dim as i32,
18874            n_head as i32,
18875            n_head_kv as i32,
18876            t as i32,
18877            t_kv as i32,
18878            causal as i32,
18879            window as i32,
18880        );
18881        let __s_b = self.gpu.stream();
18882        let mut b = __s_b.launch_builder(&f);
18883        b.arg(qb)
18884            .arg(kb)
18885            .arg(vb)
18886            .arg(o)
18887            .arg(&hd)
18888            .arg(&nh)
18889            .arg(&nhkv)
18890            .arg(&ti)
18891            .arg(&tkvi)
18892            .arg(&scale)
18893            .arg(&cz)
18894            .arg(&wi);
18895        unsafe {
18896            b.launch(cfg)?;
18897        }
18898        Ok(())
18899    }
18900
18901    /// Windowed FA prefill with the stage arm FORCED — the kernel_check bit-identity entry.
18902    #[allow(clippy::too_many_arguments)]
18903    pub fn fa_prefill_w_arm(
18904        &self,
18905        q: &CudaSlice<f32>,
18906        k: &CudaSlice<f32>,
18907        v: &CudaSlice<f32>,
18908        o: &mut CudaSlice<f32>,
18909        head_dim: usize,
18910        n_head: usize,
18911        n_head_kv: usize,
18912        t: usize,
18913        t_kv: usize,
18914        scale: f32,
18915        causal: bool,
18916        window: usize,
18917        f32_stage: bool,
18918        floor: bool,
18919    ) -> Result<(), Box<dyn std::error::Error>> {
18920        const BLOCK_Q: usize = 64;
18921        const BK: usize = 32;
18922        debug_assert_eq!(head_dim, 256, "fa_prefill_w is stamped hd256 only");
18923        // P1 (2026-07-22 engine study): per-head Br=64 stamp with the FA2 schedule (V-copy
18924        // over GEMM0, next-K over softmax+GEMM1) + boundary/interior mask split. FP order
18925        // preserved -> bit-identical (gated). MEMRA_FAW_P1=0 reverts to the g4/o2 arms.
18926        static P1_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
18927        let p1 = !floor
18928            && !f32_stage
18929            && *P1_ON.get_or_init(|| {
18930                std::env::var("MEMRA_FAW_P1")
18931                    .map(|v| v != "0")
18932                    .unwrap_or(true)
18933            });
18934        let hp =
18935            p1 && fa_f16pv_on() && faw_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
18936        if hp {
18937            const BLOCK_QH: usize = 32;
18938            let f = self.func("fa_prefill_w_bf16_p1h2");
18939            let shmem = (2 * (2 * BK * head_dim + 2 * BLOCK_QH * BK) + 4 * (2 * BLOCK_QH)) as u32;
18940            use cudarc::driver::sys::CUfunction_attribute_enum as A;
18941            f.set_attribute(
18942                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18943                shmem as i32,
18944            )?;
18945            let cfg = LaunchConfig {
18946                grid_dim: ((t as u32).div_ceil(BLOCK_QH as u32), (n_head / 2) as u32, 1),
18947                block_dim: (32, 4, 1),
18948                shared_mem_bytes: shmem,
18949            };
18950            let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
18951                head_dim as i32,
18952                n_head as i32,
18953                n_head_kv as i32,
18954                t as i32,
18955                t_kv as i32,
18956                causal as i32,
18957                window as i32,
18958            );
18959            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
18960            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
18961            let vh = self.f32_to_f16(v, t_kv * n_head_kv * head_dim)?;
18962            let __s_b = self.gpu.stream();
18963            let mut b = __s_b.launch_builder(&f);
18964            b.arg(&qb)
18965                .arg(&kb)
18966                .arg(&vh)
18967                .arg(o)
18968                .arg(&hd)
18969                .arg(&nh)
18970                .arg(&nhkv)
18971                .arg(&ti)
18972                .arg(&tkvi)
18973                .arg(&scale)
18974                .arg(&cz)
18975                .arg(&wi);
18976            unsafe {
18977                b.launch(cfg)?;
18978            }
18979            return Ok(());
18980        }
18981        if p1 {
18982            let f = self.func("fa_prefill_w_bf16_p1");
18983            let shmem =
18984                (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
18985            use cudarc::driver::sys::CUfunction_attribute_enum as A;
18986            f.set_attribute(
18987                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18988                shmem as i32,
18989            )?;
18990            let cfg = LaunchConfig {
18991                grid_dim: (
18992                    (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
18993                    n_head as u32,
18994                    1,
18995                ),
18996                block_dim: (32, 4, 1),
18997                shared_mem_bytes: shmem,
18998            };
18999            let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
19000                head_dim as i32,
19001                n_head as i32,
19002                n_head_kv as i32,
19003                t as i32,
19004                t_kv as i32,
19005                causal as i32,
19006                window as i32,
19007            );
19008            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
19009            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
19010            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
19011            let __s_b = self.gpu.stream();
19012            let mut b = __s_b.launch_builder(&f);
19013            b.arg(&qb)
19014                .arg(&kb)
19015                .arg(&vb)
19016                .arg(o)
19017                .arg(&hd)
19018                .arg(&nh)
19019                .arg(&nhkv)
19020                .arg(&ti)
19021                .arg(&tkvi)
19022                .arg(&scale)
19023                .arg(&cz)
19024                .arg(&wi);
19025            unsafe {
19026                b.launch(cfg)?;
19027            }
19028            return Ok(());
19029        }
19030        // MQA head-grouping (MEMRA_FAW_G4=0 reverts): 4 heads/CTA share the staged K/V —
19031        // per-(head,row) FP chain identical to the per-head stamp -> bit-identical (gated).
19032        static G4_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19033        let g4 = !floor
19034            && !f32_stage
19035            && n_head_kv == 1
19036            && n_head % 4 == 0
19037            && *G4_ON.get_or_init(|| {
19038                std::env::var("MEMRA_FAW_G4")
19039                    .map(|v| v != "0")
19040                    .unwrap_or(true)
19041            });
19042        if g4 {
19043            const SP_M: usize = 16;
19044            // Occupancy-2 twin (MEMRA_FAW_O2=0 reverts): one shared K/V buffer inside the dead
19045            // Q-stage region -> ~36.5KB smem, 2 CTA/SM (the llama hd256 mechanism). Bit-identical.
19046            static O2_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19047            let o2 = *O2_ON.get_or_init(|| {
19048                std::env::var("MEMRA_FAW_O2")
19049                    .map(|v| v != "0")
19050                    .unwrap_or(true)
19051            });
19052            let f = self.func(if o2 {
19053                "fa_prefill_w_bf16_g4o2"
19054            } else {
19055                "fa_prefill_w_bf16_g4"
19056            });
19057            let shmem = if o2 {
19058                (2 * (4 * SP_M * head_dim + 4 * SP_M * BK) + 4 * (4 * SP_M)) as u32
19059            } else {
19060                (2 * (2 * BK * head_dim + 4 * SP_M * head_dim + 4 * SP_M * BK) + 4 * (4 * SP_M))
19061                    as u32
19062            };
19063            use cudarc::driver::sys::CUfunction_attribute_enum as A;
19064            f.set_attribute(
19065                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19066                shmem as i32,
19067            )?;
19068            let cfg = LaunchConfig {
19069                grid_dim: ((t as u32).div_ceil(SP_M as u32), (n_head / 4) as u32, 1),
19070                block_dim: (32, 4, 1),
19071                shared_mem_bytes: shmem,
19072            };
19073            let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
19074                head_dim as i32,
19075                n_head as i32,
19076                n_head_kv as i32,
19077                t as i32,
19078                t_kv as i32,
19079                causal as i32,
19080                window as i32,
19081            );
19082            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
19083            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
19084            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
19085            let __s_b = self.gpu.stream();
19086            let mut b = __s_b.launch_builder(&f);
19087            b.arg(&qb)
19088                .arg(&kb)
19089                .arg(&vb)
19090                .arg(o)
19091                .arg(&hd)
19092                .arg(&nh)
19093                .arg(&nhkv)
19094                .arg(&ti)
19095                .arg(&tkvi)
19096                .arg(&scale)
19097                .arg(&cz)
19098                .arg(&wi);
19099            unsafe {
19100                b.launch(cfg)?;
19101            }
19102            return Ok(());
19103        }
19104        let f = self.func(if floor {
19105            "fa_prefill_w_f32"
19106        } else if f32_stage {
19107            "fa_prefill_w_f32_pp"
19108        } else {
19109            "fa_prefill_w_bf16_pp"
19110        });
19111        let shmem =
19112            (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
19113        use cudarc::driver::sys::CUfunction_attribute_enum as A;
19114        f.set_attribute(
19115            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19116            shmem as i32,
19117        )?;
19118        let cfg = LaunchConfig {
19119            grid_dim: (
19120                (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
19121                n_head as u32,
19122                1,
19123            ),
19124            block_dim: (32, 4, 1),
19125            shared_mem_bytes: shmem,
19126        };
19127        let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
19128            head_dim as i32,
19129            n_head as i32,
19130            n_head_kv as i32,
19131            t as i32,
19132            t_kv as i32,
19133            causal as i32,
19134            window as i32,
19135        );
19136        if f32_stage {
19137            let __s_b = self.gpu.stream();
19138            let mut b = __s_b.launch_builder(&f);
19139            b.arg(q)
19140                .arg(k)
19141                .arg(v)
19142                .arg(o)
19143                .arg(&hd)
19144                .arg(&nh)
19145                .arg(&nhkv)
19146                .arg(&ti)
19147                .arg(&tkvi)
19148                .arg(&scale)
19149                .arg(&cz)
19150                .arg(&wi);
19151            unsafe {
19152                b.launch(cfg)?;
19153            }
19154        } else {
19155            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
19156            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
19157            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
19158            let __s_b = self.gpu.stream();
19159            let mut b = __s_b.launch_builder(&f);
19160            b.arg(&qb)
19161                .arg(&kb)
19162                .arg(&vb)
19163                .arg(o)
19164                .arg(&hd)
19165                .arg(&nh)
19166                .arg(&nhkv)
19167                .arg(&ti)
19168                .arg(&tkvi)
19169                .arg(&scale)
19170                .arg(&cz)
19171                .arg(&wi);
19172            unsafe {
19173                b.launch(cfg)?;
19174            }
19175        }
19176        Ok(())
19177    }
19178
19179    /// hd512 FA prefill (gemma4 GLOBAL layers): BLOCK_Q=32 x 2 warps, Q staged in smem,
19180    /// grid.z = 2 O-halves (each CTA computes the full 512-dim scores, accumulates half the
19181    /// V dims). Replaces the scalar sdpa_naive on the prime path's globals.
19182    #[allow(clippy::too_many_arguments)]
19183    pub fn fa_prefill_hd512(
19184        &self,
19185        q: &CudaSlice<f32>,
19186        k: &CudaSlice<f32>,
19187        v: &CudaSlice<f32>,
19188        o: &mut CudaSlice<f32>,
19189        head_dim: usize,
19190        n_head: usize,
19191        n_head_kv: usize,
19192        t: usize,
19193        t_kv: usize,
19194        scale: f32,
19195        causal: bool,
19196    ) -> Result<(), Box<dyn std::error::Error>> {
19197        // sm_90a rides the mma twins (portable_mma_gated, 2026-07-31 — same flip as _w).
19198        if portable_mma_gated() {
19199            return self.sdpa_naive(
19200                q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
19201            );
19202        }
19203        // Default: pre-convert Q/K/V to bf16 once and stage int4 (8 bf16/copy) — at 1 CTA/SM the
19204        // synchronous stage serializes with compute and MQA re-stages the same K/V per head CTA;
19205        // pre-converting halves staged bytes and cuts stage instructions 8x. BIT-IDENTICAL to the
19206        // f32-staged kernel (the converter applies the same __float2bfloat16 the stage applied;
19207        // kernel_check gates the identity). MEMRA_FA512_STAGE=f32 = rollback to the f32 kernel.
19208        static F32_STAGE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19209        let f32_stage =
19210            *F32_STAGE.get_or_init(|| std::env::var("MEMRA_FA512_STAGE").as_deref() == Ok("f32"));
19211        // Single-pass arm (MEMRA_FA512_SP=0 reverts to the z=2 bf16 kernel): GEMM0 split-K across
19212        // the 2 warps instead of recomputed per O-half CTA — the 2026-07-22 kernel-diff excess.
19213        // Own numeric config (partial-sum order) — battery-gated.
19214        static SP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19215        let sp = !f32_stage
19216            && *SP_ON.get_or_init(|| {
19217                std::env::var("MEMRA_FA512_SP")
19218                    .map(|v| v != "0")
19219                    .unwrap_or(true)
19220            });
19221        self.fa_prefill_hd512_arm(
19222            q,
19223            k,
19224            v,
19225            o,
19226            head_dim,
19227            n_head,
19228            n_head_kv,
19229            t,
19230            t_kv,
19231            scale,
19232            causal,
19233            f32_stage,
19234            sp,
19235            sp && fa_f16pv_on(),
19236        )
19237    }
19238
19239    /// hd512 single-pass FA with PRE-CONVERTED bf16 operands (producer-emitted).
19240    #[allow(clippy::too_many_arguments)]
19241    pub fn fa_prefill_hd512_pre(
19242        &self,
19243        qb: &CudaSlice<u8>,
19244        kb: &CudaSlice<u8>,
19245        vb: &CudaSlice<u8>,
19246        o: &mut CudaSlice<f32>,
19247        head_dim: usize,
19248        n_head: usize,
19249        n_head_kv: usize,
19250        t: usize,
19251        t_kv: usize,
19252        scale: f32,
19253        causal: bool,
19254        v_f16: bool,
19255    ) -> Result<(), Box<dyn std::error::Error>> {
19256        debug_assert_eq!(head_dim, 512);
19257        const SP_M: usize = 16;
19258        const BKS: usize = 32;
19259        // f16-P/V door (MEMRA_FA_F16PV=1): P and the P@V accumulation in f16 (llama's fa=1 VKQ
19260        // class); KQ/softmax/rescale-band/final-normalize stay f32. Own numeric config —
19261        // battery-gated. V bytes must be f16 for the sp16 kernel (stage/ldmatrix are typeless).
19262        let f16pv = fa_f16pv_on();
19263        let nw = if f16pv { fa512_wide_warps() } else { 2 };
19264        let hp = f16pv && fa512_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
19265        debug_assert!(!v_f16 || f16pv, "f16 V emitted without the door on");
19266        let mut vguard = self.fa_vf16_scratch.lock().unwrap();
19267        let vref: &CudaSlice<u8> = if f16pv && !v_f16 {
19268            // Fallback re-encode (producer emitted bf16); the emit lane normally hands f16.
19269            let n = t_kv * n_head_kv * head_dim;
19270            let need = n * 2;
19271            if vguard.as_ref().map(|b| b.len() < need).unwrap_or(true) {
19272                *vguard = Some(self.alloc_uninit::<u8>(need)?);
19273            }
19274            let dst = vguard.as_mut().unwrap();
19275            self.bf16_to_f16_into(vb, n, dst)?;
19276            vguard.as_ref().unwrap()
19277        } else {
19278            vb
19279        };
19280        let f = self.func(if hp {
19281            "fa_prefill_bf16_hd512_sp16h2"
19282        } else {
19283            match (f16pv, nw) {
19284                (true, 4) => "fa_prefill_bf16_hd512_sp16w4",
19285                (true, _) => "fa_prefill_bf16_hd512_sp16",
19286                _ => "fa_prefill_bf16_hd512_sp",
19287            }
19288        });
19289        let (nwarp, npart) = if hp {
19290            (4usize, 4usize)
19291        } else if nw > 2 {
19292            (nw, nw)
19293        } else {
19294            (2, 1)
19295        };
19296        // h2 drops sQ (Q register-resident) and doubles sP/sS/sL for the head pair.
19297        let shmem = if hp {
19298            (2 * (2 * BKS * head_dim + 2 * SP_M * BKS) + 4 * (2 * npart * SP_M * BKS + 2 * SP_M))
19299                as u32
19300        } else {
19301            (2 * (SP_M * head_dim + 2 * BKS * head_dim + SP_M * BKS)
19302                + 4 * (npart * SP_M * BKS + SP_M)) as u32
19303        };
19304        use cudarc::driver::sys::CUfunction_attribute_enum as A;
19305        f.set_attribute(
19306            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19307            shmem as i32,
19308        )?;
19309        let grid_y = if hp {
19310            (n_head / 2) as u32
19311        } else {
19312            n_head as u32
19313        };
19314        let cfg = LaunchConfig {
19315            grid_dim: ((t as u32).div_ceil(SP_M as u32), grid_y, 1),
19316            block_dim: (32, nwarp as u32, 1),
19317            shared_mem_bytes: shmem,
19318        };
19319        let (hd, nh, nhkv, ti, tkvi, cz) = (
19320            head_dim as i32,
19321            n_head as i32,
19322            n_head_kv as i32,
19323            t as i32,
19324            t_kv as i32,
19325            causal as i32,
19326        );
19327        let __s_b = self.gpu.stream();
19328        let mut b = __s_b.launch_builder(&f);
19329        b.arg(qb)
19330            .arg(kb)
19331            .arg(vref)
19332            .arg(o)
19333            .arg(&hd)
19334            .arg(&nh)
19335            .arg(&nhkv)
19336            .arg(&ti)
19337            .arg(&tkvi)
19338            .arg(&scale)
19339            .arg(&cz);
19340        unsafe {
19341            b.launch(cfg)?;
19342        }
19343        Ok(())
19344    }
19345
19346    /// hd512 FA prefill with the stage/sp arms FORCED — the kernel_check gate entry
19347    /// (`fa_prefill_hd512` picks the arms from MEMRA_FA512_STAGE / MEMRA_FA512_SP).
19348    #[allow(clippy::too_many_arguments)]
19349    pub fn fa_prefill_hd512_arm(
19350        &self,
19351        q: &CudaSlice<f32>,
19352        k: &CudaSlice<f32>,
19353        v: &CudaSlice<f32>,
19354        o: &mut CudaSlice<f32>,
19355        head_dim: usize,
19356        n_head: usize,
19357        n_head_kv: usize,
19358        t: usize,
19359        t_kv: usize,
19360        scale: f32,
19361        causal: bool,
19362        f32_stage: bool,
19363        sp: bool,
19364        f16pv: bool,
19365    ) -> Result<(), Box<dyn std::error::Error>> {
19366        debug_assert_eq!(head_dim, 512, "fa_prefill_hd512 is hd512 only");
19367        if sp && !f32_stage {
19368            // Single-pass: 16 q-rows/CTA, 2 warps, grid (ceil(T/16), n_head, 1).
19369            // smem: sQ[16][512] + sK[32][512] + sV[32][512] + sP[16][32] (bf16) + sS[16][32]+sL f32.
19370            // f16pv: sp16 kernel — f16 P + f16 P@V accum, V operand encoded f16.
19371            const SP_M: usize = 16;
19372            const BKS: usize = 32;
19373            let nw = if f16pv { fa512_wide_warps() } else { 2 };
19374            let hp = f16pv && fa512_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
19375            let f = self.func(if hp {
19376                "fa_prefill_bf16_hd512_sp16h2"
19377            } else {
19378                match (f16pv, nw) {
19379                    (true, 4) => "fa_prefill_bf16_hd512_sp16w4",
19380                    (true, _) => "fa_prefill_bf16_hd512_sp16",
19381                    _ => "fa_prefill_bf16_hd512_sp",
19382                }
19383            });
19384            let (nwarp, npart) = if hp {
19385                (4usize, 4usize)
19386            } else if nw > 2 {
19387                (nw, nw)
19388            } else {
19389                (2, 1)
19390            };
19391            let shmem = if hp {
19392                (2 * (2 * BKS * head_dim + 2 * SP_M * BKS)
19393                    + 4 * (2 * npart * SP_M * BKS + 2 * SP_M)) as u32
19394            } else {
19395                (2 * (SP_M * head_dim + 2 * BKS * head_dim + SP_M * BKS)
19396                    + 4 * (npart * SP_M * BKS + SP_M)) as u32
19397            };
19398            use cudarc::driver::sys::CUfunction_attribute_enum as A;
19399            f.set_attribute(
19400                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19401                shmem as i32,
19402            )?;
19403            let grid_y = if hp {
19404                (n_head / 2) as u32
19405            } else {
19406                n_head as u32
19407            };
19408            let cfg = LaunchConfig {
19409                grid_dim: ((t as u32).div_ceil(SP_M as u32), grid_y, 1),
19410                block_dim: (32, nwarp as u32, 1),
19411                shared_mem_bytes: shmem,
19412            };
19413            let (hd, nh, nhkv, ti, tkvi, cz) = (
19414                head_dim as i32,
19415                n_head as i32,
19416                n_head_kv as i32,
19417                t as i32,
19418                t_kv as i32,
19419                causal as i32,
19420            );
19421            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
19422            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
19423            let vb = if f16pv {
19424                self.f32_to_f16(v, t_kv * n_head_kv * head_dim)?
19425            } else {
19426                self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?
19427            };
19428            let __s_b = self.gpu.stream();
19429            let mut b = __s_b.launch_builder(&f);
19430            b.arg(&qb)
19431                .arg(&kb)
19432                .arg(&vb)
19433                .arg(o)
19434                .arg(&hd)
19435                .arg(&nh)
19436                .arg(&nhkv)
19437                .arg(&ti)
19438                .arg(&tkvi)
19439                .arg(&scale)
19440                .arg(&cz);
19441            unsafe {
19442                b.launch(cfg)?;
19443            }
19444            return Ok(());
19445        }
19446        const BLOCK_Q: usize = 32;
19447        const BK: usize = 32;
19448        const HALF: usize = 256;
19449        let f = self.func(if f32_stage {
19450            "fa_prefill_f32_hd512"
19451        } else {
19452            "fa_prefill_bf16_hd512"
19453        });
19454        // sQ[32][512] + sK[BK][512] + sV[BK][256] + sP[32][BK] (bf16) + sL[32] f32
19455        let shmem = (2 * (BLOCK_Q * head_dim + BK * head_dim + BK * HALF + BLOCK_Q * BK)
19456            + 4 * BLOCK_Q) as u32;
19457        use cudarc::driver::sys::CUfunction_attribute_enum as A;
19458        f.set_attribute(
19459            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19460            shmem as i32,
19461        )?;
19462        let cfg = LaunchConfig {
19463            grid_dim: (
19464                (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
19465                n_head as u32,
19466                2,
19467            ),
19468            block_dim: (32, 2, 1),
19469            shared_mem_bytes: shmem,
19470        };
19471        let (hd, nh, nhkv, ti, tkvi, cz) = (
19472            head_dim as i32,
19473            n_head as i32,
19474            n_head_kv as i32,
19475            t as i32,
19476            t_kv as i32,
19477            causal as i32,
19478        );
19479        if f32_stage {
19480            let __s_b = self.gpu.stream();
19481            let mut b = __s_b.launch_builder(&f);
19482            b.arg(q)
19483                .arg(k)
19484                .arg(v)
19485                .arg(o)
19486                .arg(&hd)
19487                .arg(&nh)
19488                .arg(&nhkv)
19489                .arg(&ti)
19490                .arg(&tkvi)
19491                .arg(&scale)
19492                .arg(&cz);
19493            unsafe {
19494                b.launch(cfg)?;
19495            }
19496        } else {
19497            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
19498            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
19499            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
19500            let __s_b = self.gpu.stream();
19501            let mut b = __s_b.launch_builder(&f);
19502            b.arg(&qb)
19503                .arg(&kb)
19504                .arg(&vb)
19505                .arg(o)
19506                .arg(&hd)
19507                .arg(&nh)
19508                .arg(&nhkv)
19509                .arg(&ti)
19510                .arg(&tkvi)
19511                .arg(&scale)
19512                .arg(&cz);
19513            unsafe {
19514                b.launch(cfg)?;
19515            }
19516        }
19517        Ok(())
19518    }
19519
19520    /// rope_neox2 with bf16 EMIT (31B glue lane): identical rope math/stores plus the post-rope
19521    /// values written as bf16 — the FA q/k operands come from this launch (bit-identical to the
19522    /// separate f32_to_bf16 the FA entries would run).
19523    #[allow(clippy::too_many_arguments)]
19524    pub fn rope_neox2_bf16e(
19525        &self,
19526        q: &mut CudaSlice<f32>,
19527        k: &mut CudaSlice<f32>,
19528        qb: &mut CudaSlice<u8>,
19529        kb: &mut CudaSlice<u8>,
19530        pos: &CudaSlice<i32>,
19531        head_dim: usize,
19532        n_dims: usize,
19533        nh_q: usize,
19534        nh_k: usize,
19535        n_tokens: usize,
19536        base: f32,
19537        freq_scale: f32,
19538        ff: Option<&CudaSlice<f32>>,
19539    ) -> Result<(), Box<dyn std::error::Error>> {
19540        let f = self.func("rope_neox2_bf16e_f32");
19541        let rows = ((nh_q + nh_k) * n_tokens) as u32;
19542        let cfg = LaunchConfig {
19543            grid_dim: (rows, 1, 1),
19544            block_dim: ((head_dim / 2) as u32, 1, 1),
19545            shared_mem_bytes: 0,
19546        };
19547        let theta_scale = base.powf(-2.0 / n_dims as f32);
19548        let (hd, nd, nhq, nhk, nt) = (
19549            head_dim as i32,
19550            n_dims as i32,
19551            nh_q as i32,
19552            nh_k as i32,
19553            n_tokens as i32,
19554        );
19555        let __s_b = self.gpu.stream();
19556        let mut b = __s_b.launch_builder(&f);
19557        match ff {
19558            Some(t) => {
19559                b.arg(&mut *q)
19560                    .arg(&mut *k)
19561                    .arg(&mut *qb)
19562                    .arg(&mut *kb)
19563                    .arg(pos)
19564                    .arg(&hd)
19565                    .arg(&nd)
19566                    .arg(&nhq)
19567                    .arg(&nhk)
19568                    .arg(&nt)
19569                    .arg(&theta_scale)
19570                    .arg(&freq_scale)
19571                    .arg(t);
19572                unsafe {
19573                    b.launch(cfg)?;
19574                }
19575            }
19576            None => {
19577                let null: u64 = 0;
19578                b.arg(&mut *q)
19579                    .arg(&mut *k)
19580                    .arg(&mut *qb)
19581                    .arg(&mut *kb)
19582                    .arg(pos)
19583                    .arg(&hd)
19584                    .arg(&nd)
19585                    .arg(&nhq)
19586                    .arg(&nhk)
19587                    .arg(&nt)
19588                    .arg(&theta_scale)
19589                    .arg(&freq_scale)
19590                    .arg(&null);
19591                unsafe {
19592                    b.launch(cfg)?;
19593                }
19594            }
19595        }
19596        Ok(())
19597    }
19598
19599    /// Flat f32 -> bf16 conversion into a fresh scratch buffer (2 bytes/elem). `n % 4 == 0`
19600    /// (float4 in, 4x bf16 out). Feeds the bf16-staged hd512 FA prefill.
19601    pub fn f32_to_bf16(
19602        &self,
19603        x: &CudaSlice<f32>,
19604        n: usize,
19605    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
19606        assert!(n % 4 == 0, "f32_to_bf16 requires n % 4 == 0, got {n}");
19607        let mut y = self.alloc_uninit::<u8>(n * 2)?;
19608        let f = self.func("f32_to_bf16_flat");
19609        let n_i = n as i64;
19610        let cfg = LaunchConfig {
19611            grid_dim: (((n / 4) as u32).div_ceil(256), 1, 1),
19612            block_dim: (256, 1, 1),
19613            shared_mem_bytes: 0,
19614        };
19615        let __s_b = self.gpu.stream();
19616        let mut b = __s_b.launch_builder(&f);
19617        b.arg(x).arg(&mut y).arg(&n_i);
19618        unsafe {
19619            b.launch(cfg)?;
19620        }
19621        Ok(y)
19622    }
19623
19624    pub fn f32_to_f16(
19625        &self,
19626        x: &CudaSlice<f32>,
19627        n: usize,
19628    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
19629        assert!(n % 4 == 0, "f32_to_f16 requires n % 4 == 0, got {n}");
19630        let mut y = self.alloc_uninit::<u8>(n * 2)?;
19631        let f = self.func("f32_to_f16_flat");
19632        let n_i = n as i64;
19633        let cfg = LaunchConfig {
19634            grid_dim: (((n / 4) as u32).div_ceil(256), 1, 1),
19635            block_dim: (256, 1, 1),
19636            shared_mem_bytes: 0,
19637        };
19638        let __s_b = self.gpu.stream();
19639        let mut b = __s_b.launch_builder(&f);
19640        b.arg(x).arg(&mut y).arg(&n_i);
19641        unsafe {
19642            b.launch(cfg)?;
19643        }
19644        Ok(y)
19645    }
19646
19647    /// bf16 bytes -> f16 bytes, n elements (the f16-P/V door's V re-encode on the emit lane).
19648    pub fn bf16_to_f16(
19649        &self,
19650        xb: &CudaSlice<u8>,
19651        n: usize,
19652    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
19653        let mut y = self.alloc_uninit::<u8>(n * 2)?;
19654        self.bf16_to_f16_into(xb, n, &mut y)?;
19655        Ok(y)
19656    }
19657
19658    /// Same conversion into a caller-owned (pooled) buffer; `y.len() >= n*2`.
19659    pub fn bf16_to_f16_into(
19660        &self,
19661        xb: &CudaSlice<u8>,
19662        n: usize,
19663        y: &mut CudaSlice<u8>,
19664    ) -> Result<(), Box<dyn std::error::Error>> {
19665        assert!(n % 2 == 0, "bf16_to_f16 requires n % 2 == 0, got {n}");
19666        assert!(y.len() >= n * 2);
19667        let f = self.func("bf16_to_f16_flat");
19668        let n2 = (n / 2) as i64;
19669        let cfg = LaunchConfig {
19670            grid_dim: (((n / 2) as u32).div_ceil(256), 1, 1),
19671            block_dim: (256, 1, 1),
19672            shared_mem_bytes: 0,
19673        };
19674        let __s_b = self.gpu.stream();
19675        let mut b = __s_b.launch_builder(&f);
19676        b.arg(xb).arg(y).arg(&n2);
19677        unsafe {
19678            b.launch(cfg)?;
19679        }
19680        Ok(())
19681    }
19682
19683    /// task #18 (attn side): varlen FA — bf16 K/V mirrors (2 launches) + ONE
19684    /// fa_prefill_bf16kv launch for every fresh sequence. Same per-block math as the
19685    /// per-seq path (bit-gateable). Caller guarantees: fresh causal (T_kv == T),
19686    /// head_dim in {256, 128}, bf16kv lane on.
19687    #[allow(clippy::too_many_arguments)]
19688    pub fn fa_prefill_vl8(
19689        &self,
19690        seqs: &[FaSeqVl],
19691        head_dim: usize,
19692        n_head: usize,
19693        n_head_kv: usize,
19694        scale: f32,
19695    ) -> Result<(), Box<dyn std::error::Error>> {
19696        const BK: usize = 32;
19697        let b = seqs.len();
19698        assert!(b >= 1 && b <= 8);
19699        let mut packed = [FaSeqVl::default(); 8];
19700        packed[..b].copy_from_slice(seqs);
19701        let v = FaVl8(packed);
19702        let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
19703        let ept = (n_head_kv * head_dim) as i32;
19704        {
19705            let f = self.func("fa_mirror_vl");
19706            let max_n = (max_t as i64) * ept as i64;
19707            let blocks = ((max_n as u32).div_ceil(4)).div_ceil(256);
19708            for which in 0..2i32 {
19709                let cfg = LaunchConfig {
19710                    grid_dim: (blocks, 1, b as u32),
19711                    block_dim: (256, 1, 1),
19712                    shared_mem_bytes: 0,
19713                };
19714                let __s_lb = self.gpu.stream();
19715                let mut lb = __s_lb.launch_builder(&f);
19716                lb.arg(&v).arg(&ept).arg(&which);
19717                unsafe {
19718                    lb.launch(cfg)?;
19719                }
19720            }
19721        }
19722        let hd_sfx = fa_hd_suffix(head_dim)?;
19723        let f = self.func(&format!("fa_prefill_bf16kv_vl{hd_sfx}"));
19724        let block_q = 64usize;
19725        let kv_stages = 2usize;
19726        let shmem = (2 * (kv_stages * 2 * BK * head_dim + block_q * BK)
19727            + 4 * (block_q * BK + 2 * block_q)) as u32;
19728        use cudarc::driver::sys::CUfunction_attribute_enum as A;
19729        f.set_attribute(
19730            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19731            shmem as i32,
19732        )?;
19733        let cfg = LaunchConfig {
19734            grid_dim: (max_t.div_ceil(block_q as u32), n_head as u32, b as u32),
19735            block_dim: (32, 4, 1),
19736            shared_mem_bytes: shmem,
19737        };
19738        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
19739        let __s_lb = self.gpu.stream();
19740        let mut lb = __s_lb.launch_builder(&f);
19741        lb.arg(&v).arg(&hd).arg(&nh).arg(&nhkv).arg(&scale);
19742        unsafe {
19743            lb.launch(cfg)?;
19744        }
19745        Ok(())
19746    }
19747
19748    /// task #18 (attn pre-FA): varlen split + QK-norm + RoPE + KV-append — FOUR launches
19749    /// for every fresh sequence (was 6 x B, plus the q/k/v split copies which the view
19750    /// inputs remove entirely). Fresh-only (append at t0=0, RoPE pos = token index).
19751    #[allow(clippy::too_many_arguments)]
19752    pub fn attn_pre_vl8(
19753        &self,
19754        seqs: &[AttnPreVl],
19755        wq: &CudaSlice<f32>,
19756        wk: &CudaSlice<f32>,
19757        head_dim: usize,
19758        rope_dims: usize,
19759        n_head: usize,
19760        n_head_kv: usize,
19761        eps: f32,
19762        freq_base: f32,
19763        freq_scale: f32,
19764        kv_dim_k: usize,
19765        kv_dim_v: usize,
19766        k_tok_bytes: usize,
19767        v_tok_bytes: usize,
19768    ) -> Result<(), Box<dyn std::error::Error>> {
19769        let b = seqs.len();
19770        assert!(b >= 1 && b <= 8);
19771        let mut packed = [AttnPreVl::default(); 8];
19772        packed[..b].copy_from_slice(seqs);
19773        let v = AttnPreVl8(packed);
19774        let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
19775        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
19776        {
19777            let f = self.func("q_gate_split_vl");
19778            let n = max_t * (n_head * head_dim) as u32;
19779            let cfg = LaunchConfig {
19780                grid_dim: (n.div_ceil(256), 1, b as u32),
19781                block_dim: (256, 1, 1),
19782                shared_mem_bytes: 0,
19783            };
19784            let __s_lb = self.gpu.stream();
19785            let mut lb = __s_lb.launch_builder(&f);
19786            lb.arg(&v).arg(&hd).arg(&nh);
19787            unsafe {
19788                lb.launch(cfg)?;
19789            }
19790        }
19791        {
19792            let f = self.func("attn_rms_vl");
19793            let cfg = LaunchConfig {
19794                grid_dim: (max_t * n_head as u32, 2, b as u32),
19795                block_dim: (rms_block(), 1, 1),
19796                shared_mem_bytes: 0,
19797            };
19798            let __s_lb = self.gpu.stream();
19799            let mut lb = __s_lb.launch_builder(&f);
19800            lb.arg(&v)
19801                .arg(wq)
19802                .arg(wk)
19803                .arg(&hd)
19804                .arg(&nh)
19805                .arg(&nhkv)
19806                .arg(&eps);
19807            unsafe {
19808                lb.launch(cfg)?;
19809            }
19810        }
19811        {
19812            let f = self.func("attn_rope_vl");
19813            let theta_scale = freq_base.powf(-2.0 / rope_dims as f32);
19814            let nd = rope_dims as i32;
19815            let cfg = LaunchConfig {
19816                grid_dim: (max_t * n_head as u32, 2, b as u32),
19817                block_dim: ((head_dim / 2) as u32, 1, 1),
19818                shared_mem_bytes: 0,
19819            };
19820            let __s_lb = self.gpu.stream();
19821            let mut lb = __s_lb.launch_builder(&f);
19822            lb.arg(&v)
19823                .arg(&hd)
19824                .arg(&nd)
19825                .arg(&nh)
19826                .arg(&nhkv)
19827                .arg(&theta_scale)
19828                .arg(&freq_scale);
19829            unsafe {
19830                lb.launch(cfg)?;
19831            }
19832        }
19833        {
19834            let f = self.func("append_kv_vl");
19835            let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
19836            let cfg = LaunchConfig {
19837                grid_dim: (nblk, max_t, b as u32),
19838                block_dim: (32, 1, 1),
19839                shared_mem_bytes: 0,
19840            };
19841            let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
19842            let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
19843            let __s_lb = self.gpu.stream();
19844            let mut lb = __s_lb.launch_builder(&f);
19845            lb.arg(&v).arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
19846            unsafe {
19847                lb.launch(cfg)?;
19848            }
19849        }
19850        Ok(())
19851    }
19852
19853    /// FA prefill where K/V are QUANTIZED CudaViews into the resident byte KV cache (the T=K verify
19854    /// path, MTP-PLAN §D.3). Uses `fa_prefill_q` (inline-dequant during stage-to-smem). The view's
19855    /// base+offset pointer is honored; the kernel reads [0..t_kv*tok_bytes). Q is the T fresh query
19856    /// rows; t = T, t_kv = cache len. k_tok_bytes/v_tok_bytes are the per-token byte strides.
19857    pub fn fa_prefill_view(
19858        &self,
19859        q: &CudaSlice<f32>,
19860        k: &cudarc::driver::CudaView<u8>,
19861        v: &cudarc::driver::CudaView<u8>,
19862        o: &mut CudaSlice<f32>,
19863        head_dim: usize,
19864        n_head: usize,
19865        n_head_kv: usize,
19866        t: usize,
19867        t_kv: usize,
19868        scale: f32,
19869        causal: bool,
19870        k_tok_bytes: usize,
19871        v_tok_bytes: usize,
19872        g: bool,
19873    ) -> Result<(), Box<dyn std::error::Error>> {
19874        if portable_mma_gated() {
19875            return self.sdpa_naive_quantized_view(
19876                q,
19877                k,
19878                v,
19879                o,
19880                head_dim,
19881                n_head,
19882                n_head_kv,
19883                t,
19884                t_kv,
19885                scale,
19886                causal,
19887                k_tok_bytes,
19888                v_tok_bytes,
19889            );
19890        }
19891        const BLOCK_Q: usize = 64;
19892        const BK: usize = 32;
19893        // g = e4m3 cache: the kernel parses via DQ_K_ELEM/DQ_V_ELEM (format macros) — the
19894        // kf8vf8-module stamp reads fp8 with the identical MMA/softmax/PV body.
19895        let name = format!("fa_prefill_q{}", fa_hd_suffix(head_dim)?);
19896        let f = if g {
19897            self.func_g(&name)
19898        } else {
19899            self.func(&name)
19900        };
19901        let shmem =
19902            (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
19903        use cudarc::driver::sys::CUfunction_attribute_enum as A;
19904        f.set_attribute(
19905            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19906            shmem as i32,
19907        )?;
19908        let cfg = LaunchConfig {
19909            grid_dim: (
19910                (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
19911                n_head as u32,
19912                1,
19913            ),
19914            block_dim: (32, 4, 1),
19915            shared_mem_bytes: shmem,
19916        };
19917        let (hd, nh, nhkv, ti, tkvi, cz) = (
19918            head_dim as i32,
19919            n_head as i32,
19920            n_head_kv as i32,
19921            t as i32,
19922            t_kv as i32,
19923            causal as i32,
19924        );
19925        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
19926        let __s_b = self.gpu.stream();
19927        let mut b = __s_b.launch_builder(&f);
19928        b.arg(q)
19929            .arg(k)
19930            .arg(v)
19931            .arg(o)
19932            .arg(&hd)
19933            .arg(&nh)
19934            .arg(&nhkv)
19935            .arg(&ti)
19936            .arg(&tkvi)
19937            .arg(&scale)
19938            .arg(&cz)
19939            .arg(&ktb)
19940            .arg(&vtb);
19941        unsafe {
19942            b.launch(cfg)?;
19943        }
19944        Ok(())
19945    }
19946
19947    /// ARC B (2026-07-05): dequant-once chunk-prime FA. Same contract as `fa_prefill_view`, but
19948    /// instead of every (q-block, head) CTA re-dequanting the whole quantized KV stream inline
19949    /// (T/64 x n_head redundant at chunk prime — 30.5% of the 32k prime wall), dequant the full
19950    /// [t_kv, kv_dim] K and V ONCE into a resident bf16 workspace (fa_dequant_kv_ws_bf16), then
19951    /// run `fa_prefill_qw` (the bf16-workspace twin) over it. EXACT: the workspace holds the same
19952    /// __float2bfloat16(dq_*_elem(...)) values fa_prefill_q stages to smem, and the twin's MMA/
19953    /// softmax/PV code is byte-identical -> bit-identical O (kernel_check pins bitdiff=0).
19954    /// The workspace allocation is REUSED across layers/chunks (grown to the largest shape);
19955    /// contents are rewritten per call. MEMRA_PRIME_DEQW=0 falls back to fa_prefill_view (callers gate).
19956    #[allow(clippy::too_many_arguments)]
19957    pub fn fa_prefill_view_ws(
19958        &self,
19959        q: &CudaSlice<f32>,
19960        k: &cudarc::driver::CudaView<u8>,
19961        v: &cudarc::driver::CudaView<u8>,
19962        o: &mut CudaSlice<f32>,
19963        head_dim: usize,
19964        n_head: usize,
19965        n_head_kv: usize,
19966        t: usize,
19967        t_kv: usize,
19968        scale: f32,
19969        causal: bool,
19970        k_tok_bytes: usize,
19971        v_tok_bytes: usize,
19972        g: bool,
19973    ) -> Result<(), Box<dyn std::error::Error>> {
19974        if portable_mma_gated() {
19975            return self.sdpa_naive_quantized_view(
19976                q,
19977                k,
19978                v,
19979                o,
19980                head_dim,
19981                n_head,
19982                n_head_kv,
19983                t,
19984                t_kv,
19985                scale,
19986                causal,
19987                k_tok_bytes,
19988                v_tok_bytes,
19989            );
19990        }
19991        const BLOCK_Q: usize = 64;
19992        const BK: usize = 32;
19993        let kv_dim_k = n_head_kv * head_dim;
19994        let kv_dim_v = n_head_kv * head_dim;
19995        let k_ws_bytes = t_kv * kv_dim_k * 2; // bf16
19996        let v_ws_bytes = t_kv * kv_dim_v * 2;
19997        // Lock held across BOTH launches: enqueue-only (µs), all compute serializes on gpu.stream.
19998        let mut guard = self.prime_deqw_ws.lock().unwrap();
19999        let need_grow = match guard.as_ref() {
20000            Some((kw, vw)) => kw.len() < k_ws_bytes || vw.len() < v_ws_bytes,
20001            None => true,
20002        };
20003        if need_grow {
20004            let grow = |cur: usize, need: usize| if cur >= need { cur } else { need };
20005            let (ck, cv) = guard
20006                .as_ref()
20007                .map(|(a, b)| (a.len(), b.len()))
20008                .unwrap_or((0, 0));
20009            *guard = Some((
20010                self.alloc_u8(grow(ck, k_ws_bytes))?,
20011                self.alloc_u8(grow(cv, v_ws_bytes))?,
20012            ));
20013        }
20014        let (kw, vw) = guard.as_mut().unwrap();
20015        // pass 1: dequant K+V once into the bf16 workspace (grid-stride, 1 thread/elem)
20016        {
20017            // only THIS pass parses KV bytes — pass 2 reads the bf16 workspace (format-free).
20018            let f = if g {
20019                self.func_g("fa_dequant_kv_ws_bf16")
20020            } else {
20021                self.func("fa_dequant_kv_ws_bf16")
20022            };
20023            let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
20024            let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
20025            let cfg = LaunchConfig {
20026                grid_dim: (nblk.max(1), 1, 1),
20027                block_dim: (256, 1, 1),
20028                shared_mem_bytes: 0,
20029            };
20030            let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
20031            let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
20032            let __s_b = self.gpu.stream();
20033            let mut b = __s_b.launch_builder(&f);
20034            b.arg(k)
20035                .arg(v)
20036                .arg(&mut *kw)
20037                .arg(&mut *vw)
20038                .arg(&kdk)
20039                .arg(&kdv)
20040                .arg(&tkvi)
20041                .arg(&ktb)
20042                .arg(&vtb);
20043            unsafe {
20044                b.launch(cfg)?;
20045            }
20046        }
20047        // pass 2: the bf16-workspace prefill twin (same tile sizes/loop structure as fa_prefill_q).
20048        // DEFAULT: cp.async double-buffered staging twin (fa_prefill_qw_db, +32KB smem for the
20049        // second K/V tile pair, 1 CTA/SM): overlaps tile n+1's L2->smem copy with tile n's MMA.
20050        // Bit-identical output (staging is a pure byte copy; kernel_check pins bitdiff=0 under
20051        // both twins). A/B (27B rtx6000, N=3): 32k prime 17.10->16.51s, 16k 9.09->8.65s — the copy
20052        // latency hides behind the MMA pipe and beats the 2-CTA/SM occupancy of the sync twin.
20053        // MEMRA_PRIME_DEQW_DB=0 falls back to the single-buffer twin.
20054        let db = std::env::var("MEMRA_PRIME_DEQW_DB")
20055            .map(|v| v != "0")
20056            .unwrap_or(true);
20057        {
20058            let hd_sfx = fa_hd_suffix(head_dim)?;
20059            let f = self.func(&format!(
20060                "fa_prefill_qw{}{hd_sfx}",
20061                if db { "_db" } else { "" }
20062            ));
20063            let shmem = if db {
20064                // 4x KV tile buffers (bf16) + sP (bf16) + sL (f32)
20065                (2 * (4 * BK * head_dim + BLOCK_Q * BK) + 4 * BLOCK_Q) as u32
20066            } else {
20067                (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32
20068            };
20069            use cudarc::driver::sys::CUfunction_attribute_enum as A;
20070            f.set_attribute(
20071                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20072                shmem as i32,
20073            )?;
20074            let cfg = LaunchConfig {
20075                grid_dim: (
20076                    (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
20077                    n_head as u32,
20078                    1,
20079                ),
20080                block_dim: (32, 4, 1),
20081                shared_mem_bytes: shmem,
20082            };
20083            let (hd, nh, nhkv, ti, tkvi, cz) = (
20084                head_dim as i32,
20085                n_head as i32,
20086                n_head_kv as i32,
20087                t as i32,
20088                t_kv as i32,
20089                causal as i32,
20090            );
20091            let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
20092            let __s_b = self.gpu.stream();
20093            let mut b = __s_b.launch_builder(&f);
20094            b.arg(q)
20095                .arg(&*kw)
20096                .arg(&*vw)
20097                .arg(o)
20098                .arg(&hd)
20099                .arg(&nh)
20100                .arg(&nhkv)
20101                .arg(&ti)
20102                .arg(&tkvi)
20103                .arg(&scale)
20104                .arg(&cz)
20105                .arg(&kdk)
20106                .arg(&kdv);
20107            unsafe {
20108                b.launch(cfg)?;
20109            }
20110        }
20111        Ok(())
20112    }
20113
20114    /// WINDOWED `fa_prefill_view_ws` twin at head_dim 128 (lane/pp-prefill 2026-08-07):
20115    /// step35's SWA prefill (win=512, 33 of 45 layers) previously had NO windowed FA prefill
20116    /// stamp — every windowed twin was hd256-only — and took `sdpa_naive_w_quantized_view`,
20117    /// the f32 floor, at 565 ms/layer on a pp4096 where the hd128 FA family does the harder
20118    /// causal-4096 in 3.3 ms (41% of the whole prime; research/pp-prefill-20260807 anatomy).
20119    /// Same two-pass shape as the unwindowed function: dequant K/V ONCE into the resident
20120    /// bf16 workspace, then the windowed qw kernel (`fa_prefill_qw_db_w_hd128`, cp.async
20121    /// double-buffered; MEMRA_PRIME_DEQW_DB=0 selects the single-buffer twin). The window
20122    /// mask is `fa_prefill_f32_body`'s exact predicate; `window == 0` is bit-identical to
20123    /// `fa_prefill_view_ws` by construction (default-arg body). NEW NUMERIC CLASS vs the
20124    /// f32 floor on SWA rows (bf16 MMA online-softmax vs f32 serial softmax) — adoption is
20125    /// gated by the full battery, and the class must change UNIFORMLY for a whole request
20126    /// (kernel selection keys on seq_end, never per chunk — the chunkfix law).
20127    /// hd128-only deliberately: the only windowed-prefill consumer at another head_dim is
20128    /// gemma4 (hd256), which already has `fa_prefill_w_f32`.
20129    #[allow(clippy::too_many_arguments)]
20130    pub fn fa_prefill_view_ws_w_hd128(
20131        &self,
20132        q: &CudaSlice<f32>,
20133        k: &cudarc::driver::CudaView<u8>,
20134        v: &cudarc::driver::CudaView<u8>,
20135        o: &mut CudaSlice<f32>,
20136        head_dim: usize,
20137        n_head: usize,
20138        n_head_kv: usize,
20139        t: usize,
20140        t_kv: usize,
20141        scale: f32,
20142        causal: bool,
20143        window: usize,
20144        k_tok_bytes: usize,
20145        v_tok_bytes: usize,
20146    ) -> Result<(), Box<dyn std::error::Error>> {
20147        assert_eq!(
20148            head_dim, 128,
20149            "fa_prefill_view_ws_w_hd128: only the hd128 twin is stamped"
20150        );
20151        if portable_mma_gated() {
20152            return self.sdpa_naive_w_quantized_view(
20153                q,
20154                k,
20155                v,
20156                o,
20157                head_dim,
20158                n_head,
20159                n_head_kv,
20160                t,
20161                t_kv,
20162                scale,
20163                causal,
20164                window,
20165                k_tok_bytes,
20166                v_tok_bytes,
20167            );
20168        }
20169        const BLOCK_Q: usize = 64;
20170        const BK: usize = 32;
20171        let kv_dim_k = n_head_kv * head_dim;
20172        let kv_dim_v = n_head_kv * head_dim;
20173        let k_ws_bytes = t_kv * kv_dim_k * 2; // bf16
20174        let v_ws_bytes = t_kv * kv_dim_v * 2;
20175        let mut guard = self.prime_deqw_ws.lock().unwrap();
20176        let need_grow = match guard.as_ref() {
20177            Some((kw, vw)) => kw.len() < k_ws_bytes || vw.len() < v_ws_bytes,
20178            None => true,
20179        };
20180        if need_grow {
20181            let grow = |cur: usize, need: usize| if cur >= need { cur } else { need };
20182            let (ck, cv) = guard
20183                .as_ref()
20184                .map(|(a, b)| (a.len(), b.len()))
20185                .unwrap_or((0, 0));
20186            *guard = Some((
20187                self.alloc_u8(grow(ck, k_ws_bytes))?,
20188                self.alloc_u8(grow(cv, v_ws_bytes))?,
20189            ));
20190        }
20191        let (kw, vw) = guard.as_mut().unwrap();
20192        // pass 1: dequant K+V once into the bf16 workspace (identical to fa_prefill_view_ws —
20193        // the workspace bytes are the SAME __float2bfloat16(dq(...)) values either way).
20194        {
20195            let f = self.func("fa_dequant_kv_ws_bf16");
20196            let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
20197            let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
20198            let cfg = LaunchConfig {
20199                grid_dim: (nblk.max(1), 1, 1),
20200                block_dim: (256, 1, 1),
20201                shared_mem_bytes: 0,
20202            };
20203            let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
20204            let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
20205            let __s_b = self.gpu.stream();
20206            let mut b = __s_b.launch_builder(&f);
20207            b.arg(k)
20208                .arg(v)
20209                .arg(&mut *kw)
20210                .arg(&mut *vw)
20211                .arg(&kdk)
20212                .arg(&kdv)
20213                .arg(&tkvi)
20214                .arg(&ktb)
20215                .arg(&vtb);
20216            unsafe {
20217                b.launch(cfg)?;
20218            }
20219        }
20220        // pass 2: the WINDOWED qw twin (db default, same as the unwindowed wrapper).
20221        let db = std::env::var("MEMRA_PRIME_DEQW_DB")
20222            .map(|v| v != "0")
20223            .unwrap_or(true);
20224        {
20225            let f = self.func(if db {
20226                "fa_prefill_qw_db_w_hd128"
20227            } else {
20228                "fa_prefill_qw_w_hd128"
20229            });
20230            let shmem = if db {
20231                (2 * (4 * BK * head_dim + BLOCK_Q * BK) + 4 * BLOCK_Q) as u32
20232            } else {
20233                (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32
20234            };
20235            use cudarc::driver::sys::CUfunction_attribute_enum as A;
20236            f.set_attribute(
20237                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20238                shmem as i32,
20239            )?;
20240            let cfg = LaunchConfig {
20241                grid_dim: (
20242                    (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
20243                    n_head as u32,
20244                    1,
20245                ),
20246                block_dim: (32, 4, 1),
20247                shared_mem_bytes: shmem,
20248            };
20249            let (hd, nh, nhkv, ti, tkvi, cz) = (
20250                head_dim as i32,
20251                n_head as i32,
20252                n_head_kv as i32,
20253                t as i32,
20254                t_kv as i32,
20255                causal as i32,
20256            );
20257            let (kdk, kdv, wnd) = (kv_dim_k as i32, kv_dim_v as i32, window as i32);
20258            let __s_b = self.gpu.stream();
20259            let mut b = __s_b.launch_builder(&f);
20260            b.arg(q)
20261                .arg(&*kw)
20262                .arg(&*vw)
20263                .arg(o)
20264                .arg(&hd)
20265                .arg(&nh)
20266                .arg(&nhkv)
20267                .arg(&ti)
20268                .arg(&tkvi)
20269                .arg(&scale)
20270                .arg(&cz)
20271                .arg(&kdk)
20272                .arg(&kdv)
20273                .arg(&wnd);
20274            unsafe {
20275                b.launch(cfg)?;
20276            }
20277        }
20278        Ok(())
20279    }
20280
20281    /// FA decode (T=1 split-K) over the resident QUANTIZED KV cache (q8_0 K / q5_1 V) as u8 views.
20282    /// Replaces sdpa_naive_view for decode; inline-dequants per element. k_tok_bytes/v_tok_bytes are
20283    /// the per-token byte strides (differ: q8_0=34*nblk, q5_1=24*nblk per token).
20284    pub fn fa_decode(
20285        &self,
20286        q: &CudaSlice<f32>,
20287        k: &cudarc::driver::CudaView<u8>,
20288        v: &cudarc::driver::CudaView<u8>,
20289        o: &mut CudaSlice<f32>,
20290        head_dim: usize,
20291        n_head: usize,
20292        n_head_kv: usize,
20293        t_kv: usize,
20294        scale: f32,
20295        k_tok_bytes: usize,
20296        v_tok_bytes: usize,
20297    ) -> Result<(), Box<dyn std::error::Error>> {
20298        self.fa_decode_kvmod(
20299            q,
20300            k,
20301            v,
20302            o,
20303            head_dim,
20304            n_head,
20305            n_head_kv,
20306            t_kv,
20307            scale,
20308            k_tok_bytes,
20309            v_tok_bytes,
20310            false,
20311        )
20312    }
20313
20314    /// `fa_decode` with an explicit fp8-module flag (`g`): gemma windowed layers under
20315    /// MEMRA_GEMMA_WKV read an e4m3 cache — every kernel must come from the kf8vf8 module
20316    /// and the v4 lane (q8_0-hardcoded staging) is excluded.
20317    #[allow(clippy::too_many_arguments)]
20318    /// UNIFIED scalar decode launch (fa_decode_f32, nullable-ctr): ONE symbol for host-len
20319    /// (kvmod eager) and device-len (graph/stream) callers — the textually-identical f32_dc
20320    /// twin compiled apart and its ULP drift flipped 31B verify argmaxes (2026-07-12).
20321    #[allow(clippy::too_many_arguments)]
20322    #[allow(clippy::too_many_arguments)]
20323    fn fa_decode_scalar_unified(
20324        &self,
20325        q: &cudarc::driver::CudaView<f32>,
20326        k: &cudarc::driver::CudaView<u8>,
20327        v: &cudarc::driver::CudaView<u8>,
20328        o: &mut cudarc::driver::CudaViewMut<f32>,
20329        head_dim: usize,
20330        n_head: usize,
20331        n_head_kv: usize,
20332        t_kv_host: usize,
20333        t_kv_dev: Option<&CudaSlice<i32>>,
20334        scale: f32,
20335        n_splits: usize,
20336        split_keys: usize,
20337        k_tok_bytes: usize,
20338        v_tok_bytes: usize,
20339        g: bool,
20340        part_o: &mut CudaSlice<f32>,
20341        part_m: &mut CudaSlice<f32>,
20342        part_l: &mut CudaSlice<f32>,
20343        q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
20344    ) -> Result<(), Box<dyn std::error::Error>> {
20345        let f = if g {
20346            self.func_g("fa_decode_f32")
20347        } else {
20348            self.fa_func("fa_decode_f32", head_dim)
20349        };
20350        let cfg = LaunchConfig {
20351            grid_dim: (n_head as u32, n_splits as u32, 1),
20352            block_dim: (head_dim as u32, 1, 1),
20353            shared_mem_bytes: (4 * (head_dim + 32)) as u32,
20354        };
20355        let (hd, nh, nhkv, nsp) = (
20356            head_dim as i32,
20357            n_head as i32,
20358            n_head_kv as i32,
20359            n_splits as i32,
20360        );
20361        let (ktb, vtb, tkvi, ski) = (
20362            k_tok_bytes as i64,
20363            v_tok_bytes as i64,
20364            t_kv_host as i32,
20365            split_keys as i32,
20366        );
20367        let __s_b = self.gpu.stream();
20368        let mut b = __s_b.launch_builder(&f);
20369        match t_kv_dev {
20370            Some(d) => {
20371                b.arg(q)
20372                    .arg(k)
20373                    .arg(v)
20374                    .arg(&mut *part_o)
20375                    .arg(&mut *part_m)
20376                    .arg(&mut *part_l)
20377                    .arg(&hd)
20378                    .arg(&nh)
20379                    .arg(&nhkv)
20380                    .arg(&tkvi)
20381                    .arg(d)
20382                    .arg(&scale)
20383                    .arg(&nsp)
20384                    .arg(&ski)
20385                    .arg(&ktb)
20386                    .arg(&vtb);
20387                unsafe {
20388                    b.launch(cfg)?;
20389                }
20390            }
20391            None => {
20392                let null: u64 = 0;
20393                b.arg(q)
20394                    .arg(k)
20395                    .arg(v)
20396                    .arg(&mut *part_o)
20397                    .arg(&mut *part_m)
20398                    .arg(&mut *part_l)
20399                    .arg(&hd)
20400                    .arg(&nh)
20401                    .arg(&nhkv)
20402                    .arg(&tkvi)
20403                    .arg(&null)
20404                    .arg(&scale)
20405                    .arg(&nsp)
20406                    .arg(&ski)
20407                    .arg(&ktb)
20408                    .arg(&vtb);
20409                unsafe {
20410                    b.launch(cfg)?;
20411                }
20412            }
20413        }
20414        let cfg2 = LaunchConfig {
20415            grid_dim: (n_head as u32, 1, 1),
20416            block_dim: (head_dim as u32, 1, 1),
20417            shared_mem_bytes: 0,
20418        };
20419        if let Some((oq, od)) = q8_out {
20420            // wave-5b: q8-emitting combine — the wo matmul_pre consumes the pair directly.
20421            let fc = if g {
20422                self.func_g("fa_decode_combine_q8_1")
20423            } else {
20424                self.fa_func("fa_decode_combine_q8_1", head_dim)
20425            };
20426            let __s_b2 = self.gpu.stream();
20427            let mut b2 = __s_b2.launch_builder(&fc);
20428            b2.arg(&*part_o)
20429                .arg(&*part_m)
20430                .arg(&*part_l)
20431                .arg(oq)
20432                .arg(od)
20433                .arg(&hd)
20434                .arg(&nh)
20435                .arg(&nsp);
20436            unsafe {
20437                b2.launch(cfg2)?;
20438            }
20439            return Ok(());
20440        }
20441        let fc = if g {
20442            self.func_g("fa_decode_combine_f32")
20443        } else {
20444            self.fa_func("fa_decode_combine_f32", head_dim)
20445        };
20446        let __s_b2 = self.gpu.stream();
20447        let mut b2 = __s_b2.launch_builder(&fc);
20448        b2.arg(&*part_o)
20449            .arg(&*part_m)
20450            .arg(&*part_l)
20451            .arg(o)
20452            .arg(&hd)
20453            .arg(&nh)
20454            .arg(&nsp);
20455        unsafe {
20456            b2.launch(cfg2)?;
20457        }
20458        Ok(())
20459    }
20460
20461    pub fn fa_decode_kvmod(
20462        &self,
20463        q: &CudaSlice<f32>,
20464        k: &cudarc::driver::CudaView<u8>,
20465        v: &cudarc::driver::CudaView<u8>,
20466        o: &mut CudaSlice<f32>,
20467        head_dim: usize,
20468        n_head: usize,
20469        n_head_kv: usize,
20470        t_kv: usize,
20471        scale: f32,
20472        k_tok_bytes: usize,
20473        v_tok_bytes: usize,
20474        g: bool,
20475    ) -> Result<(), Box<dyn std::error::Error>> {
20476        let q_view = q.as_view();
20477        let mut o_view = o.as_view_mut();
20478        self.fa_decode_kvmod_view(
20479            &q_view,
20480            k,
20481            v,
20482            &mut o_view,
20483            head_dim,
20484            n_head,
20485            n_head_kv,
20486            t_kv,
20487            scale,
20488            k_tok_bytes,
20489            v_tok_bytes,
20490            g,
20491        )
20492    }
20493
20494    /// Row-view entry into `fa_decode_kvmod`. The kernel sees the selected Q/output rows as its
20495    /// base pointers, so the launch geometry and arithmetic are identical to the owned-slice entry.
20496    /// Batched fallback callers use this to avoid materializing rows around an otherwise unchanged
20497    /// per-session KV view and FA launch.
20498    #[allow(clippy::too_many_arguments)]
20499    pub fn fa_decode_kvmod_view(
20500        &self,
20501        q: &cudarc::driver::CudaView<f32>,
20502        k: &cudarc::driver::CudaView<u8>,
20503        v: &cudarc::driver::CudaView<u8>,
20504        o: &mut cudarc::driver::CudaViewMut<f32>,
20505        head_dim: usize,
20506        n_head: usize,
20507        n_head_kv: usize,
20508        t_kv: usize,
20509        scale: f32,
20510        k_tok_bytes: usize,
20511        v_tok_bytes: usize,
20512        g: bool,
20513    ) -> Result<(), Box<dyn std::error::Error>> {
20514        // PERF-4: the warp-per-token vec path replaces the scalar element-per-thread fa_decode_f32 —
20515        // warp-per-token fa_decode_vec_q (grid=(n_head_kv,n_splits), block=(32,gqa_ratio)).
20516        // The block dequants each KV tile ONCE into smem (bf16) and broadcasts to all gqa Q-head
20517        // warps -> each KV byte leaves HBM/L2 ~1x/group (vs 4x). ARGS identical; func/grid/block/
20518        // smem/n_splits differ. fa_decode_f32 stays the bit-reference fallback. Combine is shared.
20519        //
20520        // SPLIT-K: the scalar path has grid.x=n_head (32) blocks; the vec path only has
20521        // grid.x=n_head_kv (8). To avoid starving the GPU at mid ctx, the vec path splits MORE
20522        // aggressively (64 keys/split vs 256) so grid.y rises and 8*n_splits fills the SMs.
20523        // At VERY short ctx (t_kv<96) even 1 split can't fill the GPU from 8 KV heads, so the
20524        // broadcast can't beat the scalar path's 4x-more-blocks latency hiding — fall back to
20525        // scalar there (measured crossover: vec 0.68x at t_kv=64, 1.23x at t_kv=96, 2.2x at 256).
20526        // DEFAULT-ON (2026-06-28): clean clock-locked sweep proved vec beats scalar at every
20527        // t_kv>=96 and the gain WIDENS with ctx (graph decode: +9.5% @128, +11.6% @512, +11.8%
20528        // @2048) — the KV-byte-broadcast (4x fewer HBM reads/group) compounds as attention grows.
20529        // MEMRA_NO_FA_VEC forces the scalar bit-reference. Below FA_VEC_MIN_TKV the scalar path's
20530        // 4x-more-blocks (grid.x=n_head=32 vs n_head_kv=8) hides latency better, so keep scalar there.
20531        // g + no-v4: the g-module REGISTER twin mis-decodes the gemma windowed shape
20532        // (root-cause open, jsonl) — only reachable by forcing v4 off (MEMRA_FA_V4_MAX);
20533        // fall to the exact scalar there instead of the broken register arm.
20534        let mut fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && t_kv >= fa_vec_min_tkv();
20535        // hd256 under g: v4-or-scalar. hd128/other under g ride the REGISTER g-lane (the
20536        // dq_K_lane/dq_V_lane macros are format-aware; v3 is format-gated off, v2/smem
20537        // arms are excluded under g). Mirrored in kvmod / fa_decode_dc / fa_geom_eager.
20538        if g && head_dim == 256 && !fa_v4_at(t_kv) {
20539            fa_vec = false;
20540        }
20541        let sp = fa_split_keys(t_kv, n_head_kv);
20542        let n_splits = if fa_vec {
20543            ((t_kv + sp - 1) / sp).max(1)
20544        } else {
20545            ((t_kv + 255) / 256).max(1)
20546        };
20547        let o_len = n_head * n_splits * head_dim;
20548        let ml_len = n_head * n_splits;
20549        let mut part_guard = self.fa_part_pool.lock().unwrap();
20550        if part_guard
20551            .as_ref()
20552            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
20553            .unwrap_or(true)
20554        {
20555            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
20556            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
20557            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
20558            // later live allocations land at those addresses, and the next graph REPLAY writes
20559            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
20560            // output corruption began the burst after the trunk's t_kv growth first realloc'd
20561            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
20562            // the baked addresses alive (single-stream: eager writes the new buffers, replays
20563            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
20564            // (total retired < final size).
20565            let old = part_guard.take();
20566            let (co, cm) = old
20567                .as_ref()
20568                .map(|pp| (pp.0.len(), pp.1.len()))
20569                .unwrap_or((0, 0));
20570            if let Some(old) = old {
20571                self.fa_part_retired.lock().unwrap().push(old);
20572            }
20573            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
20574                eprintln!(
20575                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
20576                    co, o_len, cm, ml_len
20577                );
20578            }
20579            *part_guard = Some((
20580                self.alloc_uninit::<f32>(o_len.max(2 * co))?,
20581                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
20582                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
20583            ));
20584        }
20585        let pg = part_guard.as_mut().unwrap();
20586        self.gpu
20587            .stream()
20588            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
20589        self.gpu
20590            .stream()
20591            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
20592        self.gpu
20593            .stream()
20594            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
20595        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
20596        let (part_o, part_m, part_l) = (&mut *part_o, &mut *part_m, &mut *part_l);
20597        let (hd, nh, nhkv, tkvi, nsp) = (
20598            head_dim as i32,
20599            n_head as i32,
20600            n_head_kv as i32,
20601            t_kv as i32,
20602            n_splits as i32,
20603        );
20604        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
20605        // The vec kernel holds head_dim/32 register accumulators (FA_DEC_MAX_DPL=8 -> head_dim<=256).
20606        // All shipped models use head_dim=256; fall back to scalar for anything wider rather than
20607        // silently truncating the accumulator.
20608        let fa_vec = fa_vec && head_dim <= 512 && head_dim % 32 == 0;
20609        // hd-512 vec crossover (MEMRA_FA512_MIN, default 512): the DPL16 twin wins at depth
20610        // (82.5 -> vec at 1736) but the scalar's more-blocks latency hiding wins at tiny t_kv
20611        // (the same scalar-floor physics as hd256's old 96 floor; short-ctx plain regressed
20612        // 178.4 -> 173.7 when 512 rode vec unconditionally).
20613        let fa512_min = fa512_min_tkv();
20614        // FA-DEEP pick (bit-identical twins, see fa_deep_at): default module only — the
20615        // g-module keeps the v4 pick (its class is not the depth-decay class).
20616        let deep = fa_vec
20617            && head_dim == 256
20618            && fa_v4_at(t_kv)
20619            && !g
20620            && fa_deep_at(t_kv)
20621            && !matches!(fa_v4_mode(), "noB3" | "stage");
20622        let (f, cfg) = if fa_vec && head_dim == 512 && t_kv >= fa512_min {
20623            // gemma4 globals (hd 512): the DPL16 register twin (fa_decode_vec_q body with a
20624            // 16-slot accumulator ceiling). Scalar fallback measured 82.5us/layer at 1736 ctx.
20625            let gqa = (n_head / n_head_kv).max(1) as u32;
20626            let fv = self.fa_func("fa_decode_vec_q_dpl16", head_dim);
20627            (
20628                fv,
20629                LaunchConfig {
20630                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
20631                    block_dim: (32, gqa, 1),
20632                    shared_mem_bytes: 0,
20633                },
20634            )
20635        } else if fa_vec && head_dim <= 256 {
20636            let gqa = (n_head / n_head_kv).max(1) as u32;
20637            // DEEP-CTX smem twin (2026-07-05): the register-dequant path's GQA reuse rides L2,
20638            // which holds to ~8k ctx but dies at 40k (layer KV ~37MB) — the 4 GQA warps then
20639            // re-read every KV byte from DRAM (4x traffic). Above MEMRA_FA_SMEM_TKV (default
20640            // 1024 — the 2026-07-05 crossover re-sweep on real prompts: p3 spec 73.8->79.2 at
20641            // 2048, flat down to 512, p2 +5%, p1/9B unchanged; the ARC-A probe's synthetic
20642            // 2.1x smem-at-all-depths pointed here; 0=never) dispatch the smem-broadcast twin:
20643            // dequant each tile ONCE per block.
20644            // Bit-identical per (token,split): same bf16 round-trip, same accumulation order,
20645            // same partial layout -> same combine. Short/mid ctx keeps the register path (it won
20646            // there by 12x — latency, not bandwidth, rules small KV).
20647            static SMEM_TKV: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
20648            let smem_tkv = *SMEM_TKV.get_or_init(|| {
20649                std::env::var("MEMRA_FA_SMEM_TKV")
20650                    .ok()
20651                    .and_then(|v| v.parse().ok())
20652                    .unwrap_or_else(|| {
20653                        FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
20654                    })
20655            });
20656            if fa_v4_at(t_kv) && head_dim == 256 {
20657                // FA v4 lane (2026-07-10): key-per-lane score phase, zero shuffles per key.
20658                // NEW NUMERIC CONFIG (chunk-serial per-key dot) — battery-arbitrated.
20659                // g (fp8-windowed): the v4 staging is format-aware (2026-07-12) — kf8vf8 module.
20660                let v4name = match fa_v4_mode() {
20661                    "noB3" => "fa_decode_vec_q_v4_noB3", // phase probe (WRONG OUTPUT)
20662                    "stage" => "fa_decode_vec_q_v4_stage", // phase probe (WRONG OUTPUT)
20663                    _ if deep => "fa_decode_vec_q_v4_deep",
20664                    _ => "fa_decode_vec_q_v4",
20665                };
20666                let fv = if g {
20667                    self.func_g(v4name)
20668                } else {
20669                    self.func(v4name)
20670                };
20671                // fa_v4_smem (deep: fa_v4_deep_smem, +640B row pads) + sV (g: raw e4m3 sV
20672                // tile = 1B/elem — half the smem, 3->5 blocks/SM)
20673                let shmem = (if deep { 12160 } else { 11520 }
20674                    + 32 * head_dim * if g { 1 } else { 2 }) as u32;
20675                use cudarc::driver::sys::CUfunction_attribute_enum as A;
20676                fv.set_attribute(
20677                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20678                    shmem as i32,
20679                )?;
20680                (
20681                    fv,
20682                    LaunchConfig {
20683                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
20684                        block_dim: (32, gqa, 1),
20685                        shared_mem_bytes: shmem,
20686                    },
20687                )
20688            } else if fa_v3_active(head_dim) {
20689                // FA v3 lane: dp4a-K hybrid (register-quantized Q, raw q8_0 K, staged-V kept).
20690                // smem = sV only (half of v2's).
20691                let fv = if g {
20692                    self.func_g("fa_decode_vec_q_v3")
20693                } else {
20694                    self.func("fa_decode_vec_q_v3")
20695                };
20696                let shmem = (32 * head_dim * 2) as u32; // sV bf16 [FA_DEC_TILE=32][hd]
20697                (
20698                    fv,
20699                    LaunchConfig {
20700                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
20701                        block_dim: (32, gqa, 1),
20702                        shared_mem_bytes: shmem,
20703                    },
20704                )
20705            } else if fa_v2_on() {
20706                // FAVENDOR lane: llama fattn-vec tile-batched softmax + wide-load staging on
20707                // OUR smem KV broadcast. Replaces BOTH per-key twins when on; same grid/block/
20708                // partials; same 32KB sK+sV tile as the smem twin.
20709                let fv = if g {
20710                    self.func_g("fa_decode_vec_q_v2")
20711                } else {
20712                    self.func("fa_decode_vec_q_v2")
20713                };
20714                let shmem = (2 * 32 * head_dim * 2) as u32; // sK+sV bf16 [FA_DEC_TILE=32][hd]
20715                (
20716                    fv,
20717                    LaunchConfig {
20718                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
20719                        block_dim: (32, gqa, 1),
20720                        shared_mem_bytes: shmem,
20721                    },
20722                )
20723            } else if smem_tkv > 0 && t_kv >= smem_tkv && !g && !(head_dim == 512 && Self::gkv_on())
20724            {
20725                // (fp8 exclusions: the smem twin's V-stage is q5_1-hardcoded — neither the wkv
20726                // windowed layers (g) nor the gkv globals (hd512) may be forced onto it via
20727                // MEMRA_FA_SMEM_TKV; they fall through to the format-clean register/scalar arms.)
20728                let fv = if g {
20729                    self.func_g("fa_decode_vec_q_smem")
20730                } else {
20731                    self.func("fa_decode_vec_q_smem")
20732                };
20733                let shmem = (2 * 32 * head_dim * 2) as u32; // sK+sV bf16 [FA_DEC_TILE=32][hd]
20734                use cudarc::driver::sys::CUfunction_attribute_enum as A;
20735                fv.set_attribute(
20736                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20737                    shmem as i32,
20738                )?;
20739                (
20740                    fv,
20741                    LaunchConfig {
20742                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
20743                        block_dim: (32, gqa, 1),
20744                        shared_mem_bytes: shmem,
20745                    },
20746                )
20747            } else {
20748                // REGISTER-DEQUANT kernel (2026-07-03): per-warp direct q8_0/q5_1 register
20749                // dequant, zero dynamic shared memory.
20750                let fv = if g {
20751                    self.func_g("fa_decode_vec_q")
20752                } else {
20753                    self.func("fa_decode_vec_q")
20754                };
20755                (
20756                    fv,
20757                    LaunchConfig {
20758                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
20759                        block_dim: (32, gqa, 1),
20760                        shared_mem_bytes: 0,
20761                    },
20762                )
20763            }
20764        } else {
20765            // UNIFIED scalar (nullable-ctr symbol shared with graph/stream callers). The
20766            // split ladder value rides along so ns_eff reproduces THIS n_splits in-kernel.
20767            return self.fa_decode_scalar_unified(
20768                q,
20769                k,
20770                v,
20771                o,
20772                head_dim,
20773                n_head,
20774                n_head_kv,
20775                t_kv,
20776                None,
20777                scale,
20778                n_splits,
20779                if fa_vec { sp } else { 256 },
20780                k_tok_bytes,
20781                v_tok_bytes,
20782                g,
20783                part_o,
20784                part_m,
20785                part_l,
20786                None,
20787            );
20788        };
20789        let __s_b = self.gpu.stream();
20790        let mut b = __s_b.launch_builder(&f);
20791        b.arg(q)
20792            .arg(k)
20793            .arg(v)
20794            .arg(&mut *part_o)
20795            .arg(&mut *part_m)
20796            .arg(&mut *part_l)
20797            .arg(&hd)
20798            .arg(&nh)
20799            .arg(&nhkv)
20800            .arg(&tkvi)
20801            .arg(&scale)
20802            .arg(&nsp)
20803            .arg(&ktb)
20804            .arg(&vtb);
20805        unsafe {
20806            b.launch(cfg)?;
20807        }
20808        // (combine re-tile refuted in the fa-deep lane — flat/worse both shapes; the v4
20809        // combine stays for all arms. Receipts research/fa-decode-deep-20260802/.)
20810        let (fc, cfg2) = (
20811            if g {
20812                self.func_g("fa_decode_combine_f32")
20813            } else {
20814                self.fa_func("fa_decode_combine_f32", head_dim)
20815            },
20816            LaunchConfig {
20817                grid_dim: (n_head as u32, 1, 1),
20818                block_dim: (head_dim as u32, 1, 1),
20819                shared_mem_bytes: 0,
20820            },
20821        );
20822        let __s_b2 = self.gpu.stream();
20823        let mut b2 = __s_b2.launch_builder(&fc);
20824        b2.arg(&*part_o)
20825            .arg(&*part_m)
20826            .arg(&*part_l)
20827            .arg(o)
20828            .arg(&hd)
20829            .arg(&nh)
20830            .arg(&nsp);
20831        unsafe {
20832            b2.launch(cfg2)?;
20833        }
20834        Ok(())
20835    }
20836
20837    /// BATCHED-TICK increment 2: ONE fa_decode launch covering ALL B sequences of the
20838    /// batched decode step (blockIdx.z = sequence). Per-seq K/V cache bases ride a device
20839    /// pointer table (`kv_ptrs`, [2B] interleaved k0,v0,...); per-seq key bounds ride the
20840    /// tick's position table (`pos_seq`, T_kv = pos+1). v4-lane only: the CALLER
20841    /// (decode_batch) gates every row through `fa_seqs_eligible` AND one `fa_split_keys`
20842    /// rung (`split_keys`), so each sequence's split partition, key walk and combine order
20843    /// reproduce its per-seq eager v4 program exactly (kernel-check pins seqs-vs-loop bit
20844    /// identity; decode-batch-gate strict pins the whole tick vs decode_step_h).
20845    /// q is the stacked [B, n_head, head_dim] tick buffer read in place (no per-seq q
20846    /// copies); o is written [B, n_head, head_dim] in place (no per-seq a copies).
20847    #[allow(clippy::too_many_arguments)]
20848    pub fn fa_decode_batch_seqs_v4(
20849        &self,
20850        q: &CudaSlice<f32>,
20851        kv_ptrs: &cudarc::driver::CudaView<u64>,
20852        pos_seq: &CudaSlice<i32>,
20853        o: &mut CudaSlice<f32>,
20854        head_dim: usize,
20855        n_head: usize,
20856        n_head_kv: usize,
20857        b_n: usize,
20858        t_kv_max: usize,
20859        scale: f32,
20860        split_keys: usize,
20861        k_tok_bytes: usize,
20862        v_tok_bytes: usize,
20863    ) -> Result<(), Box<dyn std::error::Error>> {
20864        debug_assert!(head_dim == 256, "seqs twin is v4-stamped (hd256 only)");
20865        let n_splits_max = (t_kv_max + split_keys - 1) / split_keys;
20866        let o_len = b_n * n_head * n_splits_max * head_dim;
20867        let ml_len = b_n * n_head * n_splits_max;
20868        let mut part_guard = self.fa_part_pool.lock().unwrap();
20869        if part_guard
20870            .as_ref()
20871            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
20872            .unwrap_or(true)
20873        {
20874            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
20875            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
20876            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
20877            // later live allocations land at those addresses, and the next graph REPLAY writes
20878            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
20879            // output corruption began the burst after the trunk's t_kv growth first realloc'd
20880            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
20881            // the baked addresses alive (single-stream: eager writes the new buffers, replays
20882            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
20883            // (total retired < final size).
20884            let old = part_guard.take();
20885            let (co, cm) = old
20886                .as_ref()
20887                .map(|pp| (pp.0.len(), pp.1.len()))
20888                .unwrap_or((0, 0));
20889            if let Some(old) = old {
20890                self.fa_part_retired.lock().unwrap().push(old);
20891            }
20892            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
20893                eprintln!(
20894                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
20895                    co, o_len, cm, ml_len
20896                );
20897            }
20898            *part_guard = Some((
20899                self.alloc_uninit::<f32>(o_len.max(2 * co))?,
20900                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
20901                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
20902            ));
20903        }
20904        let pg = part_guard.as_mut().unwrap();
20905        self.gpu
20906            .stream()
20907            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
20908        self.gpu
20909            .stream()
20910            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
20911        self.gpu
20912            .stream()
20913            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
20914        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
20915        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
20916        let (nspm, spk) = (n_splits_max as i32, split_keys as i32);
20917        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
20918        let gqa = (n_head / n_head_kv).max(1) as u32;
20919        let f = self.func("fa_decode_vec_q_seqs_v4");
20920        // fa_v4_smem (11520B) + sV bf16 tile — the v4 eager arm's sizing on the default module.
20921        let shmem = (11520 + 32 * head_dim * 2) as u32;
20922        use cudarc::driver::sys::CUfunction_attribute_enum as A;
20923        f.set_attribute(
20924            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20925            shmem as i32,
20926        )?;
20927        let cfg = LaunchConfig {
20928            grid_dim: (n_head_kv as u32, n_splits_max as u32, b_n as u32),
20929            block_dim: (32, gqa, 1),
20930            shared_mem_bytes: shmem,
20931        };
20932        {
20933            let __s_b = self.gpu.stream();
20934            let mut b = __s_b.launch_builder(&f);
20935            b.arg(q)
20936                .arg(kv_ptrs)
20937                .arg(pos_seq)
20938                .arg(&mut *part_o)
20939                .arg(&mut *part_m)
20940                .arg(&mut *part_l)
20941                .arg(&hd)
20942                .arg(&nh)
20943                .arg(&nhkv)
20944                .arg(&scale)
20945                .arg(&nspm)
20946                .arg(&spk)
20947                .arg(&ktb)
20948                .arg(&vtb);
20949            unsafe {
20950                b.launch(cfg)?;
20951            }
20952        }
20953        let fc = self.func("fa_decode_combine_seqs");
20954        let cfg2 = LaunchConfig {
20955            grid_dim: (n_head as u32, b_n as u32, 1),
20956            block_dim: (head_dim as u32, 1, 1),
20957            shared_mem_bytes: 0,
20958        };
20959        let __s_b2 = self.gpu.stream();
20960        let mut b2 = __s_b2.launch_builder(&fc);
20961        b2.arg(&*part_o)
20962            .arg(&*part_m)
20963            .arg(&*part_l)
20964            .arg(o)
20965            .arg(&hd)
20966            .arg(&nh)
20967            .arg(pos_seq)
20968            .arg(&nspm)
20969            .arg(&spk);
20970        unsafe {
20971            b2.launch(cfg2)?;
20972        }
20973        Ok(())
20974    }
20975
20976    /// BATCHED-TICK increment 2: z-batched decode KV append — one launch appends this
20977    /// step's B rows, each into ITS OWN sequence cache at slot pos_seq[z], through the same
20978    /// [2B] interleaved pointer table the seqs FA reads. Each (block, z) warp executes the
20979    /// per-token appender's exact warp program on row z of the stacked [B, kv_dim] k/v —
20980    /// written cache bytes are BIT-IDENTICAL to the B per-seq calls it replaces
20981    /// (kernel-check pins the bytes). Default flash module only (callers exclude fp8-KV).
20982    #[allow(clippy::too_many_arguments)]
20983    pub fn append_kv_quantized_seqs(
20984        &self,
20985        k_rows: &CudaSlice<f32>,
20986        v_rows: &CudaSlice<f32>,
20987        kv_ptrs: &cudarc::driver::CudaView<u64>,
20988        pos_seq: &CudaSlice<i32>,
20989        b_n: usize,
20990        kv_dim_k: usize,
20991        kv_dim_v: usize,
20992        k_tok_bytes: usize,
20993        v_tok_bytes: usize,
20994    ) -> Result<(), Box<dyn std::error::Error>> {
20995        let f = self.func("append_quantize_kv_q8_0_q5_1_seqs");
20996        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
20997        let cfg = LaunchConfig {
20998            grid_dim: (nblk, b_n as u32, 1),
20999            block_dim: (32, 1, 1),
21000            shared_mem_bytes: 0,
21001        };
21002        let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
21003        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
21004        let __s_b = self.gpu.stream();
21005        let mut b = __s_b.launch_builder(&f);
21006        b.arg(k_rows)
21007            .arg(v_rows)
21008            .arg(kv_ptrs)
21009            .arg(pos_seq)
21010            .arg(&kdk)
21011            .arg(&kdv)
21012            .arg(&ktb)
21013            .arg(&vtb);
21014        unsafe {
21015            b.launch(cfg)?;
21016        }
21017        Ok(())
21018    }
21019
21020    /// True iff the MULTI-ROW verify FA (`fa_decode_rows`) is usable for a verify batch whose
21021    /// FIRST row attends `base_len + 1` keys: every row must take the SAME kernel eager decode
21022    /// would (the vec path) — mirrors fa_decode's gate exactly (MEMRA_NO_FA_VEC + FA_VEC_MIN_TKV +
21023    /// head_dim), evaluated at the MINIMUM row bound so no row could have picked scalar.
21024    /// MEMRA_FA_ROWS_OFF=1 is the A/B + fallback seam (per-row loop).
21025    pub fn fa_rows_eligible(&self, base_len: usize, head_dim: usize) -> bool {
21026        std::env::var("MEMRA_NO_FA_VEC").is_err()
21027            && std::env::var("MEMRA_FA_ROWS_OFF").is_err()
21028            && base_len + 1 >= fa_vec_min_tkv()
21029            && head_dim <= 256
21030            && head_dim % 32 == 0
21031    }
21032
21033    /// MULTI-ROW verify FA: run fa_decode_vec_q's EXACT per-row program for T causal query rows
21034    /// (row r attends keys [0..base_len+r+1)) in ONE kernel launch with grid.z = row, plus ONE
21035    /// row-batched combine. Replaces the T separate (fa_decode + combine) launches of the spec
21036    /// verify — same per-row split partition (n_splits_r = ceil(t_kv_r/split_keys), the
21037    /// fa_split_keys formula), same key-walk order, same reduce shapes => bit-identical outputs
21038    /// per row (kernel-check pins rows-vs-loop byte identity; run-spec is the end gate).
21039    /// Caller must have checked `fa_rows_eligible(base_len, head_dim)`.
21040    /// q is the verify's token-major [T, n_head, head_dim] stack; o is written [T, n_head, head_dim].
21041    #[allow(clippy::too_many_arguments)]
21042    pub fn fa_decode_rows(
21043        &self,
21044        q: &CudaSlice<f32>,
21045        k: &cudarc::driver::CudaView<u8>,
21046        v: &cudarc::driver::CudaView<u8>,
21047        o: &mut CudaSlice<f32>,
21048        head_dim: usize,
21049        n_head: usize,
21050        n_head_kv: usize,
21051        base_len: usize,
21052        t: usize,
21053        scale: f32,
21054        k_tok_bytes: usize,
21055        v_tok_bytes: usize,
21056        // hd512 dpl16 twin is DEVICE-LEN (graph arc): base_dev/plus feed the
21057        // kernel; host base_len keeps sizing the splits/partials. hd256 twins
21058        // keep the host arg. None is a bug for hd512 (asserted below).
21059        base_dev: Option<(&CudaSlice<i32>, i32)>,
21060        // K and V planes hold the same values (gemma globals, wv:=wk): pick
21061        // the _kv twin — V plane never read, value rides the q8_0 key dq.
21062        kv_shared: bool,
21063        // this layer's cache is e4m3 (gemma windowed under wkv): resolve the
21064        // hd256 rows kernel from the kf8vf8 module. PER-CALL — a global env
21065        // check here hijacked qwen/kernel-check hd256 rows (8 FAILs, 230ebbe).
21066        g: bool,
21067        // t=1 decode arm only: emit (int8, per-32 scales) from the dc combine
21068        // (hd512 path) — the standalone quantize launch folds away.
21069        mut q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
21070    ) -> Result<(), Box<dyn std::error::Error>> {
21071        debug_assert!(base_len + 1 >= fa_vec_min_tkv() && head_dim <= 512 && head_dim % 32 == 0);
21072        let t_kv_max = base_len + t; // LAST row's key bound
21073        let mut sp = fa_split_keys(t_kv_max, n_head_kv); // env/default — same value every row
21074        // hd512 split override (MEMRA_FA_SP512, 2026-07-11): gemma globals have n_head_kv=2 so
21075        // the grid is (2 x n_splits) — at depth ~29 splits = 58 blocks on 82 SMs (half idle,
21076        // rows_dpl16 8x off its byte floor). EVERY gemma hd512 caller shares THIS wrapper
21077        // (parity law), so the partition is freely tunable — verify and decode move together.
21078        if head_dim == 512 {
21079            static SP512: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
21080            // default 16 (2026-07-11 depth sweep, N=2: plain 155.4->156.5, depth spec
21081            // 236.9->250.4; 12/24/32 all worse). hd512 exists only on gemma globals.
21082            let v = *SP512.get_or_init(|| {
21083                std::env::var("MEMRA_FA_SP512")
21084                    .ok()
21085                    .and_then(|x| x.parse().ok())
21086                    .unwrap_or(0)
21087            });
21088            sp = if v >= 8 {
21089                v
21090            } else {
21091                FA_SP512_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
21092            };
21093        }
21094        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
21095        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
21096        let gqa = (n_head / n_head_kv).max(1) as u32;
21097        // LADDER-RUNG STRADDLE FIX (issue #10, 2026-07-13, rtx6000-proven): one sp for every row
21098        // diverges from eager decode when a split-ladder rung falls INSIDE the batch — row r's
21099        // eager twin used fa_split_keys(t_kv_r), the batch used fa_split_keys(t_kv_max), and
21100        // the different partition changes the combine's FP order (greedy tie flips at depth;
21101        // MEMRA_FA_SPLIT=64 pin -> PASS on the exact rtx6000 failing config). Fix: group
21102        // consecutive rows by their OWN ladder value and launch once per group — each row then
21103        // executes the exact per-row program eager ran. Rungs land once per doubling, so this
21104        // is 1 launch in the common case and 2 on a crossing round. hd512 keeps one group (its
21105        // sp override is t_kv-independent by construction).
21106        let mut groups: Vec<(usize, usize, usize)> = Vec::new(); // (row0, t_g, sp_g)
21107        if head_dim == 512 || fa_split_keys(base_len + 1, n_head_kv) == sp {
21108            groups.push((0, t, sp));
21109        } else {
21110            let mut r0 = 0usize;
21111            while r0 < t {
21112                let sp_g = fa_split_keys(base_len + r0 + 1, n_head_kv);
21113                let mut r1 = r0 + 1;
21114                while r1 < t && fa_split_keys(base_len + r1 + 1, n_head_kv) == sp_g {
21115                    r1 += 1;
21116                }
21117                groups.push((r0, r1 - r0, sp_g));
21118                r0 = r1;
21119            }
21120        }
21121        // Deep-ctx smem twin for the VERIFY rows (2026-07-05): same threshold + rationale as
21122        // fa_decode's dispatch — at 40k the register path's GQA L2-reuse premise is dead and the
21123        // verify multiplies the 4x DRAM re-read by T rows. Bit-identical per (row,token,split).
21124        static SMEM_TKV_R: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
21125        let smem_tkv = *SMEM_TKV_R.get_or_init(|| {
21126            std::env::var("MEMRA_FA_SMEM_TKV")
21127                .ok()
21128                .and_then(|v| v.parse().ok())
21129                .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
21130        });
21131        let v4 = fa_v4_at(base_len + t) && head_dim == 256;
21132        let v3 = fa_v3_active(head_dim);
21133        let smem_rows =
21134            head_dim <= 256 && !v3 && !fa_v2_on() && smem_tkv > 0 && t_kv_max >= smem_tkv;
21135        // kv_shared twin RETIRED (2026-07-11 depth run-gen gate): the wv:=wk premise fails
21136        // POST-cache — cached K is k-normed+roped, cached V is not; the twin fed roped keys
21137        // in as values. Verify/decode/stream gates were blind (both sides shared the wrong
21138        // symbol — the parity law's blind spot); only prefill-vs-decode at depth caught it.
21139        let _ = kv_shared;
21140        // i2 twin: 2-key interleaved walk (MEMRA_FA_I2=0 reverts). i4 probed NEGATIVE
21141        // (157.3 vs 161.2 depth plain — register pressure past i2's sweet spot; jsonl).
21142        let i2 = head_dim == 512 && std::env::var("MEMRA_FA_I2").as_deref() != Ok("0");
21143        // v4-hd512 (MEMRA_FA_V512=1 opt-in, 2026-07-14): the v4 key-per-lane recipe on the
21144        // globals lane (depth profile: i2 ~4.6x off its byte floor — the v3-class
21145        // reduce-per-key latency signature). NEW NUMERIC CONFIG shared by every hd512
21146        // caller (decode+verify flip together); run-gen argmax + acceptance arbitrate.
21147        // T-BATCHED hd512 (DEFAULT ON 2026-07-14, MEMRA_FA_TB512=0 seam): one block per
21148        // (kv_head, split) stages its tile once and loops the rows over it — kills the
21149        // x t DRAM re-read of the full-ctx globals (depth cell +1.4%, plain flat, N=3
21150        // interleaved). FIXED absolute partition = NEW NUMERIC for the combine order,
21151        // shared by every hd512 caller through this wrapper (decode+verify flip together;
21152        // depth stream identical, acceptance unshifted, spec 256/256 x3 models).
21153        // Requires sp <= 32 (single staged tile; acc reused per row). The z-form v4_512
21154        // sibling (in-kernel dp4a port alone) probed FLAT — hd512 was DRAM-re-read-bound,
21155        // not unpack-bound; jsonl 2026-07-14.
21156        static TB512: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
21157        // gqa <= 16 = fa_v4_smem_512's q-array capacity; past it fall to the register twins.
21158        let tb512 = head_dim == 512
21159            && sp <= 32
21160            && n_head / n_head_kv.max(1) <= 16
21161            && *TB512.get_or_init(|| std::env::var("MEMRA_FA_TB512").as_deref() != Ok("0"));
21162        let fname = if tb512 {
21163            "fa_decode_vec_q_rows_v4_512_tb"
21164        } else if i2 {
21165            "fa_decode_vec_q_rows_dpl16_i2"
21166        } else if head_dim == 512 {
21167            "fa_decode_vec_q_rows_dpl16"
21168        }
21169        // gemma globals (parity law)
21170        else if v4 {
21171            "fa_decode_vec_q_rows_v4"
21172        } else if v3 {
21173            "fa_decode_vec_q_rows_v3"
21174        } else if fa_v2_on() {
21175            "fa_decode_vec_q_rows_v2"
21176        } else if smem_rows {
21177            "fa_decode_vec_q_rows_smem"
21178        } else {
21179            "fa_decode_vec_q_rows"
21180        };
21181        let f = if head_dim == 512 {
21182            self.fa_func(fname, head_dim)
21183        } else if g {
21184            // FP8-WINDOWED: hd256 rows over an e4m3 cache — kf8vf8 module, SAME symbol
21185            // choice as decode's kvmod dispatch (parity law: excluding v4 here paired
21186            // g-module rows against decode's g-module v4 — different programs, short-VG
21187            // maxdiff 2.0 / spec stream 0/128, 2026-07-12). rows_v4 is format-aware
21188            // since fda9790; only the smem twin stays excluded (V-stage q5_1-only).
21189            // hd128 (qwen fp8-KV) lands on the base/register rows via fname — the
21190            // dq macros are format-aware.
21191            self.func_g(if smem_rows {
21192                "fa_decode_vec_q_rows"
21193            } else {
21194                fname
21195            })
21196        } else {
21197            self.func(fname)
21198        };
21199        let shmem = if tb512 {
21200            // fa_v4_smem_512 (q 9KB gqa<=16 + k tile 18KB) + sV 32*512 (e4m3 module halves it)
21201            let gk = Self::gkv_on();
21202            let sh =
21203                (8192 + 1024 + 32 * 512 + 32 * 64 + 32 * head_dim * if gk { 1 } else { 2 }) as u32;
21204            use cudarc::driver::sys::CUfunction_attribute_enum as A;
21205            f.set_attribute(
21206                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21207                sh as i32,
21208            )?;
21209            sh
21210        } else if v4 || v3 || smem_rows || fa_v2_on() {
21211            // v4: fa_v4_smem (11.5KB) + sV; v3 stages sV only; v2/smem twins stage sK+sV.
21212            let sh = (if v4 {
21213                11520 + 32 * head_dim * if g { 1 } else { 2 }
21214            } else if v3 {
21215                32 * head_dim * 2
21216            } else {
21217                2 * 32 * head_dim * 2
21218            }) as u32;
21219            use cudarc::driver::sys::CUfunction_attribute_enum as A;
21220            f.set_attribute(
21221                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21222                sh as i32,
21223            )?;
21224            sh
21225        } else {
21226            0
21227        };
21228        // Per-GROUP launches (single group in the common case — identical to the pre-fix
21229        // single launch there): each group gets its own partials (the rows kernel indexes
21230        // partials by its LOCAL grid.z row) and q/o row-offset views.
21231        for &(r0, t_g, sp_g) in &groups {
21232            let n_splits_g = (base_len + r0 + t_g).div_ceil(sp_g);
21233            let (nspm, spk) = (n_splits_g as i32, sp_g as i32);
21234            let base_i = (base_len + r0) as i32;
21235            let o_len = t_g * n_head * n_splits_g * head_dim;
21236            let ml_len = t_g * n_head * n_splits_g;
21237            let mut part_guard = self.fa_part_pool.lock().unwrap();
21238            if part_guard
21239                .as_ref()
21240                .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
21241                .unwrap_or(true)
21242            {
21243                // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
21244                // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
21245                // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
21246                // later live allocations land at those addresses, and the next graph REPLAY writes
21247                // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
21248                // output corruption began the burst after the trunk's t_kv growth first realloc'd
21249                // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
21250                // the baked addresses alive (single-stream: eager writes the new buffers, replays
21251                // touch only the old — never concurrently). Doubling growth bounds retired VRAM
21252                // (total retired < final size).
21253                let old = part_guard.take();
21254                let (co, cm) = old
21255                    .as_ref()
21256                    .map(|pp| (pp.0.len(), pp.1.len()))
21257                    .unwrap_or((0, 0));
21258                if let Some(old) = old {
21259                    self.fa_part_retired.lock().unwrap().push(old);
21260                }
21261                if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
21262                    eprintln!(
21263                        "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
21264                        co, o_len, cm, ml_len
21265                    );
21266                }
21267                *part_guard = Some((
21268                    self.alloc_uninit::<f32>(o_len.max(2 * co))?,
21269                    self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
21270                    self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
21271                ));
21272            }
21273            let pg = part_guard.as_mut().unwrap();
21274            self.gpu
21275                .stream()
21276                .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
21277            self.gpu
21278                .stream()
21279                .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
21280            self.gpu
21281                .stream()
21282                .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
21283            let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
21284            let (part_o, part_m, part_l) = (&mut *part_o, &mut *part_m, &mut *part_l);
21285            let qv = self.view(q, t * n_head * head_dim);
21286            let q_g = qv.slice(r0 * n_head * head_dim..(r0 + t_g) * n_head * head_dim);
21287            let cfg = LaunchConfig {
21288                grid_dim: (n_head_kv as u32, n_splits_g as u32, t_g as u32),
21289                block_dim: (32, gqa, 1),
21290                shared_mem_bytes: shmem,
21291            };
21292            {
21293                let __s_b = self.gpu.stream();
21294                let mut b = __s_b.launch_builder(&f);
21295                if tb512 {
21296                    // rows-inner launch: grid.z dropped, the kernel loops n_rows itself.
21297                    let (bd, plus) =
21298                        base_dev.expect("hd512 rows twin requires a device base counter");
21299                    let plus_g = plus + r0 as i32;
21300                    let nr = t_g as i32;
21301                    if Self::pdl_on() && Self::pdl_wb_on() {
21302                        // wave-B2b: flavor mirrors fa_func(fname, 512) = gkv.
21303                        use cudarc::driver::{DevicePtr, DevicePtrMut};
21304                        let s = &self.gpu.stream();
21305                        let (pq, _b0) = q_g.device_ptr(s);
21306                        let (pk, _b1) = k.device_ptr(s);
21307                        let (pv, _b2) = v.device_ptr(s);
21308                        let (po, _b3) = part_o.device_ptr_mut(s);
21309                        let (pm, _b4) = part_m.device_ptr_mut(s);
21310                        let (pl, _b5) = part_l.device_ptr_mut(s);
21311                        let (pb, _b6) = bd.device_ptr(s);
21312                        let mut ps = [
21313                            &pq as *const _ as *mut std::ffi::c_void,
21314                            &pk as *const _ as *mut _,
21315                            &pv as *const _ as *mut _,
21316                            &po as *const _ as *mut _,
21317                            &pm as *const _ as *mut _,
21318                            &pl as *const _ as *mut _,
21319                            &hd as *const _ as *mut _,
21320                            &nh as *const _ as *mut _,
21321                            &nhkv as *const _ as *mut _,
21322                            &pb as *const _ as *mut _,
21323                            &plus_g as *const _ as *mut _,
21324                            &scale as *const _ as *mut _,
21325                            &nspm as *const _ as *mut _,
21326                            &spk as *const _ as *mut _,
21327                            &ktb as *const _ as *mut _,
21328                            &vtb as *const _ as *mut _,
21329                            &nr as *const _ as *mut _,
21330                        ];
21331                        unsafe {
21332                            self.launch_pdl_flash(
21333                                Self::gkv_on(),
21334                                "fa_decode_vec_q_rows_v4_512_tb",
21335                                (n_head_kv as u32, n_splits_g as u32, 1),
21336                                (32, gqa, 1),
21337                                shmem,
21338                                &mut ps,
21339                            )?;
21340                        }
21341                    } else {
21342                        let cfg_tb = LaunchConfig {
21343                            grid_dim: (n_head_kv as u32, n_splits_g as u32, 1),
21344                            block_dim: (32, gqa, 1),
21345                            shared_mem_bytes: shmem,
21346                        };
21347                        b.arg(&q_g)
21348                            .arg(k)
21349                            .arg(v)
21350                            .arg(&mut *part_o)
21351                            .arg(&mut *part_m)
21352                            .arg(&mut *part_l)
21353                            .arg(&hd)
21354                            .arg(&nh)
21355                            .arg(&nhkv)
21356                            .arg(bd)
21357                            .arg(&plus_g)
21358                            .arg(&scale)
21359                            .arg(&nspm)
21360                            .arg(&spk)
21361                            .arg(&ktb)
21362                            .arg(&vtb)
21363                            .arg(&nr);
21364                        unsafe {
21365                            b.launch(cfg_tb)?;
21366                        }
21367                    }
21368                } else if head_dim == 512 {
21369                    let (bd, plus) =
21370                        base_dev.expect("hd512 rows twin requires a device base counter");
21371                    let plus_g = plus + r0 as i32;
21372                    b.arg(&q_g)
21373                        .arg(k)
21374                        .arg(v)
21375                        .arg(&mut *part_o)
21376                        .arg(&mut *part_m)
21377                        .arg(&mut *part_l)
21378                        .arg(&hd)
21379                        .arg(&nh)
21380                        .arg(&nhkv)
21381                        .arg(bd)
21382                        .arg(&plus_g)
21383                        .arg(&scale)
21384                        .arg(&nspm)
21385                        .arg(&spk)
21386                        .arg(&ktb)
21387                        .arg(&vtb);
21388                    unsafe {
21389                        b.launch(cfg)?;
21390                    }
21391                } else {
21392                    b.arg(&q_g)
21393                        .arg(k)
21394                        .arg(v)
21395                        .arg(&mut *part_o)
21396                        .arg(&mut *part_m)
21397                        .arg(&mut *part_l)
21398                        .arg(&hd)
21399                        .arg(&nh)
21400                        .arg(&nhkv)
21401                        .arg(&base_i)
21402                        .arg(&scale)
21403                        .arg(&nspm)
21404                        .arg(&spk)
21405                        .arg(&ktb)
21406                        .arg(&vtb);
21407                    unsafe {
21408                        b.launch(cfg)?;
21409                    }
21410                }
21411            }
21412            let cfg2 = LaunchConfig {
21413                grid_dim: (n_head as u32, t_g as u32, 1),
21414                block_dim: (head_dim as u32, 1, 1),
21415                shared_mem_bytes: 0,
21416            };
21417            let mut o_g = o.slice_mut(r0 * n_head * head_dim..(r0 + t_g) * n_head * head_dim);
21418            if head_dim == 512 {
21419                // device-len combine (shared by verify/eager/graph — parity by symbol): the
21420                // per-row n_splits derives from the SAME counter the rows kernel read.
21421                let (bd, plus) = base_dev.unwrap();
21422                let plus_g = plus + r0 as i32;
21423                if let Some((oq, od)) = q8_out.as_mut() {
21424                    // wave-5b port (2026-07-23, t=1 decode only): q8-emitting dc combine.
21425                    debug_assert!(t == 1, "rows q8 emit is a t=1 decode arm");
21426                    if Self::pdl_on() && Self::pdl_wb_on() {
21427                        // wave-B2: flavor mirrors fa_func (hd512 + gkv → kf8vf8).
21428                        use cudarc::driver::{DevicePtr, DevicePtrMut};
21429                        let s = &self.gpu.stream();
21430                        let (po, _g0) = part_o.device_ptr(s);
21431                        let (pm, _g1) = part_m.device_ptr(s);
21432                        let (pl, _g2) = part_l.device_ptr(s);
21433                        let (pq, _g3) = oq.device_ptr_mut(s);
21434                        let (pd, _g4) = od.device_ptr_mut(s);
21435                        let (pb, _g5) = bd.device_ptr(s);
21436                        let mut ps = [
21437                            &po as *const _ as *mut std::ffi::c_void,
21438                            &pm as *const _ as *mut _,
21439                            &pl as *const _ as *mut _,
21440                            &pq as *const _ as *mut _,
21441                            &pd as *const _ as *mut _,
21442                            &hd as *const _ as *mut _,
21443                            &nh as *const _ as *mut _,
21444                            &pb as *const _ as *mut _,
21445                            &plus_g as *const _ as *mut _,
21446                            &nspm as *const _ as *mut _,
21447                            &spk as *const _ as *mut _,
21448                        ];
21449                        unsafe {
21450                            self.launch_pdl_flash(
21451                                Self::gkv_on(),
21452                                "fa_decode_combine_rows_dc_q8_1",
21453                                cfg2.grid_dim,
21454                                cfg2.block_dim,
21455                                0,
21456                                &mut ps,
21457                            )?;
21458                        }
21459                        continue;
21460                    }
21461                    let fc = self.fa_func("fa_decode_combine_rows_dc_q8_1", head_dim);
21462                    let __s_b2 = self.gpu.stream();
21463                    let mut b2 = __s_b2.launch_builder(&fc);
21464                    b2.arg(&*part_o)
21465                        .arg(&*part_m)
21466                        .arg(&*part_l)
21467                        .arg(&mut **oq)
21468                        .arg(&mut **od)
21469                        .arg(&hd)
21470                        .arg(&nh)
21471                        .arg(bd)
21472                        .arg(&plus_g)
21473                        .arg(&nspm)
21474                        .arg(&spk);
21475                    unsafe {
21476                        b2.launch(cfg2)?;
21477                    }
21478                    continue;
21479                }
21480                let fc = self.fa_func("fa_decode_combine_rows_dc", head_dim);
21481                let __s_b2 = self.gpu.stream();
21482                let mut b2 = __s_b2.launch_builder(&fc);
21483                b2.arg(&*part_o)
21484                    .arg(&*part_m)
21485                    .arg(&*part_l)
21486                    .arg(&mut o_g)
21487                    .arg(&hd)
21488                    .arg(&nh)
21489                    .arg(bd)
21490                    .arg(&plus_g)
21491                    .arg(&nspm)
21492                    .arg(&spk);
21493                unsafe {
21494                    b2.launch(cfg2)?;
21495                }
21496            } else {
21497                // q8 emit is wired for the hd512 dc-combine arm only — a Some here would
21498                // leave the caller's pair unwritten (consumer would read garbage).
21499                assert!(
21500                    q8_out.is_none(),
21501                    "rows q8 emit requires the hd512 dc combine"
21502                );
21503                let fc = self.func("fa_decode_combine_rows");
21504                let __s_b2 = self.gpu.stream();
21505                let mut b2 = __s_b2.launch_builder(&fc);
21506                b2.arg(&*part_o)
21507                    .arg(&*part_m)
21508                    .arg(&*part_l)
21509                    .arg(&mut o_g)
21510                    .arg(&hd)
21511                    .arg(&nh)
21512                    .arg(&base_i)
21513                    .arg(&nspm)
21514                    .arg(&spk);
21515                unsafe {
21516                    b2.launch(cfg2)?;
21517                }
21518            }
21519        }
21520        Ok(())
21521    }
21522
21523    /// WINDOWED verify rows (gemma R6 deep-ctx): every row attends exactly `window` keys —
21524    /// bit-identical per row to the T=1 decode's fa_decode over the window VIEW. Caller gates
21525    /// base_len + 1 >= window (no under-window rows) and head_dim == 256 (v4 stamp).
21526    #[allow(clippy::too_many_arguments)]
21527    pub fn fa_decode_rows_w(
21528        &self,
21529        q: &CudaSlice<f32>,
21530        k: &cudarc::driver::CudaView<u8>,
21531        v: &cudarc::driver::CudaView<u8>,
21532        o: &mut CudaSlice<f32>,
21533        head_dim: usize,
21534        n_head: usize,
21535        n_head_kv: usize,
21536        base_dev: &CudaSlice<i32>,
21537        base_plus: i32,
21538        t: usize,
21539        scale: f32,
21540        window: usize,
21541        k_tok_bytes: usize,
21542        v_tok_bytes: usize,
21543        q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
21544    ) -> Result<(), Box<dyn std::error::Error>> {
21545        // DEVICE-LEN (graph arc step 1, 2026-07-11): the causal base rides an i32 counter
21546        // (kernel T_kv = dev[0] + base_plus + r + 1) so depth graphs can replay with len
21547        // advancing on-device. dc paths pass kvl.len_d with plus=-1; verify/eager sync the
21548        // counter with one async set_i32_one first. Partials/splits size from `window` (host).
21549        debug_assert!(head_dim == 256);
21550        // windowed split (MEMRA_FA_SPW, default 32 — re-swept 2026-07-12 under the raw-e4m3 sV
21551        // occupancy ceiling (4 blocks/SM): t=1 decode is GRID-limited (win/sp splits x nkv
21552        // blocks), so smaller splits fill the ceiling — 1.7k 174.4/174.0 vs 48's 170.7/170.3,
21553        // 4.9k 159.8 vs 157.4 (N=2 interleaved, stable window). Spec serving prefers 64
21554        // (verify t=K+1 fills the grid via grid.z=t; depth K=7 281.3 vs 249.3 at 32) — set
21555        // MEMRA_FA_SPW=64 there, same config law as MEMRA_GEMMA_GKV=0. MUST be one value for
21556        // ALL widths: a t-keyed probe broke decode-vs-verify combine order (stream 9/128).
21557        let sp = {
21558            static SPW: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
21559            let v = *SPW.get_or_init(|| {
21560                std::env::var("MEMRA_FA_SPW")
21561                    .ok()
21562                    .and_then(|x| x.parse().ok())
21563                    .unwrap_or(0)
21564            });
21565            if v >= 8 {
21566                v
21567            } else {
21568                FA_SPW_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
21569            }
21570        };
21571        let n_splits_max = (window + sp - 1) / sp;
21572        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
21573        let (nspm, spk, wini) = (n_splits_max as i32, sp as i32, window as i32);
21574        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
21575        let gqa = (n_head / n_head_kv).max(1) as u32;
21576        let o_len = t * n_head * n_splits_max * head_dim;
21577        let ml_len = t * n_head * n_splits_max;
21578        let mut part_guard = self.fa_part_pool.lock().unwrap();
21579        if part_guard
21580            .as_ref()
21581            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
21582            .unwrap_or(true)
21583        {
21584            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
21585            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
21586            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
21587            // later live allocations land at those addresses, and the next graph REPLAY writes
21588            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
21589            // output corruption began the burst after the trunk's t_kv growth first realloc'd
21590            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
21591            // the baked addresses alive (single-stream: eager writes the new buffers, replays
21592            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
21593            // (total retired < final size).
21594            let old = part_guard.take();
21595            let (co, cm) = old
21596                .as_ref()
21597                .map(|pp| (pp.0.len(), pp.1.len()))
21598                .unwrap_or((0, 0));
21599            if let Some(old) = old {
21600                self.fa_part_retired.lock().unwrap().push(old);
21601            }
21602            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
21603                eprintln!(
21604                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
21605                    co, o_len, cm, ml_len
21606                );
21607            }
21608            *part_guard = Some((
21609                self.alloc_uninit::<f32>(o_len.max(2 * co))?,
21610                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
21611                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
21612            ));
21613        }
21614        let pg = part_guard.as_mut().unwrap();
21615        self.gpu
21616            .stream()
21617            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
21618        self.gpu
21619            .stream()
21620            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
21621        self.gpu
21622            .stream()
21623            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
21624        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
21625        // Lane pick: decode AND verify both land here in the windowed regime (parity law —
21626        // hybrid_forward verify_attn), so the pick only needs internal consistency, not
21627        // clone-of-decode bit fidelity (SASS-proven impossible for textually identical
21628        // kernels, jsonl 2026-07-10). v4 under the threshold; smem twin at/above the smem
21629        // floor (deep-ctx broadcast win); register twin between.
21630        static SMEM_TKV_W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
21631        let smem_tkv = *SMEM_TKV_W.get_or_init(|| {
21632            std::env::var("MEMRA_FA_SMEM_TKV")
21633                .ok()
21634                .and_then(|v| v.parse().ok())
21635                .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
21636        });
21637        // MULTI-ROW v4: resurrected 2026-07-14 (the '33 tok/s collapse' was a paired-map
21638        // partial-write bug, not the mechanism) and falsified HONESTLY at gqa 2: bit-exact
21639        // but −1.7% on the 31B depth cell — the sp helper warp already hides staging
21640        // in-block, and mr trades L2-cheap redundant bytes for serialized per-warp gqa
21641        // score/B3 chains. Arm deleted; jsonl row 2026-07-14 is the record.
21642        use cudarc::driver::sys::CUfunction_attribute_enum as A;
21643        // FP8-WINDOWED (wkv): the v4 family is format-aware (2026-07-12 KFMT/VFMT staging
21644        // arms) — wkv rides the SAME lane logic, resolved from the kf8vf8 module. One symbol
21645        // per (lane, format-module) keeps parity structural; the old register-i2 detour
21646        // (-33%) is retired.
21647        let wg = Self::wkv_on();
21648        // STAGING-PARALLEL v4 (MEMRA_FA_SPW2, default ON at gqa==1): warp 1 = staging helper
21649        // (v4 is 61% staging); score phases identical to v4_w. Same symbol all t.
21650        let sp2 =
21651            gqa <= 4 && fa_v4_at(window) && std::env::var("MEMRA_FA_SPW2").as_deref() != Ok("0");
21652        if sp2 {
21653            let sh = (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32;
21654            if Self::pdl_on() && Self::pdl_wb_on() {
21655                // wave-B2b: flavor mirrors wg.
21656                use cudarc::driver::{DevicePtr, DevicePtrMut};
21657                let s = &self.gpu.stream();
21658                let (pq, _b0) = q.device_ptr(s);
21659                let (pk, _b1) = k.device_ptr(s);
21660                let (pv, _b2) = v.device_ptr(s);
21661                let (po, _b3) = part_o.device_ptr_mut(s);
21662                let (pm, _b4) = part_m.device_ptr_mut(s);
21663                let (pl, _b5) = part_l.device_ptr_mut(s);
21664                let (pb, _b6) = base_dev.device_ptr(s);
21665                let mut ps = [
21666                    &pq as *const _ as *mut std::ffi::c_void,
21667                    &pk as *const _ as *mut _,
21668                    &pv as *const _ as *mut _,
21669                    &po as *const _ as *mut _,
21670                    &pm as *const _ as *mut _,
21671                    &pl as *const _ as *mut _,
21672                    &hd as *const _ as *mut _,
21673                    &nh as *const _ as *mut _,
21674                    &nhkv as *const _ as *mut _,
21675                    &pb as *const _ as *mut _,
21676                    &base_plus as *const _ as *mut _,
21677                    &scale as *const _ as *mut _,
21678                    &nspm as *const _ as *mut _,
21679                    &spk as *const _ as *mut _,
21680                    &ktb as *const _ as *mut _,
21681                    &vtb as *const _ as *mut _,
21682                    &wini as *const _ as *mut _,
21683                ];
21684                unsafe {
21685                    self.launch_pdl_flash(
21686                        wg,
21687                        "fa_decode_vec_q_rows_v4_w_sp",
21688                        (n_head_kv as u32, n_splits_max as u32, t as u32),
21689                        (32, gqa + 1, 1),
21690                        sh,
21691                        &mut ps,
21692                    )?;
21693                }
21694            } else {
21695                let f = if wg {
21696                    self.func_g("fa_decode_vec_q_rows_v4_w_sp")
21697                } else {
21698                    self.func("fa_decode_vec_q_rows_v4_w_sp")
21699                };
21700                f.set_attribute(
21701                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21702                    sh as i32,
21703                )?;
21704                let cfg = LaunchConfig {
21705                    grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
21706                    block_dim: (32, gqa + 1, 1),
21707                    shared_mem_bytes: sh,
21708                };
21709                let __s_b = self.gpu.stream();
21710                let mut b = __s_b.launch_builder(&f);
21711                b.arg(q)
21712                    .arg(k)
21713                    .arg(v)
21714                    .arg(&mut *part_o)
21715                    .arg(&mut *part_m)
21716                    .arg(&mut *part_l)
21717                    .arg(&hd)
21718                    .arg(&nh)
21719                    .arg(&nhkv)
21720                    .arg(base_dev)
21721                    .arg(&base_plus)
21722                    .arg(&scale)
21723                    .arg(&nspm)
21724                    .arg(&spk)
21725                    .arg(&ktb)
21726                    .arg(&vtb)
21727                    .arg(&wini);
21728                unsafe {
21729                    b.launch(cfg)?;
21730                }
21731            }
21732        } else {
21733            if fa_v4_at(window) && Self::pdl_on() && Self::pdl_wb_on() {
21734                // wave-B2b: the v4_w pick only (smem/reg twins stay builder-launched).
21735                let sh = (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32;
21736                use cudarc::driver::{DevicePtr, DevicePtrMut};
21737                let s = &self.gpu.stream();
21738                let (pq, _b0) = q.device_ptr(s);
21739                let (pk, _b1) = k.device_ptr(s);
21740                let (pv, _b2) = v.device_ptr(s);
21741                let (po, _b3) = part_o.device_ptr_mut(s);
21742                let (pm, _b4) = part_m.device_ptr_mut(s);
21743                let (pl, _b5) = part_l.device_ptr_mut(s);
21744                let (pb, _b6) = base_dev.device_ptr(s);
21745                let mut ps = [
21746                    &pq as *const _ as *mut std::ffi::c_void,
21747                    &pk as *const _ as *mut _,
21748                    &pv as *const _ as *mut _,
21749                    &po as *const _ as *mut _,
21750                    &pm as *const _ as *mut _,
21751                    &pl as *const _ as *mut _,
21752                    &hd as *const _ as *mut _,
21753                    &nh as *const _ as *mut _,
21754                    &nhkv as *const _ as *mut _,
21755                    &pb as *const _ as *mut _,
21756                    &base_plus as *const _ as *mut _,
21757                    &scale as *const _ as *mut _,
21758                    &nspm as *const _ as *mut _,
21759                    &spk as *const _ as *mut _,
21760                    &ktb as *const _ as *mut _,
21761                    &vtb as *const _ as *mut _,
21762                    &wini as *const _ as *mut _,
21763                ];
21764                unsafe {
21765                    self.launch_pdl_flash(
21766                        wg,
21767                        "fa_decode_vec_q_rows_v4_w",
21768                        (n_head_kv as u32, n_splits_max as u32, t as u32),
21769                        (32, gqa, 1),
21770                        sh,
21771                        &mut ps,
21772                    )?;
21773                }
21774            } else {
21775                let pick = |name: &str| {
21776                    if wg {
21777                        self.func_g(name)
21778                    } else {
21779                        self.func(name)
21780                    }
21781                };
21782                let (f, sh) = if fa_v4_at(window) {
21783                    let f = pick("fa_decode_vec_q_rows_v4_w");
21784                    (f, (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32)
21785                } else if smem_tkv > 0 && window >= smem_tkv {
21786                    // NOTE: the smem twin's V-stage is still q5_1-hardcoded — unreachable under wkv
21787                    // at the gemma window (v4 covers it); revisit if the smem floor ever drops.
21788                    (
21789                        pick("fa_decode_vec_q_rows_smem_w"),
21790                        (2 * 32 * head_dim * 2) as u32,
21791                    )
21792                } else {
21793                    (pick("fa_decode_vec_q_rows_reg_w"), 0u32)
21794                };
21795                f.set_attribute(
21796                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21797                    sh as i32,
21798                )?;
21799                let cfg = LaunchConfig {
21800                    grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
21801                    block_dim: (32, gqa, 1),
21802                    shared_mem_bytes: sh,
21803                };
21804                let __s_b = self.gpu.stream();
21805                let mut b = __s_b.launch_builder(&f);
21806                b.arg(q)
21807                    .arg(k)
21808                    .arg(v)
21809                    .arg(&mut *part_o)
21810                    .arg(&mut *part_m)
21811                    .arg(&mut *part_l)
21812                    .arg(&hd)
21813                    .arg(&nh)
21814                    .arg(&nhkv)
21815                    .arg(base_dev)
21816                    .arg(&base_plus)
21817                    .arg(&scale)
21818                    .arg(&nspm)
21819                    .arg(&spk)
21820                    .arg(&ktb)
21821                    .arg(&vtb)
21822                    .arg(&wini);
21823                unsafe {
21824                    b.launch(cfg)?;
21825                }
21826            }
21827        }
21828        let cfg2 = LaunchConfig {
21829            grid_dim: (n_head as u32, t as u32, 1),
21830            block_dim: (head_dim as u32, 1, 1),
21831            shared_mem_bytes: 0,
21832        };
21833        if let Some((oq, od)) = q8_out {
21834            // wave-5b port (2026-07-23): q8-emitting combine — the t=1 decode's wo matvec
21835            // consumes the pair directly; the standalone quantize launch folds away.
21836            if Self::pdl_on() && Self::pdl_wb_on() {
21837                // wave-B2: flavor mirrors the builder's wg choice.
21838                use cudarc::driver::{DevicePtr, DevicePtrMut};
21839                let s = &self.gpu.stream();
21840                let (po, _g0) = part_o.device_ptr(s);
21841                let (pm, _g1) = part_m.device_ptr(s);
21842                let (pl, _g2) = part_l.device_ptr(s);
21843                let (pq, _g3) = oq.device_ptr_mut(s);
21844                let (pd, _g4) = od.device_ptr_mut(s);
21845                let mut ps = [
21846                    &po as *const _ as *mut std::ffi::c_void,
21847                    &pm as *const _ as *mut _,
21848                    &pl as *const _ as *mut _,
21849                    &pq as *const _ as *mut _,
21850                    &pd as *const _ as *mut _,
21851                    &hd as *const _ as *mut _,
21852                    &nh as *const _ as *mut _,
21853                    &nspm as *const _ as *mut _,
21854                    &spk as *const _ as *mut _,
21855                    &wini as *const _ as *mut _,
21856                ];
21857                unsafe {
21858                    self.launch_pdl_flash(
21859                        wg,
21860                        "fa_decode_combine_rows_w_q8_1",
21861                        cfg2.grid_dim,
21862                        cfg2.block_dim,
21863                        0,
21864                        &mut ps,
21865                    )?;
21866                }
21867                return Ok(());
21868            }
21869            let fc = if wg {
21870                self.func_g("fa_decode_combine_rows_w_q8_1")
21871            } else {
21872                self.func("fa_decode_combine_rows_w_q8_1")
21873            };
21874            let __s_b2 = self.gpu.stream();
21875            let mut b2 = __s_b2.launch_builder(&fc);
21876            b2.arg(&*part_o)
21877                .arg(&*part_m)
21878                .arg(&*part_l)
21879                .arg(oq)
21880                .arg(od)
21881                .arg(&hd)
21882                .arg(&nh)
21883                .arg(&nspm)
21884                .arg(&spk)
21885                .arg(&wini);
21886            unsafe {
21887                b2.launch(cfg2)?;
21888            }
21889            return Ok(());
21890        }
21891        let fc = if wg {
21892            self.func_g("fa_decode_combine_rows_w")
21893        } else {
21894            self.func("fa_decode_combine_rows_w")
21895        };
21896        let __s_b2 = self.gpu.stream();
21897        let mut b2 = __s_b2.launch_builder(&fc);
21898        b2.arg(&*part_o)
21899            .arg(&*part_m)
21900            .arg(&*part_l)
21901            .arg(o)
21902            .arg(&hd)
21903            .arg(&nh)
21904            .arg(&nspm)
21905            .arg(&spk)
21906            .arg(&wini);
21907        unsafe {
21908            b2.launch(cfg2)?;
21909        }
21910        Ok(())
21911    }
21912
21913    /// ROUND-STREAM stage (c): fa rows with the causal base from a device counter. Two lanes:
21914    /// v3 (qwen stream, fa_v3_active) and v4 (gemma hd256 burst — rows_v4_dc, g-module aware);
21915    /// `t_kv_upper` sizes splits/partials — the same one-sp-for-all-rows approximation class
21916    /// the host rows path already uses (battery-arbitrated); actual per-row bounds derive
21917    /// in-kernel from the counter (+ base_plus, v4 lane only — v3's kernel has no plus arg).
21918    #[allow(clippy::too_many_arguments)]
21919    pub fn fa_decode_rows_dc(
21920        &self,
21921        q: &CudaSlice<f32>,
21922        k: &cudarc::driver::CudaView<u8>,
21923        v: &cudarc::driver::CudaView<u8>,
21924        o: &mut CudaSlice<f32>,
21925        head_dim: usize,
21926        n_head: usize,
21927        n_head_kv: usize,
21928        base_dev: &CudaSlice<i32>,
21929        t_kv_upper: usize,
21930        t: usize,
21931        scale: f32,
21932        k_tok_bytes: usize,
21933        v_tok_bytes: usize,
21934        base_plus: i32,
21935        g: bool,
21936    ) -> Result<(), Box<dyn std::error::Error>> {
21937        let v4 = head_dim == 256 && fa_v4_at(t_kv_upper);
21938        assert!(
21939            v4 || fa_v3_active(head_dim),
21940            "stream fa rows requires the v3 or v4 lane"
21941        );
21942        assert!(v4 || base_plus == 0, "v3_dc kernel takes no plus arg");
21943        if v4 {
21944            let sp = fa_split_keys(t_kv_upper, n_head_kv);
21945            let n_splits_max = (t_kv_upper + sp - 1) / sp;
21946            let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
21947            let (nspm, spk) = (n_splits_max as i32, sp as i32);
21948            let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
21949            let gqa = (n_head / n_head_kv).max(1) as u32;
21950            let o_len = t * n_head * n_splits_max * head_dim;
21951            let ml_len = t * n_head * n_splits_max;
21952            let mut part_guard = self.fa_part_pool.lock().unwrap();
21953            if part_guard
21954                .as_ref()
21955                .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
21956                .unwrap_or(true)
21957            {
21958                // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
21959                // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
21960                // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
21961                // later live allocations land at those addresses, and the next graph REPLAY writes
21962                // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
21963                // output corruption began the burst after the trunk's t_kv growth first realloc'd
21964                // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
21965                // the baked addresses alive (single-stream: eager writes the new buffers, replays
21966                // touch only the old — never concurrently). Doubling growth bounds retired VRAM
21967                // (total retired < final size).
21968                let old = part_guard.take();
21969                let (co, cm) = old
21970                    .as_ref()
21971                    .map(|pp| (pp.0.len(), pp.1.len()))
21972                    .unwrap_or((0, 0));
21973                if let Some(old) = old {
21974                    self.fa_part_retired.lock().unwrap().push(old);
21975                }
21976                if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
21977                    eprintln!(
21978                        "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
21979                        co, o_len, cm, ml_len
21980                    );
21981                }
21982                *part_guard = Some((
21983                    self.alloc_uninit::<f32>(o_len.max(2 * co))?,
21984                    self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
21985                    self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
21986                ));
21987            }
21988            let pg = part_guard.as_mut().unwrap();
21989            self.gpu
21990                .stream()
21991                .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
21992            self.gpu
21993                .stream()
21994                .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
21995            self.gpu
21996                .stream()
21997                .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
21998            let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
21999            let f = if g {
22000                self.func_g("fa_decode_vec_q_rows_v4_dc")
22001            } else {
22002                self.func("fa_decode_vec_q_rows_v4_dc")
22003            };
22004            let sh = (11520 + 32 * head_dim * if g { 1 } else { 2 }) as u32;
22005            use cudarc::driver::sys::CUfunction_attribute_enum as A;
22006            f.set_attribute(
22007                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22008                sh as i32,
22009            )?;
22010            let cfg = LaunchConfig {
22011                grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
22012                block_dim: (32, gqa, 1),
22013                shared_mem_bytes: sh,
22014            };
22015            let __s_b = self.gpu.stream();
22016            let mut b = __s_b.launch_builder(&f);
22017            b.arg(q)
22018                .arg(k)
22019                .arg(v)
22020                .arg(&mut *part_o)
22021                .arg(&mut *part_m)
22022                .arg(&mut *part_l)
22023                .arg(&hd)
22024                .arg(&nh)
22025                .arg(&nhkv)
22026                .arg(base_dev)
22027                .arg(&base_plus)
22028                .arg(&scale)
22029                .arg(&nspm)
22030                .arg(&spk)
22031                .arg(&ktb)
22032                .arg(&vtb);
22033            unsafe {
22034                b.launch(cfg)?;
22035            }
22036            let fc = self.func("fa_decode_combine_rows_dc");
22037            let cfg2 = LaunchConfig {
22038                grid_dim: (n_head as u32, t as u32, 1),
22039                block_dim: (head_dim as u32, 1, 1),
22040                shared_mem_bytes: 0,
22041            };
22042            let __s_b2 = self.gpu.stream();
22043            let mut b2 = __s_b2.launch_builder(&fc);
22044            b2.arg(&*part_o)
22045                .arg(&*part_m)
22046                .arg(&*part_l)
22047                .arg(o)
22048                .arg(&hd)
22049                .arg(&nh)
22050                .arg(base_dev)
22051                .arg(&base_plus)
22052                .arg(&nspm)
22053                .arg(&spk);
22054            unsafe {
22055                b2.launch(cfg2)?;
22056            }
22057            return Ok(());
22058        }
22059        let sp = fa_split_keys(t_kv_upper, n_head_kv);
22060        let n_splits_max = (t_kv_upper + sp - 1) / sp;
22061        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
22062        let (nspm, spk) = (n_splits_max as i32, sp as i32);
22063        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
22064        let gqa = (n_head / n_head_kv).max(1) as u32;
22065        let o_len = t * n_head * n_splits_max * head_dim;
22066        let ml_len = t * n_head * n_splits_max;
22067        let mut part_guard = self.fa_part_pool.lock().unwrap();
22068        if part_guard
22069            .as_ref()
22070            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
22071            .unwrap_or(true)
22072        {
22073            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
22074            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
22075            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
22076            // later live allocations land at those addresses, and the next graph REPLAY writes
22077            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
22078            // output corruption began the burst after the trunk's t_kv growth first realloc'd
22079            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
22080            // the baked addresses alive (single-stream: eager writes the new buffers, replays
22081            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
22082            // (total retired < final size).
22083            let old = part_guard.take();
22084            let (co, cm) = old
22085                .as_ref()
22086                .map(|pp| (pp.0.len(), pp.1.len()))
22087                .unwrap_or((0, 0));
22088            if let Some(old) = old {
22089                self.fa_part_retired.lock().unwrap().push(old);
22090            }
22091            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
22092                eprintln!(
22093                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
22094                    co, o_len, cm, ml_len
22095                );
22096            }
22097            *part_guard = Some((
22098                self.alloc_uninit::<f32>(o_len.max(2 * co))?,
22099                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22100                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22101            ));
22102        }
22103        let pg = part_guard.as_mut().unwrap();
22104        self.gpu
22105            .stream()
22106            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
22107        self.gpu
22108            .stream()
22109            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
22110        self.gpu
22111            .stream()
22112            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
22113        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
22114        let f = self.func("fa_decode_vec_q_rows_v3_dc");
22115        let sh = (32 * head_dim * 2) as u32;
22116        use cudarc::driver::sys::CUfunction_attribute_enum as A;
22117        f.set_attribute(
22118            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22119            sh as i32,
22120        )?;
22121        let cfg = LaunchConfig {
22122            grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
22123            block_dim: (32, gqa, 1),
22124            shared_mem_bytes: sh,
22125        };
22126        let __s_b = self.gpu.stream();
22127        let mut b = __s_b.launch_builder(&f);
22128        b.arg(q)
22129            .arg(k)
22130            .arg(v)
22131            .arg(&mut *part_o)
22132            .arg(&mut *part_m)
22133            .arg(&mut *part_l)
22134            .arg(&hd)
22135            .arg(&nh)
22136            .arg(&nhkv)
22137            .arg(base_dev)
22138            .arg(&scale)
22139            .arg(&nspm)
22140            .arg(&spk)
22141            .arg(&ktb)
22142            .arg(&vtb);
22143        unsafe {
22144            b.launch(cfg)?;
22145        }
22146        let fc = self.func("fa_decode_combine_rows_dc");
22147        let cfg2 = LaunchConfig {
22148            grid_dim: (n_head as u32, t as u32, 1),
22149            block_dim: (head_dim as u32, 1, 1),
22150            shared_mem_bytes: 0,
22151        };
22152        let plus0 = 0i32;
22153        let __s_b2 = self.gpu.stream();
22154        let mut b2 = __s_b2.launch_builder(&fc);
22155        b2.arg(&*part_o)
22156            .arg(&*part_m)
22157            .arg(&*part_l)
22158            .arg(o)
22159            .arg(&hd)
22160            .arg(&nh)
22161            .arg(base_dev)
22162            .arg(&plus0)
22163            .arg(&nspm)
22164            .arg(&spk);
22165        unsafe {
22166            b2.launch(cfg2)?;
22167        }
22168        Ok(())
22169    }
22170
22171    /// Device-counter variant of `fa_decode` (CUDA-GRAPH-PLAN Phase 2). The sequence length is read
22172    /// from `t_kv_dev[0]` (resident device i32[1]) for the attention loop bound + per-split key range;
22173    /// the GRID `n_splits` is sized for `bucket_max` (the bucket's max t_kv — baked at capture time).
22174    /// Empty splits (key range beyond the actual t_kv) write an empty partial (m=NEG_INF) so the
22175    /// shared combine skips them -> bit-correct for ANY actual t_kv <= bucket_max.
22176    ///
22177    /// BIT-IDENTITY (the gate): pass `bucket_max == actual_t_kv` and this reproduces `fa_decode`
22178    /// EXACTLY (same n_splits, same per, same split boundaries, same combine) while reading t_kv from
22179    /// device. Bucketing (bucket_max > t_kv) is for the future captured path and changes split
22180    /// grouping (different but mathematically-equal log-sum-exp merge).
22181    pub fn fa_decode_dc(
22182        &self,
22183        q: &CudaSlice<f32>,
22184        k: &cudarc::driver::CudaView<u8>,
22185        v: &cudarc::driver::CudaView<u8>,
22186        o: &mut CudaSlice<f32>,
22187        head_dim: usize,
22188        n_head: usize,
22189        n_head_kv: usize,
22190        t_kv_dev: &CudaSlice<i32>,
22191        bucket_max: usize,
22192        scale: f32,
22193        k_tok_bytes: usize,
22194        v_tok_bytes: usize,
22195        g: bool,
22196    ) -> Result<(), Box<dyn std::error::Error>> {
22197        self.fa_decode_dc_q8(
22198            q,
22199            k,
22200            v,
22201            o,
22202            head_dim,
22203            n_head,
22204            n_head_kv,
22205            t_kv_dev,
22206            bucket_max,
22207            scale,
22208            k_tok_bytes,
22209            v_tok_bytes,
22210            g,
22211            None,
22212        )
22213    }
22214
22215    /// `fa_decode_dc` with an optional q8_1 sink (wave 5b): when `q8_out` is given the
22216    /// combine emits (int8, per-32 scales) for the wo matmul_pre and skips the f32 O write.
22217    #[allow(clippy::too_many_arguments)]
22218    pub fn fa_decode_dc_q8(
22219        &self,
22220        q: &CudaSlice<f32>,
22221        k: &cudarc::driver::CudaView<u8>,
22222        v: &cudarc::driver::CudaView<u8>,
22223        o: &mut CudaSlice<f32>,
22224        head_dim: usize,
22225        n_head: usize,
22226        n_head_kv: usize,
22227        t_kv_dev: &CudaSlice<i32>,
22228        bucket_max: usize,
22229        scale: f32,
22230        k_tok_bytes: usize,
22231        v_tok_bytes: usize,
22232        g: bool,
22233        q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
22234    ) -> Result<(), Box<dyn std::error::Error>> {
22235        // The fa_vec gate + n_splits are sized from bucket_max (host, fixed at capture). The kernel
22236        // reads the ACTUAL t_kv from t_kv_dev for the per-split bound. DEFAULT-ON to MATCH the eager
22237        // `fa_decode` gate above — graph capture must mirror eager's kernel choice or the graph-vs-eager
22238        // bit-identity gate breaks. MEMRA_NO_FA_VEC forces scalar on BOTH paths in lockstep.
22239        // `g` = this layer's cache is e4m3 (gemma windowed under wkv) — every pick below must
22240        // mirror fa_decode_kvmod's g-routing or the graph diverges from eager (short/mid 1/96,
22241        // 2026-07-12).
22242        let mut fa_vec =
22243            std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
22244        if g && head_dim == 256 && !fa_v4_at(bucket_max) {
22245            fa_vec = false;
22246        } // mirror kvmod/geom
22247        let sp = fa_split_keys(bucket_max, n_head_kv);
22248        let n_splits = if fa_vec {
22249            ((bucket_max + sp - 1) / sp).max(1)
22250        } else {
22251            ((bucket_max + 255) / 256).max(1)
22252        };
22253        let o_len = n_head * n_splits * head_dim;
22254        let ml_len = n_head * n_splits;
22255        let mut part_guard = self.fa_part_pool.lock().unwrap();
22256        if part_guard
22257            .as_ref()
22258            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
22259            .unwrap_or(true)
22260        {
22261            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
22262            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
22263            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
22264            // later live allocations land at those addresses, and the next graph REPLAY writes
22265            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
22266            // output corruption began the burst after the trunk's t_kv growth first realloc'd
22267            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
22268            // the baked addresses alive (single-stream: eager writes the new buffers, replays
22269            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
22270            // (total retired < final size).
22271            let old = part_guard.take();
22272            let (co, cm) = old
22273                .as_ref()
22274                .map(|pp| (pp.0.len(), pp.1.len()))
22275                .unwrap_or((0, 0));
22276            if let Some(old) = old {
22277                self.fa_part_retired.lock().unwrap().push(old);
22278            }
22279            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
22280                eprintln!(
22281                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
22282                    co, o_len, cm, ml_len
22283                );
22284            }
22285            *part_guard = Some((
22286                self.alloc_uninit::<f32>(o_len.max(2 * co))?,
22287                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22288                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22289            ));
22290        }
22291        let pg = part_guard.as_mut().unwrap();
22292        self.gpu
22293            .stream()
22294            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
22295        self.gpu
22296            .stream()
22297            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
22298        self.gpu
22299            .stream()
22300            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
22301        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
22302        let (hd, nh, nhkv, nsp) = (
22303            head_dim as i32,
22304            n_head as i32,
22305            n_head_kv as i32,
22306            n_splits as i32,
22307        );
22308        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
22309        let fa_vec = fa_vec && head_dim <= 512 && head_dim % 32 == 0;
22310        // FA-DEEP pick keyed on bucket_max (the fa_v4_at precedent) — bit-identical twins,
22311        // so a threshold falling between t_kv and bucket_max cannot diverge eager-vs-graph.
22312        let deep = fa_vec
22313            && head_dim == 256
22314            && fa_v4_at(bucket_max)
22315            && !g
22316            && fa_deep_at(bucket_max)
22317            && !matches!(fa_v4_mode(), "noB3" | "stage");
22318        let (f, cfg) = if fa_vec
22319            && head_dim == 512
22320            && bucket_max >= {
22321                static FA512_MIN_DC: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
22322                *FA512_MIN_DC.get_or_init(|| {
22323                    std::env::var("MEMRA_FA512_MIN")
22324                        .ok()
22325                        .and_then(|v| v.parse().ok())
22326                        .unwrap_or(512)
22327                })
22328            } {
22329            // gemma globals dc twin (mirror the eager dpl16 pick incl the crossover floor).
22330            let gqa = (n_head / n_head_kv).max(1) as u32;
22331            (
22332                self.fa_func("fa_decode_vec_q_dpl16_dc", head_dim),
22333                LaunchConfig {
22334                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22335                    block_dim: (32, gqa, 1),
22336                    shared_mem_bytes: 0,
22337                },
22338            )
22339        } else if fa_vec && head_dim == 512 {
22340            // under the 512 floor eager runs scalar — the SAME unified symbol, ctr non-null;
22341            // ns_eff in-kernel reproduces eager's ceil(t_kv/sp) partition for the LIVE len.
22342            let q_view = q.as_view();
22343            let mut o_view = o.as_view_mut();
22344            return self.fa_decode_scalar_unified(
22345                &q_view,
22346                k,
22347                v,
22348                &mut o_view,
22349                head_dim,
22350                n_head,
22351                n_head_kv,
22352                0,
22353                Some(t_kv_dev),
22354                scale,
22355                n_splits,
22356                sp,
22357                k_tok_bytes,
22358                v_tok_bytes,
22359                g,
22360                &mut *part_o,
22361                &mut *part_m,
22362                &mut *part_l,
22363                q8_out,
22364            );
22365        } else if fa_vec && head_dim == 256 && fa_v4_at(bucket_max) {
22366            // gemma/qwen v4 dc twin (eager default lane) — capture must mirror eager's pick,
22367            // incl the g-module route + raw-e4m3 sV sizing.
22368            let gqa = (n_head / n_head_kv).max(1) as u32;
22369            let fv = if g {
22370                self.func_g("fa_decode_vec_q_v4_dc")
22371            } else if deep {
22372                self.func("fa_decode_vec_q_v4_deep_dc")
22373            } else {
22374                self.func("fa_decode_vec_q_v4_dc")
22375            };
22376            let shmem =
22377                (if deep { 12160 } else { 11520 } + 32 * head_dim * if g { 1 } else { 2 }) as u32;
22378            use cudarc::driver::sys::CUfunction_attribute_enum as A;
22379            fv.set_attribute(
22380                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22381                shmem as i32,
22382            )?;
22383            (
22384                fv,
22385                LaunchConfig {
22386                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22387                    block_dim: (32, gqa, 1),
22388                    shared_mem_bytes: shmem,
22389                },
22390            )
22391        } else if fa_vec && fa_v3_active(head_dim) {
22392            // FA v3 lane _dc twin: the captured graph must run the SAME walk body as eager
22393            // under MEMRA_FA_V3=1 (eager, rows-verify and graph switch together).
22394            let gqa = (n_head / n_head_kv).max(1) as u32;
22395            let fv = if g {
22396                self.func_g("fa_decode_vec_q_v3_dc")
22397            } else {
22398                self.func("fa_decode_vec_q_v3_dc")
22399            };
22400            let shmem = (32 * head_dim * 2) as u32; // sV bf16 [FA_DEC_TILE=32][hd]
22401            (
22402                fv,
22403                LaunchConfig {
22404                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22405                    block_dim: (32, gqa, 1),
22406                    shared_mem_bytes: shmem,
22407                },
22408            )
22409        } else if fa_vec && fa_v2_on() {
22410            // FAVENDOR lane: v2 _dc twin — the captured graph must run the SAME walk body as
22411            // eager under MEMRA_FA_V2=1 or graph_decode_gate's bit-identity breaks (the flag is
22412            // a numeric config; eager, rows-verify and graph all switch together).
22413            let gqa = (n_head / n_head_kv).max(1) as u32;
22414            let fv = if g {
22415                self.func_g("fa_decode_vec_q_v2_dc")
22416            } else {
22417                self.func("fa_decode_vec_q_v2_dc")
22418            };
22419            let shmem = (2 * 32 * head_dim * 2) as u32; // sK+sV bf16 [FA_DEC_TILE=32][hd]
22420            (
22421                fv,
22422                LaunchConfig {
22423                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22424                    block_dim: (32, gqa, 1),
22425                    shared_mem_bytes: shmem,
22426                },
22427            )
22428        } else if fa_vec {
22429            let gqa = (n_head / n_head_kv).max(1) as u32;
22430            // REGISTER-DEQUANT twin: zero dynamic smem (see fa_decode above).
22431            let fv = if g {
22432                self.func_g("fa_decode_vec_q_dc")
22433            } else {
22434                self.func("fa_decode_vec_q_dc")
22435            };
22436            (
22437                fv,
22438                LaunchConfig {
22439                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22440                    block_dim: (32, gqa, 1),
22441                    shared_mem_bytes: 0,
22442                },
22443            )
22444        } else {
22445            let q_view = q.as_view();
22446            let mut o_view = o.as_view_mut();
22447            return self.fa_decode_scalar_unified(
22448                &q_view,
22449                k,
22450                v,
22451                &mut o_view,
22452                head_dim,
22453                n_head,
22454                n_head_kv,
22455                0,
22456                Some(t_kv_dev),
22457                scale,
22458                n_splits,
22459                if fa_vec { sp } else { 256 },
22460                k_tok_bytes,
22461                v_tok_bytes,
22462                g,
22463                &mut *part_o,
22464                &mut *part_m,
22465                &mut *part_l,
22466                q8_out,
22467            );
22468        };
22469        let ski = sp as i32; // one-partition law: the twins derive ns_eff from (T_kv, ski)
22470        let __s_b = self.gpu.stream();
22471        let mut b = __s_b.launch_builder(&f);
22472        b.arg(q)
22473            .arg(k)
22474            .arg(v)
22475            .arg(&mut *part_o)
22476            .arg(&mut *part_m)
22477            .arg(&mut *part_l)
22478            .arg(&hd)
22479            .arg(&nh)
22480            .arg(&nhkv)
22481            .arg(t_kv_dev)
22482            .arg(&scale)
22483            .arg(&nsp)
22484            .arg(&ski)
22485            .arg(&ktb)
22486            .arg(&vtb);
22487        unsafe {
22488            b.launch(cfg)?;
22489        }
22490        let cfg2 = LaunchConfig {
22491            grid_dim: (n_head as u32, 1, 1),
22492            block_dim: (head_dim as u32, 1, 1),
22493            shared_mem_bytes: 0,
22494        };
22495        if let Some((oq, od)) = q8_out {
22496            let fc = if g {
22497                self.func_g("fa_decode_combine_q8_1")
22498            } else {
22499                self.fa_func("fa_decode_combine_q8_1", head_dim)
22500            };
22501            let __s_b2 = self.gpu.stream();
22502            let mut b2 = __s_b2.launch_builder(&fc);
22503            b2.arg(&*part_o)
22504                .arg(&*part_m)
22505                .arg(&*part_l)
22506                .arg(oq)
22507                .arg(od)
22508                .arg(&hd)
22509                .arg(&nh)
22510                .arg(&nsp);
22511            unsafe {
22512                b2.launch(cfg2)?;
22513            }
22514            return Ok(());
22515        }
22516        let fc = if g {
22517            self.func_g("fa_decode_combine_f32")
22518        } else {
22519            self.fa_func("fa_decode_combine_f32", head_dim)
22520        };
22521        let __s_b2 = self.gpu.stream();
22522        let mut b2 = __s_b2.launch_builder(&fc);
22523        b2.arg(&*part_o)
22524            .arg(&*part_m)
22525            .arg(&*part_l)
22526            .arg(o)
22527            .arg(&hd)
22528            .arg(&nh)
22529            .arg(&nsp);
22530        unsafe {
22531            b2.launch(cfg2)?;
22532        }
22533        Ok(())
22534    }
22535
22536    /// _dcw append (t=1): physical write row = len_dev[0] - base_dev[0] in-kernel; follow
22537    /// with `inc_i32(len_dev)` on the same stream. Bit-identical bytes to the host-row append
22538    /// at equal rows.
22539    #[allow(clippy::too_many_arguments)]
22540    pub fn append_kv_quantized_dcw(
22541        &self,
22542        k_row: &CudaSlice<f32>,
22543        v_row: &CudaSlice<f32>,
22544        kc: &mut CudaSlice<u8>,
22545        vc: &mut CudaSlice<u8>,
22546        len_dev: &CudaSlice<i32>,
22547        base_dev: Option<&CudaSlice<i32>>,
22548        kv_dim_k: usize,
22549        kv_dim_v: usize,
22550        k_tok_bytes: usize,
22551        v_tok_bytes: usize,
22552    ) -> Result<(), Box<dyn std::error::Error>> {
22553        let f = self.func("append_quantize_kv_q8_0_q5_1_dcw");
22554        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
22555        let cfg = LaunchConfig {
22556            grid_dim: (nblk, 1, 1),
22557            block_dim: (32, 1, 1),
22558            shared_mem_bytes: 0,
22559        };
22560        let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
22561        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
22562        let null: u64 = 0;
22563        let __s_b = self.gpu.stream();
22564        let mut b = __s_b.launch_builder(&f);
22565        b.arg(k_row).arg(v_row).arg(kc).arg(vc).arg(len_dev);
22566        match base_dev {
22567            Some(base) => {
22568                b.arg(base);
22569            }
22570            None => {
22571                b.arg(&null);
22572            }
22573        }
22574        b.arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
22575        unsafe {
22576            b.launch(cfg)?;
22577        }
22578        Ok(())
22579    }
22580
22581    /// Increment a device i32 counter (graph-capturable; the `inc_i32` kernel).
22582    pub fn inc_i32(&self, counter: &mut CudaSlice<i32>) -> Result<(), Box<dyn std::error::Error>> {
22583        let f = self.func("inc_i32");
22584        let cfg = LaunchConfig {
22585            grid_dim: (1, 1, 1),
22586            block_dim: (1, 1, 1),
22587            shared_mem_bytes: 0,
22588        };
22589        let __s_b = self.gpu.stream();
22590        let mut b = __s_b.launch_builder(&f);
22591        b.arg(counter);
22592        unsafe {
22593            b.launch(cfg)?;
22594        }
22595        Ok(())
22596    }
22597
22598    /// Windowed device-counter fa decode (step TP graph increment A): the KV view derives
22599    /// entirely from device state — `len_dev` (staged length), `base_dev` (physical row of
22600    /// logical 0 after the last ring rebase; None reads as 0), and `window` (0 = global) — so
22601    /// a captured child replays with ZERO per-token node updates. v3-vec only (the default
22602    /// kernel class on this lane); callers keep eager below the vec floor and for any other
22603    /// class. Scratch comes from the engine's fa partial pool sized at `bucket_max` (for SWA
22604    /// layers pass min(bucket, window)); the pool's retire-on-grow keeps captured addresses
22605    /// alive across bucket growth.
22606    #[allow(clippy::too_many_arguments)]
22607    /// Retire-on-grow ensure for the fa partial pool (see the #68 comment on the eager
22608    /// twin). Split out so graph capture can pre-run it OUTSIDE the capture region — an
22609    /// alloc inside a captured section becomes a mem node, and child graphs reject those.
22610    fn fa_part_pool_grow(
22611        &self,
22612        part_guard: &mut Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>,
22613        o_len: usize,
22614        ml_len: usize,
22615    ) -> Result<(), Box<dyn std::error::Error>> {
22616        if part_guard
22617            .as_ref()
22618            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
22619            .unwrap_or(true)
22620        {
22621            let old = part_guard.take();
22622            let (co, cm) = old
22623                .as_ref()
22624                .map(|pp| (pp.0.len(), pp.1.len()))
22625                .unwrap_or((0, 0));
22626            if let Some(old) = old {
22627                self.fa_part_retired.lock().unwrap().push(old);
22628            }
22629            *part_guard = Some((
22630                self.alloc_uninit::<f32>(o_len.max(2 * co))?,
22631                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22632                self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22633            ));
22634        }
22635        Ok(())
22636    }
22637
22638    /// Pre-grow the fa partial pool for a dcw call at (n_head, bucket_max) geometry, from
22639    /// OUTSIDE any capture region. Idempotent and cheap when already big enough.
22640    pub fn fa_dcw_pool_ensure(
22641        &self,
22642        head_dim: usize,
22643        n_head: usize,
22644        n_head_kv: usize,
22645        bucket_max: usize,
22646    ) -> Result<(), Box<dyn std::error::Error>> {
22647        let sp = fa_split_keys(bucket_max, n_head_kv);
22648        let n_splits = ((bucket_max + sp - 1) / sp).max(1);
22649        let o_len = n_head * n_splits * head_dim;
22650        let ml_len = n_head * n_splits;
22651        let mut part_guard = self.fa_part_pool.lock().unwrap();
22652        self.fa_part_pool_grow(&mut part_guard, o_len, ml_len)
22653    }
22654
22655    pub fn fa_decode_dcw(
22656        &self,
22657        q: &CudaSlice<f32>,
22658        k_ring: &cudarc::driver::CudaView<u8>,
22659        v_ring: &cudarc::driver::CudaView<u8>,
22660        o: &mut CudaSlice<f32>,
22661        head_dim: usize,
22662        n_head: usize,
22663        n_head_kv: usize,
22664        len_dev: &CudaSlice<i32>,
22665        base_dev: Option<&CudaSlice<i32>>,
22666        window: usize,
22667        bucket_max: usize,
22668        scale: f32,
22669        k_tok_bytes: usize,
22670        v_tok_bytes: usize,
22671        // FUSION #2d: Some(gate_row) fuses the head gate into the combine (bit-identical,
22672        // one launch saved); `o` then receives the GATED output and the caller skips its
22673        // attn_head_gate call.
22674        fused_gate: Option<&CudaSlice<f32>>,
22675    ) -> Result<(), Box<dyn std::error::Error>> {
22676        let fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
22677        if !fa_vec || head_dim > 256 || head_dim % 32 != 0 || !fa_v3_on() {
22678            return Err("fa_decode_dcw supports the default v3-vec class only                         (bucket >= vec floor, head_dim <= 256, MEMRA_FA_V3 on);                         keep eager outside it"
22679                .into());
22680        }
22681        let sp = fa_split_keys(bucket_max, n_head_kv);
22682        let n_splits = ((bucket_max + sp - 1) / sp).max(1);
22683        let o_len = n_head * n_splits * head_dim;
22684        let ml_len = n_head * n_splits;
22685        let mut part_guard = self.fa_part_pool.lock().unwrap();
22686        Self::fa_part_pool_grow(self, &mut part_guard, o_len, ml_len)?;
22687        let pg = part_guard.as_mut().unwrap();
22688        // MEMRA_FA_DCW_MEMSET=0: skip the partial-pool zeroing — every (head, split) in
22689        // [0, nsp) writes its partial before the combine reads it (per = ceil(len/nsp), so
22690        // split s starts at s*per < len for all s < nsp), making the zeros dead stores.
22691        // Door-gated pending the identity battery; =0 saves 3 memset launches/rank/layer.
22692        static MEMSET_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
22693        // Token-graph capture ALWAYS keeps the memsets: the retarget path (increment C)
22694        // finds the attention children BY their three-memset signature and updates the
22695        // memset widths per bucket — capturing without them silently kills retargeting
22696        // (battery-v8 token drift, 2026-08-21).
22697        let memset_on = *MEMSET_ON
22698            .get_or_init(|| std::env::var("MEMRA_FA_DCW_MEMSET").as_deref() != Ok("0"))
22699            || crate::tp::token_graph_building();
22700        if memset_on {
22701            self.gpu
22702                .stream()
22703                .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
22704            self.gpu
22705                .stream()
22706                .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
22707            self.gpu
22708                .stream()
22709                .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
22710        }
22711        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
22712        let (hd, nh, nhkv, nsp) = (
22713            head_dim as i32,
22714            n_head as i32,
22715            n_head_kv as i32,
22716            n_splits as i32,
22717        );
22718        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
22719        let (ski, win) = (sp as i32, window as i32);
22720        let gqa = (n_head / n_head_kv).max(1) as u32;
22721        let smem = (32 * head_dim * 2) as u32; // sV bf16 [FA_DEC_TILE=32][hd] (v3 uses sV only)
22722        let f = self.func("fa_decode_vec_q_v3_dcw");
22723        let cfg = LaunchConfig {
22724            grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22725            block_dim: (32, gqa, 1),
22726            shared_mem_bytes: smem,
22727        };
22728        let null: u64 = 0;
22729        let __s_b = self.gpu.stream();
22730        let mut b = __s_b.launch_builder(&f);
22731        b.arg(q)
22732            .arg(k_ring)
22733            .arg(v_ring)
22734            .arg(&mut *part_o)
22735            .arg(&mut *part_m)
22736            .arg(&mut *part_l)
22737            .arg(&hd)
22738            .arg(&nh)
22739            .arg(&nhkv)
22740            .arg(len_dev);
22741        match base_dev {
22742            Some(base) => {
22743                b.arg(base);
22744            }
22745            None => {
22746                b.arg(&null);
22747            }
22748        }
22749        b.arg(&win)
22750            .arg(&scale)
22751            .arg(&nsp)
22752            .arg(&ski)
22753            .arg(&ktb)
22754            .arg(&vtb);
22755        unsafe {
22756            b.launch(cfg)?;
22757        }
22758        let fc = if fused_gate.is_some() {
22759            self.func("fa_decode_combine_gate_f32")
22760        } else {
22761            self.fa_func("fa_decode_combine_f32", head_dim)
22762        };
22763        let cfg2 = LaunchConfig {
22764            grid_dim: (n_head as u32, 1, 1),
22765            block_dim: (head_dim as u32, 1, 1),
22766            shared_mem_bytes: 0,
22767        };
22768        let __s_b2 = self.gpu.stream();
22769        let mut b2 = __s_b2.launch_builder(&fc);
22770        b2.arg(&*part_o).arg(&*part_m).arg(&*part_l);
22771        if let Some(gate_row) = fused_gate {
22772            b2.arg(gate_row);
22773        }
22774        b2.arg(o).arg(&hd).arg(&nh).arg(&nsp);
22775        unsafe {
22776            b2.launch(cfg2)?;
22777        }
22778        Ok(())
22779    }
22780
22781    /// EAGER fa_decode geometry for a given actual `t_kv` (CUDA-GRAPH-PLAN §3.3 bucketing). Returns
22782    /// `(fa_vec, n_splits)` EXACTLY as `fa_decode` computes them so the graph-capture path can key its
22783    /// bucket on the same `(kernel, n_splits)` pair and pass a `bucket_max` that reproduces eager's
22784    /// n_splits bit-for-bit. (Per = ceil(t_kv/n_splits) is then recomputed from the DEVICE t_kv inside
22785    /// the kernel and matches eager when n_splits matches — the bit-identity contract.)
22786    pub fn fa_geom_eager(
22787        &self,
22788        t_kv: usize,
22789        head_dim: usize,
22790        n_head_kv: usize,
22791        g: bool,
22792    ) -> (bool, usize) {
22793        // MUST mirror `fa_decode` / `fa_decode_dc` (default-ON 2026-06-28). This is the bucket-key
22794        // source: if it disagrees with the actual kernel pick, the graph captures the wrong path and
22795        // replay diverges from eager. All three sites read MEMRA_NO_FA_VEC in lockstep.
22796        let fa_ok = std::env::var("MEMRA_NO_FA_VEC").is_err() && t_kv >= fa_vec_min_tkv();
22797        // hd512 dpl16 vec lane (gemma globals, 2026-07-11 graph-arc fix): the original key
22798        // hardcoded vec = hd<=256, so for hd512 it bucketed by the SCALAR 256-key splits while
22799        // the dpl16/rows_dpl16 kernels split by the ladder — n_splits changed WITHIN a bucket
22800        // (mid-ctx graph mismatch at pos 19 + partials OOB at longer runs). Mirror the real
22801        // fa_decode dispatch: vec512 above the fa512 floor, vec256 as before.
22802        let vec512 = fa_ok && head_dim == 512 && t_kv >= fa512_min_tkv();
22803        let mut fa_vec = vec512 || (fa_ok && head_dim <= 256 && head_dim % 32 == 0);
22804        // g (fp8-windowed): mirror kvmod's clamp — only the v4 lane parses e4m3 in the vec
22805        // family; everything else falls to the g-module scalar.
22806        // hd256 under g: v4-or-scalar. hd128/other under g ride the REGISTER g-lane (the
22807        // dq_K_lane/dq_V_lane macros are format-aware; v3 is format-gated off, v2/smem
22808        // arms are excluded under g). Mirrored in kvmod / fa_decode_dc / fa_geom_eager.
22809        if g && head_dim == 256 && !fa_v4_at(t_kv) {
22810            fa_vec = false;
22811        }
22812        let sp = fa_split_keys(t_kv, n_head_kv);
22813        let n_splits = if fa_vec {
22814            ((t_kv + sp - 1) / sp).max(1)
22815        } else {
22816            ((t_kv + 255) / 256).max(1)
22817        };
22818        (fa_vec, n_splits)
22819    }
22820
22821    /// `bucket_max` (host t_kv to feed `fa_decode_dc` / `full_attn_decode_dc`) that makes the _dc
22822    /// kernel pick the SAME (fa_vec, n_splits) as eager would for actual `t_kv`. Because the dc
22823    /// launcher derives both from `bucket_max` via the same formulas, we just hand it `t_kv` itself:
22824    /// the n_splits is then identical, and the per-split boundaries (computed from the DEVICE t_kv in
22825    /// the kernel) match eager exactly. The bucket KEY (for the graph HashMap) is `(fa_vec, n_splits)`.
22826    pub fn fa_bucket_key(
22827        &self,
22828        t_kv: usize,
22829        head_dim: usize,
22830        n_head_kv: usize,
22831        g: bool,
22832    ) -> (bool, usize) {
22833        self.fa_geom_eager(t_kv, head_dim, n_head_kv, g)
22834    }
22835
22836    /// CUDA-graph capture wrapper (CUDA-GRAPH-PLAN §3.2, llama.cpp warmup pattern). Runs `step`
22837    /// inline TWICE (warmup — lets the caching allocator settle to stable pointers and any one-time
22838    /// kernel attribute/JIT happen outside capture), then captures a THIRD invocation on the Engine's
22839    /// decode stream (RELAXED mode) and instantiates it into a replayable `CudaGraph`. The closure
22840    /// must enqueue ONLY device work on `e.stream()` (no dtoh / no synchronize / no host branch on
22841    /// device data) — every per-step varying scalar must come from a device counter. Returns the
22842    /// instantiated graph; `CudaGraph::launch()` replays the whole step in one dispatch.
22843    /// `capture_graph` with CAPTURE-RETAIN: every Engine allocation made during the warmups
22844    /// and the capture is kept alive in the returned keeper — hold it as long as the graph
22845    /// replays (transients returning to the pool get reused by unrelated work and corrupt
22846    /// replays; the draft-graph root cause). Model-generic, next capture reuses it.
22847    pub fn capture_graph_retained<F>(
22848        &self,
22849        step: F,
22850    ) -> Result<
22851        (
22852            cudarc::driver::CudaGraph,
22853            Vec<Box<dyn std::any::Any + Send>>,
22854        ),
22855        Box<dyn std::error::Error>,
22856    >
22857    where
22858        F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
22859    {
22860        use cudarc::driver::sys::CUgraphInstantiate_flags;
22861        self.capture_graph_retained_flags(
22862            CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
22863            step,
22864        )
22865    }
22866
22867    /// Retained capture with an explicit instantiate flag. ALLOC-FREE captured graphs
22868    /// (zero mem nodes — the gemma slotted door) should pass UPLOAD instead of
22869    /// AUTO_FREE_ON_LAUNCH: the auto-free flag's launch-time mem-pool scan was measured at
22870    /// ~0.25us/node (205us on the 826-node step) even with nothing to free.
22871    pub fn capture_graph_retained_flags<F>(
22872        &self,
22873        flags: cudarc::driver::sys::CUgraphInstantiate_flags,
22874        mut step: F,
22875    ) -> Result<
22876        (
22877            cudarc::driver::CudaGraph,
22878            Vec<Box<dyn std::any::Any + Send>>,
22879        ),
22880        Box<dyn std::error::Error>,
22881    >
22882    where
22883        F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
22884    {
22885        use cudarc::driver::sys::CUstreamCaptureMode;
22886        // KEEP scope = WARMUPS ONLY (2026-07-13): keep_if_capturing retains via
22887        // CudaSlice::clone, which is a device ALLOC + D2D COPY on the stream — clones made
22888        // while the capture region is open become dead copy NODES replayed every launch
22889        // (E4B: 1440 copies = 0.74ms/token, the whole graph-vs-eager regression). The
22890        // warmup runs allocate the same transient sequence at the same pool addresses, so
22891        // retaining the warmup clones preserves the draft-graph fix without polluting the
22892        // captured graph.
22893        self.capture_keep.lock().unwrap().clear();
22894        let was_tracking = self.gpu.ctx.is_event_tracking();
22895        if was_tracking {
22896            unsafe {
22897                self.gpu.ctx.disable_event_tracking();
22898            }
22899        }
22900        let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
22901            self.capture_keep_on
22902                .store(true, std::sync::atomic::Ordering::Relaxed);
22903            let w = (|| {
22904                step(self)?;
22905                step(self)
22906            })();
22907            self.capture_keep_on
22908                .store(false, std::sync::atomic::Ordering::Relaxed);
22909            w?;
22910            self.gpu.stream().synchronize()?;
22911            self.gpu
22912                .stream()
22913                .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
22914            let r = step(self);
22915            let g = self.gpu.stream().end_capture(flags);
22916            r?;
22917            let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
22918            graph.upload()?;
22919            Ok(graph)
22920        };
22921        let result = run();
22922        self.capture_keep_on
22923            .store(false, std::sync::atomic::Ordering::Relaxed);
22924        if was_tracking {
22925            unsafe {
22926                self.gpu.ctx.enable_event_tracking();
22927            }
22928        }
22929        let keeper = std::mem::take(&mut *self.capture_keep.lock().unwrap());
22930        Ok((result?, keeper))
22931    }
22932
22933    /// Retained capture WITHOUT the two warmup executions. The warmups exist for transient
22934    /// pool-address stability (draft-graph lanes); the step TP token-graph sections are
22935    /// alloc-free with persistent operands, and their bodies carry device side effects
22936    /// (dcw KV appends + counter incs) that a warmup would REALLY EXECUTE — measured as a
22937    /// +2/rank len_d drift per bucket build that marched appends past the ring planes.
22938    pub fn capture_graph_retained_nowarm<F>(
22939        &self,
22940        mut step: F,
22941    ) -> Result<
22942        (
22943            cudarc::driver::CudaGraph,
22944            Vec<Box<dyn std::any::Any + Send>>,
22945        ),
22946        Box<dyn std::error::Error>,
22947    >
22948    where
22949        F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
22950    {
22951        use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
22952        let was_tracking = self.gpu.ctx.is_event_tracking();
22953        if was_tracking {
22954            unsafe {
22955                self.gpu.ctx.disable_event_tracking();
22956            }
22957        }
22958        let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
22959            self.gpu.stream().synchronize()?;
22960            self.gpu
22961                .stream()
22962                .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
22963            let r = step(self);
22964            let g = self.gpu.stream().end_capture(
22965                CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
22966            );
22967            r?;
22968            let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
22969            graph.upload()?;
22970            Ok(graph)
22971        };
22972        let result = run();
22973        if was_tracking {
22974            unsafe {
22975                self.gpu.ctx.enable_event_tracking();
22976            }
22977        }
22978        Ok((result?, Vec::new()))
22979    }
22980
22981    pub fn capture_graph<F>(
22982        &self,
22983        mut step: F,
22984    ) -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>>
22985    where
22986        F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
22987    {
22988        use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
22989        // EVENT TRACKING OFF for capture. The Engine creates a 2nd stream (copy_stream) so cudarc is in
22990        // multi-stream mode and, by default, records a CudaEvent per CudaSlice alloc/use to serialize
22991        // cross-stream access. Those per-buffer event waits issue stream ops that are NOT permitted
22992        // inside a capture region (CUDA_ERROR_STREAM_CAPTURE_UNSUPPORTED). The captured decode step is
22993        // strictly SINGLE-STREAM (every kernel on gpu.stream), so this synchronization is unnecessary
22994        // here — disable it for the whole warmup+capture, re-enable after. SAFETY: the decode-dc path
22995        // touches only gpu.stream; no buffer crosses to copy_stream during capture.
22996        let was_tracking = self.gpu.ctx.is_event_tracking();
22997        if was_tracking {
22998            unsafe {
22999                self.gpu.ctx.disable_event_tracking();
23000            }
23001        }
23002        // Q1 PROBE (MEMRA_GRAPH_IFLAG): the generic capture body's cuMemAllocAsync nodes are
23003        // EXACTLY BALANCED by in-graph free nodes (measured census q27: 1589 ALLOC / 1589
23004        // FREE), so AUTO_FREE_ON_LAUNCH has nothing to reclaim at launch — it only pays its
23005        // per-node launch-time mem-pool scan. `upload` / `none` select the alternatives to
23006        // measure that scan's real cost on the generic path. Diagnostic door only; the
23007        // default stays AUTO_FREE until a measured A/B justifies moving it.
23008        let iflag = {
23009            static F: std::sync::OnceLock<CUgraphInstantiate_flags> = std::sync::OnceLock::new();
23010            *F.get_or_init(|| match std::env::var("MEMRA_GRAPH_IFLAG").as_deref() {
23011                // UPLOAD = the gemma slotted door's zero-mem-node choice; PRIORITY = the flag
23012                // hybrid_forward.rs:5935 actually ships (both drop the auto-free launch scan).
23013                Ok("upload") => CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_UPLOAD,
23014                Ok("priority") => {
23015                    CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY
23016                }
23017                _ => CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
23018            })
23019        };
23020        // MEMRA_GRAPH_CAPTIME=1 (Q1 lane): phase-resolved capture cost. Recapture is paid at
23021        // every kernel-class crossing, so it — not steady-state decode — is the quantity a
23022        // mem-node reduction could plausibly shrink. Only `instantiate` (cuStreamEndCapture +
23023        // cuGraphInstantiateWithFlags) and `upload` scale with node count; the warmups are
23024        // eager step executions and are node-count-invariant. Printing the split bounds the
23025        // refactor's ceiling instead of assuming it.
23026        let ct = {
23027            static T: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
23028            *T.get_or_init(|| std::env::var("MEMRA_GRAPH_CAPTIME").as_deref() == Ok("1"))
23029        };
23030        // MEMRA_GRAPH_WARMUPS (Q1 lane; DEFAULT 1 since lane/graph-warmups 2026-08-05): the
23031        // phase split showed the eager warmups are 80% of recapture cost (q27 27.4 of 34.4 ms
23032        // pod / 42% of 52.6 ms 5090) — 3x larger than the ENTIRE mem-node ceiling the audit
23033        // chased, and node-count-invariant, so no capture-body refactor could touch it.
23034        // Warmup 2's theorized job was async-pool ADDRESS STABILITY: warmup 1's allocs may
23035        // grow/map the pool, warmup 2 re-walks the same sequence over the freed blocks so the
23036        // captured third run bakes settled addresses. That hazard is the #68 stale-baked-
23037        // address class — which the engine now guards STRUCTURALLY rather than by re-walking:
23038        // in-body transients are captured as BALANCED in-graph alloc/free node pairs (census
23039        // 1589/1589 — replays allocate for themselves; no baked transient pointers), every
23040        // externally-referenced buffer is stable-pointer by design (fa_part_pool retires-on-
23041        // grow and never frees, resident counters/scratch, cache set in place), and the
23042        // draft-graph path additionally rides capture_graph_retained (capture_keep holds all
23043        // warmup+capture allocs alive). One warmup therefore suffices for kernel-attr
23044        // settling and pool mapping. Arbitrated adversarially, not by taste:
23045        // graph-warmup-stress (pool-growth cycles large<->small x10, overlap arm, forced
23046        // recaptures over freed blocks — bit-identity vs eager + canary teeth) is GREEN at
23047        // warmups=1 on the deployment rig, plus graph-decode-gate 256-step bit-identity,
23048        // graph-session-gate, run-spec K=1..8 (receipts research/graph-warmups-5090-20260805/
23049        // + the pod's research/graph-allocfree-20260805/). Measured: recapture -38..-42% q27 /
23050        // -41% q9, decode +~1%, capture+prime -13ms. MEMRA_GRAPH_WARMUPS=2 = the rollback
23051        // seam; tools/graph-warmup-stress-gate.sh = the gate any regression re-runs.
23052        let warmups = {
23053            static W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
23054            *W.get_or_init(|| {
23055                std::env::var("MEMRA_GRAPH_WARMUPS")
23056                    .ok()
23057                    .and_then(|v| v.parse().ok())
23058                    .filter(|n| *n >= 1)
23059                    .unwrap_or(1)
23060            })
23061        };
23062        let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
23063            let t_w = std::time::Instant::now();
23064            // warmup: inline runs (no capture) so allocator pointers + kernel attrs are stable.
23065            for _ in 0..warmups {
23066                step(self)?;
23067            }
23068            self.gpu.stream().synchronize()?;
23069            let ms_warm = t_w.elapsed().as_secs_f64() * 1e3;
23070            // capture the third run.
23071            let t_c = std::time::Instant::now();
23072            self.gpu
23073                .stream()
23074                .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
23075            // If the body errors mid-capture, end the capture before propagating so the stream isn't
23076            // left in a capturing state.
23077            let r = step(self);
23078            let ms_body = t_c.elapsed().as_secs_f64() * 1e3;
23079            let t_i = std::time::Instant::now();
23080            let g = self.gpu.stream().end_capture(iflag);
23081            let ms_inst = t_i.elapsed().as_secs_f64() * 1e3;
23082            r?;
23083            let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
23084            let t_u = std::time::Instant::now();
23085            graph.upload()?;
23086            if ct {
23087                println!(
23088                    "[graph-captime] warmup2x {ms_warm:.2} ms  capture-body {ms_body:.2} ms  \
23089                          instantiate {ms_inst:.2} ms  upload {:.2} ms",
23090                    t_u.elapsed().as_secs_f64() * 1e3
23091                );
23092            }
23093            Ok(graph)
23094        };
23095        let result = run();
23096        if was_tracking {
23097            unsafe {
23098                self.gpu.ctx.enable_event_tracking();
23099            }
23100        }
23101        result
23102    }
23103
23104    /// gdn_scan variant where state_in/out are CudaViews (resident SSM state, in-place per step).
23105    pub fn gdn_scan_s128_view(
23106        &self,
23107        q: &CudaSlice<f32>,
23108        k: &CudaSlice<f32>,
23109        v: &CudaSlice<f32>,
23110        g: &CudaSlice<f32>,
23111        beta: &CudaSlice<f32>,
23112        state_in: &cudarc::driver::CudaView<f32>,
23113        state_out: &mut cudarc::driver::CudaViewMut<f32>,
23114        o: &mut CudaSlice<f32>,
23115        n_head: usize,
23116        t: usize,
23117        scale: f32,
23118    ) -> Result<(), Box<dyn std::error::Error>> {
23119        let f = self.func("gdn_scan_s128");
23120        const S_V: u32 = 128;
23121        const WARP: u32 = 32;
23122        const COLS: u32 = 4;
23123        let cfg = LaunchConfig {
23124            grid_dim: (n_head as u32, 1, S_V / COLS),
23125            block_dim: (WARP, COLS, 1),
23126            shared_mem_bytes: 0,
23127        };
23128        let (h, ti) = (n_head as i32, t as i32);
23129        let __s_b = self.gpu.stream();
23130        let mut b = __s_b.launch_builder(&f);
23131        b.arg(q)
23132            .arg(k)
23133            .arg(v)
23134            .arg(g)
23135            .arg(beta)
23136            .arg(state_in)
23137            .arg(state_out)
23138            .arg(o)
23139            .arg(&h)
23140            .arg(&ti)
23141            .arg(&scale);
23142        unsafe {
23143            b.launch(cfg)?;
23144        }
23145        Ok(())
23146    }
23147
23148    /// conv1d where the input is a CudaView (resident conv state assembled in place).
23149    pub fn ssm_conv1d_view(
23150        &self,
23151        x: &cudarc::driver::CudaView<f32>,
23152        w: &CudaSlice<f32>,
23153        y: &mut CudaSlice<f32>,
23154        conv_dim: usize,
23155        t: usize,
23156        d_conv: usize,
23157        silu: bool,
23158    ) -> Result<(), Box<dyn std::error::Error>> {
23159        let f = self.func("ssm_conv1d_silu_f32");
23160        // grid.x = channel, grid.y = T-tiles (block 256 strides over T) — parallel over both axes.
23161        let cfg = LaunchConfig {
23162            grid_dim: (conv_dim as u32, ((t as u32 + 255) / 256).max(1), 1),
23163            block_dim: (256, 1, 1),
23164            shared_mem_bytes: 0,
23165        };
23166        let (cd, ti, dc, s) = (conv_dim as i32, t as i32, d_conv as i32, silu as i32);
23167        let __s_b = self.gpu.stream();
23168        let mut b = __s_b.launch_builder(&f);
23169        b.arg(x).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc).arg(&s);
23170        unsafe {
23171            b.launch(cfg)?;
23172        }
23173        Ok(())
23174    }
23175
23176    /// Depthwise causal conv1d + optional SiLU.
23177    /// x:[conv_dim, T+d_conv-1] channel-major (first d_conv-1 cols = carried state),
23178    /// w:[d_conv, conv_dim] kernel-major, y:[conv_dim, T] channel-major.
23179    /// FUSED prefill conv (token-major input, zero left-state): replaces
23180    /// transpose + zeros + conv_left_pad + ssm_conv1d with ONE launch reading the matmul output
23181    /// directly. Output channel-major [conv_dim, T], SiLU applied. BIT-IDENTICAL accumulation.
23182    pub fn ssm_conv1d_tm(
23183        &self,
23184        qkv_tm: &CudaSlice<f32>,
23185        w: &CudaSlice<f32>,
23186        y: &mut CudaSlice<f32>,
23187        conv_dim: usize,
23188        t: usize,
23189        d_conv: usize,
23190    ) -> Result<(), Box<dyn std::error::Error>> {
23191        let f = self.func("ssm_conv1d_tm_f32");
23192        let cfg = LaunchConfig {
23193            grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
23194            block_dim: (256, 1, 1),
23195            shared_mem_bytes: 0,
23196        };
23197        let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
23198        let __s_b = self.gpu.stream();
23199        let mut b = __s_b.launch_builder(&f);
23200        b.arg(qkv_tm).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc);
23201        unsafe {
23202            b.launch(cfg)?;
23203        }
23204        Ok(())
23205    }
23206
23207    /// BATCHED verify conv (T>1, carried state): window reads the resident conv ring for
23208    /// negative rows; separate ring-update launch afterwards. BIT-IDENTICAL per value to the
23209    /// T=1 chain. T >= pad rides the pure input-column ring update (unchanged legacy path);
23210    /// T < pad (the MEMRA_SPEC_M2 t=2 verify arm) needs old-ring sources for the roll — the
23211    /// update kernel would race reading the ring it rewrites, so that arm clones the ring
23212    /// (dtod) and rolls via ssm_conv_ring_rebuild (PURE COPIES: the ring stores raw input
23213    /// columns; the final ring == what T sequential decode ring rolls leave).
23214    pub fn ssm_conv1d_tm_state(
23215        &self,
23216        qkv_tm: &CudaSlice<f32>,
23217        conv_state: &mut CudaSlice<f32>,
23218        w: &CudaSlice<f32>,
23219        y: &mut CudaSlice<f32>,
23220        conv_dim: usize,
23221        t: usize,
23222        d_conv: usize,
23223    ) -> Result<(), Box<dyn std::error::Error>> {
23224        self.ssm_conv1d_tm_state_pad(qkv_tm, conv_state, w, y, conv_dim, t, d_conv, None)
23225    }
23226
23227    /// task #14: `pad_len` = device true length for PADDED prime graphs — the ring update
23228    /// reads rows [len-pad, len) instead of the pad tail. None = the classic host-T path.
23229    #[allow(clippy::too_many_arguments)]
23230    pub fn ssm_conv1d_tm_state_pad(
23231        &self,
23232        qkv_tm: &CudaSlice<f32>,
23233        conv_state: &mut CudaSlice<f32>,
23234        w: &CudaSlice<f32>,
23235        y: &mut CudaSlice<f32>,
23236        conv_dim: usize,
23237        t: usize,
23238        d_conv: usize,
23239        pad_len: Option<&CudaSlice<i32>>,
23240    ) -> Result<(), Box<dyn std::error::Error>> {
23241        assert!(t >= 1, "ssm_conv1d_tm_state requires T >= 1");
23242        // clone BEFORE the window kernel is issued is not required (stream-ordered: the dtod and
23243        // the window kernel both read the pre-roll ring; the roll launches after both) — but
23244        // cloning first keeps the ordering trivially correct under any future stream split.
23245        let ring_old = if t < d_conv - 1 {
23246            Some(self.clone_dtod(conv_state)?)
23247        } else {
23248            None
23249        };
23250        {
23251            let f = self.func("ssm_conv1d_tm_state_f32");
23252            let cfg = LaunchConfig {
23253                grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
23254                block_dim: (256, 1, 1),
23255                shared_mem_bytes: 0,
23256            };
23257            let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
23258            let __s_b = self.gpu.stream();
23259            let mut b = __s_b.launch_builder(&f);
23260            b.arg(qkv_tm)
23261                .arg(&*conv_state)
23262                .arg(w)
23263                .arg(y)
23264                .arg(&cd)
23265                .arg(&ti)
23266                .arg(&dc);
23267            unsafe {
23268                b.launch(cfg)?;
23269            }
23270        }
23271        match (ring_old, pad_len) {
23272            (None, Some(len_d)) => {
23273                let f = self.func("ssm_conv_ring_update_dev_f32");
23274                let n = conv_dim * (d_conv - 1);
23275                let cfg = LaunchConfig::for_num_elems(n as u32);
23276                let (cd, dc) = (conv_dim as i32, d_conv as i32);
23277                let __s_b = self.gpu.stream();
23278                let mut b = __s_b.launch_builder(&f);
23279                b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
23280                unsafe {
23281                    b.launch(cfg)?;
23282                }
23283            }
23284            (None, None) => {
23285                let f = self.func("ssm_conv_ring_update_f32");
23286                let n = conv_dim * (d_conv - 1);
23287                let cfg = LaunchConfig::for_num_elems(n as u32);
23288                let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
23289                let __s_b = self.gpu.stream();
23290                let mut b = __s_b.launch_builder(&f);
23291                b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
23292                unsafe {
23293                    b.launch(cfg)?;
23294                }
23295            }
23296            (Some(old), _) => {
23297                self.ssm_conv_ring_rebuild(qkv_tm, &old, conv_state, conv_dim, t, d_conv)?
23298            }
23299        }
23300        Ok(())
23301    }
23302
23303    /// qkv-view twin (task #16): batched prime reads the concat GEMM output directly.
23304    pub fn ssm_conv1d_tm_state_pad_v(
23305        &self,
23306        qkv_tm: &cudarc::driver::CudaView<f32>,
23307        conv_state: &mut CudaSlice<f32>,
23308        w: &CudaSlice<f32>,
23309        y: &mut CudaSlice<f32>,
23310        conv_dim: usize,
23311        t: usize,
23312        d_conv: usize,
23313        pad_len: Option<&CudaSlice<i32>>,
23314    ) -> Result<(), Box<dyn std::error::Error>> {
23315        assert!(t >= 1, "ssm_conv1d_tm_state requires T >= 1");
23316        // clone BEFORE the window kernel is issued is not required (stream-ordered: the dtod and
23317        // the window kernel both read the pre-roll ring; the roll launches after both) — but
23318        // cloning first keeps the ordering trivially correct under any future stream split.
23319        let ring_old = if t < d_conv - 1 {
23320            Some(self.clone_dtod(conv_state)?)
23321        } else {
23322            None
23323        };
23324        {
23325            let f = self.func("ssm_conv1d_tm_state_f32");
23326            let cfg = LaunchConfig {
23327                grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
23328                block_dim: (256, 1, 1),
23329                shared_mem_bytes: 0,
23330            };
23331            let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
23332            let __s_b = self.gpu.stream();
23333            let mut b = __s_b.launch_builder(&f);
23334            b.arg(qkv_tm)
23335                .arg(&*conv_state)
23336                .arg(w)
23337                .arg(y)
23338                .arg(&cd)
23339                .arg(&ti)
23340                .arg(&dc);
23341            unsafe {
23342                b.launch(cfg)?;
23343            }
23344        }
23345        match (ring_old, pad_len) {
23346            (None, Some(len_d)) => {
23347                let f = self.func("ssm_conv_ring_update_dev_f32");
23348                let n = conv_dim * (d_conv - 1);
23349                let cfg = LaunchConfig::for_num_elems(n as u32);
23350                let (cd, dc) = (conv_dim as i32, d_conv as i32);
23351                let __s_b = self.gpu.stream();
23352                let mut b = __s_b.launch_builder(&f);
23353                b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
23354                unsafe {
23355                    b.launch(cfg)?;
23356                }
23357            }
23358            (None, None) => {
23359                let f = self.func("ssm_conv_ring_update_f32");
23360                let n = conv_dim * (d_conv - 1);
23361                let cfg = LaunchConfig::for_num_elems(n as u32);
23362                let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
23363                let __s_b = self.gpu.stream();
23364                let mut b = __s_b.launch_builder(&f);
23365                b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
23366                unsafe {
23367                    b.launch(cfg)?;
23368                }
23369            }
23370            (Some(_), _) => unreachable!(
23371                "ssm_conv1d_tm_state_pad_v: T < d_conv-1 has no view path (PRIME_MIN_T gates it)"
23372            ),
23373        }
23374        Ok(())
23375    }
23376
23377    /// PREFIX conv-ring rebuild (spec REPLAY-FREE partial accept): overwrite the resident ring
23378    /// with the state a T=1 chain holds after only the FIRST `tc` columns of `qkv_tm` — the last
23379    /// `pad` entries of [ring_old | cols 0..tc-1]. PURE COPIES (the ring stores raw inputs; no
23380    /// arithmetic, cannot perturb FP order). `ring_old` = the pre-round snapshot ring.
23381    pub fn ssm_conv_ring_rebuild(
23382        &self,
23383        qkv_tm: &CudaSlice<f32>,
23384        ring_old: &CudaSlice<f32>,
23385        conv_state: &mut CudaSlice<f32>,
23386        conv_dim: usize,
23387        tc: usize,
23388        d_conv: usize,
23389    ) -> Result<(), Box<dyn std::error::Error>> {
23390        let f = self.func("ssm_conv_ring_rebuild_f32");
23391        let n = conv_dim * (d_conv - 1);
23392        let cfg = LaunchConfig::for_num_elems(n as u32);
23393        let (cd, ti, dc) = (conv_dim as i32, tc as i32, d_conv as i32);
23394        let __s_b = self.gpu.stream();
23395        let mut b = __s_b.launch_builder(&f);
23396        b.arg(qkv_tm)
23397            .arg(ring_old)
23398            .arg(conv_state)
23399            .arg(&cd)
23400            .arg(&ti)
23401            .arg(&dc);
23402        unsafe {
23403            b.launch(cfg)?;
23404        }
23405        Ok(())
23406    }
23407
23408    /// FUSED decode GDN prep (T=1): repack + q/k L2-norm + beta sigmoid + g_log in one launch.
23409    /// Replaces 5 tiny serialized kernels on the decode critical path. L2 reduce runs as a 32-lane
23410    /// warp tree (vs l2_norm_f32's 256-thread two-level tree) — same math, different FP sum order;
23411    /// the argmax + run-spec gates are the authority.
23412    #[allow(clippy::too_many_arguments)]
23413    pub fn gdn_prep_decode(
23414        &self,
23415        conv_out: &CudaSlice<f32>,
23416        beta_raw: &CudaSlice<f32>,
23417        alpha: &CudaSlice<f32>,
23418        dt_bias: &CudaSlice<f32>,
23419        a: &CudaSlice<f32>,
23420        q_l2: &mut CudaSlice<f32>,
23421        k_l2: &mut CudaSlice<f32>,
23422        v_g: &mut CudaSlice<f32>,
23423        beta: &mut CudaSlice<f32>,
23424        g_log: &mut CudaSlice<f32>,
23425        d_state: usize,
23426        num_v: usize,
23427        num_k: usize,
23428        key_dim: usize,
23429        eps: f32,
23430    ) -> Result<(), Box<dyn std::error::Error>> {
23431        let f = self.func("gdn_prep_decode_f32");
23432        let cfg = LaunchConfig {
23433            grid_dim: (num_v as u32, 1, 1),
23434            block_dim: (32, 4, 1),
23435            shared_mem_bytes: 0,
23436        };
23437        let (ds, nv, nk, kd) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
23438        let __s_b = self.gpu.stream();
23439        let mut b = __s_b.launch_builder(&f);
23440        b.arg(conv_out)
23441            .arg(beta_raw)
23442            .arg(alpha)
23443            .arg(dt_bias)
23444            .arg(a)
23445            .arg(q_l2)
23446            .arg(k_l2)
23447            .arg(v_g)
23448            .arg(beta)
23449            .arg(g_log)
23450            .arg(&ds)
23451            .arg(&nv)
23452            .arg(&nk)
23453            .arg(&kd)
23454            .arg(&eps);
23455        unsafe {
23456            b.launch(cfg)?;
23457        }
23458        Ok(())
23459    }
23460
23461    /// FUSED prefill conv + GDN repack: token-major qkv -> q_g/k_g/v_g in ONE launch (no conv_out
23462    /// materialization, no qkv_to_gdn_repack pass). BIT-IDENTICAL values; scatter matches
23463    /// qkv_to_gdn_repack's modulo head-repeat mapping exactly.
23464    #[allow(clippy::too_many_arguments)]
23465    pub fn ssm_conv1d_gdn(
23466        &self,
23467        qkv_tm: &CudaSlice<f32>,
23468        w: &CudaSlice<f32>,
23469        q_g: &mut CudaSlice<f32>,
23470        k_g: &mut CudaSlice<f32>,
23471        v_g: &mut CudaSlice<f32>,
23472        conv_dim: usize,
23473        t: usize,
23474        d_conv: usize,
23475        d_state: usize,
23476        num_v: usize,
23477        num_k: usize,
23478        key_dim: usize,
23479    ) -> Result<(), Box<dyn std::error::Error>> {
23480        let f = self.func("ssm_conv1d_gdn_f32");
23481        let cfg = LaunchConfig {
23482            grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
23483            block_dim: (256, 1, 1),
23484            shared_mem_bytes: 0,
23485        };
23486        let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
23487        let (ds, nv, nk, kd) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
23488        let __s_b = self.gpu.stream();
23489        let mut b = __s_b.launch_builder(&f);
23490        b.arg(qkv_tm)
23491            .arg(w)
23492            .arg(q_g)
23493            .arg(k_g)
23494            .arg(v_g)
23495            .arg(&cd)
23496            .arg(&ti)
23497            .arg(&dc)
23498            .arg(&ds)
23499            .arg(&nv)
23500            .arg(&nk)
23501            .arg(&kd);
23502        unsafe {
23503            b.launch(cfg)?;
23504        }
23505        Ok(())
23506    }
23507
23508    pub fn ssm_conv1d(
23509        &self,
23510        x: &CudaSlice<f32>,
23511        w: &CudaSlice<f32>,
23512        y: &mut CudaSlice<f32>,
23513        conv_dim: usize,
23514        t: usize,
23515        d_conv: usize,
23516        silu: bool,
23517    ) -> Result<(), Box<dyn std::error::Error>> {
23518        let f = self.func("ssm_conv1d_silu_f32");
23519        let cfg = LaunchConfig {
23520            grid_dim: (conv_dim as u32, ((t as u32 + 255) / 256).max(1), 1),
23521            block_dim: (256, 1, 1),
23522            shared_mem_bytes: 0,
23523        };
23524        let (cd, ti, dc, s) = (conv_dim as i32, t as i32, d_conv as i32, silu as i32);
23525        let __s_b = self.gpu.stream();
23526        let mut b = __s_b.launch_builder(&f);
23527        b.arg(x).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc).arg(&s);
23528        unsafe {
23529            b.launch(cfg)?;
23530        }
23531        Ok(())
23532    }
23533
23534    /// Gated DeltaNet scan, S_v=128. q,k,v:[128,H,T]; g,beta:[H,T]; state:[128,128,H] transposed;
23535    /// o:[128,H,T]. Single sequence.
23536    pub fn gdn_scan_s128(
23537        &self,
23538        q: &CudaSlice<f32>,
23539        k: &CudaSlice<f32>,
23540        v: &CudaSlice<f32>,
23541        g: &CudaSlice<f32>,
23542        beta: &CudaSlice<f32>,
23543        state_in: &CudaSlice<f32>,
23544        state_out: &mut CudaSlice<f32>,
23545        o: &mut CudaSlice<f32>,
23546        n_head: usize,
23547        t: usize,
23548        scale: f32,
23549    ) -> Result<(), Box<dyn std::error::Error>> {
23550        let f = self.func("gdn_scan_s128");
23551        const S_V: u32 = 128;
23552        const WARP: u32 = 32;
23553        const COLS_PER_BLOCK: u32 = 4;
23554        let cfg = LaunchConfig {
23555            grid_dim: (n_head as u32, 1, S_V / COLS_PER_BLOCK),
23556            block_dim: (WARP, COLS_PER_BLOCK, 1),
23557            shared_mem_bytes: 0,
23558        };
23559        let (h, ti) = (n_head as i32, t as i32);
23560        let __s_b = self.gpu.stream();
23561        let mut b = __s_b.launch_builder(&f);
23562        b.arg(q)
23563            .arg(k)
23564            .arg(v)
23565            .arg(g)
23566            .arg(beta)
23567            .arg(state_in)
23568            .arg(state_out)
23569            .arg(o)
23570            .arg(&h)
23571            .arg(&ti)
23572            .arg(&scale);
23573        unsafe {
23574            b.launch(cfg)?;
23575        }
23576        Ok(())
23577    }
23578
23579    // ==== B2' batched decode state ops (decode_batch.rs) ====
23580    // Per-seq state pointers ride device u64 arrays (views into the per-step pointer table).
23581    // Bodies are the single-seq kernels per sequence — bit-identical per row.
23582
23583    #[allow(clippy::too_many_arguments)]
23584    pub fn ssm_conv1d_fused_decode_b(
23585        &self,
23586        qkv_cols: &CudaSlice<f32>,
23587        conv_state_ptrs: &cudarc::driver::CudaView<u64>,
23588        w: &CudaSlice<f32>,
23589        conv_outs: &mut CudaSlice<f32>,
23590        conv_dim: usize,
23591        d_conv: usize,
23592        b_n: usize,
23593    ) -> Result<(), Box<dyn std::error::Error>> {
23594        let f = self.func("ssm_conv1d_fused_decode_b_f32");
23595        let cfg = LaunchConfig {
23596            grid_dim: (((conv_dim + 255) / 256) as u32, 1, b_n as u32),
23597            block_dim: (256, 1, 1),
23598            shared_mem_bytes: 0,
23599        };
23600        let (cd, dc) = (conv_dim as i32, d_conv as i32);
23601        let __s_b = self.gpu.stream();
23602        let mut b = __s_b.launch_builder(&f);
23603        b.arg(qkv_cols)
23604            .arg(conv_state_ptrs)
23605            .arg(w)
23606            .arg(conv_outs)
23607            .arg(&cd)
23608            .arg(&dc);
23609        unsafe {
23610            b.launch(cfg)?;
23611        }
23612        Ok(())
23613    }
23614
23615    #[allow(clippy::too_many_arguments)]
23616    pub fn gdn_prep_decode_b(
23617        &self,
23618        conv_outs: &CudaSlice<f32>,
23619        beta_raws: &CudaSlice<f32>,
23620        alphas: &CudaSlice<f32>,
23621        dt_bias: &CudaSlice<f32>,
23622        a: &CudaSlice<f32>,
23623        q_l2: &mut CudaSlice<f32>,
23624        k_l2: &mut CudaSlice<f32>,
23625        v_g: &mut CudaSlice<f32>,
23626        beta: &mut CudaSlice<f32>,
23627        g_log: &mut CudaSlice<f32>,
23628        d_state: usize,
23629        num_v: usize,
23630        num_k: usize,
23631        key_dim: usize,
23632        eps: f32,
23633        conv_dim: usize,
23634        b_n: usize,
23635    ) -> Result<(), Box<dyn std::error::Error>> {
23636        let f = self.func("gdn_prep_decode_b_f32");
23637        let cfg = LaunchConfig {
23638            grid_dim: (num_v as u32, 1, b_n as u32),
23639            block_dim: (32, 4, 1),
23640            shared_mem_bytes: 0,
23641        };
23642        let (ds, nv, nk, kd, cd) = (
23643            d_state as i32,
23644            num_v as i32,
23645            num_k as i32,
23646            key_dim as i32,
23647            conv_dim as i32,
23648        );
23649        let __s_b = self.gpu.stream();
23650        let mut b = __s_b.launch_builder(&f);
23651        b.arg(conv_outs)
23652            .arg(beta_raws)
23653            .arg(alphas)
23654            .arg(dt_bias)
23655            .arg(a)
23656            .arg(q_l2)
23657            .arg(k_l2)
23658            .arg(v_g)
23659            .arg(beta)
23660            .arg(g_log)
23661            .arg(&ds)
23662            .arg(&nv)
23663            .arg(&nk)
23664            .arg(&kd)
23665            .arg(&eps)
23666            .arg(&cd);
23667        unsafe {
23668            b.launch(cfg)?;
23669        }
23670        Ok(())
23671    }
23672
23673    #[allow(clippy::too_many_arguments)]
23674    pub fn gdn_scan_s128_batched(
23675        &self,
23676        q: &CudaSlice<f32>,
23677        k: &CudaSlice<f32>,
23678        v: &CudaSlice<f32>,
23679        g: &CudaSlice<f32>,
23680        beta: &CudaSlice<f32>,
23681        state_in_ptrs: &cudarc::driver::CudaView<u64>,
23682        state_out_ptrs: &cudarc::driver::CudaView<u64>,
23683        o: &mut CudaSlice<f32>,
23684        n_head: usize,
23685        b_n: usize,
23686        scale: f32,
23687    ) -> Result<(), Box<dyn std::error::Error>> {
23688        let f = self.func("gdn_scan_s128_b");
23689        const S_V: u32 = 128;
23690        const WARP: u32 = 32;
23691        const COLS_PER_BLOCK: u32 = 4;
23692        let cfg = LaunchConfig {
23693            grid_dim: (n_head as u32, b_n as u32, S_V / COLS_PER_BLOCK),
23694            block_dim: (WARP, COLS_PER_BLOCK, 1),
23695            shared_mem_bytes: 0,
23696        };
23697        let h = n_head as i32;
23698        let __s_b = self.gpu.stream();
23699        let mut b = __s_b.launch_builder(&f);
23700        b.arg(q)
23701            .arg(k)
23702            .arg(v)
23703            .arg(g)
23704            .arg(beta)
23705            .arg(state_in_ptrs)
23706            .arg(state_out_ptrs)
23707            .arg(o)
23708            .arg(&h)
23709            .arg(&scale);
23710        unsafe {
23711            b.launch(cfg)?;
23712        }
23713        Ok(())
23714    }
23715
23716    /// VIEW twins of the three GDN decode state wrappers (lane/verify-launchslim): identical
23717    /// launches, row args as views into the caller's packed [T, ...] buffers — the t-parallel
23718    /// verify's per-row loop passes slices instead of paying an arithmetic-free dtod per row
23719    /// per kernel (48 layers x T rows x 4 copies/round on the money path). Same kernels, same
23720    /// numeric class; only the pointer arithmetic moved host-side.
23721    #[allow(clippy::too_many_arguments)]
23722    pub fn ssm_conv1d_fused_decode_b_view(
23723        &self,
23724        qkv_cols: &cudarc::driver::CudaView<f32>,
23725        conv_state_ptrs: &cudarc::driver::CudaView<u64>,
23726        w: &CudaSlice<f32>,
23727        conv_outs: &mut CudaSlice<f32>,
23728        conv_dim: usize,
23729        d_conv: usize,
23730        b_n: usize,
23731    ) -> Result<(), Box<dyn std::error::Error>> {
23732        let f = self.func("ssm_conv1d_fused_decode_b_f32");
23733        let cfg = LaunchConfig {
23734            grid_dim: (((conv_dim + 255) / 256) as u32, 1, b_n as u32),
23735            block_dim: (256, 1, 1),
23736            shared_mem_bytes: 0,
23737        };
23738        let (cd, dc) = (conv_dim as i32, d_conv as i32);
23739        let __s_b = self.gpu.stream();
23740        let mut b = __s_b.launch_builder(&f);
23741        b.arg(qkv_cols)
23742            .arg(conv_state_ptrs)
23743            .arg(w)
23744            .arg(conv_outs)
23745            .arg(&cd)
23746            .arg(&dc);
23747        unsafe {
23748            b.launch(cfg)?;
23749        }
23750        Ok(())
23751    }
23752
23753    #[allow(clippy::too_many_arguments)]
23754    pub fn gdn_prep_decode_b_view(
23755        &self,
23756        conv_outs: &CudaSlice<f32>,
23757        beta_raws: &cudarc::driver::CudaView<f32>,
23758        alphas: &cudarc::driver::CudaView<f32>,
23759        dt_bias: &CudaSlice<f32>,
23760        a: &CudaSlice<f32>,
23761        q_l2: &mut CudaSlice<f32>,
23762        k_l2: &mut CudaSlice<f32>,
23763        v_g: &mut CudaSlice<f32>,
23764        beta: &mut CudaSlice<f32>,
23765        g_log: &mut CudaSlice<f32>,
23766        d_state: usize,
23767        num_v: usize,
23768        num_k: usize,
23769        key_dim: usize,
23770        eps: f32,
23771        conv_dim: usize,
23772        b_n: usize,
23773    ) -> Result<(), Box<dyn std::error::Error>> {
23774        let f = self.func("gdn_prep_decode_b_f32");
23775        let cfg = LaunchConfig {
23776            grid_dim: (num_v as u32, 1, b_n as u32),
23777            block_dim: (32, 4, 1),
23778            shared_mem_bytes: 0,
23779        };
23780        let (ds, nv, nk, kd, cd) = (
23781            d_state as i32,
23782            num_v as i32,
23783            num_k as i32,
23784            key_dim as i32,
23785            conv_dim as i32,
23786        );
23787        let __s_b = self.gpu.stream();
23788        let mut b = __s_b.launch_builder(&f);
23789        b.arg(conv_outs)
23790            .arg(beta_raws)
23791            .arg(alphas)
23792            .arg(dt_bias)
23793            .arg(a)
23794            .arg(q_l2)
23795            .arg(k_l2)
23796            .arg(v_g)
23797            .arg(beta)
23798            .arg(g_log)
23799            .arg(&ds)
23800            .arg(&nv)
23801            .arg(&nk)
23802            .arg(&kd)
23803            .arg(&eps)
23804            .arg(&cd);
23805        unsafe {
23806            b.launch(cfg)?;
23807        }
23808        Ok(())
23809    }
23810
23811    #[allow(clippy::too_many_arguments)]
23812    pub fn gdn_scan_s128_batched_view(
23813        &self,
23814        q: &CudaSlice<f32>,
23815        k: &CudaSlice<f32>,
23816        v: &CudaSlice<f32>,
23817        g: &CudaSlice<f32>,
23818        beta: &CudaSlice<f32>,
23819        state_in_ptrs: &cudarc::driver::CudaView<u64>,
23820        state_out_ptrs: &cudarc::driver::CudaView<u64>,
23821        o: &mut cudarc::driver::CudaViewMut<f32>,
23822        n_head: usize,
23823        b_n: usize,
23824        scale: f32,
23825    ) -> Result<(), Box<dyn std::error::Error>> {
23826        let f = self.func("gdn_scan_s128_b");
23827        const S_V: u32 = 128;
23828        const WARP: u32 = 32;
23829        const COLS_PER_BLOCK: u32 = 4;
23830        let cfg = LaunchConfig {
23831            grid_dim: (n_head as u32, b_n as u32, S_V / COLS_PER_BLOCK),
23832            block_dim: (WARP, COLS_PER_BLOCK, 1),
23833            shared_mem_bytes: 0,
23834        };
23835        let h = n_head as i32;
23836        let __s_b = self.gpu.stream();
23837        let mut b = __s_b.launch_builder(&f);
23838        b.arg(q)
23839            .arg(k)
23840            .arg(v)
23841            .arg(g)
23842            .arg(beta)
23843            .arg(state_in_ptrs)
23844            .arg(state_out_ptrs)
23845            .arg(o)
23846            .arg(&h)
23847            .arg(&scale);
23848        unsafe {
23849            b.launch(cfg)?;
23850        }
23851        Ok(())
23852    }
23853
23854    /// A4 seam: chunked WY GDN prefill. DEFAULT ON (`MEMRA_GDN_CHUNKED=0` = rollback to the
23855    /// sequential scan). Flipped 2026-07-04 with the full battery green: kernel-check ALL
23856    /// GREEN x {9B, 27B} incl the f64-truth chunk gates; run-gen argmax 82==82 both models
23857    /// on AND off (24/24 sweep runs); run-spec K={1,2,3,4,6,8} PASS x {9B synth, 9B text,
23858    /// 27B p2, 27B p3}; e2e first-16-token agreement 6/6 (full-256 drifts at index 47-125
23859    /// on 5/6 prompts — accepted cache-state-FP class, batched-prime precedent).
23860    /// PREFILL-ONLY: decode + spec verify never route here (decode==verify dispatch
23861    /// identity law); prime_cache/forward/forward_last are the only callers.
23862    pub fn gdn_chunked_enabled() -> bool {
23863        static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
23864        *E.get_or_init(|| {
23865            std::env::var("MEMRA_GDN_CHUNKED")
23866                .map(|v| v != "0")
23867                .unwrap_or(true)
23868        })
23869    }
23870
23871    /// A4 chunk size (MEMRA_GDN_CHUNK, default 32 — the sweep winner: the O(T*C) chunk
23872    /// matrices grow with C while the sequential state pass is C-flat, so smaller chunks
23873    /// win; C=32/64 also get the register-history solve template). Clamped to multiples
23874    /// of 32 in [32, 128] (kernel row mappings require it).
23875    pub fn gdn_chunk_size() -> usize {
23876        static C: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
23877        *C.get_or_init(|| {
23878            let c: usize = std::env::var("MEMRA_GDN_CHUNK")
23879                .ok()
23880                .and_then(|v| v.parse().ok())
23881                .unwrap_or(32);
23882            c.clamp(32, 128) / 32 * 32
23883        })
23884    }
23885
23886    /// A4: chunked WY / blockwise-inverse GDN prefill (see cu/hybrid.cu K1-K5 header for the
23887    /// math). Same contract as `gdn_scan_s128` (layouts, state ping-pong) but chunk-parallel:
23888    /// NOT bit-identical to the sequential scan (chunked FP accumulation order); run-gen
23889    /// argmax + run-spec batteries are the accuracy authority. PREFILL callers only.
23890    #[allow(clippy::too_many_arguments)]
23891    /// task #18: K1-K3 of the chunked WY scan (shared by the per-seq path and the
23892    /// batched-prime varlen path). Returns (gcum, P, U, W); `A` is K3-internal.
23893    #[allow(clippy::too_many_arguments, clippy::type_complexity)]
23894    #[allow(clippy::too_many_arguments)]
23895    pub fn gdn_chunk_k123(
23896        &self,
23897        q: &CudaSlice<f32>,
23898        k: &CudaSlice<f32>,
23899        v: &CudaSlice<f32>,
23900        g: &CudaSlice<f32>,
23901        beta: &CudaSlice<f32>,
23902        wb16: Option<&mut CudaSlice<u8>>,
23903        n_head: usize,
23904        t: usize,
23905        c: usize,
23906        hk: usize,
23907        k2w: Option<(&CudaSlice<u8>, &CudaSlice<u8>, &mut CudaSlice<u8>)>,
23908    ) -> Result<
23909        (
23910            CudaSlice<f32>,
23911            CudaSlice<f32>,
23912            CudaSlice<f32>,
23913            CudaSlice<f32>,
23914        ),
23915        Box<dyn std::error::Error>,
23916    > {
23917        const D: usize = 128;
23918        let h = n_head;
23919        let nc = (t + c - 1) / c;
23920        let (hi, ti, ci) = (h as i32, t as i32, c as i32);
23921        let mut gcum = self.uninit(t * h)?;
23922        let mut a = self.uninit(nc * h * c * c)?;
23923        let mut p = self.uninit(nc * h * c * c)?;
23924        let mut u = self.uninit(nc * h * c * D)?;
23925        let mut w = self.uninit(nc * h * c * D)?;
23926        {
23927            // K1
23928            let f = self.func("gdn_chunk_cumgate_f32");
23929            let cfg = LaunchConfig {
23930                grid_dim: (nc as u32, h as u32, 1),
23931                block_dim: (32, 1, 1),
23932                shared_mem_bytes: 0,
23933            };
23934            let __s_b = self.gpu.stream();
23935            let mut b = __s_b.launch_builder(&f);
23936            b.arg(g).arg(&mut gcum).arg(&hi).arg(&ti).arg(&ci);
23937            unsafe {
23938                b.launch(cfg)?;
23939            }
23940        }
23941        if let Some((qb, kb, pb)) = k2w {
23942            // K2-wgmma (MEMRA_GDN_WGMMA path, c==32): A + pre-masked Pb16 in one kernel;
23943            // the P f32 buffer stays UNWRITTEN (its only wgmma-path consumer is Pb16).
23944            assert!(c == 32, "gdn_k2_wgmma is a C==32 tile");
23945            let f = self.func("gdn_k2_wgmma");
23946            let cfg = LaunchConfig {
23947                grid_dim: (nc as u32, h as u32, 1),
23948                block_dim: (128, 1, 1),
23949                shared_mem_bytes: 0,
23950            };
23951            let hki = hk as i32;
23952            let __s_b = self.gpu.stream();
23953            let mut b = __s_b.launch_builder(&f);
23954            b.arg(qb)
23955                .arg(kb)
23956                .arg(&gcum)
23957                .arg(beta)
23958                .arg(&mut a)
23959                .arg(&mut *pb)
23960                .arg(&hi)
23961                .arg(&ti)
23962                .arg(&ci)
23963                .arg(&hki);
23964            unsafe {
23965                b.launch(cfg)?;
23966            }
23967        } else if c <= 64 && !portable_mma_gated() {
23968            // K2 register-tiled (2x2 outputs/thread, whole-chunk smem k tile)
23969            let f = self.func("gdn_chunk_attn_f32");
23970            f.set_attribute(
23971                CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
23972                GDN_K2_DYNAMIC_SHARED_BYTES as i32,
23973            )?;
23974            let jt = ((c + 31) / 32) as u32;
23975            let cfg = LaunchConfig {
23976                grid_dim: (nc as u32, h as u32, jt),
23977                block_dim: (256, 1, 1),
23978                shared_mem_bytes: GDN_K2_DYNAMIC_SHARED_BYTES,
23979            };
23980            let hki = hk as i32;
23981            let __s_b = self.gpu.stream();
23982            let mut b = __s_b.launch_builder(&f);
23983            b.arg(q)
23984                .arg(k)
23985                .arg(&gcum)
23986                .arg(beta)
23987                .arg(&mut a)
23988                .arg(&mut p)
23989                .arg(&hi)
23990                .arg(&ti)
23991                .arg(&ci)
23992                .arg(&hki);
23993            unsafe {
23994                b.launch(cfg)?;
23995            }
23996        } else {
23997            // K2 generic (C = 128, or the portable target's low-smem fallback)
23998            assert!(
23999                hk == h,
24000                "generic K2 is broadcast-only (de-broadcast rides C==32)"
24001            );
24002            let f = self.func("gdn_chunk_attn_g_f32");
24003            let cfg = LaunchConfig {
24004                grid_dim: (nc as u32, h as u32, 1),
24005                block_dim: (32, 8, 1),
24006                shared_mem_bytes: 0,
24007            };
24008            let __s_b = self.gpu.stream();
24009            let mut b = __s_b.launch_builder(&f);
24010            b.arg(q)
24011                .arg(k)
24012                .arg(&gcum)
24013                .arg(beta)
24014                .arg(&mut a)
24015                .arg(&mut p)
24016                .arg(&hi)
24017                .arg(&ti)
24018                .arg(&ci);
24019            unsafe {
24020                b.launch(cfg)?;
24021            }
24022        }
24023        {
24024            // K3 (register-history templates for C=32/64; local-memory generic otherwise)
24025            let cfg = LaunchConfig {
24026                grid_dim: (nc as u32, h as u32, 1),
24027                block_dim: (256, 1, 1),
24028                shared_mem_bytes: 0,
24029            };
24030            match c {
24031                32 | 64 => {
24032                    let f = self.func(if c == 32 {
24033                        "gdn_chunk_solve32_f32"
24034                    } else {
24035                        "gdn_chunk_solve64_f32"
24036                    });
24037                    // mirror-fold: W's bf16 twin emitted on store (0 = skip)
24038                    let wb: u64 = match wb16 {
24039                        Some(d) => self.addr_u8(d),
24040                        None => 0,
24041                    };
24042                    let hki = hk as i32;
24043                    let __s_b = self.gpu.stream();
24044                    let mut b = __s_b.launch_builder(&f);
24045                    b.arg(v)
24046                        .arg(k)
24047                        .arg(&a)
24048                        .arg(&gcum)
24049                        .arg(&mut u)
24050                        .arg(&mut w)
24051                        .arg(&wb)
24052                        .arg(&hi)
24053                        .arg(&ti)
24054                        .arg(&hki);
24055                    unsafe {
24056                        b.launch(cfg)?;
24057                    }
24058                }
24059                _ => {
24060                    assert!(hk == h, "generic K3 is broadcast-only");
24061                    let f = self.func("gdn_chunk_solve_f32");
24062                    let __s_b = self.gpu.stream();
24063                    let mut b = __s_b.launch_builder(&f);
24064                    b.arg(v)
24065                        .arg(k)
24066                        .arg(&a)
24067                        .arg(&gcum)
24068                        .arg(&mut u)
24069                        .arg(&mut w)
24070                        .arg(&hi)
24071                        .arg(&ti)
24072                        .arg(&ci);
24073                    unsafe {
24074                        b.launch(cfg)?;
24075                    }
24076                }
24077            }
24078        }
24079        Ok((gcum, p, u, w))
24080    }
24081
24082    /// task #21 de-broadcast seam: q/k stored at num_k distinct GQA heads instead of
24083    /// the num_v broadcast. MEMRA_GDN_DB=0 reverts. Only the chunked prefill path
24084    /// consumes the compact layout (hk plumbed; hk == H reproduces broadcast exactly).
24085    pub fn gdn_db_on() -> bool {
24086        std::env::var("MEMRA_GDN_DB").as_deref() != Ok("0")
24087    }
24088
24089    /// Whether the K4/K5 mma pair serves at chunk size `c` (mirrors gdn_scan_chunked's
24090    /// seam read — env re-read per call ON PURPOSE, kernel-check pins both configs).
24091    /// DEFAULT ON for sm_120a builds too (lane/moeprime-nvfp4-direct, 2026-08-21): the pair
24092    /// was qualified on 90a only and left env-opt-in elsewhere; measured on Blackwell it
24093    /// wins on BOTH rigs — one RTX PRO 6000 (ornith15 pp14715 12,036 -> 12,751/12,957,
24094    /// +6-8%, both orders) and the local 5090 (q38-27b pp6435 1,397/1,429 -> 1,427/1,446,
24095    /// both orders) — with kernel-check/run-gen/margin-gate/run-spec green under the flag.
24096    /// bf16 HMMA (m16n8k16) is sm_80-class PTX; only the wgmma nest stays Hopper-gated.
24097    pub fn gdn_mma_enabled(&self, c: usize) -> bool {
24098        !portable_mma_gated()
24099            && c == 32
24100            && match std::env::var("MEMRA_GDN_MMA").as_deref() {
24101                Ok("1") => true,
24102                Ok("0") => false,
24103                _ => gdn_mma_default_on(),
24104            }
24105    }
24106
24107    /// task #22: whether the fused K4+K5 (+K2) wgmma path serves (nested inside the
24108    /// mma config; same per-call env read discipline). HARD-gated to the Hopper build:
24109    /// the wgmma asm bodies exist only at __CUDA_ARCH__ == 900 (MEMRA_K45_REAL,
24110    /// wgmma_common.cuh) — on every other arch the kernel compiles EMPTY, so an env
24111    /// force would silently produce garbage. Required since the sm_120a mma default
24112    /// flip made MEMRA_GDN_WGMMA=1 alone reach this branch there.
24113    pub fn gdn_wgmma_on(&self, c: usize) -> bool {
24114        cfg!(memra_hopper_mma)
24115            && self.gdn_mma_enabled(c)
24116            && std::env::var("MEMRA_GDN_WGMMA").as_deref() != Ok("0")
24117    }
24118
24119    /// task #18 conv-fuse: carried-ring conv + SiLU + GDN repack in ONE pass (the
24120    /// conv_out intermediate and its transposed re-read disappear — 11.8ms of the
24121    /// T=2048 prime). Ring update stays the separate follow-up launch (pad-aware).
24122    /// BIT-IDENTICAL values to ssm_conv1d_tm_state_pad + qkv_to_gdn_repack.
24123    #[allow(clippy::too_many_arguments)]
24124    pub fn ssm_conv1d_gdn_state_pad(
24125        &self,
24126        qkv_tm: &cudarc::driver::CudaView<f32>,
24127        conv_state: &mut CudaSlice<f32>,
24128        w: &CudaSlice<f32>,
24129        q_g: &mut CudaSlice<f32>,
24130        k_g: &mut CudaSlice<f32>,
24131        v_g: &mut CudaSlice<f32>,
24132        conv_dim: usize,
24133        t: usize,
24134        d_conv: usize,
24135        d_state: usize,
24136        num_v: usize,
24137        num_k: usize,
24138        key_dim: usize,
24139        hk: usize,
24140        pad_len: Option<&CudaSlice<i32>>,
24141    ) -> Result<(), Box<dyn std::error::Error>> {
24142        assert!(
24143            t >= d_conv - 1,
24144            "fused state conv requires T >= pad (PRIME_MIN_T gates)"
24145        );
24146        {
24147            let f = self.func("ssm_conv1d_gdn_state_f32");
24148            let cfg = LaunchConfig {
24149                grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
24150                block_dim: (256, 1, 1),
24151                shared_mem_bytes: 0,
24152            };
24153            let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
24154            let (ds, nv, nk, kd, hki) = (
24155                d_state as i32,
24156                num_v as i32,
24157                num_k as i32,
24158                key_dim as i32,
24159                hk as i32,
24160            );
24161            let __s_b = self.gpu.stream();
24162            let mut b = __s_b.launch_builder(&f);
24163            b.arg(qkv_tm)
24164                .arg(&*conv_state)
24165                .arg(w)
24166                .arg(q_g)
24167                .arg(k_g)
24168                .arg(v_g)
24169                .arg(&cd)
24170                .arg(&ti)
24171                .arg(&dc)
24172                .arg(&ds)
24173                .arg(&nv)
24174                .arg(&nk)
24175                .arg(&kd)
24176                .arg(&hki);
24177            unsafe {
24178                b.launch(cfg)?;
24179            }
24180        }
24181        match pad_len {
24182            Some(len_d) => {
24183                let f = self.func("ssm_conv_ring_update_dev_f32");
24184                let n = conv_dim * (d_conv - 1);
24185                let cfg = LaunchConfig::for_num_elems(n as u32);
24186                let (cd, dc) = (conv_dim as i32, d_conv as i32);
24187                let __s_b = self.gpu.stream();
24188                let mut b = __s_b.launch_builder(&f);
24189                b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
24190                unsafe {
24191                    b.launch(cfg)?;
24192                }
24193            }
24194            None => {
24195                let f = self.func("ssm_conv_ring_update_f32");
24196                let n = conv_dim * (d_conv - 1);
24197                let cfg = LaunchConfig::for_num_elems(n as u32);
24198                let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
24199                let __s_b = self.gpu.stream();
24200                let mut b = __s_b.launch_builder(&f);
24201                b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
24202                unsafe {
24203                    b.launch(cfg)?;
24204                }
24205            }
24206        }
24207        Ok(())
24208    }
24209
24210    /// task #18 increment 2: allocate ONE sequence's chunk buffers (no launches) —
24211    /// K1-K5 all run varlen afterwards. `a`/`w` become struct members so the varlen
24212    /// K2/K3 can write them.
24213    pub fn gdn_chunk_alloc(
24214        &self,
24215        n_head: usize,
24216        t: usize,
24217        c: usize,
24218        hk: usize,
24219    ) -> Result<GdnChunkBufs, Box<dyn std::error::Error>> {
24220        const D: usize = 128;
24221        assert!(
24222            c == 32,
24223            "gdn_chunk_alloc: varlen chain is the C==32 mma pair"
24224        );
24225        let h = n_head;
24226        let nc = (t + c - 1) / c;
24227        Ok(GdnChunkBufs {
24228            gcum: self.uninit(t * h)?,
24229            a: self.uninit(nc * h * c * c)?,
24230            p: self.uninit(nc * h * c * c)?,
24231            u: self.uninit(nc * h * c * D)?,
24232            w: self.uninit(nc * h * c * D)?,
24233            kb16: self.alloc_u8_uninit(t * hk * D * 2)?,
24234            wb16: self.alloc_u8_uninit(nc * h * c * D * 2)?,
24235            y16: self.alloc_u8_uninit(nc * h * c * D * 2)?,
24236            ssnap16: self.alloc_u8_uninit(nc * h * D * D * 2)?,
24237            qb16: self.alloc_u8_uninit(t * hk * D * 2)?,
24238            pb16: self.alloc_u8_uninit(nc * h * c * c * 2)?,
24239            o: self.uninit(D * h * t)?,
24240            t,
24241            nc,
24242        })
24243    }
24244
24245    /// view-source twin of f32_to_bf16 (the batched FA3 v mirror reads a concat view).
24246    pub fn f32_to_bf16_v(
24247        &self,
24248        x: &cudarc::driver::CudaView<f32>,
24249        dst: &mut CudaSlice<u8>,
24250        n: usize,
24251    ) -> Result<(), Box<dyn std::error::Error>> {
24252        let f = self.func("f32_to_bf16_bulk");
24253        let ni = n as i64;
24254        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
24255        let __s_b = self.gpu.stream();
24256        let mut b = __s_b.launch_builder(&f);
24257        b.arg(x).arg(dst).arg(&ni);
24258        unsafe {
24259            b.launch(cfg)?;
24260        }
24261        Ok(())
24262    }
24263
24264    /// f32 -> bf16 bulk mirror into a caller buffer (the K4/K5 operand mirrors).
24265    pub fn f32_to_bf16_into(
24266        &self,
24267        x: &CudaSlice<f32>,
24268        dst: &mut CudaSlice<u8>,
24269        n: usize,
24270    ) -> Result<(), Box<dyn std::error::Error>> {
24271        let f = self.func("f32_to_bf16_bulk");
24272        let ni = n as i64;
24273        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
24274        let __s_b = self.gpu.stream();
24275        let mut b = __s_b.launch_builder(&f);
24276        b.arg(x).arg(dst).arg(&ni);
24277        unsafe {
24278            b.launch(cfg)?;
24279        }
24280        Ok(())
24281    }
24282
24283    /// task #18 increment 2: varlen K1+K2+K3 — three launches run every sequence's
24284    /// cumgate/attn/solve (per-block math identical to the per-seq kernels).
24285    pub fn gdn_chunk_k123_vl8(
24286        &self,
24287        seqs: &[GdnSeqVl],
24288        n_head: usize,
24289        hk: usize,
24290        wq: Option<&GdnWVl8>,
24291    ) -> Result<(), Box<dyn std::error::Error>> {
24292        let b = seqs.len();
24293        assert!(b >= 1 && b <= 8, "gdn_chunk_k123_vl8: 1..=8 sequences");
24294        let mut packed = [GdnSeqVl::default(); 8];
24295        packed[..b].copy_from_slice(seqs);
24296        let v = GdnVl8(packed);
24297        let (hi, ci) = (n_head as i32, 32i32);
24298        let max_nc = seqs.iter().map(|a| a.nc).max().unwrap() as u32;
24299        {
24300            let f = self.func("gdn_chunk_cumgate_vl");
24301            let cfg = LaunchConfig {
24302                grid_dim: (max_nc, n_head as u32, b as u32),
24303                block_dim: (32, 1, 1),
24304                shared_mem_bytes: 0,
24305            };
24306            let __s_lb = self.gpu.stream();
24307            let mut lb = __s_lb.launch_builder(&f);
24308            lb.arg(&v).arg(&hi).arg(&ci);
24309            unsafe {
24310                lb.launch(cfg)?;
24311            }
24312        }
24313        let hki = hk as i32;
24314        if let Some(w) = wq {
24315            // K2-wgmma vl twin (writes A + pre-masked Pb16)
24316            let f = self.func("gdn_k2_wgmma_vl");
24317            let cfg = LaunchConfig {
24318                grid_dim: (max_nc, n_head as u32, b as u32),
24319                block_dim: (128, 1, 1),
24320                shared_mem_bytes: 0,
24321            };
24322            let __s_lb = self.gpu.stream();
24323            let mut lb = __s_lb.launch_builder(&f);
24324            lb.arg(&v).arg(w).arg(&hi).arg(&ci).arg(&hki);
24325            unsafe {
24326                lb.launch(cfg)?;
24327            }
24328        } else {
24329            let f = self.func("gdn_chunk_attn_vl");
24330            f.set_attribute(
24331                CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
24332                GDN_K2_DYNAMIC_SHARED_BYTES as i32,
24333            )?;
24334            let cfg = LaunchConfig {
24335                grid_dim: (max_nc, n_head as u32, b as u32),
24336                block_dim: (256, 1, 1),
24337                shared_mem_bytes: GDN_K2_DYNAMIC_SHARED_BYTES,
24338            };
24339            let __s_lb = self.gpu.stream();
24340            let mut lb = __s_lb.launch_builder(&f);
24341            lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
24342            unsafe {
24343                lb.launch(cfg)?;
24344            }
24345        }
24346        {
24347            let f = self.func("gdn_chunk_solve32_vl");
24348            let cfg = LaunchConfig {
24349                grid_dim: (max_nc, n_head as u32, b as u32),
24350                block_dim: (256, 1, 1),
24351                shared_mem_bytes: 0,
24352            };
24353            let __s_lb = self.gpu.stream();
24354            let mut lb = __s_lb.launch_builder(&f);
24355            lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
24356            unsafe {
24357                lb.launch(cfg)?;
24358            }
24359        }
24360        Ok(())
24361    }
24362
24363    /// task #18 increment 3: varlen PREP chain — conv(+ring) / repack / fused-l2 /
24364    /// fused gate-prep, 5 launches for every sequence (per-element math identical
24365    /// to the per-seq kernels; l2/gate fusions write disjoint outputs).
24366    #[allow(clippy::too_many_arguments)]
24367    pub fn gdn_prep_vl8(
24368        &self,
24369        seqs: &[GdnPrepVl],
24370        conv_w: &CudaSlice<f32>,
24371        dt_bias: &CudaSlice<f32>,
24372        a: &CudaSlice<f32>,
24373        conv_dim: usize,
24374        d_conv: usize,
24375        d_state: usize,
24376        num_v: usize,
24377        num_k: usize,
24378        key_dim: usize,
24379        hk: usize,
24380        eps: f32,
24381    ) -> Result<(), Box<dyn std::error::Error>> {
24382        let b = seqs.len();
24383        assert!(b >= 1 && b <= 8);
24384        let mut packed = [GdnPrepVl::default(); 8];
24385        packed[..b].copy_from_slice(seqs);
24386        let v = GdnPrepVl8(packed);
24387        let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
24388        let (cdi, dci) = (conv_dim as i32, d_conv as i32);
24389        let conv_fuse = std::env::var("MEMRA_CONV_FUSE").as_deref() != Ok("0");
24390        assert!(
24391            conv_fuse || hk == num_v,
24392            "de-broadcast requires the fused conv"
24393        );
24394        if conv_fuse {
24395            let f = self.func("ssm_conv1d_gdn_state_vl");
24396            let cfg = LaunchConfig {
24397                grid_dim: ((conv_dim as u32).div_ceil(256), max_t, b as u32),
24398                block_dim: (256, 1, 1),
24399                shared_mem_bytes: 0,
24400            };
24401            let (dsi, nvi, nki, kdi, hki) = (
24402                d_state as i32,
24403                num_v as i32,
24404                num_k as i32,
24405                key_dim as i32,
24406                hk as i32,
24407            );
24408            let __s_lb = self.gpu.stream();
24409            let mut lb = __s_lb.launch_builder(&f);
24410            lb.arg(&v)
24411                .arg(conv_w)
24412                .arg(&cdi)
24413                .arg(&dci)
24414                .arg(&dsi)
24415                .arg(&nvi)
24416                .arg(&nki)
24417                .arg(&kdi)
24418                .arg(&hki);
24419            unsafe {
24420                lb.launch(cfg)?;
24421            }
24422        } else {
24423            let f = self.func("ssm_conv1d_tm_state_vl");
24424            let cfg = LaunchConfig {
24425                grid_dim: ((conv_dim as u32).div_ceil(256), max_t, b as u32),
24426                block_dim: (256, 1, 1),
24427                shared_mem_bytes: 0,
24428            };
24429            let __s_lb = self.gpu.stream();
24430            let mut lb = __s_lb.launch_builder(&f);
24431            lb.arg(&v).arg(conv_w).arg(&cdi).arg(&dci);
24432            unsafe {
24433                lb.launch(cfg)?;
24434            }
24435        }
24436        {
24437            let f = self.func("ssm_conv_ring_update_vl");
24438            let n = (conv_dim * (d_conv - 1)) as u32;
24439            let cfg = LaunchConfig {
24440                grid_dim: (n.div_ceil(256), 1, b as u32),
24441                block_dim: (256, 1, 1),
24442                shared_mem_bytes: 0,
24443            };
24444            let __s_lb = self.gpu.stream();
24445            let mut lb = __s_lb.launch_builder(&f);
24446            lb.arg(&v).arg(&cdi).arg(&dci);
24447            unsafe {
24448                lb.launch(cfg)?;
24449            }
24450        }
24451        if !conv_fuse {
24452            let f = self.func("qkv_to_gdn_repack_vl");
24453            let n = max_t * (num_v * d_state) as u32;
24454            let cfg = LaunchConfig {
24455                grid_dim: (n.div_ceil(256), 1, b as u32),
24456                block_dim: (256, 1, 1),
24457                shared_mem_bytes: 0,
24458            };
24459            let (dsi, nvi, nki, kdi) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
24460            let __s_lb = self.gpu.stream();
24461            let mut lb = __s_lb.launch_builder(&f);
24462            lb.arg(&v).arg(&dsi).arg(&nvi).arg(&nki).arg(&kdi);
24463            unsafe {
24464                lb.launch(cfg)?;
24465            }
24466        }
24467        if Self::l2_v2_on(d_state) {
24468            let f = self.func("gdn_l2_v2_vl");
24469            let cfg = LaunchConfig {
24470                grid_dim: ((max_t * hk as u32).div_ceil(8), 2, b as u32),
24471                block_dim: (256, 1, 1),
24472                shared_mem_bytes: 0,
24473            };
24474            let (dsi, nvi) = (d_state as i32, hk as i32);
24475            let __s_lb = self.gpu.stream();
24476            let mut lb = __s_lb.launch_builder(&f);
24477            lb.arg(&v).arg(&dsi).arg(&nvi).arg(&eps);
24478            unsafe {
24479                lb.launch(cfg)?;
24480            }
24481        } else {
24482            let f = self.func("gdn_l2_vl");
24483            let cfg = LaunchConfig {
24484                grid_dim: (max_t * hk as u32, 2, b as u32),
24485                block_dim: (256, 1, 1),
24486                shared_mem_bytes: 0,
24487            };
24488            let (dsi, nvi) = (d_state as i32, hk as i32);
24489            let __s_lb = self.gpu.stream();
24490            let mut lb = __s_lb.launch_builder(&f);
24491            lb.arg(&v).arg(&dsi).arg(&nvi).arg(&eps);
24492            unsafe {
24493                lb.launch(cfg)?;
24494            }
24495        }
24496        {
24497            let f = self.func("gdn_gate_prep_vl");
24498            let n = max_t * num_v as u32;
24499            let cfg = LaunchConfig {
24500                grid_dim: (n.div_ceil(256), 1, b as u32),
24501                block_dim: (256, 1, 1),
24502                shared_mem_bytes: 0,
24503            };
24504            let nvi = num_v as i32;
24505            let __s_lb = self.gpu.stream();
24506            let mut lb = __s_lb.launch_builder(&f);
24507            lb.arg(&v).arg(dt_bias).arg(a).arg(&nvi);
24508            unsafe {
24509                lb.launch(cfg)?;
24510            }
24511        }
24512        Ok(())
24513    }
24514
24515    /// varlen bf16 mirrors over the gdnseq_t table (which: 0 = k_l2 -> kb16, 1 = w -> wb16).
24516    pub fn gdn_mirror_vl8(
24517        &self,
24518        seqs: &[GdnSeqVl],
24519        n_head: usize,
24520        which: i32,
24521        hk: usize,
24522    ) -> Result<(), Box<dyn std::error::Error>> {
24523        let b = seqs.len();
24524        assert!(b >= 1 && b <= 8);
24525        let mut packed = [GdnSeqVl::default(); 8];
24526        packed[..b].copy_from_slice(seqs);
24527        let v = GdnVl8(packed);
24528        let ept = (if which == 0 { hk } else { n_head } * 128) as i32;
24529        let max_n = seqs
24530            .iter()
24531            .map(|s| {
24532                if which == 0 {
24533                    s.t as i64 * ept as i64
24534                } else {
24535                    s.nc as i64 * ept as i64 * 32
24536                }
24537            })
24538            .max()
24539            .unwrap();
24540        let f = self.func("gdn_mirror_vl");
24541        let blocks = ((max_n as u32).div_ceil(4)).div_ceil(256);
24542        let cfg = LaunchConfig {
24543            grid_dim: (blocks, 1, b as u32),
24544            block_dim: (256, 1, 1),
24545            shared_mem_bytes: 0,
24546        };
24547        let __s_lb = self.gpu.stream();
24548        let mut lb = __s_lb.launch_builder(&f);
24549        lb.arg(&v).arg(&ept).arg(&which);
24550        unsafe {
24551            lb.launch(cfg)?;
24552        }
24553        Ok(())
24554    }
24555
24556    /// varlen gated-norm tail (+f16out) — one launch replaces B gated_rmsnorm calls.
24557    pub fn gdn_tail_vl8(
24558        &self,
24559        seqs: &[GdnPrepVl],
24560        norm_w: &CudaSlice<f32>,
24561        d_state: usize,
24562        num_v: usize,
24563        eps: f32,
24564    ) -> Result<(), Box<dyn std::error::Error>> {
24565        let b = seqs.len();
24566        assert!(b >= 1 && b <= 8);
24567        let mut packed = [GdnPrepVl::default(); 8];
24568        packed[..b].copy_from_slice(seqs);
24569        let v = GdnPrepVl8(packed);
24570        let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
24571        let f = self.func("gated_rmsnorm_f16out_vl");
24572        // block_dim MUST match gated_rmsnorm's (128): the reduction tree order pins the scale
24573        let cfg = LaunchConfig {
24574            grid_dim: (max_t * num_v as u32, 1, b as u32),
24575            block_dim: (128, 1, 1),
24576            shared_mem_bytes: 0,
24577        };
24578        let (dsi, nvi) = (d_state as i32, num_v as i32);
24579        let __s_lb = self.gpu.stream();
24580        let mut lb = __s_lb.launch_builder(&f);
24581        lb.arg(&v).arg(norm_w).arg(&dsi).arg(&nvi).arg(&eps);
24582        unsafe {
24583            lb.launch(cfg)?;
24584        }
24585        Ok(())
24586    }
24587
24588    /// Raw device address helpers for the varlen by-value arg struct (single-stream
24589    /// launches; every buffer outlives the call — the f16 FFI discipline).
24590    pub fn addr_f32(&self, x: &CudaSlice<f32>) -> u64 {
24591        use cudarc::driver::DevicePtr;
24592        let s = self.gpu.stream();
24593        let (p, _g) = x.device_ptr(&s);
24594        p as u64
24595    }
24596    pub fn addr_f32_mut(&self, x: &mut CudaSlice<f32>) -> u64 {
24597        use cudarc::driver::DevicePtrMut;
24598        let s = self.gpu.stream();
24599        let (p, _g) = x.device_ptr_mut(&s);
24600        p as u64
24601    }
24602    pub fn addr_f32v(&self, x: &cudarc::driver::CudaView<f32>) -> u64 {
24603        use cudarc::driver::DevicePtr;
24604        let s = self.gpu.stream();
24605        let (p, _g) = x.device_ptr(&s);
24606        p as u64
24607    }
24608    pub fn addr_u8(&self, x: &CudaSlice<u8>) -> u64 {
24609        use cudarc::driver::DevicePtr;
24610        let s = self.gpu.stream();
24611        let (p, _g) = x.device_ptr(&s);
24612        p as u64
24613    }
24614
24615    /// task #18: the varlen K4+K5 pair — TWO launches run every sequence's state pass
24616    /// and output pass (grid gains a seq dim; per-block math identical to the per-seq
24617    /// launches, so this is strictly bit-gateable against them).
24618    pub fn gdn_chunk_vl8(
24619        &self,
24620        seqs: &[GdnSeqVl],
24621        n_head: usize,
24622        scale: f32,
24623        hk: usize,
24624        wq: Option<&GdnWVl8>,
24625    ) -> Result<(), Box<dyn std::error::Error>> {
24626        const NSPLIT: u32 = 4;
24627        let b = seqs.len();
24628        assert!(b >= 1 && b <= 8, "gdn_chunk_vl8: 1..=8 sequences");
24629        let mut packed = [GdnSeqVl::default(); 8];
24630        packed[..b].copy_from_slice(seqs);
24631        let v = GdnVl8(packed);
24632        let (hi, ci) = (n_head as i32, 32i32);
24633        let max_nc = seqs.iter().map(|a| a.nc).max().unwrap() as u32;
24634        let hki = hk as i32;
24635        if let Some(w) = wq {
24636            // K4+K5 fused wgmma vl twin: one launch, Y/Ssnap never materialized.
24637            let f = self.func("gdn_k45_wgmma_vl");
24638            let cfg = LaunchConfig {
24639                grid_dim: (n_head as u32, NSPLIT, b as u32),
24640                block_dim: (256, 1, 1),
24641                shared_mem_bytes: 0,
24642            };
24643            let __s_lb = self.gpu.stream();
24644            let mut lb = __s_lb.launch_builder(&f);
24645            lb.arg(&v).arg(w).arg(&scale).arg(&hi).arg(&ci).arg(&hki);
24646            unsafe {
24647                lb.launch(cfg)?;
24648            }
24649            let _ = max_nc;
24650            return Ok(());
24651        }
24652        {
24653            let f = self.func("gdn_chunk_state_mma_vl");
24654            let cfg = LaunchConfig {
24655                grid_dim: (n_head as u32, NSPLIT, b as u32),
24656                block_dim: (256, 1, 1),
24657                shared_mem_bytes: 0,
24658            };
24659            let __s_lb = self.gpu.stream();
24660            let mut lb = __s_lb.launch_builder(&f);
24661            lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
24662            unsafe {
24663                lb.launch(cfg)?;
24664            }
24665        }
24666        {
24667            let f = self.func("gdn_chunk_output_mma_vl");
24668            let cfg = LaunchConfig {
24669                grid_dim: (max_nc, n_head as u32, b as u32),
24670                block_dim: (256, 1, 1),
24671                shared_mem_bytes: 0,
24672            };
24673            let __s_lb = self.gpu.stream();
24674            let mut lb = __s_lb.launch_builder(&f);
24675            lb.arg(&v).arg(&hi).arg(&ci).arg(&scale).arg(&hki);
24676            unsafe {
24677                lb.launch(cfg)?;
24678            }
24679        }
24680        Ok(())
24681    }
24682    pub fn gdn_scan_chunked(
24683        &self,
24684        q: &CudaSlice<f32>,
24685        k: &CudaSlice<f32>,
24686        v: &CudaSlice<f32>,
24687        g: &CudaSlice<f32>,
24688        beta: &CudaSlice<f32>,
24689        kb16_pre: Option<&CudaSlice<u8>>,
24690        qb16_pre: Option<&CudaSlice<u8>>,
24691        state_in: &CudaSlice<f32>,
24692        state_out: &mut CudaSlice<f32>,
24693        o: &mut CudaSlice<f32>,
24694        n_head: usize,
24695        t: usize,
24696        scale: f32,
24697        c: usize,
24698        hk: usize,
24699    ) -> Result<(), Box<dyn std::error::Error>> {
24700        const D: usize = 128;
24701        const NSPLIT: u32 = 4;
24702        assert!(c >= 1 && c <= 128, "gdn_scan_chunked: C must be in 1..=128");
24703        let h = n_head;
24704        let nc = (t + c - 1) / c;
24705        let (hi, ti, ci) = (h as i32, t as i32, c as i32);
24706        // mirror-fold (round 27): on the mma path W's bf16 twin is emitted by K3's store
24707        // (wb16 pre-allocated and threaded through k123) and k's by the producer l2 when
24708        // the caller hands `kb16_pre` — both standalone mirror passes disappear.
24709        // default mirrors gdn_mma_enabled (incl. the sm_120a-build ON, lane/moeprime-nvfp4-direct)
24710        let gdn_mma_pre = !portable_mma_gated()
24711            && c == 32
24712            && match std::env::var("MEMRA_GDN_MMA").as_deref() {
24713                Ok("1") => true,
24714                Ok("0") => false,
24715                _ => gdn_mma_default_on(),
24716            };
24717        let mut wb16_pre: Option<CudaSlice<u8>> = if gdn_mma_pre {
24718            Some(self.alloc_u8_uninit(nc * h * c * D * 2)?)
24719        } else {
24720            None
24721        };
24722        // K2-wgmma pre-work (MEMRA_GDN_WGMMA): the kb16/qb16 mirrors hoist ABOVE K123 so
24723        // K2 rides them via cp.async; K2 writes the pre-masked Pb16 directly (the
24724        // gdn_p_bf16_masked pass and the in-branch mirror builds disappear).
24725        // mirrors gdn_wgmma_on: hard Hopper-build gate (empty wgmma bodies elsewhere)
24726        let gdn_wgmma_pre = cfg!(memra_hopper_mma)
24727            && gdn_mma_pre
24728            && std::env::var("MEMRA_GDN_WGMMA").as_deref() != Ok("0");
24729        let nk = t * hk * D;
24730        let mut kb16_local: Option<CudaSlice<u8>> = None;
24731        if gdn_mma_pre && kb16_pre.is_none() {
24732            let mut kb = self.alloc_u8_uninit(nk * 2)?;
24733            let f = self.func("f32_to_bf16_bulk");
24734            let n2 = nk as i64;
24735            let cfg2 = LaunchConfig::for_num_elems((nk as u32).div_ceil(4));
24736            let __s_b = self.gpu.stream();
24737            let mut b = __s_b.launch_builder(&f);
24738            b.arg(k).arg(&mut kb).arg(&n2);
24739            unsafe {
24740                b.launch(cfg2)?;
24741            }
24742            kb16_local = Some(kb);
24743        }
24744        let kb16_ref0: Option<&CudaSlice<u8>> = kb16_local.as_ref().or(kb16_pre);
24745        if let Some(kb) = kb16_pre {
24746            assert!(kb.len() >= nk * 2, "kb16_pre too small");
24747        }
24748        let mut qb16: Option<CudaSlice<u8>> = None;
24749        let mut pb16: Option<CudaSlice<u8>> = None;
24750        if gdn_wgmma_pre {
24751            // mirror-fold (round 35): prep's l2 v2 emits qb16 in-epilogue (kb16 pattern);
24752            // the standalone bulk cvt only serves callers without the prep mirror.
24753            if qb16_pre.is_none() {
24754                let mut qb = self.alloc_u8_uninit(nk * 2)?;
24755                let f = self.func("f32_to_bf16_bulk");
24756                let n2 = nk as i64;
24757                let cfg2 = LaunchConfig::for_num_elems((nk as u32).div_ceil(4));
24758                let __s_b = self.gpu.stream();
24759                let mut b = __s_b.launch_builder(&f);
24760                b.arg(q).arg(&mut qb).arg(&n2);
24761                unsafe {
24762                    b.launch(cfg2)?;
24763                }
24764                qb16 = Some(qb);
24765            } else if let Some(qb) = qb16_pre {
24766                assert!(qb.len() >= nk * 2, "qb16_pre too small");
24767            }
24768            pb16 = Some(self.alloc_u8_uninit(nc * h * c * c * 2)?);
24769        }
24770        let qb16_ref0: Option<&CudaSlice<u8>> = qb16.as_ref().or(qb16_pre);
24771        let k2w = if gdn_wgmma_pre {
24772            Some((
24773                *qb16_ref0.as_ref().unwrap(),
24774                *kb16_ref0.as_ref().unwrap(),
24775                pb16.as_mut().unwrap(),
24776            ))
24777        } else {
24778            None
24779        };
24780        let (gcum, p, u, w) =
24781            self.gdn_chunk_k123(q, k, v, g, beta, wb16_pre.as_mut(), n_head, t, c, hk, k2w)?;
24782        let _ = &w;
24783        let mut y = self.uninit(nc * h * c * D)?;
24784        let mut ssnap = self.uninit(nc * h * D * D)?; // chunk-start state snapshots (K5 phase 1)
24785        // K4-MMA seam (MEMRA_GDN_MMA; harness verdict 1.75x — tools/bench_gdn_k4.cu, ledger
24786        // 2026-07-26): M in mma accumulator fragments, bf16 W/k mirrors through a cp.async
24787        // ring. C==32 only (the kernel's tile). PROMOTED default-ON on the Hopper lane
24788        // after the STATE-CARRY battery (2026-07-26): 2048-token prime (64 in-kernel state
24789        // carries) -> 256 greedy decode tokens IDENTICAL to f32 on 3 seeds, AND chunked-
24790        // continuation prime (MEMRA_PRIME_CHUNK=512, 4 cross-call carries via cache.recur)
24791        // IDENTICAL on 2 seeds; plus argmax MATCH, pp512 +3.5% (17286), oracle out
24792        // mean_rel ~1e-4. kernel-check pins BOTH configs (f32 tight band forced =0; mma
24793        // band 8e-2/8e-1 vs f64 truth). =0 reverts; portable stays f32. NOT read via
24794        // OnceLock ON PURPOSE: kernel-check toggles the env per call to pin both forms.
24795        // default mirrors gdn_mma_enabled (incl. the sm_120a-build ON) — all three read
24796        // sites must agree or the pre-work arms while the scan takes the scalar route.
24797        let gdn_mma = !portable_mma_gated()
24798            && c == 32
24799            && match std::env::var("MEMRA_GDN_MMA").as_deref() {
24800                Ok("1") => true,
24801                Ok("0") => false,
24802                _ => gdn_mma_default_on(),
24803            };
24804        if gdn_mma {
24805            let wb16 = wb16_pre
24806                .take()
24807                .expect("mma path pre-allocates wb16 (K3 store fold)");
24808            let kb16_ref: &CudaSlice<u8> = kb16_ref0.expect("mma path pre-builds kb16 above K123");
24809            // K4+K5 FUSED wgmma seam (MEMRA_GDN_WGMMA, task #22; harness verdict
24810            // tools/bench_gdn_wgmma.cu v5, ledger 1f08b997: in-band Y 1.07e-2 / state
24811            // 1.03e-2 / O 1.08e-2, 91.3us vs 70.4 K4-only at H=32 T=512). K5's output
24812            // pass runs inside the persistent-M kernel; Y and Ssnap are never
24813            // materialized. New numeric class (gk folds into k^T instead of ys) —
24814            // explicit opt-in until the state-carry battery promotes it. Env read per
24815            // call (kernel-check pins configs by toggling env, GDN_MMA precedent).
24816            // PROMOTED default-ON hopper (2026-07-27): full battery green — harness
24817            // in-band, argmax gate PASS, 3-seed greedy IDENTICAL after ~2k prime,
24818            // chunked-continuation IDENTICAL, kernel-check + decode-batch gates green,
24819            // official prefill lane +0.74% interleaved x5 (5/5 rounds). =0 reverts.
24820            if gdn_wgmma_pre {
24821                // qb16/pb16 pre-built above K123 (K2-wgmma wrote the masked Pb16).
24822                let qb16 = qb16_ref0.unwrap();
24823                let pb16 = pb16.as_ref().unwrap();
24824                {
24825                    let f = self.func("gdn_k45_wgmma");
24826                    let cfg = LaunchConfig {
24827                        grid_dim: (h as u32, 4, 1),
24828                        block_dim: (256, 1, 1),
24829                        shared_mem_bytes: 0,
24830                    };
24831                    let hki = hk as i32;
24832                    let __s_b = self.gpu.stream();
24833                    let mut b = __s_b.launch_builder(&f);
24834                    b.arg(kb16_ref)
24835                        .arg(&gcum)
24836                        .arg(beta)
24837                        .arg(&u)
24838                        .arg(&wb16)
24839                        .arg(qb16)
24840                        .arg(pb16)
24841                        .arg(o)
24842                        .arg(&scale)
24843                        .arg(state_in)
24844                        .arg(&mut *state_out)
24845                        .arg(&hi)
24846                        .arg(&ti)
24847                        .arg(&ci)
24848                        .arg(&hki);
24849                    unsafe {
24850                        b.launch(cfg)?;
24851                    }
24852                }
24853                return Ok(());
24854            }
24855            // COUPLED PAIR: K4-mma writes Y and Ssnap as bf16 (their only consumer is
24856            // K5-mma, which rounds to bf16 regardless — identical numerics, half the
24857            // traffic; harness K5 63.0 -> 35.3us). Fresh bf16 buffers replace the f32 ones.
24858            let mut y16 = self.alloc_u8_uninit(nc * h * c * D * 2)?;
24859            let mut ssnap16 = self.alloc_u8_uninit(nc * h * D * D * 2)?;
24860            {
24861                let f = self.func("gdn_chunk_state_mma");
24862                let cfg = LaunchConfig {
24863                    grid_dim: (h as u32, NSPLIT, 1),
24864                    block_dim: (256, 1, 1),
24865                    shared_mem_bytes: 0,
24866                };
24867                let hki = hk as i32;
24868                let __s_b = self.gpu.stream();
24869                let mut b = __s_b.launch_builder(&f);
24870                b.arg(kb16_ref)
24871                    .arg(&gcum)
24872                    .arg(beta)
24873                    .arg(&u)
24874                    .arg(&wb16)
24875                    .arg(&mut y16)
24876                    .arg(&mut ssnap16)
24877                    .arg(state_in)
24878                    .arg(&mut *state_out)
24879                    .arg(&hi)
24880                    .arg(&ti)
24881                    .arg(&ci)
24882                    .arg(&hki);
24883                unsafe {
24884                    b.launch(cfg)?;
24885                }
24886            }
24887            {
24888                // K5-mma (bf16 St/Y consumers)
24889                let f = self.func("gdn_chunk_output_mma");
24890                let jt = ((c + 31) / 32) as u32;
24891                let cfg = LaunchConfig {
24892                    grid_dim: (nc as u32, h as u32, jt),
24893                    block_dim: (256, 1, 1),
24894                    shared_mem_bytes: 0,
24895                };
24896                let hki = hk as i32;
24897                let __s_b = self.gpu.stream();
24898                let mut b = __s_b.launch_builder(&f);
24899                b.arg(q)
24900                    .arg(&gcum)
24901                    .arg(&p)
24902                    .arg(&y16)
24903                    .arg(&ssnap16)
24904                    .arg(o)
24905                    .arg(&hi)
24906                    .arg(&ti)
24907                    .arg(&ci)
24908                    .arg(&scale)
24909                    .arg(&hki);
24910                unsafe {
24911                    b.launch(cfg)?;
24912                }
24913            }
24914            return Ok(());
24915        }
24916        {
24917            // K4 (sequential over chunks inside; blocks col-partition the state)
24918            let f = self.func("gdn_chunk_state_f32");
24919            let cfg = LaunchConfig {
24920                grid_dim: (h as u32, NSPLIT, 1),
24921                block_dim: (256, 1, 1),
24922                shared_mem_bytes: 0,
24923            };
24924            let __s_b = self.gpu.stream();
24925            let mut b = __s_b.launch_builder(&f);
24926            b.arg(k)
24927                .arg(&gcum)
24928                .arg(beta)
24929                .arg(&u)
24930                .arg(&w)
24931                .arg(&mut y)
24932                .arg(&mut ssnap)
24933                .arg(state_in)
24934                .arg(&mut *state_out)
24935                .arg(&hi)
24936                .arg(&ti)
24937                .arg(&ci);
24938            unsafe {
24939                b.launch(cfg)?;
24940            }
24941        }
24942        {
24943            // K5 (j-blocked: grid.z = 32-row output blocks per chunk; writes o fully)
24944            let f = self.func("gdn_chunk_output_f32");
24945            let jt = ((c + 31) / 32) as u32;
24946            let cfg = LaunchConfig {
24947                grid_dim: (nc as u32, h as u32, jt),
24948                block_dim: (256, 1, 1),
24949                shared_mem_bytes: 0,
24950            };
24951            let __s_b = self.gpu.stream();
24952            let mut b = __s_b.launch_builder(&f);
24953            b.arg(q)
24954                .arg(&gcum)
24955                .arg(&p)
24956                .arg(&y)
24957                .arg(&ssnap)
24958                .arg(o)
24959                .arg(&hi)
24960                .arg(&ti)
24961                .arg(&ci)
24962                .arg(&scale);
24963            unsafe {
24964                b.launch(cfg)?;
24965            }
24966        }
24967        Ok(())
24968    }
24969
24970    /// PREFILL GDN scan dispatch (the A4 seam): chunked WY form when enabled and T is in the
24971    /// batched-prefill regime, else the sequential scan. Callers: hybrid_forward::linear_attn
24972    /// (forward/forward_last) + linear_attn_prime (prime_cache). Decode (T=1) and the spec
24973    /// verify call `gdn_scan_s128` DIRECTLY — the decode==verify dispatch identity is untouched.
24974    ///
24975    /// MEMRA_GDN_DIFF=1: numerical-oracle mode — runs BOTH forms on the same inputs, prints the
24976    /// per-call (== per-layer, in call order) output/state error distribution, and keeps the
24977    /// SEQUENTIAL results so the run stays on the shipped path (stage-1 prototype evidence).
24978    #[allow(clippy::too_many_arguments)]
24979    #[allow(clippy::too_many_arguments)]
24980    pub fn gdn_scan_prefill(
24981        &self,
24982        q: &CudaSlice<f32>,
24983        k: &CudaSlice<f32>,
24984        v: &CudaSlice<f32>,
24985        g: &CudaSlice<f32>,
24986        beta: &CudaSlice<f32>,
24987        kb16_pre: Option<&CudaSlice<u8>>,
24988        qb16_pre: Option<&CudaSlice<u8>>,
24989        state_in: &CudaSlice<f32>,
24990        state_out: &mut CudaSlice<f32>,
24991        o: &mut CudaSlice<f32>,
24992        n_head: usize,
24993        t: usize,
24994        scale: f32,
24995        hk: usize,
24996    ) -> Result<(), Box<dyn std::error::Error>> {
24997        if std::env::var("MEMRA_GDN_DIFF").is_ok() && t >= 16 {
24998            assert!(hk == n_head, "GDN_DIFF oracle is broadcast-only");
24999            return self.gdn_scan_diff(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale);
25000        }
25001        if Self::gdn_chunked_enabled() && t >= 16 {
25002            self.gdn_scan_chunked(
25003                q,
25004                k,
25005                v,
25006                g,
25007                beta,
25008                kb16_pre,
25009                qb16_pre,
25010                state_in,
25011                state_out,
25012                o,
25013                n_head,
25014                t,
25015                scale,
25016                Self::gdn_chunk_size(),
25017                hk,
25018            )
25019        } else {
25020            assert!(
25021                hk == n_head,
25022                "s128 scan is broadcast-only (prep guarantees by predicate)"
25023            );
25024            self.gdn_scan_s128(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale)
25025        }
25026    }
25027
25028    /// Stage-1 oracle: run sequential AND chunked, report per-call error stats, keep sequential.
25029    #[allow(clippy::too_many_arguments)]
25030    fn gdn_scan_diff(
25031        &self,
25032        q: &CudaSlice<f32>,
25033        k: &CudaSlice<f32>,
25034        v: &CudaSlice<f32>,
25035        g: &CudaSlice<f32>,
25036        beta: &CudaSlice<f32>,
25037        state_in: &CudaSlice<f32>,
25038        state_out: &mut CudaSlice<f32>,
25039        o: &mut CudaSlice<f32>,
25040        n_head: usize,
25041        t: usize,
25042        scale: f32,
25043    ) -> Result<(), Box<dyn std::error::Error>> {
25044        static CALL: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
25045        let call = CALL.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
25046        let mut o_c = self.uninit(o.len())?;
25047        let mut st_c = self.uninit(state_out.len())?;
25048        self.gdn_scan_chunked(
25049            q,
25050            k,
25051            v,
25052            g,
25053            beta,
25054            None,
25055            None,
25056            state_in,
25057            &mut st_c,
25058            &mut o_c,
25059            n_head,
25060            t,
25061            scale,
25062            Self::gdn_chunk_size(),
25063            n_head,
25064        )?;
25065        self.gdn_scan_s128(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale)?;
25066        let (oh_s, oh_c) = (self.dtoh(o)?, self.dtoh(&o_c)?);
25067        let (sh_s, sh_c) = (self.dtoh(state_out)?, self.dtoh(&st_c)?);
25068        let stats = |a: &[f32], b: &[f32]| -> (f32, f32, f64) {
25069            let mut max_abs = 0f32;
25070            let mut max_rel = 0f32;
25071            let mut sum_rel = 0f64;
25072            for (x, y) in a.iter().zip(b) {
25073                let ad = (x - y).abs();
25074                let rel = ad / x.abs().max(y.abs()).max(1e-3);
25075                if ad > max_abs {
25076                    max_abs = ad;
25077                }
25078                if rel > max_rel {
25079                    max_rel = rel;
25080                }
25081                sum_rel += rel as f64;
25082            }
25083            (max_abs, max_rel, sum_rel / a.len() as f64)
25084        };
25085        let (o_ma, o_mr, o_mean) = stats(&oh_s, &oh_c);
25086        let (s_ma, s_mr, s_mean) = stats(&sh_s, &sh_c);
25087        println!(
25088            "[gdn-diff call {call:3} T={t} C={}] out: max_abs={o_ma:.3e} max_rel={o_mr:.3e} mean_rel={o_mean:.3e} | \
25089                  state: max_abs={s_ma:.3e} max_rel={s_mr:.3e} mean_rel={s_mean:.3e}",
25090            Self::gdn_chunk_size()
25091        );
25092        Ok(())
25093    }
25094
25095    /// softplus-based g_log: g_log[h,t] = a[h] * softplus(alpha[h,t] + dt_bias[h]). a pre-negated.
25096    pub fn gdn_glog(
25097        &self,
25098        alpha: &CudaSlice<f32>,
25099        dt_bias: &CudaSlice<f32>,
25100        a: &CudaSlice<f32>,
25101        g_log: &mut CudaSlice<f32>,
25102        n_head: usize,
25103        t: usize,
25104    ) -> Result<(), Box<dyn std::error::Error>> {
25105        let f = self.func("gdn_glog_f32");
25106        let cfg = LaunchConfig::for_num_elems((n_head * t) as u32);
25107        let (h, ti) = (n_head as i32, t as i32);
25108        let __s_b = self.gpu.stream();
25109        let mut b = __s_b.launch_builder(&f);
25110        b.arg(alpha).arg(dt_bias).arg(a).arg(g_log).arg(&h).arg(&ti);
25111        unsafe {
25112            b.launch(cfg)?;
25113        }
25114        Ok(())
25115    }
25116
25117    /// view twins (task #16): the batched prime's GDN core reads the CONCAT projection
25118    /// buffers at row offsets (CudaView) — same kernels, same values, no split copies.
25119    pub fn sigmoid_v(
25120        &self,
25121        x: &cudarc::driver::CudaView<f32>,
25122        y: &mut CudaSlice<f32>,
25123        n: usize,
25124    ) -> Result<(), Box<dyn std::error::Error>> {
25125        let f = self.func("sigmoid_f32");
25126        let cfg = LaunchConfig::for_num_elems(n as u32);
25127        let ni = n as i32;
25128        let __s_b = self.gpu.stream();
25129        let mut b = __s_b.launch_builder(&f);
25130        b.arg(x).arg(y).arg(&ni);
25131        unsafe {
25132            b.launch(cfg)?;
25133        }
25134        Ok(())
25135    }
25136
25137    pub fn gdn_glog_v(
25138        &self,
25139        alpha: &cudarc::driver::CudaView<f32>,
25140        dt_bias: &CudaSlice<f32>,
25141        a: &CudaSlice<f32>,
25142        g_log: &mut CudaSlice<f32>,
25143        n_head: usize,
25144        t: usize,
25145    ) -> Result<(), Box<dyn std::error::Error>> {
25146        let f = self.func("gdn_glog_f32");
25147        let cfg = LaunchConfig::for_num_elems((n_head * t) as u32);
25148        let (h, ti) = (n_head as i32, t as i32);
25149        let __s_b = self.gpu.stream();
25150        let mut b = __s_b.launch_builder(&f);
25151        b.arg(alpha).arg(dt_bias).arg(a).arg(g_log).arg(&h).arg(&ti);
25152        unsafe {
25153            b.launch(cfg)?;
25154        }
25155        Ok(())
25156    }
25157
25158    pub fn sigmoid(
25159        &self,
25160        x: &CudaSlice<f32>,
25161        y: &mut CudaSlice<f32>,
25162        n: usize,
25163    ) -> Result<(), Box<dyn std::error::Error>> {
25164        let f = self.func("sigmoid_f32");
25165        let cfg = LaunchConfig::for_num_elems(n as u32);
25166        let ni = n as i32;
25167        let __s_b = self.gpu.stream();
25168        let mut b = __s_b.launch_builder(&f);
25169        b.arg(x).arg(y).arg(&ni);
25170        unsafe {
25171            b.launch(cfg)?;
25172        }
25173        Ok(())
25174    }
25175
25176    /// attn out-gate fused epilogue (task #17): dst = a * sigmoid(g) + fp16 twin, one launch
25177    /// (replaces sigmoid + mul + convert). Bit-identical class.
25178    pub fn sig_mul_f16out(
25179        &self,
25180        a: &CudaSlice<f32>,
25181        g: &CudaSlice<f32>,
25182        dst: &mut CudaSlice<f32>,
25183        dst16: &mut CudaSlice<u8>,
25184        n: usize,
25185    ) -> Result<(), Box<dyn std::error::Error>> {
25186        let f = self.func("sig_mul_f16out_f32");
25187        let cfg = LaunchConfig::for_num_elems(n as u32);
25188        let ni = n as i32;
25189        let __s_b = self.gpu.stream();
25190        let mut b = __s_b.launch_builder(&f);
25191        b.arg(a).arg(g).arg(dst).arg(dst16).arg(&ni);
25192        unsafe {
25193            b.launch(cfg)?;
25194        }
25195        Ok(())
25196    }
25197
25198    /// step35 (Step-3.7-Flash) SEPARATE head-wise attention gate: one scalar per query head,
25199    /// broadcast over head_dim. `dst = a * sigmoid(g)` where `a`/`dst` are `[head_dim, n_head, T]`
25200    /// (the `q_gate_split` layout) and `g` is the PRE-sigmoid `attn_gate` projection output in
25201    /// token-major `[T, n_head]`. `dst16` is the optional fp16 operand for wo (None -> skipped).
25202    ///
25203    /// NOT interchangeable with `sig_mul_f16out`, which gates FULL WIDTH (qwen35 packs one gate
25204    /// value per (head, dim) element inside wq). Using this for that, or that for this, silently
25205    /// applies the wrong number of distinct gate values.
25206    #[allow(clippy::too_many_arguments)]
25207    pub fn attn_head_gate(
25208        &self,
25209        a: &CudaSlice<f32>,
25210        g: &CudaSlice<f32>,
25211        dst: &mut CudaSlice<f32>,
25212        dst16: Option<&mut CudaSlice<u8>>,
25213        head_dim: usize,
25214        n_head: usize,
25215        t: usize,
25216    ) -> Result<(), Box<dyn std::error::Error>> {
25217        let f = self.func("attn_head_gate_f32");
25218        let cfg = LaunchConfig::for_num_elems((head_dim * n_head * t) as u32);
25219        let (hd, nh, ti) = (head_dim as i32, n_head as i32, t as i32);
25220        // nullable device pointer by value (0 = skip), same convention as `l2_norm_pp`.
25221        let d16: u64 = match dst16 {
25222            Some(d) => self.addr_u8(d),
25223            None => 0,
25224        };
25225        let __s_b = self.gpu.stream();
25226        let mut b = __s_b.launch_builder(&f);
25227        b.arg(a)
25228            .arg(g)
25229            .arg(dst)
25230            .arg(&d16)
25231            .arg(&hd)
25232            .arg(&nh)
25233            .arg(&ti);
25234        unsafe {
25235            b.launch(cfg)?;
25236        }
25237        Ok(())
25238    }
25239
25240    /// step35 CLAMPED SwiGLU: `dst = min(silu(gate*gs), limit) * clamp(up*us, +-limit)`.
25241    /// Verbatim from llama.cpp `llama-graph.cpp:2146-2165` (routed, `swiglu_clamp_exp`) and
25242    /// `:1751-1770` (shared, `swiglu_clamp_shexp`), non-DEEPSEEK4 branch.
25243    ///
25244    /// This is NOT `swigluoai_mul_scaled`: that one clamps the gate BEFORE swish and multiplies by
25245    /// `(1 + clamp(up))`. Caller MUST check `limit > 1e-6` (upstream's eps gate) and use the plain
25246    /// `silu_mul_scaled` path otherwise — at limit=0 this kernel would clamp every positive
25247    /// activation to zero. On Step-3.7-Flash only layers 43 (7.0) and 44 (16.0) have a live limit.
25248    #[allow(clippy::too_many_arguments)]
25249    pub fn swiglu_clamped_mul_scaled(
25250        &self,
25251        gate: &CudaSlice<f32>,
25252        up: &CudaSlice<f32>,
25253        gs: f32,
25254        us: f32,
25255        limit: f32,
25256        dst: &mut CudaSlice<f32>,
25257        n: usize,
25258    ) -> Result<(), Box<dyn std::error::Error>> {
25259        debug_assert!(
25260            limit > 1e-6,
25261            "swiglu_clamped needs a live limit; use silu_mul_scaled"
25262        );
25263        let f = self.func("swiglu_clamped_mul_scaled_f32");
25264        let cfg = LaunchConfig::for_num_elems(n as u32);
25265        let ni = n as i32;
25266        let __s_b = self.gpu.stream();
25267        let mut b = __s_b.launch_builder(&f);
25268        b.arg(gate)
25269            .arg(up)
25270            .arg(&gs)
25271            .arg(&us)
25272            .arg(&limit)
25273            .arg(dst)
25274            .arg(&ni);
25275        unsafe {
25276            b.launch(cfg)?;
25277        }
25278        Ok(())
25279    }
25280
25281    /// gated RMSNorm: dst = RMSNorm(o, w[ncols]) * silu(z), per row of ncols. nrows blocks.
25282    pub fn gated_rmsnorm(
25283        &self,
25284        o: &CudaSlice<f32>,
25285        w: &CudaSlice<f32>,
25286        z: &CudaSlice<f32>,
25287        dst: &mut CudaSlice<f32>,
25288        ncols: usize,
25289        nrows: usize,
25290        eps: f32,
25291    ) -> Result<(), Box<dyn std::error::Error>> {
25292        let f = self.func("gated_rmsnorm_f32");
25293        let cfg = LaunchConfig {
25294            grid_dim: (nrows as u32, 1, 1),
25295            block_dim: (128, 1, 1),
25296            shared_mem_bytes: 0,
25297        };
25298        let (nc, e) = (ncols as i32, eps);
25299        let __s_b = self.gpu.stream();
25300        let mut b = __s_b.launch_builder(&f);
25301        b.arg(o).arg(w).arg(z).arg(dst).arg(&nc).arg(&e);
25302        unsafe {
25303            b.launch(cfg)?;
25304        }
25305        Ok(())
25306    }
25307
25308    /// f16out twin of `gated_rmsnorm` (task #17): epilogue also emits the fp16 operand for
25309    /// the ssm_out GEMM. Bit-identical class (same floats + the cvt kernel's __float2half).
25310    pub fn gated_rmsnorm_f16out(
25311        &self,
25312        o: &CudaSlice<f32>,
25313        w: &CudaSlice<f32>,
25314        z: &CudaSlice<f32>,
25315        dst: &mut CudaSlice<f32>,
25316        dst16: &mut CudaSlice<u8>,
25317        ncols: usize,
25318        nrows: usize,
25319        eps: f32,
25320    ) -> Result<(), Box<dyn std::error::Error>> {
25321        let f = self.func("gated_rmsnorm_f16out_f32");
25322        // block_dim MUST match gated_rmsnorm's (128): the reduction tree order pins the scale
25323        let cfg = LaunchConfig {
25324            grid_dim: (nrows as u32, 1, 1),
25325            block_dim: (128, 1, 1),
25326            shared_mem_bytes: 0,
25327        };
25328        let (nc, e) = (ncols as i32, eps);
25329        let __s_b = self.gpu.stream();
25330        let mut b = __s_b.launch_builder(&f);
25331        b.arg(o).arg(w).arg(z).arg(dst).arg(dst16).arg(&nc).arg(&e);
25332        unsafe {
25333            b.launch(cfg)?;
25334        }
25335        Ok(())
25336    }
25337
25338    /// add+RMSNorm emitting the f32 normed row AND its q8_1 quantization in one launch (the MoE
25339    /// layer input: z feeds the router matmul as f32, the expert dp4a as q8_1). BIT-IDENTICAL to
25340    /// add_rms_norm + quantize_q8_1. Returns (q, d) alongside the caller-provided res/z buffers.
25341    #[allow(clippy::too_many_arguments)]
25342    pub fn add_rms_norm_zq8(
25343        &self,
25344        a: &CudaSlice<f32>,
25345        b_in: &CudaSlice<f32>,
25346        w: &CudaSlice<f32>,
25347        res: &mut CudaSlice<f32>,
25348        z: &mut CudaSlice<f32>,
25349        ncols: usize,
25350        nrows: usize,
25351        eps: f32,
25352    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
25353        assert!(ncols % 32 == 0);
25354        let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
25355        let mut d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
25356        let f = self.func("add_rms_norm_zq8");
25357        let cfg = LaunchConfig {
25358            grid_dim: (nrows as u32, 1, 1),
25359            block_dim: (1024, 1, 1),
25360            shared_mem_bytes: 0,
25361        };
25362        let (nc, ep) = (ncols as i32, eps);
25363        let __s_b = self.gpu.stream();
25364        let mut b = __s_b.launch_builder(&f);
25365        b.arg(a)
25366            .arg(b_in)
25367            .arg(w)
25368            .arg(res)
25369            .arg(z)
25370            .arg(&mut q)
25371            .arg(&mut d)
25372            .arg(&nc)
25373            .arg(&ep);
25374        unsafe {
25375            b.launch(cfg)?;
25376        }
25377        Ok((q, d))
25378    }
25379
25380    /// gated RMSNorm emitting q8_1 directly (fused quantize epilogue) — the ssm_out matvec input.
25381    /// BIT-IDENTICAL bytes to gated_rmsnorm + quantize_q8_1 (ncols % 32 == 0; blocks never straddle
25382    /// rows). Saves one launch per linear-attn layer (36/token on the 9B).
25383    /// z-view twins of gated_rmsnorm(+f16out) — task #16 batched-prime split removal.
25384    pub fn gated_rmsnorm_zv(
25385        &self,
25386        o: &CudaSlice<f32>,
25387        w: &CudaSlice<f32>,
25388        z: &cudarc::driver::CudaView<f32>,
25389        dst: &mut CudaSlice<f32>,
25390        ncols: usize,
25391        nrows: usize,
25392        eps: f32,
25393    ) -> Result<(), Box<dyn std::error::Error>> {
25394        let f = self.func("gated_rmsnorm_f32");
25395        let cfg = LaunchConfig {
25396            grid_dim: (nrows as u32, 1, 1),
25397            block_dim: (128, 1, 1),
25398            shared_mem_bytes: 0,
25399        };
25400        let (nc, e) = (ncols as i32, eps);
25401        let __s_b = self.gpu.stream();
25402        let mut b = __s_b.launch_builder(&f);
25403        b.arg(o).arg(w).arg(z).arg(dst).arg(&nc).arg(&e);
25404        unsafe {
25405            b.launch(cfg)?;
25406        }
25407        Ok(())
25408    }
25409
25410    pub fn gated_rmsnorm_f16out_zv(
25411        &self,
25412        o: &CudaSlice<f32>,
25413        w: &CudaSlice<f32>,
25414        z: &cudarc::driver::CudaView<f32>,
25415        dst: &mut CudaSlice<f32>,
25416        dst16: &mut CudaSlice<u8>,
25417        ncols: usize,
25418        nrows: usize,
25419        eps: f32,
25420    ) -> Result<(), Box<dyn std::error::Error>> {
25421        let f = self.func("gated_rmsnorm_f16out_f32");
25422        // block_dim MUST match gated_rmsnorm's (128): the reduction tree order pins the scale
25423        let cfg = LaunchConfig {
25424            grid_dim: (nrows as u32, 1, 1),
25425            block_dim: (128, 1, 1),
25426            shared_mem_bytes: 0,
25427        };
25428        let (nc, e) = (ncols as i32, eps);
25429        let __s_b = self.gpu.stream();
25430        let mut b = __s_b.launch_builder(&f);
25431        b.arg(o).arg(w).arg(z).arg(dst).arg(dst16).arg(&nc).arg(&e);
25432        unsafe {
25433            b.launch(cfg)?;
25434        }
25435        Ok(())
25436    }
25437
25438    pub fn gated_rmsnorm_q8_1(
25439        &self,
25440        o: &CudaSlice<f32>,
25441        w: &CudaSlice<f32>,
25442        z: &CudaSlice<f32>,
25443        ncols: usize,
25444        nrows: usize,
25445        eps: f32,
25446    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
25447        assert!(ncols % 32 == 0);
25448        let f = self.func("gated_rmsnorm_q8_1");
25449        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
25450        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
25451        let cfg = LaunchConfig {
25452            grid_dim: (nrows as u32, 1, 1),
25453            block_dim: (128, 1, 1),
25454            shared_mem_bytes: 0,
25455        };
25456        let (nc, ep) = (ncols as i32, eps);
25457        let __s_b = self.gpu.stream();
25458        let mut b = __s_b.launch_builder(&f);
25459        b.arg(o)
25460            .arg(w)
25461            .arg(z)
25462            .arg(&mut out_q)
25463            .arg(&mut out_d)
25464            .arg(&nc)
25465            .arg(&ep);
25466        unsafe {
25467            b.launch(cfg)?;
25468        }
25469        Ok((out_q, out_d))
25470    }
25471
25472    /// transpose [rows,cols] row-major -> [cols,rows] row-major.
25473    pub fn transpose(
25474        &self,
25475        inp: &CudaSlice<f32>,
25476        rows: usize,
25477        cols: usize,
25478    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
25479        let f = self.func("transpose_f32");
25480        let mut out = self.zeros(rows * cols)?;
25481        let cfg = LaunchConfig::for_num_elems((rows * cols) as u32);
25482        let (r, c) = (rows as i32, cols as i32);
25483        let __s_b = self.gpu.stream();
25484        let mut b = __s_b.launch_builder(&f);
25485        b.arg(inp).arg(&mut out).arg(&r).arg(&c);
25486        unsafe {
25487            b.launch(cfg)?;
25488        }
25489        Ok(out)
25490    }
25491
25492    /// repeat-interleave heads: in[head_dim,n_in,T] -> out[head_dim,n_out,T].
25493    pub fn repeat_heads(
25494        &self,
25495        inp: &CudaSlice<f32>,
25496        out: &mut CudaSlice<f32>,
25497        head_dim: usize,
25498        n_in: usize,
25499        n_out: usize,
25500        t: usize,
25501    ) -> Result<(), Box<dyn std::error::Error>> {
25502        let f = self.func("repeat_heads_f32");
25503        let cfg = LaunchConfig::for_num_elems((head_dim * n_out * t) as u32);
25504        let (hd, ni, no, ti) = (head_dim as i32, n_in as i32, n_out as i32, t as i32);
25505        let __s_b = self.gpu.stream();
25506        let mut b = __s_b.launch_builder(&f);
25507        b.arg(inp).arg(out).arg(&hd).arg(&ni).arg(&no).arg(&ti);
25508        unsafe {
25509            b.launch(cfg)?;
25510        }
25511        Ok(())
25512    }
25513
25514    /// q|gate split (on-device). qf:[T, n_head*2*head_dim] -> q_out,gate_out:[head_dim,n_head,T].
25515    /// Replaces the dtoh->host-double-loop->htod in full_attn / full_attn_decode.
25516    ///
25517    /// Bounds-checked against `qf` before the launch: the kernel reads `2*head_dim*n_head*t`
25518    /// floats, and running it on a `wq` output that carries no fused gate reads 2x off the end
25519    /// (silently, on the device). A layout mismatch is a typed `FusedQGateExtent` here instead.
25520    pub fn q_gate_split(
25521        &self,
25522        qf: &CudaSlice<f32>,
25523        q_out: &mut CudaSlice<f32>,
25524        gate_out: &mut CudaSlice<f32>,
25525        head_dim: usize,
25526        n_head: usize,
25527        t: usize,
25528    ) -> Result<(), Box<dyn std::error::Error>> {
25529        memra_gguf::config::check_fused_q_gate_extent(qf.len(), head_dim, n_head, t)?;
25530        let out_need = head_dim * n_head * t;
25531        if q_out.len() < out_need || gate_out.len() < out_need {
25532            return Err(format!(
25533                "q_gate_split destinations too small: need {out_need} each, have q={} gate={}",
25534                q_out.len(),
25535                gate_out.len()
25536            )
25537            .into());
25538        }
25539        let f = self.func("q_gate_split_f32");
25540        let cfg = LaunchConfig::for_num_elems((head_dim * n_head * t) as u32);
25541        let (hd, nh, ti) = (head_dim as i32, n_head as i32, t as i32);
25542        let __s_b = self.gpu.stream();
25543        let mut b = __s_b.launch_builder(&f);
25544        b.arg(qf)
25545            .arg(q_out)
25546            .arg(gate_out)
25547            .arg(&hd)
25548            .arg(&nh)
25549            .arg(&ti);
25550        unsafe {
25551            b.launch(cfg)?;
25552        }
25553        Ok(())
25554    }
25555
25556    /// qkv->GDN repack (on-device). conv_out:[conv_dim,T] channel-major ->
25557    /// q_g/k_g/v_g:[d_state,num_v,T] with q/k head-repeat kh = vh % num_k (validated modulo mapping).
25558    /// Replaces the dtoh->host-q/k/v-repack->3x-htod in linear_attn / linear_attn_decode.
25559    pub fn qkv_to_gdn_repack(
25560        &self,
25561        conv_out: &CudaSlice<f32>,
25562        q_g: &mut CudaSlice<f32>,
25563        k_g: &mut CudaSlice<f32>,
25564        v_g: &mut CudaSlice<f32>,
25565        d_state: usize,
25566        num_v: usize,
25567        num_k: usize,
25568        key_dim: usize,
25569        t: usize,
25570    ) -> Result<(), Box<dyn std::error::Error>> {
25571        let f = self.func("qkv_to_gdn_repack_f32");
25572        let cfg = LaunchConfig::for_num_elems((d_state * num_v * t) as u32);
25573        let (ds, nv, nk, kd, ti) = (
25574            d_state as i32,
25575            num_v as i32,
25576            num_k as i32,
25577            key_dim as i32,
25578            t as i32,
25579        );
25580        let __s_b = self.gpu.stream();
25581        let mut b = __s_b.launch_builder(&f);
25582        b.arg(conv_out)
25583            .arg(q_g)
25584            .arg(k_g)
25585            .arg(v_g)
25586            .arg(&ds)
25587            .arg(&nv)
25588            .arg(&nk)
25589            .arg(&kd)
25590            .arg(&ti);
25591        unsafe {
25592            b.launch(cfg)?;
25593        }
25594        Ok(())
25595    }
25596
25597    /// conv left zero-pad (prefill from zero state). src:[conv_dim,T] -> dst:[conv_dim,T+pad],
25598    /// cols 0..pad = 0, cols pad..pad+T = src. `dst` MUST be pre-zeroed. No dtoh/host-loop/htod.
25599    pub fn conv_left_pad(
25600        &self,
25601        src: &CudaSlice<f32>,
25602        dst: &mut CudaSlice<f32>,
25603        conv_dim: usize,
25604        t: usize,
25605        pad: usize,
25606    ) -> Result<(), Box<dyn std::error::Error>> {
25607        let f = self.func("conv_left_pad_f32");
25608        let cfg = LaunchConfig::for_num_elems((conv_dim * t) as u32);
25609        let (cd, ti, p) = (conv_dim as i32, t as i32, pad as i32);
25610        let __s_b = self.gpu.stream();
25611        let mut b = __s_b.launch_builder(&f);
25612        b.arg(src).arg(dst).arg(&cd).arg(&ti).arg(&p);
25613        unsafe {
25614            b.launch(cfg)?;
25615        }
25616        Ok(())
25617    }
25618
25619    /// conv-state assemble + ring roll (decode T=1). conv_state:[conv_dim,pad] (resident),
25620    /// qkv_col:[conv_dim] -> conv_in:[conv_dim,pad+1]; AND rolls conv_state (keep last pad cols).
25621    /// Replaces the dtoh->host-conv-ring-assemble->ring-update->htod in linear_attn_decode.
25622    pub fn conv_assemble_and_roll(
25623        &self,
25624        qkv_col: &CudaSlice<f32>,
25625        conv_state: &mut CudaSlice<f32>,
25626        conv_in: &mut CudaSlice<f32>,
25627        conv_dim: usize,
25628        pad: usize,
25629    ) -> Result<(), Box<dyn std::error::Error>> {
25630        let f = self.func("conv_assemble_and_roll_f32");
25631        let cfg = LaunchConfig::for_num_elems(conv_dim as u32);
25632        let (cd, p) = (conv_dim as i32, pad as i32);
25633        let __s_b = self.gpu.stream();
25634        let mut b = __s_b.launch_builder(&f);
25635        b.arg(qkv_col).arg(conv_state).arg(conv_in).arg(&cd).arg(&p);
25636        unsafe {
25637            b.launch(cfg)?;
25638        }
25639        Ok(())
25640    }
25641
25642    /// RANK3 LEVER (conv fuse, T=1 DECODE): fused conv_assemble_and_roll + ssm_conv1d_silu in ONE
25643    /// launch. Assembles the conv window [conv_state | qkv_col] in registers, computes the depthwise
25644    /// causal conv + SiLU into `conv_out`, and rolls the ring — never materializing conv_in to HBM.
25645    /// Replaces e.conv_assemble_and_roll(...) + e.ssm_conv1d(...). BIT-IDENTICAL to that two-kernel
25646    /// sequence (same 8-wide accumulation order, same SiLU). `conv_out` is [conv_dim] (T=1).
25647    pub fn ssm_conv1d_fused_decode(
25648        &self,
25649        qkv_col: &CudaSlice<f32>,
25650        conv_state: &mut CudaSlice<f32>,
25651        w: &CudaSlice<f32>,
25652        conv_out: &mut CudaSlice<f32>,
25653        conv_dim: usize,
25654        d_conv: usize,
25655    ) -> Result<(), Box<dyn std::error::Error>> {
25656        let f = self.func("ssm_conv1d_fused_decode_f32");
25657        let cfg = LaunchConfig::for_num_elems(conv_dim as u32);
25658        let (cd, dc) = (conv_dim as i32, d_conv as i32);
25659        let __s_b = self.gpu.stream();
25660        let mut b = __s_b.launch_builder(&f);
25661        b.arg(qkv_col)
25662            .arg(conv_state)
25663            .arg(w)
25664            .arg(conv_out)
25665            .arg(&cd)
25666            .arg(&dc);
25667        unsafe {
25668            b.launch(cfg)?;
25669        }
25670        Ok(())
25671    }
25672
25673    /// Copy a contiguous range [start, start+len) out of src into a fresh slice (device→device via host).
25674    /// Used for qkv split views. Small/rare; not perf-critical in Stage 1.
25675    pub fn slice_range(
25676        &self,
25677        src: &CudaSlice<f32>,
25678        start: usize,
25679        len: usize,
25680    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
25681        let host = self.gpu.stream().clone_dtoh(src)?;
25682        self.gpu.stream().synchronize()?;
25683        Ok(self.htod(&host[start..start + len])?)
25684    }
25685}
25686
25687#[cfg(test)]
25688mod target_dispatch_tests {
25689    use super::legacy_quant_gemm_allowed;
25690
25691    #[test]
25692    fn legacy_quant_gemm_arch_policy_honors_the_escape_hatch() {
25693        // sm_120a native lane
25694        assert!(legacy_quant_gemm_allowed(false, false, false));
25695        assert!(!legacy_quant_gemm_allowed(false, false, true));
25696        // pure portable lane (sm_89): gated
25697        assert!(!legacy_quant_gemm_allowed(true, false, false));
25698        assert!(!legacy_quant_gemm_allowed(true, false, true));
25699        // Hopper-MMA lane (sm_90a): portable build, int8-MMA GEMM re-admitted
25700        assert!(legacy_quant_gemm_allowed(true, true, false));
25701        assert!(!legacy_quant_gemm_allowed(true, true, true));
25702    }
25703
25704    #[cfg(all(memra_portable_cuda, not(memra_hopper_mma)))]
25705    #[test]
25706    fn portable_build_disables_legacy_quant_gemm_without_an_env_override() {
25707        assert!(!legacy_quant_gemm_allowed(
25708            cfg!(memra_portable_cuda),
25709            cfg!(memra_hopper_mma),
25710            false
25711        ));
25712    }
25713
25714    #[cfg(memra_hopper_mma)]
25715    #[test]
25716    fn hopper_mma_build_re_admits_legacy_quant_gemm() {
25717        assert!(legacy_quant_gemm_allowed(
25718            cfg!(memra_portable_cuda),
25719            cfg!(memra_hopper_mma),
25720            false
25721        ));
25722        assert!(super::portable_mma_gated() == false);
25723    }
25724}
25725
25726/// The memra-kv device seam (Phase D): the cache's 7 ops delegate to the engine's
25727/// inherent methods (inherent methods win name resolution, so no recursion).
25728impl memra_kv::KvDev for Engine {
25729    fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
25730        Engine::zeros(self, n)
25731    }
25732    fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
25733        Engine::uninit(self, n)
25734    }
25735    fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
25736        Engine::alloc_u8(self, n)
25737    }
25738    fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
25739        Engine::htod_i32(self, v)
25740    }
25741    fn clone_dtod(
25742        &self,
25743        src: &CudaSlice<f32>,
25744    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
25745        Engine::clone_dtod(self, src)
25746    }
25747    fn copy_into(
25748        &self,
25749        dst: &mut CudaSlice<f32>,
25750        off: usize,
25751        src: &CudaSlice<f32>,
25752        len: usize,
25753    ) -> Result<(), Box<dyn std::error::Error>> {
25754        Engine::copy_into(self, dst, off, src, len)
25755    }
25756    fn set_i32_one(
25757        &self,
25758        d: &mut CudaSlice<i32>,
25759        v: i32,
25760    ) -> Result<(), Box<dyn std::error::Error>> {
25761        Engine::set_i32_one(self, d, v)
25762    }
25763}
25764
25765#[cfg(test)]
25766mod fused_gate_bounds_tests {
25767    use super::*;
25768
25769    /// The fused `[q|gate]` split's read-site guard, on the device.
25770    ///
25771    /// `q_gate_split_f32` reads `2*head_dim*n_head*T` floats out of `qf`. A checkpoint whose gate
25772    /// is a SEPARATE tensor produces a `wq` output of exactly half that, so before 2026-08-19 the
25773    /// kernel launched and read 2x past the end of the allocation — an out-of-bounds DEVICE read:
25774    /// no panic, no error, just whatever memory follows. The guard turns it into a typed
25775    /// `FusedQGateExtent` before the launch.
25776    ///
25777    /// Catch demonstration for this test (guard temporarily removed, then restored):
25778    /// `compute-sanitizer --tool memcheck` on the half-width case reported invalid `__global__`
25779    /// reads of size 4 in `q_gate_split_f32`; with the guard in place the same run is clean and
25780    /// the call returns `Err`. Receipt in the lane report.
25781    #[test]
25782    #[ignore = "requires a CUDA GPU"]
25783    fn q_gate_split_refuses_a_separate_gate_wq_instead_of_reading_past_it() {
25784        let e = Engine::new(0).unwrap();
25785        let (head_dim, n_head, t) = (8usize, 4usize, 2usize);
25786        let fused = 2 * head_dim * n_head * t;
25787        let out_n = head_dim * n_head * t;
25788
25789        // half-width `qf` = the separate-gate / ungated layout. MUST be refused.
25790        let narrow = e.htod(&vec![1.0f32; out_n]).unwrap();
25791        let mut q = e.uninit(out_n).unwrap();
25792        let mut gate = e.uninit(out_n).unwrap();
25793        let err = e
25794            .q_gate_split(&narrow, &mut q, &mut gate, head_dim, n_head, t)
25795            .expect_err("half-width wq must be refused, not read past")
25796            .to_string();
25797        assert!(err.contains("NO fused gate"), "{err}");
25798        assert!(err.contains(&format!("{fused}")), "{err}");
25799
25800        // full-width `qf` = a real qwen3.5 fused layout. MUST still run, and split correctly:
25801        // per head hh the block is [q(head_dim) | gate(head_dim)] at stride 2*head_dim.
25802        let host: Vec<f32> = (0..fused).map(|i| i as f32).collect();
25803        let wide = e.htod(&host).unwrap();
25804        e.q_gate_split(&wide, &mut q, &mut gate, head_dim, n_head, t)
25805            .expect("full-width wq splits");
25806        let (qh, gh) = (e.dtoh(&q).unwrap(), e.dtoh(&gate).unwrap());
25807        for tok in 0..t {
25808            for hh in 0..n_head {
25809                for d in 0..head_dim {
25810                    let base = tok * (n_head * 2 * head_dim) + hh * (2 * head_dim);
25811                    let idx = tok * (n_head * head_dim) + hh * head_dim + d;
25812                    assert_eq!(qh[idx], host[base + d], "q t{tok} h{hh} d{d}");
25813                    assert_eq!(gh[idx], host[base + head_dim + d], "gate t{tok} h{hh} d{d}");
25814                }
25815            }
25816        }
25817
25818        // undersized destinations are refused too (the other half of the extent contract)
25819        let mut small = e.uninit(out_n - 1).unwrap();
25820        assert!(
25821            e.q_gate_split(&wide, &mut small, &mut gate, head_dim, n_head, t)
25822                .is_err()
25823        );
25824    }
25825}
25826
25827/// FULL-WIDTH-ROPE CONTRACT on the fused rms_norm+qkv+rope kernels
25828/// (lane/graph-s-key-exactness-20260819, probe O-10 follow-up). CPU-only: the guard runs before
25829/// any launch, so the refusal is testable without a device.
25830#[cfg(test)]
25831mod fused_rope_width_tests {
25832    use super::Engine;
25833
25834    /// gemma-4: rotary width == head width on both classes (GGUF 256/256 and 512/512, and the
25835    /// safetensors route derives the same), which is why the fusion is legal there today.
25836    #[test]
25837    fn full_width_is_accepted() {
25838        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 256, 256).is_ok());
25839        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_cat", 512, 512).is_ok());
25840        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append", 128, 128).is_ok());
25841    }
25842
25843    /// The widths the gemma-4 31B OFFICIAL artifact declares, read from its own GGUF header
25844    /// (`gemma-4-31B-it-official-Q8_0-MTP.gguf`, box3, 2026-08-19):
25845    ///
25846    /// ```text
25847    /// attention.key_length     512   rope.dimension_count     512   (global class)
25848    /// attention.key_length_swa 256   rope.dimension_count_swa 256   (SWA class)
25849    /// ```
25850    ///
25851    /// Both classes satisfy `n_rot == head_dim`, which is why the fusion is legal for gemma and
25852    /// why `HybridModel::gemma4_rope_dims` can feed this guard without refusing what we serve.
25853    /// An artifact that ever declares otherwise gets a loud refusal at the first fused launch
25854    /// instead of a silently over-rotated head.
25855    #[test]
25856    fn gemma4_official_artifact_widths_pass() {
25857        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 512, 512).is_ok());
25858        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append_dc", 256, 256).is_ok());
25859    }
25860
25861    /// The bug this guard exists to make impossible: a partial-rotary arch fused onto a kernel
25862    /// with no `n_dims`, silently rotating the pass-through band.
25863    #[test]
25864    fn partial_rotary_is_refused_with_the_geometry_named() {
25865        // qwen3.5: n_rot 64 of head_dim 256 (the shape probe O-10 pinned in the split path).
25866        let err = Engine::full_width_rope_only("rms_norm_qkv_rope", 64, 256)
25867            .expect_err("partial rotary must refuse");
25868        let msg = err.to_string();
25869        assert!(msg.contains("PARTIAL ROTARY REFUSED"), "{msg}");
25870        assert!(msg.contains("n_rot 64"), "{msg}");
25871        assert!(msg.contains("head_dim 256"), "{msg}");
25872        assert!(
25873            msg.contains("64..256"),
25874            "names the band it would corrupt: {msg}"
25875        );
25876        // step35 full-attn: 64 of 128 (upstream halves n_rot_full).
25877        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append_dc", 64, 128).is_err());
25878        // and the reverse mismatch (a wider rope than the head) is not "close enough" either.
25879        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 256, 128).is_err());
25880    }
25881}